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

@@ -340,7 +340,11 @@ 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):
@@ -352,13 +356,12 @@ def _generate_spatial_fluid(self, field, ngz):
                         rv = self.ds.arr(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)
+                    data = gz[field][ngz:-ngz, ngz:-ngz, ngz:-ngz]
+                    ind += wogz.select(
+                        self.selector,
+                        source=data,
+                        dest=rv,
+                        offset=ind)
         if accumulate:
             rv = uconcatenate(outputs)
         return rv

yt/data_objects/octree_subset.py

@@ -35,11 +35,12 @@ 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
@@ -455,8 +456,8 @@ 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)
+        if bs.octree_subset._block_order:
+            rv = rv.copy(order=bs.octree_subset._block_order)
         return rv
 
     @property

yt/fields/fluid_fields.py

@@ -65,7 +65,7 @@ def _cell_mass(field, data):
     def _sound_speed(field, data):
         tr = data.ds.gamma * data[ftype, "pressure"] / data[ftype, "density"]
         return np.sqrt(tr)
-    
+
     registry.add_field((ftype, "sound_speed"),
                        sampling_type="local",
                        function=_sound_speed,
@@ -75,7 +75,7 @@ def _radial_mach_number(field, data):
         """ Radial component of M{|v|/c_sound} """
         tr = data[ftype, "radial_velocity"] / data[ftype, "sound_speed"]
         return np.abs(tr)
-    
+
     registry.add_field((ftype, "radial_mach_number"),
                        sampling_type="local",
                        function=_radial_mach_number,
@@ -93,7 +93,7 @@ def _kin_energy(field, data):
     def _mach_number(field, data):
         """ M{|v|/c_sound} """
         return data[ftype, "velocity_magnitude"] / data[ftype, "sound_speed"]
-    
+
     registry.add_field((ftype, "mach_number"),
                        sampling_type="local",
                        function=_mach_number,
@@ -189,7 +189,7 @@ 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)
+        mylog.warning("In %s geometry, gradient fields may contain artifacts near cartesian axes.", geom)
 
     assert(isinstance(grad_field, tuple))
     ftype, fname = grad_field
@@ -205,16 +205,27 @@ def grad_func(axi, ax):
         slice_3dl = slice_3d[:axi] + (sl_left,) + slice_3d[axi+1:]
         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

@@ -93,10 +93,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"),
                        sampling_type="cell",

yt/frontends/ramses/data_structures.py

@@ -1,6 +1,7 @@
 import os
 import numpy as np
 import weakref
+import warnings
 from collections import defaultdict
 from glob import glob
 
@@ -9,6 +10,7 @@
     setdefaultattr
 from yt.geometry.oct_geometry_handler import \
     OctreeIndex
+
 from yt.geometry.geometry_handler import \
     YTDataChunk
 from yt.data_objects.static_output import \
@@ -179,9 +181,24 @@ 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):
+                warnings.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
+
+    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.
@@ -197,13 +214,85 @@ def fill(self, fd, fields, selector, file_handler):
         # Initializing data container
         for field in fields:
             tr[field] = np.zeros(cell_count, 'float64')
-
         fill_hydro(fd, file_handler.offset,
-                   file_handler.level_count, levels, cell_inds,
+                   file_handler.level_count,
+                   [self.domain_id-1],
+                   levels, cell_inds,
                    file_inds, ndim, all_fields, fields, tr,
                    oct_handler)
         return tr
 
+    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')
+        fill_hydro(fd, file_handler.offset,
+                   file_handler.level_count,
+                   list(range(ncpu)),
+                   levels, cell_inds,
+                   file_inds, ndim, all_fields, 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 = 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)
+        oinds, cinds = oh.fill_octcellindex_neighbours(self.selector)
+
+        oinds = oinds.reshape(-1, 64)
+        cinds = cinds.reshape(-1, 64)
+
+        inds = indices[oinds, cinds]
+
+        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'):
@@ -260,7 +349,9 @@ def _identify_base_chunk(self, dobj):
             base_region = getattr(dobj, "base_region", dobj)
             if len(domains) > 1:
                 mylog.debug("Identified %s intersecting domains", len(domains))
-            subsets = [RAMSESDomainSubset(base_region, domain, self.dataset)
+            subsets = [RAMSESDomainSubset(
+                            base_region, domain, self.dataset,
+                            num_ghost_zones=dobj._num_ghost_zones)
                        for domain in domains]
             dobj._chunk_info = subsets
         dobj._current_chunk = list(self._chunk_all(dobj))[0]