Octree ghost zones — reloaded

tests/tests.yaml

@@ -142,7 +142,7 @@ answer_tests:
     - yt/frontends/boxlib/tests/test_outputs.py:test_NyxDataset
     - yt/frontends/boxlib/tests/test_outputs.py:test_WarpXDataset
 
-  local_ramses_002:
+  local_ramses_003:
     - yt/frontends/ramses/tests/test_outputs.py:test_output_00080
 
   local_ytdata_007:

yt/data_objects/data_containers.py

@@ -353,7 +353,9 @@ def _generate_spatial_fluid(self, field, ngz):
                         outputs.append(rv)
                         ind = 0  # Does this work with mesh?
                     with o._activate_cache():
-                        ind += o.select(self.selector, self[field], rv, ind)
+                        ind += o.select(
+                            self.selector, source=self[field], dest=rv, offset=ind
+                        )
         else:
             chunks = self.index._chunk(self, "spatial", ngz=ngz)
             for i, chunk in enumerate(chunks):
@@ -366,15 +368,12 @@ def _generate_spatial_fluid(self, field, ngz):
                             np.empty(wogz.ires.size, dtype="float64"), units
                         )
                         outputs.append(rv)
-                    if gz._type_name == "octree_subset":
-                        raise NotImplementedError
-                    else:
-                        ind += wogz.select(
-                            self.selector,
-                            gz[field][ngz:-ngz, ngz:-ngz, ngz:-ngz],
-                            rv,
-                            ind,
-                        )
+                    ind += wogz.select(
+                        self.selector,
+                        source=gz[field][ngz:-ngz, ngz:-ngz, ngz:-ngz],
+                        dest=rv,
+                        offset=ind,
+                    )
         if accumulate:
             rv = uconcatenate(outputs)
         return rv

yt/data_objects/octree_subset.py

@@ -37,11 +37,14 @@ class OctreeSubset(YTSelectionContainer):
     _con_args = ("base_region", "domain", "ds")
     _domain_offset = 0
     _cell_count = -1
-    _block_reorder = None
+    _block_order = "C"
 
-    def __init__(self, base_region, domain, ds, over_refine_factor=1):
+    def __init__(
+        self, base_region, domain, ds, over_refine_factor=1, num_ghost_zones=0
+    ):
         super(OctreeSubset, self).__init__(ds, None)
         self._num_zones = 1 << (over_refine_factor)
+        self._num_ghost_zones = num_ghost_zones
         self._oref = over_refine_factor
         self.domain = domain
         self.domain_id = domain.domain_id
@@ -528,9 +531,10 @@ def __init__(self, ind, block_slice):
 
     def __getitem__(self, key):
         bs = self.block_slice
-        rv = bs.octree_subset[key][:, :, :, self.ind].T
-        if bs.octree_subset._block_reorder:
-            rv = rv.copy(order=bs.octree_subset._block_reorder)
+        rv = np.require(
+            bs.octree_subset[key][:, :, :, self.ind].T,
+            requirements=bs.octree_subset._block_order,
+        )
         return rv
 
     @property

yt/fields/fluid_fields.py

@@ -221,8 +221,8 @@ def setup_gradient_fields(registry, grad_field, field_units, slice_info=None):
     geom = registry.ds.geometry
     if is_curvilinear(geom):
         mylog.warning(
-            "In %s geometry, gradient fields may contain artifacts near cartesian axes."
-            % geom
+            "In %s geometry, gradient fields may contain artifacts near cartesian axes.",
+            geom,
         )
 
     assert isinstance(grad_field, tuple)
@@ -240,16 +240,27 @@ def grad_func(axi, ax):
         slice_3dr = slice_3d[:axi] + (sl_right,) + slice_3d[axi + 1 :]
 
         def func(field, data):
-            ds = div_fac * data[ftype, "d%s" % ax]
+            block_order = getattr(data, "_block_order", "C")
+            if block_order == "F":
+                # Fortran-ordering: we need to swap axes here and
+                # reswap below
+                field_data = data[grad_field].swapaxes(0, 2)
+            else:
+                field_data = data[grad_field]
+            dx = div_fac * data[ftype, f"d{ax}"]
             if ax == "theta":
-                ds *= data[ftype, "r"]
+                dx *= data[ftype, "r"]
             if ax == "phi":
-                ds *= data[ftype, "r"] * np.sin(data[ftype, "theta"])
-            f = data[grad_field][slice_3dr] / ds[slice_3d]
-            f -= data[grad_field][slice_3dl] / ds[slice_3d]
+                dx *= data[ftype, "r"] * np.sin(data[ftype, "theta"])
+            f = field_data[slice_3dr] / dx[slice_3d]
+            f -= field_data[slice_3dl] / dx[slice_3d]
             new_field = np.zeros_like(data[grad_field], dtype=np.float64)
-            new_field = data.ds.arr(new_field, f.units)
+            new_field = data.ds.arr(new_field, field_data.units / dx.units)
             new_field[slice_3d] = f
+
+            if block_order == "F":
+                new_field = new_field.swapaxes(0, 2)
+
             return new_field
 
         return func

yt/fields/geometric_fields.py

@@ -101,10 +101,11 @@ def _zeros(field, data):
 
     def _ones(field, data):
         """Returns one for all cells"""
-        arr = np.ones(data.ires.shape, dtype="float64")
+        tmp = np.ones(data.ires.shape, dtype="float64")
+        arr = data.apply_units(tmp, field.units)
         if data._spatial:
             return data._reshape_vals(arr)
-        return data.apply_units(arr, field.units)
+        return arr
 
     registry.add_field(
         ("index", "ones"),

yt/frontends/ramses/data_structures.py

@@ -185,16 +185,47 @@ def included(self, selector):
 class RAMSESDomainSubset(OctreeSubset):
 
     _domain_offset = 1
-    _block_reorder = "F"
+    _block_order = "F"
 
-    def fill(self, fd, fields, selector, file_handler):
+    _base_domain = None
+
+    def __init__(
+        self,
+        base_region,
+        domain,
+        ds,
+        over_refine_factor=1,
+        num_ghost_zones=0,
+        base_grid=None,
+    ):
+        super(RAMSESDomainSubset, self).__init__(
+            base_region, domain, ds, over_refine_factor, num_ghost_zones
+        )
+
+        self._base_grid = base_grid
+
+        if num_ghost_zones > 0:
+            if not all(ds.periodicity):
+                mylog.warn("Ghost zones will wrongly assume the domain to be periodic.")
+            # Create a base domain *with no self._base_domain.fwidth
+            base_domain = RAMSESDomainSubset(
+                ds.all_data(), domain, ds, over_refine_factor
+            )
+            self._base_domain = base_domain
+        elif num_ghost_zones < 0:
+            raise RuntimeError(
+                "Cannot initialize a domain subset with a negative number of ghost zones,"
+                " was called with num_ghost_zones=%s" % num_ghost_zones
+            )
+
+    def _fill_no_ghostzones(self, fd, fields, selector, file_handler):
         ndim = self.ds.dimensionality
         # Here we get a copy of the file, which we skip through and read the
         # bits we want.
         oct_handler = self.oct_handler
         all_fields = [f for ft, f in file_handler.field_list]
         fields = [f for ft, f in fields]
-        tr = {}
+        data = {}
         cell_count = selector.count_oct_cells(self.oct_handler, self.domain_id)
 
         levels, cell_inds, file_inds = self.oct_handler.file_index_octs(
@@ -203,12 +234,59 @@ def fill(self, fd, fields, selector, file_handler):
 
         # Initializing data container
         for field in fields:
-            tr[field] = np.zeros(cell_count, "float64")
+            data[field] = np.zeros(cell_count, "float64")
+
+        cpu_list = [self.domain_id - 1]
+        fill_hydro(
+            fd,
+            file_handler.offset,
+            file_handler.level_count,
+            cpu_list,
+            levels,
+            cell_inds,
+            file_inds,
+            ndim,
+            all_fields,
+            fields,
+            data,
+            oct_handler,
+        )
+        return data
+
+    def _fill_with_ghostzones(
+        self, fd, fields, selector, file_handler, num_ghost_zones
+    ):
+        ndim = self.ds.dimensionality
+        ncpu = self.ds.parameters["ncpu"]
+        # Here we get a copy of the file, which we skip through and read the
+        # bits we want.
+        oct_handler = self.oct_handler
+        all_fields = [f for ft, f in file_handler.field_list]
+        fields = [f for ft, f in fields]
+        tr = {}
+
+        cell_count = (
+            selector.count_octs(self.oct_handler, self.domain_id) * self.nz ** ndim
+        )
+
+        (
+            levels,
+            cell_inds,
+            file_inds,
+            domains,
+        ) = self.oct_handler.file_index_octs_with_ghost_zones(
+            selector, self.domain_id, cell_count
+        )
 
+        # Initializing data container
+        for field in fields:
+            tr[field] = np.zeros(cell_count, "float64")
+        cpu_list = list(range(ncpu))
         fill_hydro(
             fd,
             file_handler.offset,
             file_handler.level_count,
+            cpu_list,
             levels,
             cell_inds,
             file_inds,
@@ -217,9 +295,59 @@ def fill(self, fd, fields, selector, file_handler):
             fields,
             tr,
             oct_handler,
+            domains=domains,
         )
         return tr
 
+    @property
+    def fwidth(self):
+        fwidth = super(RAMSESDomainSubset, self).fwidth
+        if self._num_ghost_zones > 0:
+            fwidth = fwidth.reshape(-1, 8, 3)
+            n_oct = fwidth.shape[0]
+            # new_fwidth contains the fwidth of the oct+ghost zones
+            # this is a constant array in each oct, so we simply copy
+            # the oct value using numpy fancy-indexing
+            new_fwidth = np.zeros((n_oct, self.nz ** 3, 3), dtype=fwidth.dtype)
+            new_fwidth[:, :, :] = fwidth[:, 0:1, :]
+            fwidth = new_fwidth.reshape(-1, 3)
+        return fwidth
+
+    @property
+    def fcoords(self):
+        num_ghost_zones = self._num_ghost_zones
+        if num_ghost_zones == 0:
+            return super(RAMSESDomainSubset, self).fcoords
+
+        oh = self.oct_handler
+
+        indices = oh.fill_index(self.selector).reshape(-1, 8)
+        oct_inds, cell_inds = oh.fill_octcellindex_neighbours(self.selector)
+
+        oct_inds = oct_inds.reshape(-1, 64)
+        cell_inds = cell_inds.reshape(-1, 64)
+
+        inds = indices[oct_inds, cell_inds]
+
+        fcoords = self.ds.arr(oh.fcoords(self.selector)[inds].reshape(-1, 3), "unitary")
+
+        return fcoords
+
+    def fill(self, fd, fields, selector, file_handler):
+        if self._num_ghost_zones == 0:
+            return self._fill_no_ghostzones(fd, fields, selector, file_handler)
+        else:
+            return self._fill_with_ghostzones(
+                fd, fields, selector, file_handler, self._num_ghost_zones
+            )
+
+    def retrieve_ghost_zones(self, ngz, fields, smoothed=False):
+        new_subset = RAMSESDomainSubset(
+            self.base_region, self.domain, self.ds, num_ghost_zones=ngz, base_grid=self
+        )
+
+        return new_subset
+
 
 class RAMSESIndex(OctreeIndex):
     def __init__(self, ds, dataset_type="ramses"):
@@ -275,7 +403,12 @@ def _identify_base_chunk(self, dobj):
             if len(domains) > 1:
                 mylog.debug("Identified %s intersecting domains", len(domains))
             subsets = [
-                RAMSESDomainSubset(base_region, domain, self.dataset)
+                RAMSESDomainSubset(
+                    base_region,
+                    domain,
+                    self.dataset,
+                    num_ghost_zones=dobj._num_ghost_zones,
+                )
                 for domain in domains
             ]
             dobj._chunk_info = subsets