Support (N, M, L) per-dimension zones in octree containers

Author: cphyc

yt/data_objects/index_subobjects/octree_subset.py

@@ -46,7 +46,10 @@ class OctreeSubset(YTSelectionContainer, abc.ABC):
 
     def __init__(self, base_region, domain, ds, num_zones=2, num_ghost_zones=0):
         super().__init__(ds, None)
-        self._num_zones = num_zones
+        if hasattr(num_zones, "__len__"):
+            self._num_zones = np.array(num_zones, dtype="int64")
+        else:
+            self._num_zones = np.array([num_zones, num_zones, num_zones], dtype="int64")
         self._num_ghost_zones = num_ghost_zones
         self.domain = domain
         self.domain_id = domain.domain_id
@@ -80,23 +83,28 @@ def __getitem__(self, key):
 
     @property
     def nz(self):
-        return self._num_zones + 2 * self._num_ghost_zones
+        nz = self._num_zones + 2 * self._num_ghost_zones
+        if hasattr(nz, "__len__"):
+            return nz
+        return np.array([nz, nz, nz], dtype="int64")
 
     def get_bbox(self):
         return self.base_region.get_bbox()
 
     def _reshape_vals(self, arr):
         nz = self.nz
+        nzx, nzy, nzz = nz[0], nz[1], nz[2]
+        nzones = nzx * nzy * nzz
         if len(arr.shape) <= 2:
-            n_oct = arr.shape[0] // (nz**3)
+            n_oct = arr.shape[0] // nzones
         elif arr.shape[-1] == 3:
             n_oct = arr.shape[-2]
         else:
             n_oct = arr.shape[-1]
-        if arr.size == nz * nz * nz * n_oct:
-            new_shape = (nz, nz, nz, n_oct)
-        elif arr.size == nz * nz * nz * n_oct * 3:
-            new_shape = (nz, nz, nz, n_oct, 3)
+        if arr.size == nzones * n_oct:
+            new_shape = (nzx, nzy, nzz, n_oct)
+        elif arr.size == nzones * n_oct * 3:
+            new_shape = (nzx, nzy, nzz, n_oct, 3)
         else:
             raise RuntimeError
         # Note that if arr is already F-contiguous, this *shouldn't* copy the
@@ -172,7 +180,7 @@ def deposit(self, positions, fields=None, method=None, kernel_name="cubic"):
         if cls is None:
             raise YTParticleDepositionNotImplemented(method)
         nz = self.nz
-        nvals = (nz, nz, nz, (self.domain_ind >= 0).sum())
+        nvals = (int(nz[0]), int(nz[1]), int(nz[2]), (self.domain_ind >= 0).sum())
         if np.max(self.domain_ind) >= nvals[-1]:
             print(
                 f"nocts, domain_ind >= 0, max {self.oct_handler.nocts} {nvals[-1]} {np.max(self.domain_ind)}"
@@ -335,7 +343,7 @@ def smooth(
                 [1, 1, 1],
                 self.ds.domain_left_edge,
                 self.ds.domain_right_edge,
-                num_zones=self._nz,
+                num_zones=self._num_zones,
             )
             # This should ensure we get everything within one neighbor of home.
             particle_octree.n_ref = nneighbors * 2
@@ -354,7 +362,7 @@ def smooth(
             raise YTParticleDepositionNotImplemented(method)
         nz = self.nz
         mdom_ind = self.domain_ind
-        nvals = (nz, nz, nz, (mdom_ind >= 0).sum())
+        nvals = (int(nz[0]), int(nz[1]), int(nz[2]), (mdom_ind >= 0).sum())
         op = cls(nvals, len(fields), nneighbors, kernel_name)
         op.initialize()
         mylog.debug(
@@ -455,7 +463,7 @@ def particle_operation(
             raise YTParticleDepositionNotImplemented(method)
         nz = self.nz
         mdom_ind = self.domain_ind
-        nvals = (nz, nz, nz, (mdom_ind >= 0).sum())
+        nvals = (int(nz[0]), int(nz[1]), int(nz[2]), (mdom_ind >= 0).sum())
         op = cls(nvals, len(fields), nneighbors, kernel_name)
         op.initialize()
         mylog.debug(
@@ -548,7 +556,7 @@ def __init__(self, ind, block_slice):
         self.ind = ind
         self.block_slice = block_slice
         nz = self.block_slice.octree_subset.nz
-        self.ActiveDimensions = np.array([nz, nz, nz], dtype="int64")
+        self.ActiveDimensions = np.array([nz[0], nz[1], nz[2]], dtype="int64")
         self.ds = block_slice.ds
 
     def __getitem__(self, key):

yt/data_objects/tests/test_octree.py

@@ -1,6 +1,7 @@
 import numpy as np
 from numpy.testing import assert_almost_equal, assert_equal
 
+from yt.geometry.oct_container import OctreeContainer
 from yt.testing import fake_sph_grid_ds
 
 n_ref = 4
@@ -118,3 +119,26 @@ def test_octree_properties():
     refined = octree["index", "refined"]
     refined_ans = np.array([True] + [False] * 7 + [True] + [False] * 8, dtype=np.bool_)
     assert_equal(refined, refined_ans)
+
+
+def test_num_zones_tuple():
+    """
+    Test that OctreeContainer accepts num_zones as a scalar or a tuple (N, M, L).
+    Both should correctly set per-dimension zone counts.
+    """
+    # Scalar: all dimensions equal
+    oct_scalar = OctreeContainer(
+        [1, 1, 1], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], num_zones=2
+    )
+    # Tuple: potentially different per-dimension
+    oct_tuple = OctreeContainer(
+        [1, 1, 1], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], num_zones=(2, 2, 2)
+    )
+    # Non-uniform tuple
+    oct_nonuniform = OctreeContainer(
+        [1, 1, 1], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], num_zones=(2, 3, 4)
+    )
+    # Verify that creating these containers doesn't raise exceptions
+    assert oct_scalar is not None
+    assert oct_tuple is not None
+    assert oct_nonuniform is not None

yt/frontends/artio/data_structures.py

@@ -128,7 +128,7 @@ def deposit(self, positions, fields=None, method=None, kernel_name="cubic"):
         if cls is None:
             raise YTParticleDepositionNotImplemented(method)
         nz = self.nz
-        nvals = (nz, nz, nz, self.ires.size)
+        nvals = (int(nz[0]), int(nz[1]), int(nz[2]), self.ires.size)
         # We allocate number of zones, not number of octs
         op = cls(nvals, kernel_name)
         op.initialize()
@@ -241,6 +241,7 @@ def _identify_base_chunk(self, dobj):
             sfc_start = getattr(dobj, "sfc_start", None)
             sfc_end = getattr(dobj, "sfc_end", None)
             nz = getattr(dobj, "_num_zones", 0)
+            nz_scalar = int(np.asarray(nz).flat[0]) if hasattr(nz, "__len__") else nz
             if all_data:
                 mylog.debug("Selecting entire artio domain")
                 list_sfc_ranges = self.ds._handle.root_sfc_ranges_all(
@@ -271,7 +272,7 @@ def _identify_base_chunk(self, dobj):
                     )
                     range_handler.construct_mesh()
                     self.range_handlers[start, end] = range_handler
-                if nz != 2:
+                if nz_scalar != 2:
                     ci.append(
                         ARTIORootMeshSubset(
                             base_region,
@@ -281,7 +282,7 @@ def _identify_base_chunk(self, dobj):
                             self.ds,
                         )
                     )
-                if nz != 1 and range_handler.total_octs > 0:
+                if nz_scalar != 1 and range_handler.total_octs > 0:
                     ci.append(
                         ARTIOOctreeSubset(
                             base_region,

yt/frontends/ramses/data_structures.py

@@ -467,8 +467,8 @@ def _fill_with_ghostzones(
         fields = [f for ft, f in fields]
         tr = {}
 
-        cell_count = (
-            selector.count_octs(self.oct_handler, self.domain_id) * self.nz**ndim
+        cell_count = selector.count_octs(self.oct_handler, self.domain_id) * int(
+            np.prod(self.nz[:ndim])
         )
 
         # Initializing data container
@@ -520,7 +520,7 @@ def fwidth(self):
             # 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 = np.zeros((n_oct, int(np.prod(self.nz)), 3), dtype=fwidth.dtype)
             new_fwidth[:, :, :] = fwidth[:, 0:1, :]
             fwidth = new_fwidth.reshape(-1, 3)
         return fwidth
@@ -538,7 +538,7 @@ def fcoords(self):
             self.selector, self._num_ghost_zones
         )
 
-        N_per_oct = self.nz**3
+        N_per_oct = int(np.prod(self.nz))
         oct_inds = oct_inds.reshape(-1, N_per_oct)
         cell_inds = cell_inds.reshape(-1, N_per_oct)
 

yt/frontends/stream/data_structures.py

@@ -832,8 +832,8 @@ def _fill_no_ghostzones(self, content, dest, selector, offset):
     def _fill_with_ghostzones(self, content, dest, selector, offset):
         oct_handler = self.oct_handler
         ndim = self.ds.dimensionality
-        cell_count = (
-            selector.count_octs(self.oct_handler, self.domain_id) * self.nz**ndim
+        cell_count = selector.count_octs(self.oct_handler, self.domain_id) * int(
+            np.prod(self.nz[:ndim])
         )
 
         gz_cache = getattr(self, "_ghost_zone_cache", None)

yt/geometry/_selection_routines/selector_object.pxi

@@ -159,7 +159,7 @@ cdef class SelectorObject:
                             visitor.pos[1] = (visitor.pos[1] >> 1)
                             visitor.pos[2] = (visitor.pos[2] >> 1)
                             visitor.level -= 1
-                        elif this_level == 1 and visitor.nz > 1:
+                        elif this_level == 1 and (visitor.nz[0] > 1 or visitor.nz[1] > 1 or visitor.nz[2] > 1):
                             visitor.global_index += increment
                             increment = 0
                             self.visit_oct_cells(root, ch, spos, sdds,
@@ -178,10 +178,10 @@ cdef class SelectorObject:
     cdef void visit_oct_cells(self, Oct *root, Oct *ch,
                               np.float64_t spos[3], np.float64_t sdds[3],
                               OctVisitor visitor, int i, int j, int k):
-        # We can short-circuit the whole process if data.nz == 2.
+        # We can short-circuit the whole process if data.nz == 2 in all dims.
         # This saves us some funny-business.
         cdef int selected
-        if visitor.nz == 2:
+        if visitor.nz[0] == 2 and visitor.nz[1] == 2 and visitor.nz[2] == 2:
             selected = self.select_cell(spos, sdds)
             if ch != NULL:
                 selected *= self.overlap_cells
@@ -199,22 +199,25 @@ cdef class SelectorObject:
         cdef np.float64_t dds[3]
         cdef np.float64_t pos[3]
         cdef int ci, cj, ck
-        cdef int nr = (visitor.nz >> 1)
+        cdef int nr[3]
+        nr[0] = visitor.nz[0] >> 1
+        nr[1] = visitor.nz[1] >> 1
+        nr[2] = visitor.nz[2] >> 1
         for ci in range(3):
-            dds[ci] = sdds[ci] / nr
+            dds[ci] = sdds[ci] / nr[ci]
         # Boot strap at the first index.
         pos[0] = (spos[0] - sdds[0]/2.0) + dds[0] * 0.5
-        for ci in range(nr):
+        for ci in range(nr[0]):
             pos[1] = (spos[1] - sdds[1]/2.0) + dds[1] * 0.5
-            for cj in range(nr):
+            for cj in range(nr[1]):
                 pos[2] = (spos[2] - sdds[2]/2.0) + dds[2] * 0.5
-                for ck in range(nr):
+                for ck in range(nr[2]):
                     selected = self.select_cell(pos, dds)
                     if ch != NULL:
                         selected *= self.overlap_cells
-                    visitor.ind[0] = ci + i * nr
-                    visitor.ind[1] = cj + j * nr
-                    visitor.ind[2] = ck + k * nr
+                    visitor.ind[0] = ci + i * nr[0]
+                    visitor.ind[1] = cj + j * nr[1]
+                    visitor.ind[2] = ck + k * nr[2]
                     visitor.visit(root, selected)
                     pos[2] += dds[2]
                 pos[1] += dds[1]

yt/geometry/oct_container.pxd

@@ -57,7 +57,7 @@ cdef class OctreeContainer:
     cdef int partial_coverage
     cdef int level_offset
     cdef int nn[3]
-    cdef np.uint8_t nz
+    cdef np.uint8_t nz[3]
     cdef np.float64_t DLE[3]
     cdef np.float64_t DRE[3]
     cdef public np.int64_t nocts