Support vertex centred data

yt/data_objects/octree_subset.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from itertools import product, repeat
 
 import numpy as np
 
@@ -104,7 +105,7 @@ def mask_refinement(self, selector):
 
     def select_blocks(self, selector):
         mask = self.oct_handler.mask(selector, domain_id=self.domain_id)
-        slicer = OctreeSubsetBlockSlice(self)
+        slicer = OctreeSubsetBlockSlice(self, self.ds)
         for i, sl in slicer:
             yield sl, np.atleast_3d(mask[i, ...])
 
@@ -528,6 +529,7 @@ def __init__(self, ind, block_slice):
         self.block_slice = block_slice
         nz = self.block_slice.octree_subset.nz
         self.ActiveDimensions = np.array([nz, nz, nz], dtype="int64")
+        self.ds = block_slice.ds
 
     def __getitem__(self, key):
         bs = self.block_slice
@@ -572,22 +574,64 @@ def dds(self):
     def clear_data(self):
         pass
 
-    def get_vertex_centered_data(self, *args, **kwargs):
-        raise NotImplementedError
+    def get_vertex_centered_data(self, fields, smoothed=False, no_ghost=False):
+        field = fields[0]
+        new_field = self.block_slice.get_vertex_centered_data(fields)[field]
+        return {field: new_field[..., self.ind]}
 
     @contextmanager
     def _field_parameter_state(self, field_parameters):
         yield self.block_slice.octree_subset._field_parameter_state(field_parameters)
 
 
 class OctreeSubsetBlockSlice:
-    def __init__(self, octree_subset):
+    def __init__(self, octree_subset, ds):
         self.octree_subset = octree_subset
+        self.ds = ds
+        self._vertex_centered_data = {}
         # Cache some attributes
         for attr in ["ires", "icoords", "fcoords", "fwidth"]:
             v = getattr(octree_subset, attr)
             setattr(self, f"_{attr}", octree_subset._reshape_vals(v))
 
+    @property
+    def octree_subset_with_gz(self):
+        subset_with_gz = getattr(self, "_octree_subset_with_gz", None)
+        if not subset_with_gz:
+            self._octree_subset_with_gz = self.octree_subset.retrieve_ghost_zones(1, [])
+        return self._octree_subset_with_gz
+
+    def get_vertex_centered_data(self, fields, smoothed=False, no_ghost=False):
+        if no_ghost is True:
+            raise NotImplementedError(
+                "get_vertex_centered_data without ghost zones for oct-based datasets has not been implemented."
+            )
+
+        # Make sure the field list has only unique entries
+        fields = list(set(fields))
+        new_fields = {}
+        cg = self.octree_subset_with_gz
+        for field in fields:
+            if field in self._vertex_centered_data:
+                new_fields[field] = self._vertex_centered_data[field]
+            else:
+                finfo = self.ds._get_field_info(field)
+                orig_field = cg[field]
+                nocts = orig_field.shape[-1]
+                new_field = np.zeros((3, 3, 3, nocts), order="F")
+
+                # Compute vertex-centred data as mean of 8 neighbours cell data
+                slices = (slice(1, None), slice(None, -1))
+                for slx, sly, slz in product(*repeat(slices, 3)):
+                    new_field += orig_field[slx, sly, slz]
+                new_field *= 0.125
+
+                new_fields[field] = self.ds.arr(new_field, finfo.output_units)
+
+                self._vertex_centered_data[field] = new_fields[field]
+
+        return new_fields
+
     def __iter__(self):
         for i in range(self._ires.shape[-1]):
             yield i, OctreeSubsetBlockSlicePosition(i, self)

yt/frontends/ramses/data_structures.py

@@ -349,9 +349,31 @@ def fill(self, fd, fields, selector, file_handler):
             )
 
     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
-        )
+        if smoothed:
+            mylog.warning(
+                f"{self}.retrieve_ghost_zones was called with the "
+                f"`smoothed` argument set to True. This is not supported, "
+                "ignoring it."
+            )
+            smoothed = False
+
+        try:
+            new_subset = self._subset_with_gz
+            mylog.debug(
+                "Reusing previous subset with ghost zone for domain %s", self.domain_id
+            )
+        except AttributeError:
+            new_subset = RAMSESDomainSubset(
+                self.base_region,
+                self.domain,
+                self.ds,
+                num_ghost_zones=ngz,
+                base_grid=self,
+            )
+            self._subset_with_gz = new_subset
+
+            # Cache the fields
+            new_subset.get_data(fields)
 
         return new_subset