Implement fcoords etc. using subclass of OctreeSubset

Author: cphyc

yt/frontends/dyablo/data_structures.py

@@ -8,6 +8,7 @@
 import h5py
 import numpy as np
 
+from yt.data_objects.index_subobjects.octree_subset import OctreeSubset
 from yt.data_objects.static_output import Dataset
 from yt.geometry.geometry_handler import YTDataChunk
 from yt.geometry.oct_container import OctreeContainer
@@ -23,6 +24,8 @@
 class DyabloOctreeIndex(OctreeIndex):
     """Octree Index for Dyablo data with leaf-only blocks."""
 
+    domain_id = 1  # Dyablo uses a single domain
+
     def __init__(self, ds, dataset_type):
         self.dataset_type = dataset_type
         self.dataset = ds
@@ -55,7 +58,8 @@ def _initialize_oct_handler(self):
         if n_cells % cells_per_block != 0:
             mylog.warning(
                 f"Number of cells ({n_cells}) is not divisible by "
-                f"cells per block ({cells_per_block}). Some cells will be ignored."
+                f"cells per block ({cells_per_block}). "
+                "Some cells will be ignored."
             )
 
         # NOTE: we're storing an octree of blocks, so we insert
@@ -82,7 +86,8 @@ def _initialize_oct_handler(self):
             block_corners = coordinates[block_connectivity, :]
 
             # Compute block bounds from all cell corners
-            all_corners = block_corners.reshape(-1, 3)  # (cells_per_block * 8, 3)
+            # (cells_per_block * 8, 3)
+            all_corners = block_corners.reshape(-1, 3)
             block_min = np.min(all_corners, axis=0)
             block_max = np.max(all_corners, axis=0)
             block_width = block_max - block_min
@@ -104,7 +109,8 @@ def _initialize_oct_handler(self):
             ]  # (cells_per_block, 8, 3)
 
             # Compute block center: mean of all cell centers
-            cell_centers = np.mean(block_corners, axis=1)  # (cells_per_block, 3)
+            # (cells_per_block, 3)
+            cell_centers = np.mean(block_corners, axis=1)
             block_centers[block_id] = np.mean(cell_centers, axis=0)
 
             # Compute block size and level
@@ -116,9 +122,8 @@ def _initialize_oct_handler(self):
             # Level is determined by block size relative to coarsest block
             max_dim_width = np.max(block_width)
             refinement_ratio = max_block_width / max_dim_width
-            block_levels[block_id] = np.round(np.log2(refinement_ratio)).astype(
-                np.uint64
-            )
+            level = np.round(np.log2(refinement_ratio)).astype(np.uint64)
+            block_levels[block_id] = level
 
         # Count number of blocks - note that a block at level l
         # requires its parent blocks at levels < l to be present
@@ -213,18 +218,175 @@ def _chunk_io(self, dobj, cache=True, local_only=False):
             yield YTDataChunk(dobj, "io", [subset], None, cache=cache)
 
 
-class DyabloOctreeSubset:
-    """Octree subset for Dyablo data."""
+class DyabloOctreeSubset(OctreeSubset):
+    """Octree subset for Dyablo data with NxMxL blocks.
+
+    Unlike standard octrees where each oct contains 2x2x2 cells,
+    Dyablo uses blocks as octs, where each block contains NxMxL cells.
+    This class overrides coordinate methods to properly handle this
+    non-standard structure.
+    """
+
+    _domain_offset = 1
+    _block_order = "C"
+
+    def __init__(self, base_region, domain, ds):
+        # Get block size for num_zones parameter
+        N, M, L = ds.block_size
+        # Use the maximum dimension for num_zones
+        num_zones = max(N, M, L)
+
+        super().__init__(base_region, domain, ds, num_zones=num_zones)
+
+        self._current_particle_type = "io"
+        self._current_fluid_type = self.ds.default_fluid_type
+
+    @property
+    def oct_handler(self):
+        return self.domain.oct_handler
+
+    def _get_selected_cell_info(self, selector):
+        """
+        Get information about selected cells using two-stage selection.
+
+        This method:
+        1. Uses octree to select blocks via domain_ind
+        2. Computes cell centers from connectivity data
+        3. Applies selector at cell level to filter cells
+
+        Parameters
+        ----------
+        selector : SelectorObject
+            Selector object for filtering cells
+
+        Returns
+        -------
+        block_inds : ndarray
+            Array of selected block indices
+            (oct indices where domain_ind != -1)
+        cell_centers : ndarray
+            Array of cell center positions (n_cells_in_blocks, 3)
+        mask : ndarray
+            Boolean mask for cells passing selector (n_cells_in_blocks,)
+        """
+        # Stage 1: Use octree to get selected block indices
+        domain_ind_array = self.oct_handler.domain_ind(selector)
+        block_inds = np.where(domain_ind_array >= 0)[0]
+
+        if len(block_inds) == 0:
+            empty_centers = np.array([]).reshape(0, 3)
+            empty_mask = np.array([], dtype=bool)
+            return block_inds, empty_centers, empty_mask
+
+        # Stage 2: Compute cell centers from connectivity
+        block_size = self.ds.block_size
+        coordinates = self.domain.coordinates
+        connectivity = self.domain.connectivity.reshape(-1, *block_size, 8)
 
-    def __init__(self, base_region, index, ds):
-        self.base_region = base_region
-        self.index = index
-        self.ds = ds
-        self.oct_handler = index.oct_handler
-        self.domain_id = 1  # Dyablo uses domain 1
+        # Get connectivity for selected blocks (n_blocks, N, M, L, 8)
+        selected_block_connectivity = connectivity[block_inds]
+        # Reshape to (n_blocks * N * M * L, 8)
+        selected_cell_connectivity = selected_block_connectivity.reshape(-1, 8)
 
-    def __getitem__(self, key):
-        return self.base_region[key]
+        # Compute cell centers: average of 8 corner coordinates
+        cell_centers = np.mean(
+            coordinates[selected_cell_connectivity], axis=1
+        )  # (n_cells, 3)
+
+        # Stage 3: Apply selector at cell level
+        radius = 0
+        mask = selector.select_points(*cell_centers.T, radius)
+
+        mylog.debug(
+            "Selected %d cells from %d blocks (%d total cells checked)",
+            mask.sum(),
+            len(block_inds),
+            len(mask),
+        )
+
+        return block_inds, cell_centers, mask
+
+    def select_icoords(self, dobj):
+        """
+        Return integer coordinates of selected cells.
+
+        Integer coordinates represent the cell index on a uniform grid
+        at the appropriate refinement level.
+        """
+        block_inds, cell_centers, mask = self._get_selected_cell_info(dobj.selector)
+
+        if len(block_inds) == 0:
+            return np.array([], dtype="int64").reshape(0, 3)
+
+        # Calculate integer coordinates from cell positions
+        # For uniform grid (level 0), cell width is domain_width / total_cells
+        block_size = self.ds.block_size
+        N, M, L = block_size
+
+        # Domain dimensions represent number of blocks in each direction
+        domain_dims = self.ds.domain_dimensions
+
+        # Total number of cells in each direction
+        total_cells = domain_dims * np.array([N, M, L])
+
+        # Convert positions to integer grid indices
+        domain_width = self.ds.domain_right_edge - self.ds.domain_left_edge
+        cell_width = domain_width / total_cells
+        rel_pos = cell_centers - self.ds.domain_left_edge
+        icoords = (rel_pos / cell_width).astype(np.int64)
+
+        # Clip to valid range [0, total_cells-1] to handle cells at boundaries
+        for i in range(3):
+            icoords[:, i] = np.clip(icoords[:, i], 0, total_cells[i] - 1)
+
+        block_inds, cell_centers, mask = self._get_selected_cell_info(dobj.selector)
+
+        if len(block_inds) == 0:
+            return self.ds.arr(np.array([]).reshape(0, 3), "code_length")
+
+        return self.ds.arr(cell_centers[mask], "code_length")
+
+    def select_fwidth(self, dobj):
+        """
+        Return cell widths of selected cells.
+        """
+        block_inds, cell_centers, mask = self._get_selected_cell_info(dobj.selector)
+
+        if len(block_inds) == 0:
+            return self.ds.arr(np.array([]).reshape(0, 3), "code_length")
+
+        # Compute cell widths from connectivity corner positions
+        block_size = self.ds.block_size
+        coordinates = self.domain.coordinates
+        connectivity = self.domain.connectivity.reshape(-1, *block_size, 8)
+
+        # Get connectivity for selected blocks
+        selected_block_connectivity = connectivity[block_inds]
+        selected_cell_connectivity = selected_block_connectivity.reshape(-1, 8)
+
+        # Width in each dimension: max - min of corner coordinates
+        # (n_cells, 8, 3)
+        cell_coords = coordinates[selected_cell_connectivity]
+        # (n_cells, 3)
+        cell_widths = cell_coords.max(axis=1) - cell_coords.min(axis=1)
+
+        return self.ds.arr(cell_widths[mask], "code_length")
+
+    def select_ires(self, dobj):
+        """
+        Return refinement level (resolution) of selected cells.
+
+        For Dyablo uniform grid data, all cells are at level 0.
+        """
+        block_inds, cell_centers, mask = self._get_selected_cell_info(dobj.selector)
+
+        if len(block_inds) == 0:
+            return np.array([], dtype="int64")
+
+        # For uniform grid, all blocks are at level 0
+        # Could use self.domain.block_levels if AMR support is added later
+        n_selected = mask.sum()
+        return np.zeros(n_selected, dtype="int64")
 
     def fill(self, fields, selector):
         """
@@ -246,9 +408,8 @@ def fill(self, fields, selector):
         dict
             Dictionary mapping (ftype, fname) to selected data arrays
         """
-        # Get the indices of the blocks to read from
-        block_inds = self.oct_handler.domain_ind(selector)
-        # Get the file indices (block indices) that contain selected cells
+        # Use helper to get selected blocks and cell mask
+        block_inds, cell_centers, mask = self._get_selected_cell_info(selector)
 
         if len(block_inds) == 0:
             return {field: np.array([]) for field in fields}
@@ -265,29 +426,14 @@ def fill(self, fields, selector):
                     # Read only the cells we need
                     full_data = f[field_path][:].reshape(-1, *block_size)
 
-                    # Only keep those octs we selected
+                    # Only keep those blocks we selected
                     data_subset = full_data[block_inds].flatten()
-                    result[ftype, fname] = data_subset
+
+                    # Apply cell-level mask
+                    result[ftype, fname] = data_subset[mask]
                 else:
                     result[ftype, fname] = np.array([])
 
-        # Also read coordinates for selector filtering
-        coordinates = self.ds._coordinates
-        connectivity = self.ds._connectivity.reshape(-1, *block_size, 8)
-        selected_block_connectivity = connectivity[block_inds].reshape(-1, 8)
-
-        # Compute cell centers
-        cell_centers = np.mean(
-            coordinates[selected_block_connectivity], axis=1
-        )  # (n_selected_cells, 3)
-
-        # Apply selector to filter cells
-        radius = 0
-        mask = selector.select_points(*cell_centers.T, radius)
-        mylog.debug("Selected %d cells within %d blocks", mask.sum(), len(block_inds))
-        for key in result:
-            result[key] = result[key][mask]
-
         return result