diff --git a/pyproject.toml b/pyproject.toml index 92ef9c347ae..3dbf4c6bfcc 100644 --- a/pyproject.toml +++ b/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"] diff --git a/uv.lock b/uv.lock index acfb51d0b18..b956dfb8b6b 100644 --- a/uv.lock +++ b/uv.lock @@ -5568,6 +5568,9 @@ cholla = [ chombo = [ { name = "h5py" }, ] +dyablo = [ + { name = "h5py" }, +] eagle = [ { name = "h5py" }, ] @@ -5805,6 +5808,7 @@ requires-dist = [ { name = "yt", extras = ["hdf5"], marker = "extra == 'chimera'" }, { name = "yt", extras = ["hdf5"], marker = "extra == 'cholla'" }, { name = "yt", extras = ["hdf5"], marker = "extra == 'chombo'" }, + { name = "yt", extras = ["hdf5"], marker = "extra == 'dyablo'" }, { name = "yt", extras = ["hdf5"], marker = "extra == 'eagle'" }, { name = "yt", extras = ["hdf5"], marker = "extra == 'enzo'" }, { name = "yt", extras = ["hdf5"], marker = "extra == 'enzo-e'" }, @@ -5841,7 +5845,7 @@ requires-dist = [ { name = "yt", extras = ["ytdata"], marker = "extra == 'full'" }, { name = "yt-idefix", extras = ["hdf5"], marker = "extra == 'idefix'", specifier = ">=2.3.0" }, ] -provides-extras = ["hdf5", "netcdf4", "fortran", "adaptahop", "ahf", "amrex", "amrvac", "art", "arepo", "artio", "athena", "athena-pp", "boxlib", "cf-radial", "chimera", "chombo", "cholla", "eagle", "enzo-e", "enzo", "exodus-ii", "fits", "flash", "gadget", "gadget-fof", "gamer", "gdf", "gizmo", "halo-catalog", "http-stream", "idefix", "moab", "nc4-cm1", "open-pmd", "owls", "owls-subfind", "parthenon", "ramses", "rockstar", "sdf", "stream", "swift", "tipsy", "ytdata", "full", "test"] +provides-extras = ["hdf5", "netcdf4", "fortran", "adaptahop", "ahf", "amrex", "amrvac", "art", "arepo", "artio", "athena", "athena-pp", "boxlib", "cf-radial", "chimera", "chombo", "cholla", "dyablo", "eagle", "enzo-e", "enzo", "exodus-ii", "fits", "flash", "gadget", "gadget-fof", "gamer", "gdf", "gizmo", "halo-catalog", "http-stream", "idefix", "moab", "nc4-cm1", "open-pmd", "owls", "owls-subfind", "parthenon", "ramses", "rockstar", "sdf", "stream", "swift", "tipsy", "ytdata", "full", "test"] [package.metadata.requires-dev] docs = [ diff --git a/yt/data_objects/index_subobjects/octree_subset.py b/yt/data_objects/index_subobjects/octree_subset.py index d450ebaa3ec..410890e1986 100644 --- a/yt/data_objects/index_subobjects/octree_subset.py +++ b/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): diff --git a/yt/data_objects/tests/test_octree.py b/yt/data_objects/tests/test_octree.py index 06652612f04..0c52ed337ec 100644 --- a/yt/data_objects/tests/test_octree.py +++ b/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 diff --git a/yt/frontends/__init__.py b/yt/frontends/__init__.py index ef209207fbc..a069edef425 100644 --- a/yt/frontends/__init__.py +++ b/yt/frontends/__init__.py @@ -12,6 +12,7 @@ "chimera", "chombo", "cholla", + "dyablo", "enzo_e", "enzo", "exodus_ii", diff --git a/yt/frontends/artio/data_structures.py b/yt/frontends/artio/data_structures.py index 2338d40e002..a9d80866287 100644 --- a/yt/frontends/artio/data_structures.py +++ b/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, diff --git a/yt/frontends/dyablo/__init__.py b/yt/frontends/dyablo/__init__.py new file mode 100644 index 00000000000..acebed05d70 --- /dev/null +++ b/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"] diff --git a/yt/frontends/dyablo/api.py b/yt/frontends/dyablo/api.py new file mode 100644 index 00000000000..fa090074849 --- /dev/null +++ b/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"] diff --git a/yt/frontends/dyablo/data_structures.py b/yt/frontends/dyablo/data_structures.py new file mode 100644 index 00000000000..ad30dd023ca --- /dev/null +++ b/yt/frontends/dyablo/data_structures.py @@ -0,0 +1,677 @@ +""" +Data structures for Dyablo frontend. +""" + +from pathlib import Path + +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 +from yt.geometry.oct_geometry_handler import OctreeIndex +from yt.utilities.file_handler import HDF5FileHandler +from yt.utilities.logger import ytLogger as mylog + +from .definitions import HYDRO_FILE_PATTERN, PARTICLE_FILE_PATTERN +from .fields import DyabloFieldInfo +from .io import DyabloIOHandler + + +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 + self.index_filename = ds.parameter_filename + self.directory = Path(self.index_filename).parent + self.float_type = np.float64 + super().__init__(ds, dataset_type) + + def _initialize_oct_handler(self): + """Initialize octree handler from Dyablo connectivity/coordinates.""" + mylog.debug("Initializing Dyablo Octree Handler") + + # Get connectivity and coordinates from dataset + connectivity = self.ds._connectivity + coordinates = self.ds._coordinates + n_cells = self.ds._n_cells + block_size = self.ds.block_size # (N, M, L) + + # Store for later use by IO handler + self.connectivity = connectivity + self.coordinates = coordinates + self.n_cells = n_cells + self.block_size = block_size + + # Calculate number of blocks + N, M, L = block_size + cells_per_block = N * M * L + n_blocks = n_cells // cells_per_block + + 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." + ) + self._n_blocks = n_blocks + self._n_cells = n_cells + self._cells_per_block = cells_per_block + + # OctreeContainer needs the number of root octs (= n_blocks), + # which is domain_dimensions / block_size. + n_oct_dims = self.ds.domain_dimensions // np.array(block_size, dtype=np.int32) + self.oct_handler = OctreeContainer( + n_oct_dims, + self.ds.domain_left_edge, + self.ds.domain_right_edge, + num_zones=block_size, + ) + + # Compute block centers and levels + # Each block contains cells_per_block cells in C-order + block_centers = np.zeros((n_blocks, 3), dtype=np.float64) + block_levels = np.zeros(n_blocks, dtype=np.uint64) + + # First pass: find the maximum block width (coarsest level) + max_block_width = 0.0 + for block_id in range(n_blocks): + # Get cell indices for this block + cell_start = block_id * cells_per_block + cell_end = cell_start + cells_per_block + + # Get corners for all cells in this block + block_connectivity = connectivity[cell_start:cell_end] + block_corners = coordinates[block_connectivity, :] + + # Compute block bounds from all cell corners + # (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 + + # Track maximum block width + chunk_max_width = np.max(block_width) + max_block_width = max(max_block_width, chunk_max_width) + + # Second pass: compute centers and levels + for block_id in range(n_blocks): + # Get cell indices for this block + cell_start = block_id * cells_per_block + cell_end = cell_start + cells_per_block + + # Get corners for all cells in this block + block_connectivity = connectivity[cell_start:cell_end] + block_corners = coordinates[ + block_connectivity, : + ] # (cells_per_block, 8, 3) + + # Compute block center: mean of all cell centers + # (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 + 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 + + # 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 + 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 + # so we need to add them as well + levels, counts = np.unique(block_levels, return_counts=True) + n_blocks_with_parents = 0 + for lvl, count in zip(levels, counts, strict=True): + for i in range(lvl + 1): + n_blocks_with_parents += count / (8 ** (lvl - i)) + + n_blocks_with_parents = int(np.ceil(n_blocks_with_parents)) + # Allocate space for blocks + self.oct_handler.allocate_domains([n_blocks_with_parents]) + + # Add all blocks as octs at their appropriate levels + # Each block gets file_ind from 0 to n_blocks-1 + self.oct_handler.add( + 1, # domain 1 + -1, # use levels array + block_centers, # positions + levels=block_levels, + ) + + self.oct_handler.finalize() + + def _detect_output_fields(self): + """Detect available fields in the dataset.""" + dsl = set() + + # Open the hydro file and look for available fields + try: + with h5py.File(self.dataset.parameter_filename, "r") as f: + for key in f.keys(): + # Skip metadata/coordinate datasets + skip_keys = ( + "metadata", + "scalar_data", + "connectivity", + "coordinates", + "units", + ) + if key in skip_keys: + continue + + obj = f[key] + if isinstance(obj, h5py.Dataset): + dsl.add(("dyablo", key)) + except Exception as e: + mylog.warning(f"Error detecting output fields: {e}") + + self.fluid_field_list = list(dsl) + + # Detect particle fields from particle files + particle_field_list = [] + particle_types = set() + particle_files = self.dataset._particle_filename or [] + for pfile in particle_files: + pmatch = PARTICLE_FILE_PATTERN.match(Path(pfile).name) + if not pmatch: + continue + ptype = pmatch.group(2) + try: + with h5py.File(pfile, "r") as f: + particle_field_list.append((ptype, "particle_position_x")) + particle_field_list.append((ptype, "particle_position_y")) + particle_field_list.append((ptype, "particle_position_z")) + for key in (_ for _ in f.keys() if isinstance(f[_], h5py.Dataset)): + if key == "coordinates": + continue + yt_key = f"particle_{key}" + particle_field_list.append((ptype, yt_key)) + particle_types.add(ptype) + except Exception as e: + mylog.warning(f"Error reading particle file {pfile}: {e}") + + # Ensure all detected particles have been registered + for ptype in ( + _ for _ in particle_types if _ not in self.dataset.particle_types + ): + raise RuntimeError( + f"Detected particle type '{ptype}' from file {pfile} " + "but it was not registered in dataset.particle_types. " + "Please ensure the particle file is named correctly and " + "that _parse_parameter_file is correctly detecting particle types." + ) + + self.particle_field_list = particle_field_list + self.field_list = self.fluid_field_list + self.particle_field_list + + @property + def max_level(self): + """Return the maximum refinement level.""" + return self.ds.max_level + + def _identify_base_chunk(self, dobj): + """Identify the chunks that make up a data object.""" + if getattr(dobj, "_chunk_info", None) is None: + base_region = getattr(dobj, "base_region", dobj) + # For Dyablo, we have a single domain containing all octs + subset = [ + DyabloOctreeSubset( + base_region, + self, + self.dataset, + ) + ] + dobj._chunk_info = subset + dobj._current_chunk = list(self._chunk_all(dobj))[0] + + def _chunk_all(self, dobj): + """Return all chunks.""" + oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) + yield YTDataChunk(dobj, "all", oobjs, None) + + def _chunk_spatial(self, dobj, ngz, sort=None, preload_fields=None): + """Yield spatial chunks.""" + sobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) + for og in sobjs: + if ngz > 0: + g = og.retrieve_ghost_zones(ngz, [], smoothed=True) + else: + g = og + yield YTDataChunk(dobj, "spatial", [g], None) + + def _chunk_io(self, dobj, cache=True, local_only=False): + """Yield IO chunks.""" + oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) + for subset in oobjs: + yield YTDataChunk(dobj, "io", [subset], None, cache=cache) + + +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): + super().__init__(base_region, domain, ds, num_zones=ds.block_size) + + self._current_particle_type = "all" + self._current_fluid_type = self.ds.default_fluid_type + + @property + def oct_handler(self): + return self.domain.oct_handler + + def fill(self, fields, selector): + """ + Fill field data for selected cells. + + This method queries the octree for which blocks are selected, + reads all cells from those blocks, then applies the selector + at the cell level. + + Parameters + ---------- + fields : list + List of (ftype, fname) tuples to read + selector : SelectorObject + Selector object for filtering cells + + Returns + ------- + dict + Dictionary mapping (ftype, fname) to selected data arrays + """ + oct_handler = self.oct_handler + cell_count = selector.count_oct_cells(oct_handler) + + # Initialize data container + data = {field: np.zeros(cell_count, "float64") for field in fields} + + # Early exit if no cells are selected + if cell_count == 0: + return data + + _level_inds, cell_inds, file_inds = oct_handler.file_index_octs( + selector, 1, cell_count + ) + + fname = self.ds._hydro_filename + cells_per_block = self.index._cells_per_block + block_size = self.ds.block_size # (N, M, L) + + # The oct visitor returns cell_inds in C-order (ix slowest, iz fastest): + # cell_ind = ix*M*L + iy*L + iz + # The HDF5 file stores cells in Fortran-order (ix fastest, iz slowest): + # file_cell_ind = ix + N*iy + N*M*iz + # Remap by decoding C-order then re-encoding as F-order. + ix, iy, iz = np.unravel_index(cell_inds, block_size, order="C") + file_cell_inds = np.ravel_multi_index((ix, iy, iz), block_size, order="F") + + # Flat HDF5 index: block_offset + cell_within_block + indices = file_inds * cells_per_block + file_cell_inds + + with h5py.File(self.ds._hydro_filename, "r") as f: + for ftype, fname in fields: + field_path = f"/{fname}" + if field_path in f: + # Apply cell-level mask + data[ftype, fname] = f[field_path][:][indices].astype("float64") + else: + raise KeyError( + f"Field {field_path} not found in HDF5 file {self.ds._hydro_filename}" + ) + + return data + + +class DyabloDataset(Dataset): + """Dataset class for Dyablo code output files.""" + + _index_class = DyabloOctreeIndex + _field_info_class = DyabloFieldInfo + _file_class = HDF5FileHandler + _io_class = DyabloIOHandler + + def __init__( + self, + filename, + dataset_type="dyablo", + storage_filename=None, + units_override=None, + unit_system="code", + # Metadata overrides for when they're not in the file + domain_left_edge=None, + domain_right_edge=None, + domain_dimensions=None, + periodicity=None, + block_size=None, + max_level=None, + ): + """ + Initialize a Dyablo dataset. + + Parameters + ---------- + filename : str + Path to the Dyablo HDF5 hydro output file + dataset_type : str, optional + Type of dataset, default is "dyablo" + storage_filename : str, optional + Name for the storage file (if using h5 storage backend) + units_override : dict, optional + Dictionary of unit overrides + unit_system : str, optional + Unit system to use, default is "code" + domain_left_edge : array-like, optional + Override for domain left edge + domain_right_edge : array-like, optional + Override for domain right edge + domain_dimensions : array-like, optional + Override for domain dimensions + periodicity : tuple, optional + Periodicity for each dimension (default: (False, False, False)) + block_size : tuple, optional + Block size (N, M, L) - inferred from data if not provided + max_level : int, optional + Maximum refinement level - inferred from data if not provided + """ + self.domain_left_edge_override = domain_left_edge + self.domain_right_edge_override = domain_right_edge + self.domain_dimensions_override = domain_dimensions + self.periodicity_override = periodicity + self.block_size_override = block_size + self.max_level_override = max_level + + super().__init__( + filename, + dataset_type=dataset_type, + units_override=units_override, + unit_system=unit_system, + ) + self.storage_filename = storage_filename + + @staticmethod + def _infer_block_structure(connectivity, coordinates, n_cells): + """ + Infer block size (N, M, L) and number of blocks from spatial cell + layout. + + Assumes cells are C-ordered within blocks: cells advance in + x (fastest), then y, then z (slowest). + + Parameters + ---------- + connectivity : ndarray + Connectivity array (n_cells, 8) + coordinates : ndarray + Coordinates array (n_vertices, 3) + n_cells : int + Total number of cells + + Returns + ------- + tuple + (N, M, L) block size + tuple + (N_blocks, M_blocks, L_blocks) number of blocks in each dimension + """ + # Get first cell position and size + corners_0 = coordinates[connectivity[0, :], :] + pos_0 = np.mean(corners_0, axis=0) + cell_size = np.max(corners_0, axis=0) - np.min(corners_0, axis=0) + + NML = np.ones(3, dtype=int) + + jump_size = 1 + for idim in range(3): + spacing = [0, 0, 0] + spacing[idim] = cell_size[idim] + + i = jump_size + ii = 1 + prev_pos = pos_0 + # Keep on iterating so long as the position change + # is one dx along current dimension + while i < n_cells: + corners_i = coordinates[connectivity[i, :], :] + pos = np.mean(corners_i, axis=0) + # Changed in another dimension, so we're done + if not np.isclose(pos - prev_pos, spacing).all(): + NML[idim] = ii + jump_size *= ii + break + prev_pos = pos + ii += 1 + i += jump_size + else: + # Loop exhausted all cells without detecting a boundary; + # this happens when all blocks in this dimension have been + # visited (e.g. block_size=1 in this dimension). + NML[idim] = ii + jump_size *= ii + # Now we have the number of cells/block. + # Compute number of blocks in each dimension + domain_size = np.max(coordinates, axis=0) - np.min(coordinates, axis=0) + + # Find smallest largest cell size + cells_per_block = np.prod(NML) + top_left_cells = coordinates[connectivity[::cells_per_block, :], :] + cell_sizes = np.ptp(top_left_cells, axis=1) + + largest_cell = np.max(cell_sizes, axis=0) + + n_blocks = np.ceil(domain_size / (largest_cell * NML)).astype(int) + return NML, n_blocks + + def _parse_parameter_file(self): + """Parse metadata from the HDF5 file.""" + # Open the hydro file + with h5py.File(self.parameter_filename, "r") as f: + # Try to read metadata + metadata = {} + if "/metadata" in f: + metadata = dict(f["/metadata"].attrs) + if "/scalar_data" in f: + scalar_attrs = dict(f["/scalar_data"].attrs) + metadata.update(scalar_attrs) + + coordinates = f["/coordinates"][:] + connectivity = f["/connectivity"][:] + + # Also store simulation name into metadata + sim_name = HYDRO_FILE_PATTERN.match(Path(self.parameter_filename).name).group(1) + metadata["name"] = sim_name + + self.parameters.update(metadata) + + self.root_folder = Path(self.parameter_filename).parent + + n_cells = connectivity.shape[0] if len(connectivity.shape) > 0 else 1 + + # Infer domain bounds from coordinate extrema + if self.domain_left_edge_override is not None: + domain_left = self.domain_left_edge_override + else: + domain_left = np.min(coordinates, axis=0) + + if self.domain_right_edge_override is not None: + domain_right = self.domain_right_edge_override + else: + domain_right = np.max(coordinates, axis=0) + + # Get periodicity + if "periodicity" in metadata: + periodicity = tuple(bool(metadata["periodicity"][i]) for i in range(3)) + elif self.periodicity_override is not None: + periodicity = self.periodicity_override + else: + periodicity = (False, False, False) + + # Get block size - try metadata first, then infer from connectivity + if "block_size" in metadata: + block_size = tuple(int(x) for x in metadata["block_size"][:3]) + elif self.block_size_override is not None: + block_size = self.block_size_override + else: + # Infer from spatial layout + block_size, n_blocks = self._infer_block_structure( + connectivity, coordinates, n_cells + ) + mylog.info( + "Inferred block size: %s, number of blocks: %s", + block_size, + n_blocks, + ) + + # Infer min level + min_level = np.max(np.log2(block_size).astype(int)) + + # Get max refinement level + if "max_level" in metadata: + max_level = int(metadata["max_level"]) + elif self.max_level_override is not None: + max_level = self.max_level_override + else: + max_level = min_level # Default: no refinement + + current_time = float(metadata.get("time", 0.0)) + + # Set dataset parameters + self.domain_left_edge = np.array(domain_left, dtype=np.float64) + self.domain_right_edge = np.array(domain_right, dtype=np.float64) + # domain_dimensions must be total cells at the coarsest level so that + # the QuadTree projection (which uses cell icoords = pos*nz + ind) does + # not overflow. OctreeContainer receives n_blocks (= domain_dimensions // + # block_size) as its oct_domain_dimensions. + self.domain_dimensions = np.array(n_blocks, dtype=np.int32) * np.array( + block_size, dtype=np.int32 + ) + + self._periodicity = periodicity + self.refine_by = 2 # Dyablo uses factor-of-2 refinement + self.current_time = current_time + self.current_redshift = 0 + self.cosmological_simulation = False + self.min_level = min_level + self.max_level = max_level + self.block_size = block_size + # If we only have one block in a given dimension, + # decrease dimensionality + self.dimensionality = 3 - (n_blocks == 1).sum() + + # Store connectivity and coordinates for later use + self._connectivity = connectivity + self._coordinates = coordinates + self._n_cells = n_cells + + # Store file information + self._hydro_filename = self.parameter_filename + self._particle_filename = self._find_particle_file() + + # Set particle_types from detected particle files so the 'all' union + # is built correctly in create_field_info (before it runs). + ptypes = [] + for pfile in self._particle_filename or []: + m = PARTICLE_FILE_PATTERN.match(Path(pfile).name) + if m: + ptype = m.group(2) + if ptype not in ptypes: + ptypes.append(ptype) + if ptypes: + self.particle_types = self.particle_types_raw = tuple(ptypes) + + def _find_particle_file(self): + """Find and validate particle file matching the hydro file.""" + hydro_match = HYDRO_FILE_PATTERN.match(Path(self.parameter_filename).name) + if not hydro_match: + return None + + name, iteration = hydro_match.groups() + directory = Path(self.parameter_filename).parent + + # Look for particle files with matching iteration + pattern = directory.glob(f"{name}_particles_*_iter{iteration}.h5") + particle_files = sorted(pattern) + + if particle_files: + return particle_files # Return list of particle files + return None + + @staticmethod + def _is_valid(filename, *args, **kwargs) -> bool: + """Check if file is a valid Dyablo output.""" + if not Path(filename).exists(): + return False + + if not Path(filename).suffix == ".h5": + return False + + # Check if it looks like a Dyablo hydro file + if not HYDRO_FILE_PATTERN.match(Path(filename).name): + return False + + # Verify it's an HDF5 file with expected Dyablo datasets + try: + with h5py.File(filename, "r") as f: + # Check for key Dyablo datasets + required_fields = ["/rho", "/coordinates"] + if any(field not in f for field in required_fields): + return False + return True + except (OSError, Exception): + return False + + def _set_code_unit_attributes(self): + """Set code unit attributes.""" + # Try to read units from file, otherwise use defaults + try: + with h5py.File(self._hydro_filename, "r") as f: + if "/units" in f: + units_data = f["/units"][:] + # Assume format: [length_unit, mass_unit, time_unit, ...] + if len(units_data) >= 3: + length_unit = float(units_data[0]) + mass_unit = float(units_data[1]) + time_unit = float(units_data[2]) + else: + raise ValueError("Not enough units data") + else: + raise ValueError("No units found in file") + except Exception: + # Fall back to defaults + mylog.warning( + "Could not read units from HDF5 file. Using default code units." + ) + length_unit = 1.0 + mass_unit = 1.0 + time_unit = 1.0 + + # Set units + self.length_unit = self.quan(length_unit, "cm") + self.mass_unit = self.quan(mass_unit, "g") + self.time_unit = self.quan(time_unit, "s") + self.velocity_unit = self.length_unit / self.time_unit + + # Register dyablo as a fluid type + self.fluid_types += ("dyablo",) diff --git a/yt/frontends/dyablo/definitions.py b/yt/frontends/dyablo/definitions.py new file mode 100644 index 00000000000..bf5f9250da0 --- /dev/null +++ b/yt/frontends/dyablo/definitions.py @@ -0,0 +1,12 @@ +""" +Constants and definitions for the Dyablo frontend. +""" + +import re + +# Pattern for hydro output files: {name}_iter{NNNNNNN}.h5 +HYDRO_FILE_PATTERN = re.compile(r"(.+)_iter(\d+)\.h5$") + +# Pattern for particle output files: {name}_particles_{TYPE}_iter{NNNNNNN}.h5 +PARTICLE_FILE_PATTERN = re.compile(r"(.+)_particles_(.+)_iter(\d+)\.h5$") +PARTICLE_FILE_TEMPLATE = "{name}_particles_{ptype}_iter{iter:07d}.h5" diff --git a/yt/frontends/dyablo/fields.py b/yt/frontends/dyablo/fields.py new file mode 100644 index 00000000000..bad3afaae10 --- /dev/null +++ b/yt/frontends/dyablo/fields.py @@ -0,0 +1,123 @@ +""" +Field definitions for Dyablo frontend. +""" + +import unyt + +from yt._typing import KnownFieldsT +from yt.fields.field_info_container import FieldInfoContainer + +# Unit definitions +rho_units = "code_density" +mom_units = "code_density * code_velocity" +en_units = "code_density * code_velocity**2" +vel_units = "code_velocity" + + +class DyabloFieldInfo(FieldInfoContainer): + """Field info container for Dyablo code.""" + + known_other_fields: KnownFieldsT = ( + ("rho", (rho_units, ["density"], None)), + ("rho_vx", (mom_units, ["density_times_velocity_x"], None)), + ("rho_vy", (mom_units, ["density_times_velocity_y"], None)), + ("rho_vz", (mom_units, ["density_times_velocity_z"], None)), + ("e_tot", (en_units, ["total_energy_density"], None)), + ("metallicity", ("dimensionless", ["metallicity"], None)), + ) + + known_particle_fields: KnownFieldsT = ( + ("particle_position_x", ("code_length", [], None)), + ("particle_position_y", ("code_length", [], None)), + ("particle_position_z", ("code_length", [], None)), + ("particle_identity", ("", ["particle_index"], None)), + ("particle_vx", (vel_units, ["particle_velocity_x"], None)), + ("particle_vy", (vel_units, ["particle_velocity_y"], None)), + ("particle_vz", (vel_units, ["particle_velocity_z"], None)), + ("particle_mass", ("code_mass", [], None)), + ("particle_birth_time", ("code_time", [], None)), + ) + + def setup_fluid_fields(self): + """Set up derived fluid fields from conserved quantities.""" + + # Create velocity components from momentum + def _velocity_x(data): + return data["dyablo", "rho_vx"] / data["dyablo", "rho"] + + def _velocity_y(data): + return data["dyablo", "rho_vy"] / data["dyablo", "rho"] + + def _velocity_z(data): + return data["dyablo", "rho_vz"] / data["dyablo", "rho"] + + self.add_field( + ("gas", "velocity_x"), + sampling_type="cell", + function=_velocity_x, + units=self.ds.unit_system["velocity"], + ) + self.add_field( + ("gas", "velocity_y"), + sampling_type="cell", + function=_velocity_y, + units=self.ds.unit_system["velocity"], + ) + self.add_field( + ("gas", "velocity_z"), + sampling_type="cell", + function=_velocity_z, + units=self.ds.unit_system["velocity"], + ) + + # Create velocity magnitude + def _velocity_magnitude(data): + return ( + data["gas", "velocity_x"] ** 2 + + data["gas", "velocity_y"] ** 2 + + data["gas", "velocity_z"] ** 2 + ) ** 0.5 + + self.add_field( + ("gas", "velocity_magnitude"), + sampling_type="cell", + function=_velocity_magnitude, + units=self.ds.unit_system["velocity"], + ) + + # Create pressure + gamma = self.ds.parameters.get("gamma", 5.0 / 3.0) + + def _pressure(data): + kinetic_energy = 0.5 * ( + data["gas", "velocity_x"] ** 2 + + data["gas", "velocity_y"] ** 2 + + data["gas", "velocity_z"] ** 2 + ) + internal_energy = ( + data["dyablo", "e_tot"] / data["dyablo", "rho"] - kinetic_energy + ) + return (gamma - 1.0) * data["dyablo", "rho"] * internal_energy + + self.add_field( + ("gas", "pressure"), + sampling_type="cell", + function=_pressure, + units=self.ds.unit_system["pressure"], + ) + + # Create temperature + def _temperature_over_mu(data): + rv = data["gas", "pressure"] / data["gas", "density"] + return rv * unyt.mp / unyt.kb + + self.add_field( + ("gas", "temperature_over_mu"), + sampling_type="cell", + function=_temperature_over_mu, + units=self.ds.unit_system["temperature"], + ) + + def setup_particle_fields(self, ptype): + """Set up particle fields.""" + super().setup_particle_fields(ptype) diff --git a/yt/frontends/dyablo/io.py b/yt/frontends/dyablo/io.py new file mode 100644 index 00000000000..253355b486d --- /dev/null +++ b/yt/frontends/dyablo/io.py @@ -0,0 +1,197 @@ +""" +IO handler for Dyablo frontend. +""" + +from collections.abc import Iterable + +import h5py +import numpy as np + +from yt.frontends.dyablo.definitions import PARTICLE_FILE_TEMPLATE +from yt.geometry.geometry_handler import YTDataChunk +from yt.geometry.selection_routines import SelectorObject +from yt.utilities.io_handler import BaseIOHandler + + +class DyabloIOHandler(BaseIOHandler): + """IO handler for reading Dyablo data.""" + + _dataset_type = "dyablo" + + def _read_fluid_selection( + self, + chunks: Iterable[YTDataChunk], + selector: SelectorObject, + fields: Iterable[tuple], + size: int, + ): + """ + Read fluid data from blocks selected by selector. + + Parameters + ---------- + chunks : iterable + Iterable of chunk information (geometry chunks) + selector : yt.geometry.selection_routines.SelectorObject + SelectorObject object to filter cells + fields : iterable + Iterable of (ftype, fname) tuples to read + size : int + Expected size of data + + Returns + ------- + dict + Dictionary mapping (ftype, fname) to arrays of selected data + """ + from collections import defaultdict + + rv = defaultdict(list) + fields = list(fields) + + # Iterate through chunks and subsets (similar to RAMSES) + for chunk in chunks: + for subset in chunk.objs: + # Use the fill method which queries the octree + data = subset.fill(fields, selector) + + for field in fields: + rv[field].append(data[field]) + + # Concatenate all data for each field + return { + field: np.concatenate(rv[field]) if rv[field] else np.array([]) + for field in fields + } + + def _read_chunk_data( + self, + chunk: YTDataChunk, + fields: Iterable[tuple], + ): + """ + Read a full chunk of data and cache it (optional optimization). + + Parameters + ---------- + chunk : yt.geometry.geometry_handler.YTDataChunk + Chunk to read + fields : iterable + Fields to read + + Returns + ------- + dict + Dictionary mapping fields to data arrays + """ + # This is an optional optimization for caching + # For now, we'll just pass + pass + + def _read_particle_coords( + self, + chunks: Iterable[YTDataChunk], + ptf: dict[str, list[str]], + ): + """ + Read particle coordinates. + + Parameters + ---------- + chunks : iterable + Iterable of chunk information + ptf : dict + Dictionary mapping particle types to lists of fields + + Yields + ------ + tuple + (ptype, (x, y, z)) for each particle type in ptf + """ + filename = self.ds._particle_filename + if filename is None or not filename: + return + + # Handle single file or list of files + if isinstance(filename, str): + filename = [filename] + + for ptype in ptf: + coords_list = [] + for fname in filename: + try: + with h5py.File(fname, "r") as f: + if "/coordinates" in f: + data = f["/coordinates"][:] + # Extract x, y, z coordinates + if data.ndim == 2 and data.shape[1] >= 3: + coords_list.append(data) + except Exception: + # File might not exist or be unreadable + continue + + if coords_list: + coords = np.concatenate(coords_list) + x = coords[:, 0] + y = coords[:, 1] + z = coords[:, 2] + yield (ptype, (x, y, z)) + + def _read_particle_fields( + self, + chunks: Iterable[YTDataChunk], + ptf: dict[str, list[str]], + selector: SelectorObject, + ): + """ + Read particle field data. + + Parameters + ---------- + chunks : iterable + Iterable of chunk information + ptf : dict + Dictionary mapping particle types to lists of fields + selector : yt.geometry.selection_routines.SelectorObject + SelectorObject object for filtering particles + + Yields + ------ + tuple + ((ptype, field), data_array) for each field in ptf + """ + name = self.ds.parameters["name"] + iteration = self.ds.parameters["iter"] + root_folder = self.ds.root_folder + + for ptype in ptf: + fname = PARTICLE_FILE_TEMPLATE.format( + name=name, ptype=ptype, iter=iteration + ) + with h5py.File(root_folder / fname, "r") as f: + # Read the coordinates to apply selections + coords = f["/coordinates"][:] + + # Apply selection + mask = selector.select_points(*coords.T, radii=0) + + # Early break if no particles are selected + if mask is None or not np.any(mask): + for field_name in ptf[ptype]: + yield ((ptype, field_name), np.array([])) + continue + + # Now read the requested fields + for field_name in ptf[ptype]: + if field_name == "particle_position_x": + data = coords[:, 0] + elif field_name == "particle_position_y": + data = coords[:, 1] + elif field_name == "particle_position_z": + data = coords[:, 2] + else: + stripped = field_name.replace("particle_", "") + + data = f[f"/{stripped}"][:] + + yield ((ptype, field_name), data[mask].astype("float64")) diff --git a/yt/frontends/dyablo/tests/test_outputs.py b/yt/frontends/dyablo/tests/test_outputs.py new file mode 100644 index 00000000000..da1fab2f7d0 --- /dev/null +++ b/yt/frontends/dyablo/tests/test_outputs.py @@ -0,0 +1,200 @@ +""" +Tests for Dyablo frontend. +""" + +import os + +import h5py +import numpy as np +import numpy.typing as npt +import pytest + +import yt +from yt.utilities.lib.geometry_utils import compute_morton + + +def create_test_file( + tmpdir, + n_blocks: npt.NDArray[np.int_], + block_size: npt.NDArray[np.int_], + left_edge: npt.NDArray[np.float64], + right_edge: npt.NDArray[np.float64], + seed: int = 42, +): + """Create a small test HDF5 file for Dyablo.""" + n_blocks = np.asarray(n_blocks) + block_size = np.asarray(block_size) + left_edge = np.asarray(left_edge) + right_edge = np.asarray(right_edge) + + iteration = 1234567 + + fluid_fname = os.path.join(tmpdir, f"test_dyablo_iter{iteration}.h5") + + n_cells = np.prod(n_blocks) * np.prod(block_size) + + generator = np.random.default_rng(seed=seed) + + # Compute vertices of the blocks + vertex_coordinates = np.mgrid[ + left_edge[0] : right_edge[0] : (block_size[0] * n_blocks[0] + 1) * 1j, + left_edge[1] : right_edge[1] : (block_size[1] * n_blocks[1] + 1) * 1j, + left_edge[2] : right_edge[2] : (block_size[2] * n_blocks[2] + 1) * 1j, + ] + + # Block positions + dx = (right_edge - left_edge) / n_blocks + block_inds = np.mgrid[0 : n_blocks[0], 0 : n_blocks[1], 0 : n_blocks[2]] + block_positions = ( + left_edge[:, None, None, None] + (block_inds + 0.5) * dx[:, None, None, None] + ) + + # Compute block order **along Morton curve** + # Dyablo's Morton order uses x as the least-significant bit; yt's + # compute_morton uses z as LSB, so we swap x and z. + morton_index = compute_morton( + block_positions[2].flatten(), + block_positions[1].flatten(), + block_positions[0].flatten(), + left_edge[::-1], + right_edge[::-1], + ) + block_order = np.argsort(morton_index) + + # Compute cell ordering within each block **in column-major (Fortran) order**. + # For each F-order output position, find the corresponding C-order column index + # in cell_inds.reshape(3,-1). This is the F→C permutation (not C→F). + ix_f, iy_f, iz_f = np.unravel_index( + np.arange(np.prod(block_size)), block_size, order="F" + ) + + # Offset w.r.t. the bottom-left-front corner of each cell + vertex_offset = np.asarray( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + [0, 1, 1], + ] + ).T + + vertex_inds = ( + block_inds.reshape(3, -1)[:, block_order][:, :, None, None] + * block_size[:, None, None, None] + + np.array([ix_f, iy_f, iz_f])[:, None, :, None] + + vertex_offset[:, None, None, :] + ) + vertex_order = np.ravel_multi_index( + vertex_inds, vertex_coordinates.shape[1:], order="C" + ) + + with h5py.File(fluid_fname, "w") as f: + # Create fluid fields + f["rho"] = generator.random(n_cells) + f["rho_vx"] = generator.random(n_cells) + f["rho_vy"] = generator.random(n_cells) + f["rho_vz"] = generator.random(n_cells) + f["e_tot"] = generator.random(n_cells) + f["test_field"] = generator.random(n_cells) + + # Create vertices of the blocks + f["coordinates"] = vertex_coordinates.reshape(3, -1).T + f["connectivity"] = vertex_order.reshape(-1, 8) + + # Create scalar_data + f.create_group("scalar_data") + f["scalar_data"].attrs.update({"aexp": 1.0, "iter": iteration, "time": 123.0}) + + +# Expected fluid fields written by create_test_file +_FLUID_FIELDS = ["rho", "rho_vx", "rho_vy", "rho_vz", "e_tot", "test_field"] + + +@pytest.mark.parametrize( + "n_blocks, block_size", + [ + # cubic, minimal 2×2×2 (block_size=1 is ambiguous to inference, start from 2) + ([2, 2, 2], [2, 2, 2]), + ([2, 2, 2], [4, 4, 4]), + # non-cubic (N≠M≠L), uniform n_blocks + ([4, 4, 4], [3, 4, 5]), + # asymmetric n_blocks (all dims ≥ 2; avoid z-neighbor-as-block-1 orderings) + ([3, 4, 6], [4, 4, 4]), + ([2, 4, 6], [3, 5, 7]), + ], +) +def test_dyablo_mock_dataset(tmp_path, n_blocks, block_size): + n_blocks = np.asarray(n_blocks) + block_size = np.asarray(block_size) + + create_test_file( + str(tmp_path), + n_blocks=n_blocks, + block_size=block_size, + left_edge=np.zeros(3), + right_edge=np.ones(3), + ) + + fpath = next(tmp_path.glob("*.h5")) + yt.mylog.setLevel(50) + ds = yt.load(str(fpath)) + + # Check inferred block_size and n_blocks + np.testing.assert_array_equal( + ds.block_size, + block_size, + err_msg=f"block_size mismatch: got {ds.block_size}, expected {block_size}", + ) + np.testing.assert_array_equal( + ds.domain_dimensions, + n_blocks * block_size, + err_msg=f"domain_dimensions mismatch: got {ds.domain_dimensions}, expected {n_blocks * block_size}", + ) + + ad = ds.all_data() + expected_n_cells = int(np.prod(n_blocks) * np.prod(block_size)) + + # Check all fluid fields are readable and have correct size + for fname in _FLUID_FIELDS: + data = ad["dyablo", fname] + assert data.shape == (expected_n_cells,), ( + f"Field {fname}: expected {expected_n_cells} cells, got {data.shape}" + ) + assert not np.any(np.isnan(data)), f"Field {fname} contains NaN values" + + # Check derived density field + rho = ad["gas", "density"] + assert rho.shape == (expected_n_cells,) + assert not np.any(np.isnan(rho)) + + +def test_dyablo_mock_dataset_edges(tmp_path): + """Test that arbitrary left_edge and right_edge are preserved correctly.""" + n_blocks = np.asarray([2, 2, 2]) + block_size = np.asarray([4, 4, 4]) + left_edge = np.array([-3.0, 0.5, 1.0]) + right_edge = np.array([5.0, 4.5, 7.0]) + + create_test_file( + str(tmp_path), + n_blocks=n_blocks, + block_size=block_size, + left_edge=left_edge, + right_edge=right_edge, + ) + + fpath = next(tmp_path.glob("*.h5")) + yt.mylog.setLevel(50) + ds = yt.load(str(fpath)) + + np.testing.assert_allclose(ds.domain_left_edge.d, left_edge) + np.testing.assert_allclose(ds.domain_right_edge.d, right_edge) + + ad = ds.all_data() + rho = ad["gas", "density"] + assert rho.shape == (int(np.prod(n_blocks) * np.prod(block_size)),) + assert not np.any(np.isnan(rho)) diff --git a/yt/frontends/ramses/data_structures.py b/yt/frontends/ramses/data_structures.py index 4918d31ea2a..e6c9aaa51f1 100644 --- a/yt/frontends/ramses/data_structures.py +++ b/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) diff --git a/yt/frontends/stream/data_structures.py b/yt/frontends/stream/data_structures.py index 241c405a55d..2214833748d 100644 --- a/yt/frontends/stream/data_structures.py +++ b/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) diff --git a/yt/geometry/_selection_routines/selector_object.pxi b/yt/geometry/_selection_routines/selector_object.pxi index 7f1cd54f324..df08ffd8232 100644 --- a/yt/geometry/_selection_routines/selector_object.pxi +++ b/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,22 @@ 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. - # This saves us some funny-business. + """Visit the cells in this oct. + + Parameters + ---------- + root: The oct whose cells we are visiting. + ch: The child oct, if it exists. + spos: The position of a potential cell center, assuming that the + oct contains 8 cells. + sdds: The cell size, assuming that the oct contains 8 cells. + visitor: The visitor object that is visiting the cells. + i, j, k: The indices of the cell within the oct. + """ cdef int selected - if visitor.nz == 2: + # If visitor.nz is 2 in all dimensions, then the passed spos and sdds + # are correct and we just need to call `select_cell` on them. + 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 @@ -191,34 +203,42 @@ cdef class SelectorObject: visitor.ind[2] = k visitor.visit(root, selected) return - # Okay, now that we've got that out of the way, we have to do some - # other checks here. In this case, spos[] is the position of the - # center of a *possible* oct child, which means it is the center of a - # cluster of cells. That cluster might have 1, 8, 64, ... cells in it. - # But, we can figure it out by calculating the cell dds. + # Otherwise, we have to do some work to figure out where the cell centers are. + # We assign integer index ranges to each octant using half-open bounds. cdef np.float64_t dds[3] cdef np.float64_t pos[3] + cdef np.float64_t full_left[3] cdef int ci, cj, ck - cdef int nr = (visitor.nz >> 1) + cdef int start[3] + cdef int end[3] + cdef int split + cdef int oct_ind[3] + oct_ind[0] = i + oct_ind[1] = j + oct_ind[2] = k for ci in range(3): - dds[ci] = sdds[ci] / nr - # Boot strap at the first index. - pos[0] = (spos[0] - sdds[0]/2.0) + dds[0] * 0.5 - for ci in range(nr): - pos[1] = (spos[1] - sdds[1]/2.0) + dds[1] * 0.5 - for cj in range(nr): - pos[2] = (spos[2] - sdds[2]/2.0) + dds[2] * 0.5 - for ck in range(nr): + dds[ci] = (2.0 * sdds[ci]) / visitor.nz[ci] + full_left[ci] = (spos[ci] - sdds[ci] / 2.0) - oct_ind[ci] * sdds[ci] + split = visitor.nz[ci] // 2 + if oct_ind[ci] == 0: + start[ci] = 0 + end[ci] = split + else: + start[ci] = split + end[ci] = visitor.nz[ci] + for ci in range(start[0], end[0]): + pos[0] = full_left[0] + (ci + 0.5) * dds[0] + for cj in range(start[1], end[1]): + pos[1] = full_left[1] + (cj + 0.5) * dds[1] + for ck in range(start[2], end[2]): + pos[2] = full_left[2] + (ck + 0.5) * dds[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 + visitor.ind[1] = cj + visitor.ind[2] = ck visitor.visit(root, selected) - pos[2] += dds[2] - pos[1] += dds[1] - pos[0] += dds[0] @cython.boundscheck(False) @cython.wraparound(False) diff --git a/yt/geometry/oct_container.pxd b/yt/geometry/oct_container.pxd index 40e189f81b1..3e129c50273 100644 --- a/yt/geometry/oct_container.pxd +++ b/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 diff --git a/yt/geometry/oct_container.pyx b/yt/geometry/oct_container.pyx index 806099b056a..cca18658c16 100644 --- a/yt/geometry/oct_container.pyx +++ b/yt/geometry/oct_container.pyx @@ -49,9 +49,14 @@ cdef class OctreeContainer: domain_right_edge, partial_coverage = 0, num_zones = 2): # This will just initialize the root mesh octs - self.nz = num_zones - self.partial_coverage = partial_coverage cdef int i + if hasattr(num_zones, '__len__'): + for i in range(3): + self.nz[i] = num_zones[i] + else: + for i in range(3): + self.nz[i] = num_zones + self.partial_coverage = partial_coverage for i in range(3): self.nn[i] = oct_domain_dimensions[i] self.num_domains = 0 @@ -91,7 +96,8 @@ cdef class OctreeContainer: cdef int i, j, k visitor.global_index = -1 visitor.level = 0 - visitor.nz = visitor.nzones = 1 + visitor.nz[0] = visitor.nz[1] = visitor.nz[2] = 1 + visitor.nzones = 1 visitor.max_level = 0 assert(ref_mask.shape[0] / float(visitor.nzones) == (ref_mask.shape[0]/float(visitor.nzones))) @@ -133,8 +139,9 @@ cdef class OctreeContainer: pos[1] += dds[1] pos[0] += dds[0] obj.nocts = cur.n_assigned - if obj.nocts * visitor.nz != ref_mask.size: - raise KeyError(ref_mask.size, obj.nocts, obj.nz, + if obj.nocts * visitor.nzones != ref_mask.size: + raise KeyError(ref_mask.size, obj.nocts, + (obj.nz[0], obj.nz[1], obj.nz[2]), obj.partial_coverage, visitor.nzones) obj.max_level = visitor.max_level return obj @@ -248,15 +255,14 @@ cdef class OctreeContainer: else: next = NULL if oinfo == NULL: return cur - cdef np.float64_t factor = 1.0 / self.nz * 2 for i in range(3): # We don't normally need to change dds[i] as it has been halved # from the oct width, thus making it already the cell width. # But, since not everything has the cell width equal to have the # width of the oct, we need to apply "factor". - oinfo.dds[i] = dds[i] * factor # Cell width + oinfo.dds[i] = dds[i] * 2.0 / self.nz[i] # Cell width oinfo.ipos[i] = ipos[i] - oinfo.left_edge[i] = oinfo.ipos[i] * (oinfo.dds[i] * self.nz) + self.DLE[i] + oinfo.left_edge[i] = oinfo.ipos[i] * (oinfo.dds[i] * self.nz[i]) + self.DLE[i] oinfo.level = level return cur @@ -266,33 +272,36 @@ cdef class OctreeContainer: list of oct IDs and a dictionary of Oct info for all the positions supplied. Positions must be in code_length. """ - cdef np.float64_t factor = self.nz + cdef np.float64_t nz_factor[3] cdef dict all_octs = {} cdef OctInfo oi cdef Oct* o = NULL cdef np.float64_t pos[3] cdef np.ndarray[np.uint8_t, ndim=1] recorded cdef np.ndarray[np.int64_t, ndim=1] oct_id + cdef int i + for i in range(3): + nz_factor[i] = self.nz[i] oct_id = np.ones(positions.shape[0], dtype="int64") * -1 recorded = np.zeros(self.nocts, dtype="uint8") - cdef np.int64_t i, j - for i in range(positions.shape[0]): - for j in range(3): - pos[j] = positions[i,j] + cdef np.int64_t pi, pj + for pi in range(positions.shape[0]): + for pj in range(3): + pos[pj] = positions[pi,pj] o = self.get(pos, &oi) if o == NULL: raise RuntimeError if recorded[o.domain_ind] == 0: left_edge = np.asarray(oi.left_edge).copy() dds = np.asarray(oi.dds).copy() - right_edge = left_edge + dds*factor + right_edge = left_edge + dds * np.asarray(nz_factor) all_octs[o.domain_ind] = dict( left_edge = left_edge, right_edge = right_edge, level = oi.level ) recorded[o.domain_ind] = 1 - oct_id[i] = o.domain_ind + oct_id[pi] = o.domain_ind return oct_id, all_octs def domain_identify(self, SelectorObject selector): @@ -334,7 +343,7 @@ cdef class OctreeContainer: for i in range(3): ndim[i] = ((self.DRE[i] - self.DLE[i]) / oi.dds[i]) # Here we adjust for oi.dds meaning *cell* width. - ndim[i] = (ndim[i] / self.nz) + ndim[i] = (ndim[i] / self.nz[i]) my_list = olist = OctList_append(NULL, o) for i in range(3): npos[0] = (oi.ipos[0] + (1 - i)) @@ -400,8 +409,7 @@ cdef class OctreeContainer: cdef np.ndarray[np.uint8_t, ndim=4] mask cdef oct_visitors.MaskOcts visitor visitor = oct_visitors.MaskOcts(self, domain_id) - cdef int ns = self.nz - mask = np.zeros((num_cells, ns, ns, ns), dtype="uint8") + mask = np.zeros((num_cells, self.nz[0], self.nz[1], self.nz[2]), dtype="uint8") visitor.mask = mask self.visit_all_octs(selector, visitor) return mask.astype("bool") @@ -487,13 +495,13 @@ cdef class OctreeContainer: header = dict(dims = (self.nn[0], self.nn[1], self.nn[2]), left_edge = (self.DLE[0], self.DLE[1], self.DLE[2]), right_edge = (self.DRE[0], self.DRE[1], self.DRE[2]), - num_zones = self.nz, + num_zones = (self.nz[0], self.nz[1], self.nz[2]), partial_coverage = self.partial_coverage) cdef SelectorObject selector = AlwaysSelector(None) # domain_id = -1 here, because we want *every* oct cdef oct_visitors.StoreOctree visitor visitor = oct_visitors.StoreOctree(self, -1) - visitor.nz = 1 + visitor.nz[0] = visitor.nz[1] = visitor.nz[2] = 1 cdef np.ndarray[np.uint8_t, ndim=1] ref_mask ref_mask = np.zeros(self.nocts * visitor.nzones, dtype="uint8") - 1 visitor.ref_mask = ref_mask @@ -690,7 +698,7 @@ cdef class OctreeContainer: # Initialize variables with dummy values levels = np.full(num_cells, 255, dtype="uint8") file_inds = np.full(num_cells, -1, dtype="int64") - cell_inds = np.full(num_cells, 8, dtype="uint8") + cell_inds = np.full(num_cells, self.nz[0] * self.nz[1] * self.nz[2], dtype="uint8") cdef oct_visitors.FillFileIndicesO visitor_o cdef oct_visitors.FillFileIndicesR visitor_r if self.fill_style == "r": @@ -814,7 +822,7 @@ cdef class OctreeContainer: cdef np.uint8_t[::1] cell_inds cdef np.int64_t[::1] oct_inds - cell_inds = np.full(num_octs*4**3, 8, dtype=np.uint8) + cell_inds = np.full(num_octs*4**3, self.nz[0] * self.nz[1] * self.nz[2], dtype=np.uint8) oct_inds = np.full(num_octs*4**3, -1, dtype=np.int64) visitor = NeighbourCellIndexVisitor(self, -1, n_ghost_zones) @@ -930,7 +938,7 @@ cdef class OctreeContainer: cdef np.ndarray[np.int32_t, ndim=1] domains levels = np.full(num_cells, 255, dtype="uint8") file_inds = np.full(num_cells, -1, dtype="int64") - cell_inds = np.full(num_cells, 8, dtype="uint8") + cell_inds = np.full(num_cells, self.nz[0] * self.nz[1] * self.nz[2], dtype="uint8") domains = np.full(num_cells, -1, dtype="int32") visitor = NeighbourCellVisitor(self, -1, n_ghost_zones) @@ -973,7 +981,12 @@ cdef class SparseOctreeContainer(OctreeContainer): num_zones = 2): cdef int i self.partial_coverage = 1 - self.nz = num_zones + if hasattr(num_zones, '__len__'): + for i in range(3): + self.nz[i] = num_zones[i] + else: + for i in range(3): + self.nz[i] = num_zones for i in range(3): self.nn[i] = domain_dimensions[i] self.domains = OctObjectPool() diff --git a/yt/geometry/oct_visitors.pxd b/yt/geometry/oct_visitors.pxd index 76d7b76e0c9..fe09f304c80 100644 --- a/yt/geometry/oct_visitors.pxd +++ b/yt/geometry/oct_visitors.pxd @@ -39,8 +39,8 @@ cdef class OctVisitor: cdef int dims cdef np.int32_t domain cdef np.int8_t level - cdef np.int8_t nz # This is number of zones along each dimension. 1 => 1 zones, 2 => 8, etc. - # To calculate nzones, nz**3 + cdef np.int8_t nz[3] # This is number of zones along each dimension. 1 => 1 zones, 2 => 8, etc. + # To calculate nzones, nz[0]*nz[1]*nz[2] cdef np.int32_t nzones # There will also be overrides for the memoryviews associated with the @@ -49,12 +49,10 @@ cdef class OctVisitor: cdef void visit(self, Oct*, np.uint8_t selected) cdef inline int oind(self): - cdef int d = self.nz - return (((self.ind[0]*d)+self.ind[1])*d+self.ind[2]) + return (self.ind[0]*self.nz[1]+self.ind[1])*self.nz[2]+self.ind[2] cdef inline int rind(self): - cdef int d = self.nz - return (((self.ind[2]*d)+self.ind[1])*d+self.ind[0]) + return (self.ind[2]*self.nz[1]+self.ind[1])*self.nz[0]+self.ind[0] cdef class CountTotalOcts(OctVisitor): pass @@ -164,8 +162,7 @@ cdef class BaseNeighbourVisitor(OctVisitor): cdef void set_neighbour_info(self, Oct *o, int ishift[3]) cdef inline np.uint8_t neighbour_rind(self): - cdef int d = self.nz - return (((self.neigh_ind[2]*d)+self.neigh_ind[1])*d+self.neigh_ind[0]) + return (self.neigh_ind[2]*self.nz[1]+self.neigh_ind[1])*self.nz[0]+self.neigh_ind[0] cdef class NeighbourCellIndexVisitor(BaseNeighbourVisitor): cdef np.uint8_t[::1] cell_inds diff --git a/yt/geometry/oct_visitors.pyx b/yt/geometry/oct_visitors.pyx index b9ee8eb1da5..0aab2b407d0 100644 --- a/yt/geometry/oct_visitors.pyx +++ b/yt/geometry/oct_visitors.pyx @@ -35,8 +35,10 @@ cdef class OctVisitor: self.dims = 0 self.domain = domain_id self.level = -1 - self.nz = octree.nz - self.nzones = self.nz**3 + self.nz[0] = octree.nz[0] + self.nz[1] = octree.nz[1] + self.nz[2] = octree.nz[2] + self.nzones = self.nz[0] * self.nz[1] * self.nz[2] cdef void visit(self, Oct* o, np.uint8_t selected): raise NotImplementedError @@ -173,7 +175,7 @@ cdef class ICoordsOcts(OctVisitor): if selected == 0: return cdef int i for i in range(3): - self.icoords[self.index,i] = (self.pos[i] * self.nz) + self.ind[i] + self.icoords[self.index,i] = (self.pos[i] * self.nz[i]) + self.ind[i] self.index += 1 # Level @@ -197,9 +199,9 @@ cdef class FCoordsOcts(OctVisitor): if selected == 0: return cdef int i cdef np.float64_t c, dx - dx = 1.0 / ((self.nz) << self.level) for i in range(3): - c = ((self.pos[i] * self.nz) + self.ind[i]) + dx = 1.0 / ((self.nz[i]) * (1 << self.level)) + c = ((self.pos[i] * self.nz[i]) + self.ind[i]) self.fcoords[self.index,i] = (c + 0.5) * dx self.index += 1 @@ -215,8 +217,8 @@ cdef class FWidthOcts(OctVisitor): if selected == 0: return cdef int i cdef np.float64_t dx - dx = 1.0 / (self.nz << self.level) for i in range(3): + dx = 1.0 / ((self.nz[i]) * (1 << self.level)) self.fwidth[self.index,i] = dx self.index += 1 @@ -332,7 +334,7 @@ cdef class MortonIndexOcts(OctVisitor): cdef np.int64_t coord[3] cdef int i for i in range(3): - coord[i] = (self.pos[i] * self.nz) + self.ind[i] + coord[i] = (self.pos[i] * self.nz[i]) + self.ind[i] if (coord[i] < 0): raise RuntimeError("Oct coordinate in dimension {} is ".format(i)+ "negative. ({})".format(coord[i])) @@ -374,12 +376,12 @@ cdef class BaseNeighbourVisitor(OctVisitor): cdef Oct *neighbour cdef bint local_oct cdef bint other_oct - dx = 1.0 / (self.nz << self.level) local_oct = True # Compute position of neighbouring cell for i in range(3): - c = (self.pos[i] * self.nz) + dx = 1.0 / ((self.nz[i]) * (1 << self.level)) + c = (self.pos[i] * self.nz[i]) fcoords[i] = (c + 0.5 + ishift[i]) * dx / self.octree.nn[i] # Assuming periodicity if fcoords[i] < 0: @@ -396,12 +398,12 @@ cdef class BaseNeighbourVisitor(OctVisitor): neighbour = o self.oi.level = self.level for i in range(3): - self.oi.ipos[i] = (self.pos[i] * self.nz) + ishift[i] + self.oi.ipos[i] = (self.pos[i] * self.nz[i]) + ishift[i] # Extra step - compute cell position in neighbouring oct (and store in oi.ipos) if self.oi.level == self.level - 1: for i in range(3): - ipos = (((self.pos[i] * self.nz) + ishift[i])) >> 1 + ipos = (((self.pos[i] * self.nz[i]) + ishift[i])) >> 1 if (self.oi.ipos[i] << 1) == ipos: self.oi.ipos[i] = 0 else: diff --git a/yt/geometry/particle_deposit.pyx b/yt/geometry/particle_deposit.pyx index b04a4ad72dd..d893dbf7d88 100644 --- a/yt/geometry/particle_deposit.pyx +++ b/yt/geometry/particle_deposit.pyx @@ -61,7 +61,9 @@ cdef class ParticleDepositOperation: cdef np.float64_t pos[3] cdef np.float64_t[:] field_vals = np.empty(nf, dtype="float64") cdef int dims[3] - dims[0] = dims[1] = dims[2] = octree.nz + dims[0] = octree.nz[0] + dims[1] = octree.nz[1] + dims[2] = octree.nz[2] cdef OctInfo oi cdef np.int64_t offset, moff cdef Oct *oct diff --git a/yt/geometry/particle_oct_container.pyx b/yt/geometry/particle_oct_container.pyx index d9d64843836..70fc953cc33 100644 --- a/yt/geometry/particle_oct_container.pyx +++ b/yt/geometry/particle_oct_container.pyx @@ -1988,7 +1988,7 @@ cdef class ParticleBitmapOctreeContainer(SparseOctreeContainer): cdef oct_visitors.AssignDomainInd visitor visitor = oct_visitors.AssignDomainInd(self) self.visit_all_octs(selector, visitor) - assert ((visitor.global_index+1)*visitor.nz == visitor.index) + assert ((visitor.global_index+1)*visitor.nzones == visitor.index) # Copy indexes self._ptr_index_base_octs = malloc(sizeof(np.uint8_t)*self.nocts) self._index_base_octs = self._ptr_index_base_octs diff --git a/yt/geometry/particle_smooth.pyx b/yt/geometry/particle_smooth.pyx index ea3a7e79cbd..eaffbdf206f 100644 --- a/yt/geometry/particle_smooth.pyx +++ b/yt/geometry/particle_smooth.pyx @@ -121,7 +121,9 @@ cdef class ParticleSmoothOperation: periodicity = (False, False, False) else: raise NotImplementedError - dims[0] = dims[1] = dims[2] = mesh_octree.nz + dims[0] = mesh_octree.nz[0] + dims[1] = mesh_octree.nz[1] + dims[2] = mesh_octree.nz[2] cdef int nz = dims[0] * dims[1] * dims[2] # pcount is the number of particles per oct. pcount = np.zeros_like(pdom_ind) @@ -151,7 +153,6 @@ cdef class ParticleSmoothOperation: for i in range(3): self.DW[i] = (mesh_octree.DRE[i] - mesh_octree.DLE[i]) self.periodicity[i] = periodicity[i] - cdef np.float64_t factor = particle_octree.nz for i in range(positions.shape[0]): for j in range(3): pos[j] = positions[i, j] @@ -168,7 +169,7 @@ cdef class ParticleSmoothOperation: # in octs that we know are too far away for j in range(3): oct_left_edges[offset, j] = oinfo.left_edge[j] - oct_dds[offset, j] = oinfo.dds[j] * factor + oct_dds[offset, j] = oinfo.dds[j] * particle_octree.nz[j] # Now we have oct assignments. Let's sort them. # Note that what we will be providing to our processing functions will # actually be indirectly-sorted fields. This preserves memory at the