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

@@ -352,13 +352,11 @@ 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,
+                        data,
+                        rv, ind)
         if accumulate:
             rv = uconcatenate(outputs)
         return rv

yt/data_objects/octree_subset.py

@@ -37,9 +37,10 @@ class OctreeSubset(YTSelectionContainer):
     _cell_count = -1
     _block_reorder = None
 
-    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

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,25 @@ 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):
+            block_reorder = getattr(data, '_block_reorder', None)
+            if block_reorder == 'F':
+                field_data = data[grad_field].swapaxes(0, 2)
+            else:
+                field_data = data[grad_field]
             ds = div_fac * data[ftype, "d%s" % ax]
             if ax == "theta":
                 ds *= 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]
+            f  = field_data[slice_3dr]/ds[slice_3d]
+            f -= field_data[slice_3dl]/ds[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 / ds.units)
             new_field[slice_3d] = f
+
+            if block_reorder == 'F':
+                new_field = new_field.swapaxes(0, 2)
+
             return new_field
         return func
 

yt/fields/geometric_fields.py

@@ -94,9 +94,10 @@ def _zeros(field, data):
     def _ones(field, data):
         """Returns one for all cells"""
         arr = np.ones(data.ires.shape, dtype="float64")
+        tmp = data.apply_units(arr, field.units)
         if data._spatial:
-            return data._reshape_vals(arr)
-        return data.apply_units(arr, field.units)
+            return data._reshape_vals(tmp)
+        return tmp
 
     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 \
@@ -181,7 +183,22 @@ class RAMSESDomainSubset(OctreeSubset):
     _domain_offset = 1
     _block_reorder = "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]

yt/frontends/ramses/io_utils.pyx

@@ -83,7 +83,7 @@ def read_amr(FortranFile f, dict headers,
 @cython.nonecheck(False)
 cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t nvar, dict headers):
 
-    cdef np.ndarray[np.int64_t, ndim=1] offset, level_count
+    cdef np.ndarray[np.int64_t, ndim=2] offset, level_count
     cdef INT64_t ndim, twotondim, nlevelmax, n_levels, nboundary, ncpu, ncpu_and_bound
     cdef INT64_t ilevel, icpu, skip_len
     cdef INT32_t file_ilevel, file_ncache
@@ -101,12 +101,11 @@ cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t n
     skip_len = twotondim * nvar
 
     # It goes: level, CPU, 8-variable (1 oct)
-    offset = np.zeros(n_levels, dtype=np.int64)
-    offset -= 1
-    level_count = np.zeros(n_levels, dtype=np.int64)
+    offset = np.full((ncpu_and_bound, n_levels), -1, dtype=np.int64)
+    level_count = np.zeros((ncpu_and_bound, n_levels), dtype=np.int64)
 
-    cdef np.int64_t[:] level_count_view = level_count
-    cdef np.int64_t[:] offset_view = offset
+    cdef np.int64_t[:,:] level_count_view = level_count
+    cdef np.int64_t[:,:] offset_view = offset
 
     for ilevel in range(nlevelmax):
         for icpu in range(ncpu_and_bound):
@@ -121,9 +120,9 @@ cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t n
                     'from data (%s) is not coherent with the expected (%s)',
                     f.name, file_ilevel, ilevel)
 
-            if icpu + 1 == domain_id and ilevel >= min_level:
-                offset[ilevel - min_level] = f.tell()
-                level_count[ilevel - min_level] = <INT64_t> file_ncache
+            if ilevel >= min_level:
+                offset_view[icpu, ilevel - min_level] = f.tell()
+                level_count_view[icpu, ilevel - min_level] = <INT64_t> file_ncache
             f.skip(skip_len)
 
     return offset, level_count
@@ -133,45 +132,57 @@ cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t n
 @cython.cdivision(True)
 @cython.nonecheck(False)
 def fill_hydro(FortranFile f,
-               np.ndarray[np.int64_t, ndim=1] offsets,
-               np.ndarray[np.int64_t, ndim=1] level_count,
+               np.ndarray[np.int64_t, ndim=2] offsets,
+               np.ndarray[np.int64_t, ndim=2] level_count,
+               list cpu_enumerator,
                np.ndarray[np.uint8_t, ndim=1] levels,
                np.ndarray[np.uint8_t, ndim=1] cell_inds,
                np.ndarray[np.int64_t, ndim=1] file_inds,
                INT64_t ndim, list all_fields, list fields,
                dict tr,
-               RAMSESOctreeContainer oct_handler):
-    cdef INT64_t ilevel, ifield, nfields, noffset
+               RAMSESOctreeContainer oct_handler,
+               np.ndarray[np.int32_t, ndim=1] domains=np.array([], dtype='int32')):
     cdef INT64_t offset
     cdef dict tmp
     cdef str field
-    cdef INT64_t twotondim, i
+    cdef INT64_t twotondim
+    cdef int ilevel, icpu, ifield, nfields, nlevels, nc, ncpu_selected
     cdef np.ndarray[np.uint8_t, ndim=1] mask
 
     twotondim = 2**ndim
     nfields = len(all_fields)
-    noffset = len(offsets)
+    ncpu = offsets.shape[0]
+    nlevels = offsets.shape[1]
+    ncpu_selected = len(cpu_enumerator)
 
     mask = np.array([(field in fields) for field in all_fields], dtype=np.uint8)
 
     # Loop over levels
-    for ilevel in range(noffset):
-        offset = offsets[ilevel]
-        if offset == -1:
-            continue
-        f.seek(offset)
-        nc = level_count[ilevel]
-        tmp = {}
-        # Initalize temporary data container for io
-        for field in all_fields:
-            tmp[field] = np.empty((nc, twotondim), dtype="float64")
-
-        for i in range(twotondim):
-            # Read the selected fields
-            for ifield in range(nfields):
-                if not mask[ifield]:
-                    f.skip()
-                else:
-                    tmp[all_fields[ifield]][:, i] = f.read_vector('d') # i-th cell
-
-        oct_handler.fill_level(ilevel, levels, cell_inds, file_inds, tr, tmp)
+    for ilevel in range(nlevels):
+        # Loop over cpu domains
+        for icpu in cpu_enumerator:
+            nc = level_count[icpu, ilevel]
+            if nc == 0:
+                continue
+            offset = offsets[icpu, ilevel]
+            if offset == -1:
+                continue
+            f.seek(offset)
+            tmp = {}
+            # Initalize temporary data container for io
+            for field in all_fields:
+                tmp[field] = np.empty((nc, twotondim), dtype="float64", order='F')
+
+            for i in range(twotondim):
+                # Read the selected fields
+                for ifield in range(nfields):
+                    if not mask[ifield]:
+                        f.skip()
+                    else:
+                        tmp[all_fields[ifield]][:, i] = f.read_vector('d') # i-th cell
+            if ncpu_selected > 1:
+                oct_handler.fill_level_with_domain(
+                    ilevel, levels, cell_inds, file_inds, domains, tr, tmp, domain=icpu+1)
+            else:
+                oct_handler.fill_level(
+                    ilevel, levels, cell_inds, file_inds, tr, tmp)