Dyablo frontend

pyproject.toml

@@ -114,6 +114,7 @@ cf-radial = ["xarray>=0.16.1", "arm-pyart>=1.19.2",]
 chimera = ["yt[HDF5]"]
 chombo = ["yt[HDF5]"]
 cholla = ["yt[HDF5]"]
+dyablo = ["yt[HDF5]"]
 eagle = ["yt[HDF5]"]
 enzo-e = ["yt[HDF5]", "libconf>=1.0.1"]
 enzo = ["yt[HDF5]", "libconf>=1.0.1"]

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/__init__.py

@@ -12,6 +12,7 @@
     "chimera",
     "chombo",
     "cholla",
+    "dyablo",
     "enzo_e",
     "enzo",
     "exodus_ii",

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/dyablo/__init__.py

@@ -0,0 +1,10 @@
+"""
+Dyablo frontend module for yt.
+
+Dyablo is a regular-sized-block AMR hydrodynamical simulation code
+with outputs in HDF5 format, ordered along a Morton curve.
+"""
+
+from .api import DyabloDataset
+
+__all__ = ["DyabloDataset"]

yt/frontends/dyablo/api.py

@@ -0,0 +1,9 @@
+"""
+API for Dyablo frontend.
+"""
+
+from .data_structures import DyabloDataset, DyabloOctreeIndex
+from .fields import DyabloFieldInfo
+from .io import DyabloIOHandler
+
+__all__ = ["DyabloDataset", "DyabloOctreeIndex", "DyabloFieldInfo", "DyabloIOHandler"]