diff --git a/yt/frontends/__init__.py b/yt/frontends/__init__.py index ef209207fbc..b57350d6cf0 100644 --- a/yt/frontends/__init__.py +++ b/yt/frontends/__init__.py @@ -36,6 +36,7 @@ "owls", "owls_subfind", "parthenon", + "mini_ramses", "ramses", "rockstar", "sdf", diff --git a/yt/frontends/mini_ramses/__init__.py b/yt/frontends/mini_ramses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/yt/frontends/mini_ramses/api.py b/yt/frontends/mini_ramses/api.py new file mode 100644 index 00000000000..4117615ebdb --- /dev/null +++ b/yt/frontends/mini_ramses/api.py @@ -0,0 +1,5 @@ +from . import tests +from .data_structures import MiniRAMSESDataset +from .definitions import field_aliases +from .fields import MiniRAMSESFieldInfo +from .io import IOHandlerMiniRAMSES diff --git a/yt/frontends/mini_ramses/data_structures.py b/yt/frontends/mini_ramses/data_structures.py new file mode 100644 index 00000000000..6008f8165d8 --- /dev/null +++ b/yt/frontends/mini_ramses/data_structures.py @@ -0,0 +1,890 @@ +import os +import struct +from pathlib import Path + +import numpy as np + +from yt.data_objects.index_subobjects.octree_subset import OctreeSubset +from yt.data_objects.static_output import Dataset +from yt.funcs import mylog, setdefaultattr +from yt.geometry.geometry_handler import YTDataChunk +from yt.geometry.oct_container import RAMSESOctreeContainer +from yt.geometry.oct_geometry_handler import OctreeIndex + +from .definitions import ( + MINI_RAMSES_FILE_RE, + OUTPUT_DIR_RE, +) +from .fields import MiniRAMSESFieldInfo + + +class MiniRAMSESFileSanitizer: + """Handle the different files that can be passed and associated + safely to a mini-ramses output.""" + + root_folder = None + info_fname = None + + def __init__(self, filename): + paths_to_try = (Path(filename), Path(filename).resolve()) + + self.original_filename = filename + self.is_valid = False + + for path in paths_to_try: + if self._test_path(path): + break + + def _test_path(self, path): + """Test if path leads to a valid mini-ramses output.""" + # If it's a file, check if it's a known mini-ramses file + if path.is_file(): + if MINI_RAMSES_FILE_RE.match(path.name): + return self._test_parent_folder(path.parent) + # Could be the info.txt file directly + if path.name == "info.txt": + return self._test_parent_folder(path.parent) + return False + # If it's a directory, check if it's an output directory + if path.is_dir(): + return self._test_parent_folder(path) + return False + + def _test_parent_folder(self, folder): + """Check if folder is a valid mini-ramses output directory.""" + match = OUTPUT_DIR_RE.match(folder.name) + if not match: + return False + + info_path = folder / "info.txt" + if not info_path.exists(): + return False + + # Verify this is mini-ramses (not RAMSES) by checking for + # mini-ramses-specific file patterns (e.g., amr.NNNNN not amr_NNNNN.outNNNNN) + # and info.txt format (has "nfile" as first line) + if not self._is_mini_ramses_info(info_path): + return False + + self.root_folder = folder + self.info_fname = info_path + self.is_valid = True + return True + + @staticmethod + def _is_mini_ramses_info(info_path): + """Check if this is a mini-ramses info file (has nfile as first field).""" + try: + with open(info_path) as f: + first_line = f.readline().strip() + return first_line.startswith("nfile") + except (OSError, UnicodeDecodeError): + return False + + def validate(self): + """Raise an error if the file is not valid.""" + if not self.is_valid: + raise ValueError( + f"Cannot identify {self.original_filename} as a " + "mini-ramses output" + ) + + +class MiniRAMSESDomainFile: + """Manage a single domain (CPU file) of a mini-ramses output.""" + + _last_mask = None + _last_selector_id = None + + def __init__(self, ds, domain_id): + self.ds = ds + self.domain_id = domain_id + self._level_count = None + self._level_offsets = None + self._octree = None + + iout = int(ds.basename.split("_")[1]) + basename = ds.root_folder + + # File paths for this domain + icpu_str = f"{domain_id:05d}" + self.amr_fn = os.path.join(basename, f"amr.{icpu_str}") + self.hydro_fn = os.path.join(basename, f"hydro.{icpu_str}") + self.grav_fn = os.path.join(basename, f"grav.{icpu_str}") + + @property + def level_count(self): + if self._level_count is not None: + return self._level_count + self._read_amr_header() + return self._level_count + + @property + def level_offsets(self): + """Return cumulative oct count up to (but not including) each level.""" + if self._level_offsets is not None: + return self._level_offsets + self._read_amr_header() + return self._level_offsets + + def _read_amr_header(self): + """Read the AMR header from the stream-based binary file.""" + if not os.path.exists(self.amr_fn): + self._level_count = np.zeros( + self.ds.max_level - self.ds.min_level + 1, dtype="int64" + ) + self._level_offsets = np.zeros_like(self._level_count) + return + + with open(self.amr_fn, "rb") as f: + ndim = struct.unpack("i", f.read(4))[0] + levelmin = struct.unpack("i", f.read(4))[0] + nlevelmax = struct.unpack("i", f.read(4))[0] + + noct = np.zeros(nlevelmax, dtype="int32") + for ilevel in range(levelmin - 1, nlevelmax): + noct[ilevel] = struct.unpack("i", f.read(4))[0] + + nlevels = self.ds.max_level - self.ds.min_level + 1 + self._level_count = np.zeros(nlevels, dtype="int64") + for ilevel in range(levelmin - 1, nlevelmax): + idx = ilevel - (self.ds.min_level - 1) + if 0 <= idx < nlevels: + self._level_count[idx] = noct[ilevel] + + # Compute cumulative offsets: offset[i] = sum of octs at levels < i + self._level_offsets = np.zeros(nlevels, dtype="int64") + cumsum = 0 + for i in range(nlevels): + self._level_offsets[i] = cumsum + cumsum += self._level_count[i] + + def _read_hydro_header(self): + """Read the hydro header to get nvar.""" + if not os.path.exists(self.hydro_fn): + return 0 + with open(self.hydro_fn, "rb") as f: + ndim = struct.unpack("i", f.read(4))[0] + nvar = struct.unpack("i", f.read(4))[0] + return nvar + + def _read_grav_header(self): + """Read the gravity header to get nvar.""" + if not os.path.exists(self.grav_fn): + return 0 + with open(self.grav_fn, "rb") as f: + ndim = struct.unpack("i", f.read(4))[0] + nvar = struct.unpack("i", f.read(4))[0] + return nvar + + +class MiniRAMSESDomainSubset(OctreeSubset): + _domain_offset = 1 + _block_order = "F" + + _base_domain = None + + def __init__( + self, base_region, domain, ds, over_refine_factor=1, num_ghost_zones=0 + ): + super().__init__( + base_region, domain, ds, over_refine_factor, num_ghost_zones + ) + self._base_domain = domain + + @property + def oct_handler(self): + return self.ds.index.oct_handler + + def fill(self, fn, fields, selector, field_list): + """Read and fill field data from a mini-ramses binary file. + + Parameters + ---------- + fn : str + Filename to read from + fields : list + List of (ftype, fname) tuples to read + selector : SelectorObject + The selector determining which cells to include + field_list : list + List of all fields in the file (for determining indices) + + Returns + ------- + data : dict + Dictionary of {fname: array} for each requested field + """ + ndim = self.ds.dimensionality + twotondim = 2**ndim + oct_handler = self.oct_handler + domain = self._base_domain + + # Count selected cells + cell_count = selector.count_oct_cells(oct_handler, self.domain_id) + + # Get field names only (without ftype) + field_names = [f for ft, f in fields] + + # Initialize output data + data = {} + for fname in field_names: + data[fname] = np.zeros(cell_count, dtype="float64") + + if cell_count == 0: + return data + + # Get indices of selected cells + # level_inds: which level (0-indexed relative to min_level) + # cell_inds: which cell within the oct (0-7) + # file_inds: which oct within that level + level_inds, cell_inds, file_inds = oct_handler.file_index_octs( + selector, self.domain_id, cell_count + ) + + # Read all data from file into a flat array + all_data = self._read_all_cell_data(fn, field_list) + + if all_data is None: + return data + + # Get level offsets to compute global oct index + level_offsets = domain.level_offsets + + # Compute global oct index: offset for this level + oct index within level + # file_inds gives oct index within each level + # level_offsets[level] gives cumulative count of octs at levels < level + global_oct_inds = np.zeros(cell_count, dtype="int64") + for i in range(cell_count): + level = level_inds[i] + if level < len(level_offsets): + global_oct_inds[i] = level_offsets[level] + file_inds[i] + else: + global_oct_inds[i] = file_inds[i] + + # Cell index in flat array: oct_idx * 8 + cell_idx + indices = global_oct_inds * twotondim + cell_inds + + # Select the requested cells + for ftype, fname in fields: + if fname in all_data: + data[fname] = all_data[fname][indices] + + return data + + def _read_all_cell_data(self, fn, field_list): + """Read all cell data from file into arrays indexed by oct. + + Returns dict with field name -> flat array where index = oct_idx * 8 + cell_idx + """ + if not os.path.exists(fn): + return None + + ndim = self.ds.dimensionality + twotondim = 2**ndim + + with open(fn, "rb") as f: + # Read header + f_ndim = struct.unpack("i", f.read(4))[0] + f_nvar = struct.unpack("i", f.read(4))[0] + levelmin = struct.unpack("i", f.read(4))[0] + nlevelmax = struct.unpack("i", f.read(4))[0] + + noct = np.zeros(nlevelmax, dtype="int32") + for ilevel in range(levelmin - 1, nlevelmax): + noct[ilevel] = struct.unpack("i", f.read(4))[0] + + total_octs = noct.sum() + total_cells = total_octs * twotondim + + # Build field index map + all_field_names = [f for ft, f in field_list] + field_indices = {} + for ftype, fname in field_list: + if fname in all_field_names: + try: + fidx = all_field_names.index(fname) + if fidx < f_nvar: + field_indices[fname] = fidx + except ValueError: + pass + + # Pre-allocate arrays + raw_data = {fname: np.zeros(total_cells, dtype="float64") + for fname in field_indices} + + # Read all data + # Layout is: for each oct, nvar blocks of 8 cells each (field-major within oct) + cell_offset = 0 + for ilevel in range(levelmin - 1, nlevelmax): + ncache = noct[ilevel] + if ncache == 0: + continue + + for igrid in range(ncache): + # Each oct: nvar fields * twotondim cells * 4 bytes + # Layout is (nvar, twotondim) - field-major + oct_data = np.frombuffer( + f.read(4 * twotondim * f_nvar), + dtype=" 0: + g = og.retrieve_ghost_zones(ngz, []) + else: + g = og + yield YTDataChunk(dobj, "spatial", [g], None) + + def _chunk_io(self, dobj, cache=True, local_only=False): + oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) + for subset in oobjs: + yield YTDataChunk(dobj, "io", [subset], None) + + +class MiniRAMSESDataset(Dataset): + _index_class = MiniRAMSESIndex + _field_info_class = MiniRAMSESFieldInfo + _dataset_type = "mini_ramses" + + # Attributes set during parsing + gamma = 1.4 + nfile = 1 + + def __init__( + self, + filename, + dataset_type="mini_ramses", + fields=None, + storage_filename=None, + units_override=None, + unit_system="cgs", + default_species_fields=None, + ): + self.fluid_types += ("mini-ramses", "gravity") + self._fields_in_file = [] + self._gravity_fields_in_file = [] + self.storage_filename = storage_filename + self.default_species_fields = default_species_fields + + # Sanitize input filename + sanitizer = MiniRAMSESFileSanitizer(filename) + sanitizer.validate() + self.root_folder = str(sanitizer.root_folder) + + super().__init__( + str(sanitizer.info_fname), + dataset_type, + units_override=units_override, + unit_system=unit_system, + ) + + @property + def basename(self): + return os.path.basename(self.root_folder) + + def _set_code_unit_attributes(self): + """Set code unit attributes from info file parameters.""" + setdefaultattr(self, "length_unit", self.quan(self.unit_l, "cm")) + setdefaultattr(self, "mass_unit", self.quan( + self.unit_d * self.unit_l**3, "g" + )) + setdefaultattr(self, "time_unit", self.quan(self.unit_t, "s")) + setdefaultattr( + self, "velocity_unit", + self.quan(self.unit_l / self.unit_t, "cm/s"), + ) + setdefaultattr( + self, + "magnetic_unit", + self.quan( + np.sqrt(4.0 * np.pi * self.unit_d) + * self.unit_l + / self.unit_t, + "gauss", + ), + ) + + def _parse_parameter_file(self): + """Parse the mini-ramses info.txt file.""" + info_path = os.path.join(self.root_folder, "info.txt") + + # Read all key=value pairs from info.txt + params = {} + with open(info_path) as f: + for line in f: + line = line.strip() + if "=" in line: + key, val = line.split("=", 1) + key = key.strip() + val = val.strip() + try: + if "." in val or "E" in val.upper(): + params[key] = float(val) + else: + params[key] = int(val) + except ValueError: + params[key] = val + + # Extract parameters + # mini-ramses info.txt has: nfile, ncpu, ndim, levelmin, levelmax, + # ngridmax, nstep_coarse, boxlen, time, texp, aexp, H0, + # omega_m, omega_l, omega_k, omega_b, gamma, unit_l, unit_d, unit_t + self.nfile = params.get("nfile", 1) + self.dimensionality = params.get("ndim", 3) + self.min_level = params.get("levelmin", 1) + self.max_level = params.get("levelmax", 1) + self.domain_dimensions = np.ones(3, dtype="int64") * 2 ** ( + self.min_level + ) + + self.gamma = params.get("gamma", 1.4) + boxlen = params.get("boxlen", 1.0) + + self.domain_left_edge = np.zeros(3, dtype="float64") + self.domain_right_edge = np.ones(3, dtype="float64") * boxlen + + # Physical parameters + self.unit_l = params.get("unit_l", 1.0) + self.unit_d = params.get("unit_d", 1.0) + self.unit_t = params.get("unit_t", 1.0) + + # Cosmology parameters + self.omega_matter = params.get("omega_m", 0.0) + self.omega_lambda = params.get("omega_l", 0.0) + self.omega_radiation = 0.0 + self.hubble_constant = params.get("H0", 0.0) / 100.0 # Convert to H100 + current_time = params.get("time", 0.0) + self.current_time = current_time + + aexp = params.get("aexp", 1.0) + h0 = params.get("H0", 0.0) + + # Determine if this is a cosmological simulation + # Same logic as RAMSES: if time >= 0 and H0 == 1 and aexp == 1 + # then it's NOT cosmological (all three must be true) + is_cosmological = not ( + current_time >= 0 and h0 == 1.0 and aexp == 1.0 + ) + + if is_cosmological: + self.cosmological_simulation = True + self.current_redshift = 1.0 / aexp - 1.0 + else: + self.cosmological_simulation = False + self.current_redshift = 0.0 + + # Store all parameters + self.parameters = params + self.parameters["boxlen"] = boxlen + self.parameters["time"] = current_time + self.parameters["gamma"] = self.gamma + + self.unique_identifier = int( + os.stat(self.parameter_filename).st_ctime + ) + + # Periodicity + self._periodicity = (True, True, True) + + # Refine by 2 in each dimension + self.refine_by = 2 + + # Detect fluid fields (hydro and gravity) + self._detect_fluid_fields() + self._detect_gravity_fields() + + # Detect particle types + self._detect_particle_types() + + def _detect_fluid_fields(self): + """Detect available fluid fields from hydro header or first hydro file.""" + # Try reading hydro_header.txt first + header_fn = os.path.join(self.root_folder, "hydro_header.txt") + if os.path.exists(header_fn): + self._fields_in_file = self._parse_hydro_header(header_fn) + return + + # Try reading the first hydro file header + hydro_fn = os.path.join(self.root_folder, "hydro.00001") + if os.path.exists(hydro_fn): + with open(hydro_fn, "rb") as f: + ndim = struct.unpack("i", f.read(4))[0] + nvar = struct.unpack("i", f.read(4))[0] + self._fields_in_file = self._default_hydro_fields(nvar) + return + + self._fields_in_file = [] + + def _parse_hydro_header(self, header_fn): + """Parse the mini-ramses hydro_header.txt file.""" + fields = [] + field_map = { + "density": ("mini-ramses", "Density"), + "velocity_x": ("mini-ramses", "x-velocity"), + "velocity_y": ("mini-ramses", "y-velocity"), + "velocity_z": ("mini-ramses", "z-velocity"), + "thermal_pressure": ("mini-ramses", "Pressure"), + "magnetic_field_x": ("mini-ramses", "B_x"), + "magnetic_field_y": ("mini-ramses", "B_y"), + "magnetic_field_z": ("mini-ramses", "B_z"), + "metal_mass_fraction": ("mini-ramses", "Metallicity"), + } + with open(header_fn) as f: + for line in f: + line = line.strip() + if line.startswith("nvar"): + continue + if line.startswith("variable"): + # Parse "variable # N: name" + parts = line.split(":") + if len(parts) >= 2: + fname = parts[1].strip() + if fname in field_map: + fields.append(field_map[fname]) + else: + fields.append(("mini-ramses", fname)) + return fields + + def _default_hydro_fields(self, nvar): + """Return default field names based on nvar count.""" + ndim = self.dimensionality + fields = [("mini-ramses", "Density")] + for ax in "xyz"[:ndim]: + fields.append(("mini-ramses", f"{ax}-velocity")) + fields.append(("mini-ramses", "Pressure")) + # MHD fields + if nvar > ndim + 2: + remaining = nvar - (ndim + 2) + if remaining >= 3: + for ax in "xyz": + fields.append(("mini-ramses", f"B_{ax}")) + remaining -= 3 + for i in range(remaining): + fields.append(("mini-ramses", f"scalar_{i + 1}")) + return fields + + def _detect_gravity_fields(self): + """Detect available gravity fields from grav header or first grav file.""" + # Try reading grav_header.txt first + header_fn = os.path.join(self.root_folder, "grav_header.txt") + if os.path.exists(header_fn): + self._gravity_fields_in_file = self._parse_grav_header(header_fn) + return + + # Try reading the first grav file header + grav_fn = os.path.join(self.root_folder, "grav.00001") + if os.path.exists(grav_fn): + with open(grav_fn, "rb") as f: + ndim = struct.unpack("i", f.read(4))[0] + nvar = struct.unpack("i", f.read(4))[0] + self._gravity_fields_in_file = self._default_gravity_fields(nvar) + return + + self._gravity_fields_in_file = [] + + def _parse_grav_header(self, header_fn): + """Parse the mini-ramses grav_header.txt file.""" + fields = [] + field_map = { + "potential": ("gravity", "Potential"), + "accel_x": ("gravity", "x-acceleration"), + "accel_y": ("gravity", "y-acceleration"), + "accel_z": ("gravity", "z-acceleration"), + } + with open(header_fn) as f: + for line in f: + line = line.strip() + if line.startswith("nvar"): + continue + if line.startswith("variable"): + # Parse "variable # N: name" + parts = line.split(":") + if len(parts) >= 2: + fname = parts[1].strip() + if fname in field_map: + fields.append(field_map[fname]) + else: + fields.append(("gravity", fname)) + return fields + + def _default_gravity_fields(self, nvar): + """Return default gravity field names based on nvar count.""" + fields = [("gravity", "Potential")] + for ax in "xyz"[:min(nvar - 1, 3)]: + fields.append(("gravity", f"{ax}-acceleration")) + return fields + + def _detect_particle_types(self): + """Detect available particle types from output files.""" + particle_types = [] + ptype_map = { + "part": "io", + "star": "star", + "sink": "sink", + } + for prefix, ptype in ptype_map.items(): + fn = os.path.join(self.root_folder, f"{prefix}.00001") + if os.path.exists(fn): + particle_types.append(ptype) + + if particle_types: + self.particle_types = tuple(particle_types) + self.particle_types_raw = tuple(particle_types) + else: + self.particle_types = () + self.particle_types_raw = () + + @classmethod + def _is_valid(cls, filename, *args, **kwargs): + if cls._missing_load_requirements(): + return False + return MiniRAMSESFileSanitizer(filename).is_valid + + def __str__(self): + return self.basename diff --git a/yt/frontends/mini_ramses/definitions.py b/yt/frontends/mini_ramses/definitions.py new file mode 100644 index 00000000000..37d1c88cbec --- /dev/null +++ b/yt/frontends/mini_ramses/definitions.py @@ -0,0 +1,50 @@ +import re + +# mini-ramses uses different file naming than RAMSES: +# - output_XXXXX/info.txt (no number in info filename) +# - output_XXXXX/amr.NNNNN (dot-separated, no "out" suffix) +# - output_XXXXX/hydro.NNNNN +# - output_XXXXX/part.NNNNN, star.NNNNN, sink.NNNNN, tree.NNNNN +# - output_XXXXX/grav.NNNNN +# - output_XXXXX/params.bin +# - output_XXXXX/hydro_header.txt, part_header.txt, etc. + +OUTPUT_DIR_EXP = r"output_(\d{5})" +OUTPUT_DIR_RE = re.compile(OUTPUT_DIR_EXP) + +# Matches mini-ramses file patterns (dot-separated, no "out" suffix) +# e.g. amr.00001, hydro.00001, part.00001, info.txt +MINI_RAMSES_FILE_RE = re.compile( + r"((amr|hydro|part|grav|star|sink|tree|trac|rt)\.\d{5}" + r"|info\.txt|params\.bin)" +) + +# Regular expressions used to parse file descriptors +VERSION_RE = re.compile(r"# version: *(\d+)") +VAR_DESC_RE = re.compile(r"\s*([^\s]+),\s*([^\s]+),\s*([^\s]+)") + +field_aliases = { + "standard_five": ( + "Density", + "x-velocity", + "y-velocity", + "z-velocity", + "Pressure", + ), + "standard_six": ( + "Density", + "x-velocity", + "y-velocity", + "z-velocity", + "Pressure", + "Metallicity", + ), +} + +# mini-ramses uses separate files for each particle type, +# so the family concept is simpler +particle_families = { + "DM": 1, + "star": 2, + "cloud": 3, +} diff --git a/yt/frontends/mini_ramses/fields.py b/yt/frontends/mini_ramses/fields.py new file mode 100644 index 00000000000..ca6d1f67bd9 --- /dev/null +++ b/yt/frontends/mini_ramses/fields.py @@ -0,0 +1,62 @@ +from yt._typing import KnownFieldsT +from yt.fields.field_info_container import FieldInfoContainer +from yt.utilities.physical_constants import ( + boltzmann_constant_cgs, + mass_hydrogen_cgs, +) + +b_units = "code_magnetic" +ra_units = "code_length / code_time**2" +rho_units = "code_density" +vel_units = "code_velocity" +pressure_units = "code_pressure" +ener_units = "code_mass * code_velocity**2" +specific_ener_units = "code_velocity**2" +ang_mom_units = "code_mass * code_velocity * code_length" + + +class MiniRAMSESFieldInfo(FieldInfoContainer): + known_other_fields: KnownFieldsT = ( + ("Density", (rho_units, ["density"], None)), + ("x-velocity", (vel_units, ["velocity_x"], None)), + ("y-velocity", (vel_units, ["velocity_y"], None)), + ("z-velocity", (vel_units, ["velocity_z"], None)), + ("Pressure", (pressure_units, ["pressure"], None)), + ("Metallicity", ("", ["metallicity"], None)), + ("x-acceleration", (ra_units, ["acceleration_x"], None)), + ("y-acceleration", (ra_units, ["acceleration_y"], None)), + ("z-acceleration", (ra_units, ["acceleration_z"], None)), + ("Potential", (specific_ener_units, ["potential"], None)), + ("B_x", (b_units, ["magnetic_field_x"], None)), + ("B_y", (b_units, ["magnetic_field_y"], None)), + ("B_z", (b_units, ["magnetic_field_z"], None)), + ) + known_particle_fields: KnownFieldsT = ( + ("particle_position_x", ("code_length", [], None)), + ("particle_position_y", ("code_length", [], None)), + ("particle_position_z", ("code_length", [], None)), + ("particle_velocity_x", (vel_units, [], None)), + ("particle_velocity_y", (vel_units, [], None)), + ("particle_velocity_z", (vel_units, [], None)), + ("particle_mass", ("code_mass", [], None)), + ("particle_identity", ("", ["particle_index"], None)), + ("particle_refinement_level", ("", [], None)), + ("particle_birth_time", ("code_time", ["age"], None)), + ("particle_metallicity", ("", [], None)), + ) + + def setup_particle_fields(self, ptype): + super().setup_particle_fields(ptype) + + def setup_fluid_fields(self): + def _temperature(data): + rv = data["gas", "pressure"] / data["gas", "density"] + rv *= mass_hydrogen_cgs / boltzmann_constant_cgs + return rv + + self.add_field( + ("gas", "temperature"), + sampling_type="cell", + function=_temperature, + units=self.ds.unit_system["temperature"], + ) diff --git a/yt/frontends/mini_ramses/io.py b/yt/frontends/mini_ramses/io.py new file mode 100644 index 00000000000..76df3a644c0 --- /dev/null +++ b/yt/frontends/mini_ramses/io.py @@ -0,0 +1,300 @@ +import os +import struct +from collections import defaultdict + +import numpy as np + +from yt.utilities.io_handler import BaseIOHandler +from yt.utilities.logger import ytLogger as mylog + + +class IOHandlerMiniRAMSES(BaseIOHandler): + _dataset_type = "mini_ramses" + _particle_reader = True + + def _read_fluid_selection(self, chunks, selector, fields, size): + # Read fluid (hydro and gravity) data from mini-ramses output files + # Use the selector to determine which cells to read + tr = defaultdict(list) + + # Separate fields by type + ftypes = {f[0] for f in fields} + + for chunk in chunks: + for ft in ftypes: + # Get all fields of this type + field_subs = [f for f in fields if f[0] == ft] + + for subset in chunk.objs: + domain = subset.domain + ds = subset.ds + + # Determine file and field list based on field type + if ft == "mini-ramses": + fn = domain.hydro_fn + field_list = ds._fields_in_file + elif ft == "gravity": + fn = domain.grav_fn + field_list = ds._gravity_fields_in_file + else: + continue + + if not os.path.exists(fn): + continue + + # Use subset.fill to read with selection + rv = subset.fill(fn, field_subs, selector, field_list) + + for ftype, fname in field_subs: + d = rv.get(fname, np.empty(0, dtype="float64")) + if d.size == 0: + continue + mylog.debug( + "Filling %s with %s (%0.3e %0.3e) (%s zones)", + fname, + d.size, + d.min(), + d.max(), + d.size, + ) + tr[ftype, fname].append(d) + + # Concatenate results + rv = {} + for field in fields: + tmp = tr.pop(field, None) + rv[field] = np.concatenate(tmp) if tmp else np.empty(0, dtype="float64") + + return rv + + def _read_particle_coords(self, chunks, ptf): + for chunk in chunks: + for subset in chunk.objs: + domain = subset.domain + ds = subset.ds + ndim = ds.dimensionality + boxlen = float(ds.domain_right_edge[0]) + + for ptype in sorted(ptf): + prefix = self._ptype_to_prefix(ptype) + fn = os.path.join( + ds.root_folder, + f"{prefix}.{domain.domain_id:05d}", + ) + if not os.path.exists(fn): + continue + + pdata = self._read_particle_file(fn, ndim, prefix) + if pdata is None: + continue + + npart = pdata["npart"] + pos = np.zeros((npart, 3), dtype="float64") + for i, ax in enumerate("xyz"[:ndim]): + pos[:, i] = pdata[f"pos_{ax}"] + + # Normalize positions to code units + yield ptype, ( + pos[:, 0] / boxlen, + pos[:, 1] / boxlen, + pos[:, 2] / boxlen, + ), 0.0 + + def _read_particle_fields(self, chunks, ptf, selector): + for chunk in chunks: + for subset in chunk.objs: + domain = subset.domain + ds = subset.ds + ndim = ds.dimensionality + boxlen = float(ds.domain_right_edge[0]) + + for ptype, field_list in sorted(ptf.items()): + prefix = self._ptype_to_prefix(ptype) + fn = os.path.join( + ds.root_folder, + f"{prefix}.{domain.domain_id:05d}", + ) + if not os.path.exists(fn): + continue + + pdata = self._read_particle_file(fn, ndim, prefix) + if pdata is None: + continue + + npart = pdata["npart"] + pos = np.zeros((npart, 3), dtype="float64") + for i, ax in enumerate("xyz"[:ndim]): + pos[:, i] = pdata[f"pos_{ax}"] + + mask = selector.select_points( + pos[:, 0] / boxlen, + pos[:, 1] / boxlen, + pos[:, 2] / boxlen, + 0.0, + ) + + if mask is None: + continue + + for fname in field_list: + data = self._get_particle_field_data( + pdata, fname, ndim, boxlen + ) + if data is not None: + yield (ptype, fname), data[mask] + + def _get_particle_field_data(self, pdata, fname, ndim, boxlen): + """Map a yt particle field name to data from the particle file.""" + field_map = { + "particle_position_x": "pos_x", + "particle_position_y": "pos_y", + "particle_position_z": "pos_z", + "particle_velocity_x": "vel_x", + "particle_velocity_y": "vel_y", + "particle_velocity_z": "vel_z", + "particle_mass": "mass", + "particle_refinement_level": "level", + "particle_identity": "birth_id", + "particle_metallicity": "metallicity", + "particle_birth_time": "birth_date", + } + + if fname in field_map: + key = field_map[fname] + if key in pdata: + return pdata[key].astype("float64") + return None + + @staticmethod + def _ptype_to_prefix(ptype): + """Map yt particle type name to mini-ramses file prefix.""" + ptype_map = { + "io": "part", + "star": "star", + "sink": "sink", + "tree": "tree", + "trac": "trac", + } + return ptype_map.get(ptype, "part") + + @staticmethod + def _read_particle_file(fn, ndim, prefix): + """Read a mini-ramses particle file (stream binary format). + + Mini-ramses output particle files use stream I/O with: + - Header: ndim (int32) + npart (int32) = 8 bytes + - Then contiguous arrays of float32 for positions, velocities, mass + - Then int32 for level, int32 for birth_id + """ + try: + with open(fn, "rb") as f: + file_ndim = struct.unpack("i", f.read(4))[0] + npart = struct.unpack("i", f.read(4))[0] + + if npart == 0: + return None + + pdata = {"npart": npart} + + # Positions (float32) + for ax in "xyz"[:ndim]: + pdata[f"pos_{ax}"] = np.frombuffer( + f.read(4 * npart), dtype="= opt_size_per * npart: + pdata[opt_name] = np.frombuffer( + f.read(opt_size_per * npart), dtype=opt_dtype + ).copy() + remaining -= opt_size_per * npart + + # If no header was found, try to read level and birth_id + if not opt_fields: + if remaining >= 4 * npart: + pdata["level"] = np.frombuffer( + f.read(4 * npart), dtype="= 4 * npart: + pdata["birth_id"] = np.frombuffer( + f.read(4 * npart), dtype="11s}\n") + f.write(f"ncpu ={'1':>11s}\n") + f.write(f"ndim ={ndim:>11d}\n") + f.write(f"levelmin ={levelmin:>11d}\n") + f.write(f"levelmax ={nlevelmax:>11d}\n") + f.write(f"ngridmax ={'100':>11s}\n") + f.write(f"nstep_coarse={'0':>11s}\n") + f.write("\n") + f.write(f"boxlen ={boxlen:>23.15E}\n") + f.write(f"time ={0.0:>23.15E}\n") + f.write(f"texp ={0.0:>23.15E}\n") + f.write(f"aexp ={1.0:>23.15E}\n") + f.write(f"H0 ={1.0:>23.15E}\n") + f.write(f"omega_m ={0.0:>23.15E}\n") + f.write(f"omega_l ={0.0:>23.15E}\n") + f.write(f"omega_k ={0.0:>23.15E}\n") + f.write(f"omega_b ={0.0:>23.15E}\n") + f.write(f"gamma ={gamma:>23.15E}\n") + f.write(f"unit_l ={3.08568e+21:>23.15E}\n") + f.write(f"unit_d ={1.6726e-24:>23.15E}\n") + f.write(f"unit_t ={3.1557e+13:>23.15E}\n") + f.write("\n") + + # Write amr.00001 (stream binary) + twotondim = 2**ndim + # Place a few octs at levelmin + nocts_per_level = np.zeros(nlevelmax, dtype="int32") + n_base_octs = 2 # number of octs at the base level + nocts_per_level[levelmin - 1] = n_base_octs + + with open(os.path.join(outdir, "amr.00001"), "wb") as f: + f.write(struct.pack("i", ndim)) + f.write(struct.pack("i", levelmin)) + f.write(struct.pack("i", nlevelmax)) + for ilevel in range(levelmin - 1, nlevelmax): + f.write(struct.pack("i", nocts_per_level[ilevel])) + + # Write grid data + for ilevel in range(levelmin - 1, nlevelmax): + for igrid in range(nocts_per_level[ilevel]): + # Cartesian key (ndim int32 values) + ckey = [igrid] * ndim + f.write(struct.pack(f"{ndim}i", *ckey)) + # Refinement map (twotondim int32 values) - not refined + refined = [0] * twotondim + f.write(struct.pack(f"{twotondim}i", *refined)) + + # Write hydro.00001 (stream binary) + with open(os.path.join(outdir, "hydro.00001"), "wb") as f: + f.write(struct.pack("i", ndim)) + f.write(struct.pack("i", nvar)) + f.write(struct.pack("i", levelmin)) + f.write(struct.pack("i", nlevelmax)) + for ilevel in range(levelmin - 1, nlevelmax): + f.write(struct.pack("i", nocts_per_level[ilevel])) + + # Write cell data per level + for ilevel in range(levelmin - 1, nlevelmax): + ncache = nocts_per_level[ilevel] + for igrid in range(ncache): + # qout(twotondim, nvar) as float32 + qout = np.random.rand(twotondim, nvar).astype("float32") + # Set density to be positive + qout[:, 0] = np.abs(qout[:, 0]) + 0.1 + # Set pressure to be positive + if nvar >= 5: + qout[:, 4] = np.abs(qout[:, 4]) + 0.01 + f.write(qout.tobytes()) + + # Write hydro_header.txt + with open(os.path.join(outdir, "hydro_header.txt"), "w") as f: + f.write(f"nvar ={nvar:>11d}\n") + f.write("variable # 1: density\n") + f.write("variable # 2: velocity_x\n") + f.write("variable # 3: velocity_y\n") + f.write("variable # 4: velocity_z\n") + f.write("variable # 5: thermal_pressure\n") + + # Write part.00001 (stream binary) - particle data + with open(os.path.join(outdir, "part.00001"), "wb") as f: + f.write(struct.pack("i", ndim)) + f.write(struct.pack("i", npart)) + + # Positions (float32) + for ax in range(ndim): + pos = np.random.rand(npart).astype("float32") * boxlen + f.write(pos.tobytes()) + + # Velocities (float32) + for ax in range(ndim): + vel = (np.random.rand(npart).astype("float32") - 0.5) * 100 + f.write(vel.tobytes()) + + # Mass (float32) + mass = np.random.rand(npart).astype("float32") * 1e-3 + f.write(mass.tobytes()) + + # Level (int32) + levels = np.ones(npart, dtype="int32") * levelmin + f.write(levels.tobytes()) + + # Birth ID (int32) + ids = np.arange(1, npart + 1, dtype="int32") + f.write(ids.tobytes()) + + # Write part_header.txt + with open(os.path.join(outdir, "part_header.txt"), "w") as f: + f.write("Total number of particles\n") + f.write(f"{npart}\n") + f.write("Total number of files\n") + f.write("1\n") + f.write("Particle fields\n") + f.write("pos vel mass level birth_id \n") + + return outdir + + +class TestMiniRAMSESFileSanitizer: + def test_valid_info_file(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + info_path = os.path.join(outdir, "info.txt") + sanitizer = MiniRAMSESFileSanitizer(info_path) + assert sanitizer.is_valid + + def test_valid_directory(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + sanitizer = MiniRAMSESFileSanitizer(outdir) + assert sanitizer.is_valid + + def test_valid_amr_file(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + amr_path = os.path.join(outdir, "amr.00001") + sanitizer = MiniRAMSESFileSanitizer(amr_path) + assert sanitizer.is_valid + + def test_invalid_file(self, tmp_path): + sanitizer = MiniRAMSESFileSanitizer(str(tmp_path / "nonexistent")) + assert not sanitizer.is_valid + + def test_ramses_info_not_valid(self, tmp_path): + """Ensure RAMSES info files (ncpu first, not nfile) are rejected.""" + outdir = tmp_path / "output_00001" + outdir.mkdir() + info = outdir / "info.txt" + info.write_text("ncpu = 1\nndim = 3\n") + sanitizer = MiniRAMSESFileSanitizer(str(info)) + assert not sanitizer.is_valid + + +class TestMiniRAMSESDataset: + def test_is_valid(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + assert MiniRAMSESDataset._is_valid(outdir) + + def test_is_valid_info_file(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + info_path = os.path.join(outdir, "info.txt") + assert MiniRAMSESDataset._is_valid(info_path) + + def test_not_valid_for_ramses(self, tmp_path): + """Ensure standard RAMSES outputs are NOT detected as mini-ramses.""" + outdir = tmp_path / "output_00001" + outdir.mkdir() + info = outdir / "info.txt" + info.write_text( + "ncpu = 1\n" + "ndim = 3\n" + "levelmin = 3\n" + ) + assert not MiniRAMSESDataset._is_valid(str(outdir)) + + def test_load_dataset(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + ds = MiniRAMSESDataset(outdir) + assert ds is not None + assert ds.dimensionality == 3 + assert ds.min_level == 3 + assert ds.max_level == 5 + assert ds.gamma == 1.4 + assert ds.cosmological_simulation is False + + def test_domain_dimensions(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path), levelmin=3) + ds = MiniRAMSESDataset(outdir) + # domain_dimensions should be 2^levelmin in each direction + assert np.all(ds.domain_dimensions == 8) + + def test_unit_attributes(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + ds = MiniRAMSESDataset(outdir) + assert ds.length_unit is not None + assert ds.mass_unit is not None + assert ds.time_unit is not None + + def test_fluid_fields_detected(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + ds = MiniRAMSESDataset(outdir) + assert len(ds._fields_in_file) > 0 + field_names = [f[1] for f in ds._fields_in_file] + assert "Density" in field_names + + def test_particle_types_detected(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + ds = MiniRAMSESDataset(outdir) + assert "io" in ds.particle_types + + def test_str_representation(self, tmp_path): + outdir = _create_mini_ramses_output(str(tmp_path)) + ds = MiniRAMSESDataset(outdir) + assert str(ds) == "output_00001"