From 6ab7bcb6399bc00bb41496dd66ca07c285b5a6ef Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Fri, 6 Mar 2026 09:02:43 +0100 Subject: [PATCH 1/9] Support new format --- yt_derived_fields/cutouts/loader.py | 356 +++++++++++++++++++--------- 1 file changed, 244 insertions(+), 112 deletions(-) diff --git a/yt_derived_fields/cutouts/loader.py b/yt_derived_fields/cutouts/loader.py index fb00fdd..c7b572f 100644 --- a/yt_derived_fields/cutouts/loader.py +++ b/yt_derived_fields/cutouts/loader.py @@ -17,116 +17,234 @@ class Scale(Enum): LOG = 1 -header: list[tuple[str, Scale, str]] = [ - ("redshift", Scale.LINEAR, "1"), - ("dx", Scale.LOG, "cm"), - ("x", Scale.LINEAR, "Mpccm/h"), - ("y", Scale.LINEAR, "Mpccm/h"), - ("z", Scale.LINEAR, "Mpccm/h"), - ("vx", Scale.LINEAR, "cm/s"), - ("vy", Scale.LINEAR, "cm/s"), - ("vz", Scale.LINEAR, "cm/s"), - ("density", Scale.LOG, "mp/cm**3"), - ("temperature", Scale.LOG, "K"), - ("pressure", Scale.LINEAR, "dyne/cm**2"), - ("iron_number_density", Scale.LOG, "1/cm**3"), - ("oxygen_number_density", Scale.LOG, "1/cm**3"), - ("nitrogen_number_density", Scale.LOG, "1/cm**3"), - ("magnesium_number_density", Scale.LOG, "1/cm**3"), - ("neon_number_density", Scale.LOG, "1/cm**3"), - ("silicon_number_density", Scale.LOG, "1/cm**3"), - ("calcium_number_density", Scale.LOG, "1/cm**3"), - ("carbon_number_density", Scale.LOG, "1/cm**3"), - ("sulfur_number_density", Scale.LOG, "1/cm**3"), - ("carbon_monoxide_number_density", Scale.LOG, "1/cm**3"), - ("oxygen_01", Scale.LINEAR, "1"), - ("oxygen_02", Scale.LINEAR, "1"), - ("oxygen_03", Scale.LINEAR, "1"), - ("oxygen_04", Scale.LINEAR, "1"), - ("oxygen_05", Scale.LINEAR, "1"), - ("oxygen_06", Scale.LINEAR, "1"), - ("oxygen_07", Scale.LINEAR, "1"), - ("oxygen_08", Scale.LINEAR, "1"), - ("nitrogen_01", Scale.LINEAR, "1"), - ("nitrogen_02", Scale.LINEAR, "1"), - ("nitrogen_03", Scale.LINEAR, "1"), - ("nitrogen_04", Scale.LINEAR, "1"), - ("nitrogen_05", Scale.LINEAR, "1"), - ("nitrogen_06", Scale.LINEAR, "1"), - ("nitrogen_07", Scale.LINEAR, "1"), - ("carbon_01", Scale.LINEAR, "1"), - ("carbon_02", Scale.LINEAR, "1"), - ("carbon_03", Scale.LINEAR, "1"), - ("carbon_04", Scale.LINEAR, "1"), - ("carbon_05", Scale.LINEAR, "1"), - ("carbon_06", Scale.LINEAR, "1"), - ("magnesium_01", Scale.LINEAR, "1"), - ("magnesium_02", Scale.LINEAR, "1"), - ("magnesium_03", Scale.LINEAR, "1"), - ("magnesium_04", Scale.LINEAR, "1"), - ("magnesium_05", Scale.LINEAR, "1"), - ("magnesium_06", Scale.LINEAR, "1"), - ("magnesium_07", Scale.LINEAR, "1"), - ("magnesium_08", Scale.LINEAR, "1"), - ("magnesium_09", Scale.LINEAR, "1"), - ("magnesium_10", Scale.LINEAR, "1"), - ("silicon_01", Scale.LINEAR, "1"), - ("silicon_02", Scale.LINEAR, "1"), - ("silicon_03", Scale.LINEAR, "1"), - ("silicon_04", Scale.LINEAR, "1"), - ("silicon_05", Scale.LINEAR, "1"), - ("silicon_06", Scale.LINEAR, "1"), - ("silicon_07", Scale.LINEAR, "1"), - ("silicon_08", Scale.LINEAR, "1"), - ("silicon_09", Scale.LINEAR, "1"), - ("silicon_10", Scale.LINEAR, "1"), - ("silicon_11", Scale.LINEAR, "1"), - ("sulfur_01", Scale.LINEAR, "1"), - ("sulfur_02", Scale.LINEAR, "1"), - ("sulfur_03", Scale.LINEAR, "1"), - ("sulfur_04", Scale.LINEAR, "1"), - ("sulfur_05", Scale.LINEAR, "1"), - ("sulfur_06", Scale.LINEAR, "1"), - ("sulfur_07", Scale.LINEAR, "1"), - ("sulfur_08", Scale.LINEAR, "1"), - ("sulfur_09", Scale.LINEAR, "1"), - ("sulfur_10", Scale.LINEAR, "1"), - ("sulfur_11", Scale.LINEAR, "1"), - ("iron_01", Scale.LINEAR, "1"), - ("iron_02", Scale.LINEAR, "1"), - ("iron_03", Scale.LINEAR, "1"), - ("iron_04", Scale.LINEAR, "1"), - ("iron_05", Scale.LINEAR, "1"), - ("iron_06", Scale.LINEAR, "1"), - ("iron_07", Scale.LINEAR, "1"), - ("iron_08", Scale.LINEAR, "1"), - ("iron_09", Scale.LINEAR, "1"), - ("iron_10", Scale.LINEAR, "1"), - ("iron_11", Scale.LINEAR, "1"), - ("neon_01", Scale.LINEAR, "1"), - ("neon_02", Scale.LINEAR, "1"), - ("neon_03", Scale.LINEAR, "1"), - ("neon_04", Scale.LINEAR, "1"), - ("neon_05", Scale.LINEAR, "1"), - ("neon_06", Scale.LINEAR, "1"), - ("neon_07", Scale.LINEAR, "1"), - ("neon_08", Scale.LINEAR, "1"), - ("neon_09", Scale.LINEAR, "1"), - ("neon_10", Scale.LINEAR, "1"), - ("hydrogen_01", Scale.LINEAR, "1"), - ("hydrogen_02", Scale.LINEAR, "1"), - ("helium_02", Scale.LINEAR, "1"), - ("helium_03", Scale.LINEAR, "1"), - ("Habing", Scale.LOG, "erg/s/cm**2"), - ("Lyman_Werner", Scale.LOG, "erg/s/cm**2"), - ("HI_Ionising", Scale.LOG, "erg/s/cm**2"), - ("H2_Ionising", Scale.LOG, "erg/s/cm**2"), - ("HeI_Ionising", Scale.LOG, "erg/s/cm**2"), - ("HeII_ionising", Scale.LOG, "erg/s/cm**2"), -] - - -def load_cutout(filename: str | Path, boxsize: float = 50, h0: float = 0.6727, verbose: bool = True): +headers: dict[int, list[tuple[str, Scale, str, str]]] = { + 1: [ + ("redshift", Scale.LINEAR, "1", "d"), + ("dx", Scale.LOG, "cm", "d"), + ("x", Scale.LINEAR, "Mpccm/h", "d"), + ("y", Scale.LINEAR, "Mpccm/h", "d"), + ("z", Scale.LINEAR, "Mpccm/h", "d"), + ("vx", Scale.LINEAR, "cm/s", "d"), + ("vy", Scale.LINEAR, "cm/s", "d"), + ("vz", Scale.LINEAR, "cm/s", "d"), + ("density", Scale.LOG, "mp/cm**3", "d"), + ("temperature", Scale.LOG, "K", "d"), + ("pressure", Scale.LINEAR, "dyne/cm**2", "d"), + ("iron_number_density", Scale.LOG, "1/cm**3", "d"), + ("oxygen_number_density", Scale.LOG, "1/cm**3", "d"), + ("nitrogen_number_density", Scale.LOG, "1/cm**3", "d"), + ("magnesium_number_density", Scale.LOG, "1/cm**3", "d"), + ("neon_number_density", Scale.LOG, "1/cm**3", "d"), + ("silicon_number_density", Scale.LOG, "1/cm**3", "d"), + ("calcium_number_density", Scale.LOG, "1/cm**3", "d"), + ("carbon_number_density", Scale.LOG, "1/cm**3", "d"), + ("sulfur_number_density", Scale.LOG, "1/cm**3", "d"), + ("carbon_monoxide_number_density", Scale.LOG, "1/cm**3", "d"), + ("oxygen_01", Scale.LINEAR, "1", "d"), + ("oxygen_02", Scale.LINEAR, "1", "d"), + ("oxygen_03", Scale.LINEAR, "1", "d"), + ("oxygen_04", Scale.LINEAR, "1", "d"), + ("oxygen_05", Scale.LINEAR, "1", "d"), + ("oxygen_06", Scale.LINEAR, "1", "d"), + ("oxygen_07", Scale.LINEAR, "1", "d"), + ("oxygen_08", Scale.LINEAR, "1", "d"), + ("nitrogen_01", Scale.LINEAR, "1", "d"), + ("nitrogen_02", Scale.LINEAR, "1", "d"), + ("nitrogen_03", Scale.LINEAR, "1", "d"), + ("nitrogen_04", Scale.LINEAR, "1", "d"), + ("nitrogen_05", Scale.LINEAR, "1", "d"), + ("nitrogen_06", Scale.LINEAR, "1", "d"), + ("nitrogen_07", Scale.LINEAR, "1", "d"), + ("carbon_01", Scale.LINEAR, "1", "d"), + ("carbon_02", Scale.LINEAR, "1", "d"), + ("carbon_03", Scale.LINEAR, "1", "d"), + ("carbon_04", Scale.LINEAR, "1", "d"), + ("carbon_05", Scale.LINEAR, "1", "d"), + ("carbon_06", Scale.LINEAR, "1", "d"), + ("magnesium_01", Scale.LINEAR, "1", "d"), + ("magnesium_02", Scale.LINEAR, "1", "d"), + ("magnesium_03", Scale.LINEAR, "1", "d"), + ("magnesium_04", Scale.LINEAR, "1", "d"), + ("magnesium_05", Scale.LINEAR, "1", "d"), + ("magnesium_06", Scale.LINEAR, "1", "d"), + ("magnesium_07", Scale.LINEAR, "1", "d"), + ("magnesium_08", Scale.LINEAR, "1", "d"), + ("magnesium_09", Scale.LINEAR, "1", "d"), + ("magnesium_10", Scale.LINEAR, "1", "d"), + ("silicon_01", Scale.LINEAR, "1", "d"), + ("silicon_02", Scale.LINEAR, "1", "d"), + ("silicon_03", Scale.LINEAR, "1", "d"), + ("silicon_04", Scale.LINEAR, "1", "d"), + ("silicon_05", Scale.LINEAR, "1", "d"), + ("silicon_06", Scale.LINEAR, "1", "d"), + ("silicon_07", Scale.LINEAR, "1", "d"), + ("silicon_08", Scale.LINEAR, "1", "d"), + ("silicon_09", Scale.LINEAR, "1", "d"), + ("silicon_10", Scale.LINEAR, "1", "d"), + ("silicon_11", Scale.LINEAR, "1", "d"), + ("sulfur_01", Scale.LINEAR, "1", "d"), + ("sulfur_02", Scale.LINEAR, "1", "d"), + ("sulfur_03", Scale.LINEAR, "1", "d"), + ("sulfur_04", Scale.LINEAR, "1", "d"), + ("sulfur_05", Scale.LINEAR, "1", "d"), + ("sulfur_06", Scale.LINEAR, "1", "d"), + ("sulfur_07", Scale.LINEAR, "1", "d"), + ("sulfur_08", Scale.LINEAR, "1", "d"), + ("sulfur_09", Scale.LINEAR, "1", "d"), + ("sulfur_10", Scale.LINEAR, "1", "d"), + ("sulfur_11", Scale.LINEAR, "1", "d"), + ("iron_01", Scale.LINEAR, "1", "d"), + ("iron_02", Scale.LINEAR, "1", "d"), + ("iron_03", Scale.LINEAR, "1", "d"), + ("iron_04", Scale.LINEAR, "1", "d"), + ("iron_05", Scale.LINEAR, "1", "d"), + ("iron_06", Scale.LINEAR, "1", "d"), + ("iron_07", Scale.LINEAR, "1", "d"), + ("iron_08", Scale.LINEAR, "1", "d"), + ("iron_09", Scale.LINEAR, "1", "d"), + ("iron_10", Scale.LINEAR, "1", "d"), + ("iron_11", Scale.LINEAR, "1", "d"), + ("neon_01", Scale.LINEAR, "1", "d"), + ("neon_02", Scale.LINEAR, "1", "d"), + ("neon_03", Scale.LINEAR, "1", "d"), + ("neon_04", Scale.LINEAR, "1", "d"), + ("neon_05", Scale.LINEAR, "1", "d"), + ("neon_06", Scale.LINEAR, "1", "d"), + ("neon_07", Scale.LINEAR, "1", "d"), + ("neon_08", Scale.LINEAR, "1", "d"), + ("neon_09", Scale.LINEAR, "1", "d"), + ("neon_10", Scale.LINEAR, "1", "d"), + ("hydrogen_01", Scale.LINEAR, "1", "d"), + ("hydrogen_02", Scale.LINEAR, "1", "d"), + ("helium_02", Scale.LINEAR, "1", "d"), + ("helium_03", Scale.LINEAR, "1", "d"), + ("Habing", Scale.LOG, "erg/s/cm**2", "d"), + ("Lyman_Werner", Scale.LOG, "erg/s/cm**2", "d"), + ("HI_Ionising", Scale.LOG, "erg/s/cm**2", "d"), + ("H2_Ionising", Scale.LOG, "erg/s/cm**2", "d"), + ("HeI_Ionising", Scale.LOG, "erg/s/cm**2", "d"), + ("HeII_ionising", Scale.LOG, "erg/s/cm**2", "d"), + ], + 2: [ + ("redshift", Scale.LINEAR, "1", "f"), + ("dx", Scale.LOG, "cm", "f"), + ("x", Scale.LINEAR, "Mpccm/h", "d"), + ("y", Scale.LINEAR, "Mpccm/h", "d"), + ("z", Scale.LINEAR, "Mpccm/h", "d"), + ("vx", Scale.LINEAR, "cm/s", "f"), + ("vy", Scale.LINEAR, "cm/s", "f"), + ("vz", Scale.LINEAR, "cm/s", "f"), + ("density", Scale.LOG, "mp/cm**3", "f"), + ("hydrogen_density", Scale.LOG, "1/cm**3", "f"), + ("temperature", Scale.LOG, "K", "f"), + ("pressure", Scale.LOG, "dyne/cm**2", "f"), + ("iron_number_density", Scale.LOG, "1/cm**3", "f"), + ("oxygen_number_density", Scale.LOG, "1/cm**3", "f"), + ("nitrogen_number_density", Scale.LOG, "1/cm**3", "f"), + ("magnesium_number_density", Scale.LOG, "1/cm**3", "f"), + ("neon_number_density", Scale.LOG, "1/cm**3", "f"), + ("silicon_number_density", Scale.LOG, "1/cm**3", "f"), + ("calcium_number_density", Scale.LOG, "1/cm**3", "f"), + ("carbon_number_density", Scale.LOG, "1/cm**3", "f"), + ("sulfur_number_density", Scale.LOG, "1/cm**3", "f"), + ("carbon_monoxide_number_density", Scale.LOG, "1/cm**3", "f"), + ("oxygen_01", Scale.LOG, "1", "f"), + ("oxygen_02", Scale.LOG, "1", "f"), + ("oxygen_03", Scale.LOG, "1", "f"), + ("oxygen_04", Scale.LOG, "1", "f"), + ("oxygen_05", Scale.LOG, "1", "f"), + ("oxygen_06", Scale.LOG, "1", "f"), + ("oxygen_07", Scale.LOG, "1", "f"), + ("oxygen_08", Scale.LOG, "1", "f"), + ("nitrogen_01", Scale.LOG, "1", "f"), + ("nitrogen_02", Scale.LOG, "1", "f"), + ("nitrogen_03", Scale.LOG, "1", "f"), + ("nitrogen_04", Scale.LOG, "1", "f"), + ("nitrogen_05", Scale.LOG, "1", "f"), + ("nitrogen_06", Scale.LOG, "1", "f"), + ("nitrogen_07", Scale.LOG, "1", "f"), + ("carbon_01", Scale.LOG, "1", "f"), + ("carbon_02", Scale.LOG, "1", "f"), + ("carbon_03", Scale.LOG, "1", "f"), + ("carbon_04", Scale.LOG, "1", "f"), + ("carbon_05", Scale.LOG, "1", "f"), + ("carbon_06", Scale.LOG, "1", "f"), + ("magnesium_01", Scale.LOG, "1", "f"), + ("magnesium_02", Scale.LOG, "1", "f"), + ("magnesium_03", Scale.LOG, "1", "f"), + ("magnesium_04", Scale.LOG, "1", "f"), + ("magnesium_05", Scale.LOG, "1", "f"), + ("magnesium_06", Scale.LOG, "1", "f"), + ("magnesium_07", Scale.LOG, "1", "f"), + ("magnesium_08", Scale.LOG, "1", "f"), + ("magnesium_09", Scale.LOG, "1", "f"), + ("magnesium_10", Scale.LOG, "1", "f"), + ("silicon_01", Scale.LOG, "1", "f"), + ("silicon_02", Scale.LOG, "1", "f"), + ("silicon_03", Scale.LOG, "1", "f"), + ("silicon_04", Scale.LOG, "1", "f"), + ("silicon_05", Scale.LOG, "1", "f"), + ("silicon_06", Scale.LOG, "1", "f"), + ("silicon_07", Scale.LOG, "1", "f"), + ("silicon_08", Scale.LOG, "1", "f"), + ("silicon_09", Scale.LOG, "1", "f"), + ("silicon_10", Scale.LOG, "1", "f"), + ("silicon_11", Scale.LOG, "1", "f"), + ("sulfur_01", Scale.LOG, "1", "f"), + ("sulfur_02", Scale.LOG, "1", "f"), + ("sulfur_03", Scale.LOG, "1", "f"), + ("sulfur_04", Scale.LOG, "1", "f"), + ("sulfur_05", Scale.LOG, "1", "f"), + ("sulfur_06", Scale.LOG, "1", "f"), + ("sulfur_07", Scale.LOG, "1", "f"), + ("sulfur_08", Scale.LOG, "1", "f"), + ("sulfur_09", Scale.LOG, "1", "f"), + ("sulfur_10", Scale.LOG, "1", "f"), + ("sulfur_11", Scale.LOG, "1", "f"), + ("iron_01", Scale.LOG, "1", "f"), + ("iron_02", Scale.LOG, "1", "f"), + ("iron_03", Scale.LOG, "1", "f"), + ("iron_04", Scale.LOG, "1", "f"), + ("iron_05", Scale.LOG, "1", "f"), + ("iron_06", Scale.LOG, "1", "f"), + ("iron_07", Scale.LOG, "1", "f"), + ("iron_08", Scale.LOG, "1", "f"), + ("iron_09", Scale.LOG, "1", "f"), + ("iron_10", Scale.LOG, "1", "f"), + ("iron_11", Scale.LOG, "1", "f"), + ("neon_01", Scale.LOG, "1", "f"), + ("neon_02", Scale.LOG, "1", "f"), + ("neon_03", Scale.LOG, "1", "f"), + ("neon_04", Scale.LOG, "1", "f"), + ("neon_05", Scale.LOG, "1", "f"), + ("neon_06", Scale.LOG, "1", "f"), + ("neon_07", Scale.LOG, "1", "f"), + ("neon_08", Scale.LOG, "1", "f"), + ("neon_09", Scale.LOG, "1", "f"), + ("neon_10", Scale.LOG, "1", "f"), + ("hydrogen_01", Scale.LOG, "1", "f"), + ("hydrogen_02", Scale.LOG, "1", "f"), + ("helium_02", Scale.LOG, "1", "f"), + ("helium_03", Scale.LOG, "1", "f"), + ("Habing", Scale.LOG, "erg/s/cm**2", "f"), + ("Lyman_Werner", Scale.LOG, "erg/s/cm**2", "f"), + ("HI_Ionising", Scale.LOG, "erg/s/cm**2", "f"), + ("H2_Ionising", Scale.LOG, "erg/s/cm**2", "f"), + ("HeI_Ionising", Scale.LOG, "erg/s/cm**2", "f"), + ("HeII_ionising", Scale.LOG, "erg/s/cm**2", "f"), + ("heating_rate", Scale.LOG, "erg/s", "f"), + ("cooling_rate", Scale.LOG, "erg/s", "f"), + ], +} + + +def load_cutout( + filename: str | Path, + boxsize: float = 50, + h0: float = 0.6727, + verbose: bool = True, + version: int | list[tuple[str, Scale, str, str]] = 1, +): """Load a Megatron cutout file as a yt dataset. Parameters @@ -139,6 +257,15 @@ def load_cutout(filename: str | Path, boxsize: float = 50, h0: float = 0.6727, v The Hubble constant of the original simulation. Default is 0.6727. verbose : bool Whether to show a progress bar when loading the data. Default is True. + version : int or list of (name, scale, unit) tuples + The version of the cutout format to load. If an int, it must be a key + in the `headers` dict. If a list, it should be a custom header + specification. Default is 2. + + Returns + ------- + yt.Dataset + The loaded yt dataset. """ original_path = path = Path(filename) try: @@ -151,12 +278,17 @@ def load_cutout(filename: str | Path, boxsize: float = 50, h0: float = 0.6727, v if not path.exists(): path = Path(pooch.retrieve(str(filename), known_hash=None)) + if isinstance(version, int): + header = headers[version] + else: + header = version + data = {} with FortranFile(path, "r") as ff: prog = tqdm if verbose and yt.is_root() else lambda x, *args, **kwargs: x - for name, scale, _unit in prog(header, desc="Loading cutout"): + for name, scale, _unit, dtype in prog(header, desc="Loading cutout"): # Read in the quantity - raw_data = ff.read_reals("float64") + raw_data = ff.read_reals(dtype) if scale == Scale.LOG: if use_numexpr: ne.evaluate("10 ** raw_data", out=raw_data) From 02b7d59494984643048e0c054ea9d3ba189a4e9d Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 21 Apr 2026 12:26:51 +0200 Subject: [PATCH 2/9] Adapt to Martin's format FIXME: star_ds should have the right units, it doesn not presently --- yt_derived_fields/cutouts/loader.py | 152 +++++++++++++++++++++------- 1 file changed, 117 insertions(+), 35 deletions(-) diff --git a/yt_derived_fields/cutouts/loader.py b/yt_derived_fields/cutouts/loader.py index c7b572f..cde10f1 100644 --- a/yt_derived_fields/cutouts/loader.py +++ b/yt_derived_fields/cutouts/loader.py @@ -4,19 +4,89 @@ import numpy as np import pooch import unyt +from unyt.dimensions import length + import yt -from scipy.io import FortranFile +from cython_fortran_file import FortranFile from tqdm import tqdm from yt_experiments.octree.converter import OctTree +from unyt import unyt_array from yt_derived_fields.megatron_derived_fields.chemistry_derived_fields import metal_data +try: + import numexpr as ne + USE_NUMEXPR = True +except ImportError: + USE_NUMEXPR = False + + class Scale(Enum): LINEAR = 0 LOG = 1 + + +star_headers: dict[int, list[tuple[str, Scale, str, str]]] = { + 1: [], + 2: [ + ("particle_position_x", Scale.LINEAR, "unitary", "d"), + ("particle_position_y", Scale.LINEAR, "unitary", "d"), + ("particle_position_z", Scale.LINEAR, "unitary", "d"), + ("particle_velocity_x", Scale.LINEAR, "cm/s", "f"), + ("particle_velocity_y", Scale.LINEAR, "cm/s", "f"), + ("particle_velocity_z", Scale.LINEAR, "cm/s", "f"), + ("age", Scale.LOG, "Myr", "f"), + ("iron_mass_fraction", Scale.LOG, "1", "f"), + ("oxygen_mass_fraction", Scale.LOG, "1", "f"), + ("nitrogen_mass_fraction", Scale.LOG, "1", "f"), + ("magnesium_mass_fraction", Scale.LOG, "1", "f"), + ("neon_mass_fraction", Scale.LOG, "1", "f"), + ("silicon_mass_fraction", Scale.LOG, "1", "f"), + ("calcium_mass_fraction", Scale.LOG, "1", "f"), + ("carbon_mass_fraction", Scale.LOG, "1", "f"), + ("sulfur_mass_fraction", Scale.LOG, "1", "f"), + ("initial_mass", Scale.LOG, "Msun", "f"), + ("mass", Scale.LOG, "Msun", "f"), + ] +} + +def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source = None): + if isinstance(fname, str): + fname = Path(fname) + + registry = data_source.ds.unit_registry if data_source is not None else unyt.UnitRegistry() + + with FortranFile(fname, "r") as ff: + ff.seek(0, 2) + endpos = ff.tell() + ff.seek(0) + + header = star_headers[2] + + nstars = ff.read_int() + + data = {} + + for name, scale, unit, dtype in tqdm(header, desc="Loading star cutout"): + raw_data = ff.read_vector(dtype).astype("d") + if scale == Scale.LOG: + if USE_NUMEXPR: + ne.evaluate("10 ** raw_data", out=raw_data) + else: + raw_data = 10**raw_data + + assert len(raw_data) == nstars, f"Expected {nstars} stars but got {len(raw_data)} for field {name}" + + data[name] = (raw_data, unit) + + assert ff.tell() == endpos, "Did not read entire file" + + return yt.load_particles(data, data_source=data_source) + + headers: dict[int, list[tuple[str, Scale, str, str]]] = { 1: [ ("redshift", Scale.LINEAR, "1", "d"), @@ -128,14 +198,14 @@ class Scale(Enum): 2: [ ("redshift", Scale.LINEAR, "1", "f"), ("dx", Scale.LOG, "cm", "f"), - ("x", Scale.LINEAR, "Mpccm/h", "d"), - ("y", Scale.LINEAR, "Mpccm/h", "d"), - ("z", Scale.LINEAR, "Mpccm/h", "d"), + ("x", Scale.LINEAR, "unitary", "d"), + ("y", Scale.LINEAR, "unitary", "d"), + ("z", Scale.LINEAR, "unitary", "d"), ("vx", Scale.LINEAR, "cm/s", "f"), ("vy", Scale.LINEAR, "cm/s", "f"), ("vz", Scale.LINEAR, "cm/s", "f"), - ("density", Scale.LOG, "mp/cm**3", "f"), - ("hydrogen_density", Scale.LOG, "1/cm**3", "f"), + ("density", Scale.LOG, "g/cm**3", "f"), + ("hydrogen_number_density", Scale.LOG, "1/cm**3", "f"), ("temperature", Scale.LOG, "K", "f"), ("pressure", Scale.LOG, "dyne/cm**2", "f"), ("iron_number_density", Scale.LOG, "1/cm**3", "f"), @@ -240,10 +310,13 @@ class Scale(Enum): def load_cutout( filename: str | Path, - boxsize: float = 50, - h0: float = 0.6727, verbose: bool = True, version: int | list[tuple[str, Scale, str, str]] = 1, + h0 = 0.672699966430664, + boxsize = 50., + omega_m = 0.313899993896484, + omega_l = 0.686094999313354, + omega_b = 0.4916, ): """Load a Megatron cutout file as a yt dataset. @@ -253,14 +326,20 @@ def load_cutout( Path to the cutout file. If a URL, it will be downloaded using pooch. boxsize : boxsize in Mpccm/h The boxsize of the original simulation in comoving Mpc/h. Default is 50. - h0 : float - The Hubble constant of the original simulation. Default is 0.6727. verbose : bool Whether to show a progress bar when loading the data. Default is True. version : int or list of (name, scale, unit) tuples The version of the cutout format to load. If an int, it must be a key in the `headers` dict. If a list, it should be a custom header specification. Default is 2. + h0 : float + The Hubble constant of the original simulation. + omega_m : float + The matter density parameter of the original simulation. + omega_l : float + The dark energy density parameter of the original simulation. + omega_b : float + The baryon density parameter of the original simulation. Returns ------- @@ -268,12 +347,6 @@ def load_cutout( The loaded yt dataset. """ original_path = path = Path(filename) - try: - import numexpr as ne - - use_numexpr = True - except ImportError: - use_numexpr = False if not path.exists(): path = Path(pooch.retrieve(str(filename), known_hash=None)) @@ -285,40 +358,41 @@ def load_cutout( data = {} with FortranFile(path, "r") as ff: + # Seek to end to compute file size + ff.seek(0, 2) + endpos = ff.tell() + ff.seek(0) + prog = tqdm if verbose and yt.is_root() else lambda x, *args, **kwargs: x - for name, scale, _unit, dtype in prog(header, desc="Loading cutout"): - # Read in the quantity - raw_data = ff.read_reals(dtype) - if scale == Scale.LOG: - if use_numexpr: - ne.evaluate("10 ** raw_data", out=raw_data) - else: - raw_data = 10**raw_data + for name, scale, _unit, dtype in prog(header, desc="Loading gas cutout"): + data[name] = ff.read_vector(dtype) - if name == "density": - if use_numexpr: - ne.evaluate("raw_data / 0.76", out=raw_data) # Convert from nH to rho - else: - raw_data = raw_data / 0.76 - data[name] = raw_data + if scale == Scale.LOG and USE_NUMEXPR: + ne.evaluate("10 ** data", out=data[name]) + elif scale == Scale.LOG: + data[name] = 10**data[name] + + # Make sure we read the entire file + assert ff.tell() == endpos, "Did not read entire file" - redshift = data.pop("redshift")[0] + redshift = data.pop("redshift")()[0][0] aexp = 1 / (1 + redshift) # Create a unyt registry boxsize_physical = boxsize * unyt.Mpc * aexp / h0 registry = unyt.UnitRegistry() + registry.add("unitary", float(boxsize_physical.to("m")), length) # Get xc (no need for unit conversion thus) - xc = np.stack([data.pop(_) for _ in "xyz"], axis=-1) + xc = np.stack([data.pop(_)()[0] for _ in "xyz"], axis=-1) center = (xc.max(axis=0) + xc.min(axis=0)) / 2 # Special case for dx (needs precise conversion from pc) - dx = data.pop("dx") / 3.08e18 * unyt.pc / boxsize_physical + dx = data.pop("dx")()[0] / 3.08e18 * unyt.pc / boxsize_physical # Convert everything else - for name, _, unit in header: + for name, _, unit, dtype in header: if name not in data: continue data[name] = unyt.unyt_array(data[name], unit, registry=registry) @@ -368,6 +442,9 @@ def reorder(dt): "cosmological_simulation": True, "current_redshift": redshift, "hubble_constant": h0, + "omega_matter": omega_m, + "omega_lambda": omega_l, + "omega_baryon": omega_b, } yt.mylog.debug("Loading octree dataset") @@ -380,6 +457,7 @@ def reorder(dt): parameters=params, length_unit=boxsize_physical, ) + ds.domain_center = ds.arr(center, "code_length") yt.mylog.debug("---------------------------------------------") @@ -414,4 +492,8 @@ def _metal_mass_fraction(field, data): for element in metal_data.keys(): create_density(element) - return ds + # Now try creating star dataset if possible + star_fname = ds.filename.replace("gas", "stars") + star_ds = load_star_cutout(original_path.parent / star_fname, boxsize, h0, aexp, ds.all_data()) + + return ds, star_ds From ceb00f43f658c37880e124f554c6b0e86d2755e7 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 21 Apr 2026 13:27:08 +0200 Subject: [PATCH 3/9] Use lazy loading for performances --- yt_derived_fields/cutouts/loader.py | 118 +++++++++++++++++++++------- 1 file changed, 88 insertions(+), 30 deletions(-) diff --git a/yt_derived_fields/cutouts/loader.py b/yt_derived_fields/cutouts/loader.py index cde10f1..b275d7d 100644 --- a/yt_derived_fields/cutouts/loader.py +++ b/yt_derived_fields/cutouts/loader.py @@ -2,10 +2,12 @@ from pathlib import Path import numpy as np +import numpy.typing as npt import pooch import unyt from unyt.dimensions import length - +from dataclasses import dataclass, field +from functools import partial import yt from cython_fortran_file import FortranFile from tqdm import tqdm @@ -16,19 +18,17 @@ try: import numexpr as ne + USE_NUMEXPR = True except ImportError: USE_NUMEXPR = False - class Scale(Enum): LINEAR = 0 LOG = 1 - - star_headers: dict[int, list[tuple[str, Scale, str, str]]] = { 1: [], 2: [ @@ -50,10 +50,11 @@ class Scale(Enum): ("sulfur_mass_fraction", Scale.LOG, "1", "f"), ("initial_mass", Scale.LOG, "Msun", "f"), ("mass", Scale.LOG, "Msun", "f"), - ] + ], } -def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source = None): + +def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source=None): if isinstance(fname, str): fname = Path(fname) @@ -308,15 +309,55 @@ def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source = None): } +@dataclass +class IOHandler: + filename: str | Path + unyt_registry: unyt.UnitRegistry + metadata: dict[str, tuple[int, str, Scale, str, str]] = field(default_factory=dict) + fp: FortranFile = field(init=False, repr=False) + + def __post_init__(self): + self.fp = FortranFile(self.filename) + + def __delete__(self, instance): + self.fp.close() + + def read(self, name: str): + pos, name, scale, _unit, dtype = self.metadata[name] + self.fp.seek(pos) + raw_data = self.fp.read_vector(dtype) + if scale == Scale.LOG and USE_NUMEXPR: + ne.evaluate("10 ** raw_data", out=raw_data) + elif scale == Scale.LOG: + raw_data = 10**raw_data + + return raw_data + + def read_with_order(self, name: str, order: npt.NDArray[int], nan_mask: npt.NDArray): + dt = self.read(name) + tmp = dt[order] * nan_mask + return tmp[:, None] + + def get_data_object(self, ptype: str, leaf_order, nan_mask): + """Return a dictionary of field name to delayed read functions.""" + data = {} + for k, (_pos, _name, _scale, unit, _dtype) in self.metadata.items(): + data[ptype, k] = ( + partial(self.read_with_order, k, leaf_order, nan_mask), + unit, + ) + return data + + def load_cutout( filename: str | Path, verbose: bool = True, version: int | list[tuple[str, Scale, str, str]] = 1, - h0 = 0.672699966430664, - boxsize = 50., - omega_m = 0.313899993896484, - omega_l = 0.686094999313354, - omega_b = 0.4916, + h0=0.672699966430664, + boxsize=50.0, + omega_m=0.313899993896484, + omega_l=0.686094999313354, + omega_b=0.4916, ): """Load a Megatron cutout file as a yt dataset. @@ -356,6 +397,13 @@ def load_cutout( else: header = version + # Create unyt registry + registry = unyt.UnitRegistry() + io_handler = IOHandler( + filename=path, + unyt_registry=registry, + ) + data = {} with FortranFile(path, "r") as ff: # Seek to end to compute file size @@ -364,18 +412,25 @@ def load_cutout( ff.seek(0) prog = tqdm if verbose and yt.is_root() else lambda x, *args, **kwargs: x - for name, scale, _unit, dtype in prog(header, desc="Loading gas cutout"): - data[name] = ff.read_vector(dtype) + for name, scale, unit, dtype in prog(header, desc="Loading gas cutout"): + if name in ("x", "y", "z", "dx", "redshift"): + raw_data = ff.read_vector(dtype) + + if scale == Scale.LOG and USE_NUMEXPR: + ne.evaluate("10 ** raw_data", out=raw_data) + elif scale == Scale.LOG: + raw_data = 10**raw_data - if scale == Scale.LOG and USE_NUMEXPR: - ne.evaluate("10 ** data", out=data[name]) - elif scale == Scale.LOG: - data[name] = 10**data[name] + data[name] = raw_data + else: + print(name, scale, unit, dtype) + io_handler.metadata[name] = ff.tell(), name, scale, unit, dtype + ff.skip() # Make sure we read the entire file assert ff.tell() == endpos, "Did not read entire file" - redshift = data.pop("redshift")()[0][0] + redshift = data.pop("redshift")[0] aexp = 1 / (1 + redshift) # Create a unyt registry @@ -384,18 +439,18 @@ def load_cutout( registry.add("unitary", float(boxsize_physical.to("m")), length) # Get xc (no need for unit conversion thus) - xc = np.stack([data.pop(_)()[0] for _ in "xyz"], axis=-1) + xc = np.stack([data.pop(_) for _ in "xyz"], axis=-1) center = (xc.max(axis=0) + xc.min(axis=0)) / 2 # Special case for dx (needs precise conversion from pc) - dx = data.pop("dx")()[0] / 3.08e18 * unyt.pc / boxsize_physical + dx = data.pop("dx") / 3.08e18 * unyt.pc / boxsize_physical - # Convert everything else - for name, _, unit, dtype in header: - if name not in data: - continue - data[name] = unyt.unyt_array(data[name], unit, registry=registry) + # # Convert everything else + # for name, _, unit, dtype in header: + # if name not in data: + # continue + # data[name] = unyt.unyt_array(data[name], unit, registry=registry) # Get level level = np.round(np.log2(1 / dx)).astype(int) @@ -431,12 +486,15 @@ def load_cutout( nan_mask = np.where(leaf_order < 0, np.nan, 1) - def reorder(dt): - tmp = dt[leaf_order] * nan_mask - return tmp[:, None] + data = io_handler.get_data_object("gas", leaf_order, nan_mask) + + # def reorder(field): + # dt = io_handler.read(field) + # tmp = dt[leaf_order] * nan_mask + # return tmp[:, None] - yt.mylog.debug("Reordering data according to octree leaf order") - data = {("gas", k): reorder(v) for k, v in data.items()} + # yt.mylog.debug("Reordering data according to octree leaf order") + # data = {("gas", k): lambda: reorder(k) for k, v in data.items()} params = { "cosmological_simulation": True, From 5a5c3a6c756a7679b00e5ae648dd24b1cb31180d Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 21 Apr 2026 13:39:11 +0200 Subject: [PATCH 4/9] Pass quantities as tuple (value, unit) This prevents yt from getting confused when converting length_unit into code_length. --- yt_derived_fields/cutouts/loader.py | 26 +++----------------------- 1 file changed, 3 insertions(+), 23 deletions(-) diff --git a/yt_derived_fields/cutouts/loader.py b/yt_derived_fields/cutouts/loader.py index b275d7d..05eec81 100644 --- a/yt_derived_fields/cutouts/loader.py +++ b/yt_derived_fields/cutouts/loader.py @@ -58,8 +58,6 @@ def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source=None): if isinstance(fname, str): fname = Path(fname) - registry = data_source.ds.unit_registry if data_source is not None else unyt.UnitRegistry() - with FortranFile(fname, "r") as ff: ff.seek(0, 2) endpos = ff.tell() @@ -312,7 +310,6 @@ def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source=None): @dataclass class IOHandler: filename: str | Path - unyt_registry: unyt.UnitRegistry metadata: dict[str, tuple[int, str, Scale, str, str]] = field(default_factory=dict) fp: FortranFile = field(init=False, repr=False) @@ -398,10 +395,8 @@ def load_cutout( header = version # Create unyt registry - registry = unyt.UnitRegistry() io_handler = IOHandler( filename=path, - unyt_registry=registry, ) data = {} @@ -423,7 +418,6 @@ def load_cutout( data[name] = raw_data else: - print(name, scale, unit, dtype) io_handler.metadata[name] = ff.tell(), name, scale, unit, dtype ff.skip() @@ -435,8 +429,6 @@ def load_cutout( # Create a unyt registry boxsize_physical = boxsize * unyt.Mpc * aexp / h0 - registry = unyt.UnitRegistry() - registry.add("unitary", float(boxsize_physical.to("m")), length) # Get xc (no need for unit conversion thus) xc = np.stack([data.pop(_) for _ in "xyz"], axis=-1) @@ -446,12 +438,6 @@ def load_cutout( # Special case for dx (needs precise conversion from pc) dx = data.pop("dx") / 3.08e18 * unyt.pc / boxsize_physical - # # Convert everything else - # for name, _, unit, dtype in header: - # if name not in data: - # continue - # data[name] = unyt.unyt_array(data[name], unit, registry=registry) - # Get level level = np.round(np.log2(1 / dx)).astype(int) @@ -488,14 +474,6 @@ def load_cutout( data = io_handler.get_data_object("gas", leaf_order, nan_mask) - # def reorder(field): - # dt = io_handler.read(field) - # tmp = dt[leaf_order] * nan_mask - # return tmp[:, None] - - # yt.mylog.debug("Reordering data according to octree leaf order") - # data = {("gas", k): lambda: reorder(k) for k, v in data.items()} - params = { "cosmological_simulation": True, "current_redshift": redshift, @@ -513,7 +491,9 @@ def load_cutout( num_zones=1, dataset_name=original_path.name, parameters=params, - length_unit=boxsize_physical, + length_unit=(boxsize_physical.value, str(boxsize_physical.units)), + mass_unit=(1, "Msun"), + time_unit=(1, "Gyr"), ) ds.domain_center = ds.arr(center, "code_length") From 06f242851d5dec214329b2cdd97e5e3c5cb1ad48 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Wed, 22 Apr 2026 12:05:33 +0200 Subject: [PATCH 5/9] feat: Stars are now part of the dataset Note that this requires the dev branch of yt --- pyproject.toml | 2 +- yt_derived_fields/cutouts/loader.py | 134 +++++++++++++++------------- 2 files changed, 75 insertions(+), 61 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b172dca..f22b153 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ classifiers = [ "Programming Language :: Python :: 3 :: Only", ] dependencies = [ - "yt>=4.3.0", + "yt>=4.5.dev0", "yt_experiments>=0.3.0", "roman", "pyneb", diff --git a/yt_derived_fields/cutouts/loader.py b/yt_derived_fields/cutouts/loader.py index 05eec81..1765d2c 100644 --- a/yt_derived_fields/cutouts/loader.py +++ b/yt_derived_fields/cutouts/loader.py @@ -1,5 +1,6 @@ from enum import Enum from pathlib import Path +from typing import Union import numpy as np import numpy.typing as npt @@ -32,9 +33,9 @@ class Scale(Enum): star_headers: dict[int, list[tuple[str, Scale, str, str]]] = { 1: [], 2: [ - ("particle_position_x", Scale.LINEAR, "unitary", "d"), - ("particle_position_y", Scale.LINEAR, "unitary", "d"), - ("particle_position_z", Scale.LINEAR, "unitary", "d"), + ("particle_position_x", Scale.LINEAR, "code_length", "d"), + ("particle_position_y", Scale.LINEAR, "code_length", "d"), + ("particle_position_z", Scale.LINEAR, "code_length", "d"), ("particle_velocity_x", Scale.LINEAR, "cm/s", "f"), ("particle_velocity_y", Scale.LINEAR, "cm/s", "f"), ("particle_velocity_z", Scale.LINEAR, "cm/s", "f"), @@ -53,39 +54,6 @@ class Scale(Enum): ], } - -def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source=None): - if isinstance(fname, str): - fname = Path(fname) - - with FortranFile(fname, "r") as ff: - ff.seek(0, 2) - endpos = ff.tell() - ff.seek(0) - - header = star_headers[2] - - nstars = ff.read_int() - - data = {} - - for name, scale, unit, dtype in tqdm(header, desc="Loading star cutout"): - raw_data = ff.read_vector(dtype).astype("d") - if scale == Scale.LOG: - if USE_NUMEXPR: - ne.evaluate("10 ** raw_data", out=raw_data) - else: - raw_data = 10**raw_data - - assert len(raw_data) == nstars, f"Expected {nstars} stars but got {len(raw_data)} for field {name}" - - data[name] = (raw_data, unit) - - assert ff.tell() == endpos, "Did not read entire file" - - return yt.load_particles(data, data_source=data_source) - - headers: dict[int, list[tuple[str, Scale, str, str]]] = { 1: [ ("redshift", Scale.LINEAR, "1", "d"), @@ -309,20 +277,32 @@ def load_star_cutout(fname: str | Path, boxsize, h0, aexp, data_source=None): @dataclass class IOHandler: - filename: str | Path - metadata: dict[str, tuple[int, str, Scale, str, str]] = field(default_factory=dict) - fp: FortranFile = field(init=False, repr=False) + field_filename: str | Path + particle_filename: str | Path + field_metadata: dict[str, tuple[int, str, Scale, str, str]] = field(default_factory=dict) + particle_metadata: dict[str, tuple[int, str, Scale, str, str]] = field(default_factory=dict) + + field_fp: FortranFile = field(init=False, repr=False) def __post_init__(self): - self.fp = FortranFile(self.filename) + self.field_fp = FortranFile(self.field_filename) + self.particles_fp = FortranFile(self.particle_filename) def __delete__(self, instance): - self.fp.close() - - def read(self, name: str): - pos, name, scale, _unit, dtype = self.metadata[name] - self.fp.seek(pos) - raw_data = self.fp.read_vector(dtype) + try: + self.field_fp.close() + except Exception: + pass + try: + self.particles_fp.close() + except Exception: + pass + + @classmethod + def read_data(cls, fp: FortranFile, metadata: dict, key: str): + pos, name, scale, _unit, dtype = metadata[key] + fp.seek(pos) + raw_data = fp.read_vector(dtype) if scale == Scale.LOG and USE_NUMEXPR: ne.evaluate("10 ** raw_data", out=raw_data) elif scale == Scale.LOG: @@ -330,17 +310,38 @@ def read(self, name: str): return raw_data - def read_with_order(self, name: str, order: npt.NDArray[int], nan_mask: npt.NDArray): - dt = self.read(name) + def read_field(self, name: str) -> npt.NDArray: + return self.read_data(self.field_fp, self.field_metadata, name) + + def read_field_in_order(self, name: str, order: npt.NDArray[int], nan_mask: npt.NDArray) -> npt.NDArray: + dt = self.read_field(name) tmp = dt[order] * nan_mask return tmp[:, None] - def get_data_object(self, ptype: str, leaf_order, nan_mask): + def read_particle(self, name: str) -> npt.NDArray: + return self.read_data(self.particles_fp, self.particle_metadata, name) + + def get_data_object( + self, + *, + ftype: str, + ptype: str, + leaf_order: npt.NDArray[int], + nan_mask: npt.NDArray, + ) -> dict[tuple[str, str], tuple[Union[callable, npt.NDArray], str]]: """Return a dictionary of field name to delayed read functions.""" data = {} - for k, (_pos, _name, _scale, unit, _dtype) in self.metadata.items(): + # Loader for field values + for k, (_pos, _name, _scale, unit, _dtype) in self.field_metadata.items(): + data[ftype, k] = ( + partial(self.read_field_in_order, k, leaf_order, nan_mask), + unit, + ) + + # Loader for particles + for k, (_pos, _name, _scale, unit, _dtype) in self.particle_metadata.items(): data[ptype, k] = ( - partial(self.read_with_order, k, leaf_order, nan_mask), + self.read_particle(k), unit, ) return data @@ -366,7 +367,7 @@ def load_cutout( The boxsize of the original simulation in comoving Mpc/h. Default is 50. verbose : bool Whether to show a progress bar when loading the data. Default is True. - version : int or list of (name, scale, unit) tuples + version : int or list of (name, scale, unit, dtype) tuples The version of the cutout format to load. If an int, it must be a key in the `headers` dict. If a list, it should be a custom header specification. Default is 2. @@ -385,21 +386,26 @@ def load_cutout( The loaded yt dataset. """ original_path = path = Path(filename) + star_path = path.parent / path.name.replace("gas", "stars") if not path.exists(): path = Path(pooch.retrieve(str(filename), known_hash=None)) if isinstance(version, int): header = headers[version] + particle_header = star_headers[version] else: header = version + particle_header = [] # Create unyt registry io_handler = IOHandler( - filename=path, + field_filename=path, + particle_filename=star_path, ) data = {} + # Handle fields with FortranFile(path, "r") as ff: # Seek to end to compute file size ff.seek(0, 2) @@ -418,12 +424,19 @@ def load_cutout( data[name] = raw_data else: - io_handler.metadata[name] = ff.tell(), name, scale, unit, dtype + io_handler.field_metadata[name] = ff.tell(), name, scale, unit, dtype ff.skip() # Make sure we read the entire file assert ff.tell() == endpos, "Did not read entire file" + if star_path.exists(): + with FortranFile(star_path, "r") as ff: + _nstar = ff.read_int() + for name, scale, unit, dtype in particle_header: + io_handler.particle_metadata[name] = ff.tell(), name, scale, unit, dtype + ff.skip() + redshift = data.pop("redshift")[0] aexp = 1 / (1 + redshift) @@ -472,7 +485,12 @@ def load_cutout( nan_mask = np.where(leaf_order < 0, np.nan, 1) - data = io_handler.get_data_object("gas", leaf_order, nan_mask) + data = io_handler.get_data_object( + ftype="gas", + ptype="star", + leaf_order=leaf_order, + nan_mask=nan_mask, + ) params = { "cosmological_simulation": True, @@ -530,8 +548,4 @@ def _metal_mass_fraction(field, data): for element in metal_data.keys(): create_density(element) - # Now try creating star dataset if possible - star_fname = ds.filename.replace("gas", "stars") - star_ds = load_star_cutout(original_path.parent / star_fname, boxsize, h0, aexp, ds.all_data()) - - return ds, star_ds + return ds From a2eb17a1190c49fedf59253f18b731881403bed7 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 12 May 2026 18:00:52 +0200 Subject: [PATCH 6/9] Add cosmology informations --- yt_derived_fields/cutouts/loader.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/yt_derived_fields/cutouts/loader.py b/yt_derived_fields/cutouts/loader.py index 1765d2c..fbcfd17 100644 --- a/yt_derived_fields/cutouts/loader.py +++ b/yt_derived_fields/cutouts/loader.py @@ -1,19 +1,19 @@ +from dataclasses import dataclass, field from enum import Enum +from functools import partial from pathlib import Path from typing import Union +import astropy import numpy as np import numpy.typing as npt import pooch import unyt -from unyt.dimensions import length -from dataclasses import dataclass, field -from functools import partial import yt from cython_fortran_file import FortranFile from tqdm import tqdm +from yt.utilities.cosmology import Cosmology from yt_experiments.octree.converter import OctTree -from unyt import unyt_array from yt_derived_fields.megatron_derived_fields.chemistry_derived_fields import metal_data @@ -355,7 +355,7 @@ def load_cutout( boxsize=50.0, omega_m=0.313899993896484, omega_l=0.686094999313354, - omega_b=0.4916, + omega_b=0.04916, ): """Load a Megatron cutout file as a yt dataset. @@ -492,9 +492,13 @@ def load_cutout( nan_mask=nan_mask, ) + cosmo = astropy.cosmology.FlatLambdaCDM(H0=h0 * 100, Om0=omega_m, Ob0=omega_b) + current_time = cosmo.age(redshift).to("Gyr").value + params = { "cosmological_simulation": True, "current_redshift": redshift, + "current_time": current_time, "hubble_constant": h0, "omega_matter": omega_m, "omega_lambda": omega_l, @@ -512,6 +516,14 @@ def load_cutout( length_unit=(boxsize_physical.value, str(boxsize_physical.units)), mass_unit=(1, "Msun"), time_unit=(1, "Gyr"), + sim_time = current_time, + ) + + ds.cosmology = Cosmology( + hubble_constant=ds.hubble_constant, + omega_matter=ds.omega_matter, + omega_lambda=ds.omega_lambda, + unit_registry=ds.unit_registry, ) ds.domain_center = ds.arr(center, "code_length") From 83ef5694057bcd661b229ee06883d2558a5789b3 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 12 May 2026 18:01:11 +0200 Subject: [PATCH 7/9] FEAT: add integrated luminosity --- .../stars_derived_fields.py | 30 +++++++++++++++++++ .../spectral_utils/pop2_stellar_spectra.py | 14 +++++++++ 2 files changed, 44 insertions(+) diff --git a/yt_derived_fields/megatron_derived_fields/stars_derived_fields.py b/yt_derived_fields/megatron_derived_fields/stars_derived_fields.py index a44d735..aed6039 100644 --- a/yt_derived_fields/megatron_derived_fields/stars_derived_fields.py +++ b/yt_derived_fields/megatron_derived_fields/stars_derived_fields.py @@ -203,6 +203,36 @@ def _pop2_spectra_combined(field, data): display_name="Pop. II Combined Star Spectra", ) + def _pop2_spectra_wavelengths(field, data): + + pop2_spec_wavelengths = pop2_stellar_spectra.get_pop_2_spectrum(data, wavelength_only=True) + + return pop2_spec_wavelengths + + ds.add_field( + name=("pop2", "spectra_wavelengths"), + function=_pop2_spectra_wavelengths, + #force_override=True, + units="Angstrom", + sampling_type="particle", + vector_field=True, + display_name="Pop. II Star Spectra Wavelengths", + ) + + def _pop2_luminosity(field, data): + pop2_lum = pop2_stellar_spectra.get_pop_2_spectrum(data, total_luminosity=True, parallel=parallel) + + return pop2_lum + + ds.add_field( + name=("pop2", "luminosity"), + function=_pop2_luminosity, + #force_override=True, + units="erg/s", + sampling_type="particle", + display_name="Pop. II Star Total Luminosity", + ) + # An attempt to generate a spectra per cell. The idea is to check the stars in # each cell and assign them the summed stellar spectra. In practice, one would diff --git a/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py b/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py index 01459d4..68772f2 100644 --- a/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py +++ b/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py @@ -144,6 +144,7 @@ def generate_pop_II_spec_interp( def get_pop_2_spectrum( data, combined: bool = False, + total_luminosity: bool = False, lmin: int = 1150, lmax: int = 10000, downsample: bool = True, @@ -152,6 +153,7 @@ def get_pop_2_spectrum( n_batch: int = 5000, ncpu_max: int = 10, data_dir: Optional[str] = None, + wavelength_only: bool = False, ): """ Calculates the Population II spectrum (BPASS v2.2.1). @@ -165,6 +167,7 @@ def get_pop_2_spectrum( - data["pop2", "age"] - data["pop2", "particle_initial_mass"] combined: sum spectra over all particles if True + total_luminosity: if True, return the total luminosity (integrated over wavelength) instead of the spectrum lmin, lmax: wavelength range in Å downsample: block-mean spectra and wavelengths by ds_nwv ds_nwv: integer downsampling factor @@ -172,12 +175,19 @@ def get_pop_2_spectrum( ncpu_max: max CPUs for joblib parallelization data_dir: directory containing reduced_spectra-bin-imf_chab300.*.dat.npy files progress: if True, show a tqdm progress bar when parallelizing + wavelength_only: if True, skip interpolation and return the wavelength array corresponding to the spectra bins Returns: - If combined=True: 1D unyt array (n_wvl or n_wvl_ds,) with erg/s - If combined=False: 2D array (N_pop2, n_wvl or n_wvl_ds) with erg/s """ + if wavelength_only: + return wavelength_space(lmin, lmax, downsample, ds_nwv) * u.angstrom + + if total_luminosity and combined: + raise ValueError("total_luminosity=True is not compatible with combined=True; total luminosity is a single number, not a spectrum.") + N_pop2 = int(np.sum(data["pop2", "particle_ones"])) if N_pop2 == 0: # No Pop II stars; return zeroed spectrum (no contribution) @@ -214,6 +224,8 @@ def get_pop_2_spectrum( p2_spec = spec_interp_p2(to_interp) * initial_masses[:, None] if combined: p2_spec = p2_spec.sum(axis=0) + elif total_luminosity: + p2_spec = p2_spec.sum(axis=1) return p2_spec * u.erg / u.s # Chunk the data for efficient parallelization @@ -239,5 +251,7 @@ def batch_interp(c1, c2): p2_spec = np.array(results) if combined: p2_spec = p2_spec.sum(axis=0) + elif total_luminosity: + p2_spec = p2_spec.sum(axis=1) return p2_spec * u.erg / u.s From 06c716977f5121ba7a1166689156f4e42d24e843 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 12 May 2026 18:02:33 +0200 Subject: [PATCH 8/9] Fix path on infinty --- yt_derived_fields/spectral_utils/generate_atomic_grids.py | 4 ++-- yt_derived_fields/spectral_utils/pop2_stellar_spectra.py | 2 +- yt_derived_fields/spectral_utils/pop3_stellar_spectra.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/yt_derived_fields/spectral_utils/generate_atomic_grids.py b/yt_derived_fields/spectral_utils/generate_atomic_grids.py index dd21646..7897c3f 100755 --- a/yt_derived_fields/spectral_utils/generate_atomic_grids.py +++ b/yt_derived_fields/spectral_utils/generate_atomic_grids.py @@ -24,7 +24,7 @@ candidates.append(Path(data_dir)) # Fallback onto known paths (glamdring, infinity) candidates.append(Path("/mnt/glacier/chianti/")) -candidates.append(Path("/data100/cadiou/Megatron/chianti/")) +candidates.append(Path("/data122/cadiou/Megatron/chianti/")) chianti_path = None for base in candidates: @@ -37,7 +37,7 @@ "Give a valid data directory, or place files under one of the known paths." ) -os.environ["XUVTOP"] = chianti_path +os.environ["XUVTOP"] = str(chianti_path) print("setting environment variable XUVTOP to", os.environ["XUVTOP"]) import ChiantiPy.core as ch diff --git a/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py b/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py index 68772f2..ba35fcd 100644 --- a/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py +++ b/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py @@ -46,7 +46,7 @@ def _resolve_bpass_dir(data_dir: Optional[str]) -> Path: candidates.append(Path(data_dir)) # Fallback onto known paths (glamdring, infinity) candidates.append(Path("/mnt/glacier/DATA/bpass_v2.2.1_imf_chab300")) - candidates.append(Path("/data100/cadiou/Megatron/DATA/bpass_v2.2.1_imf_chab300")) + candidates.append(Path("/data122/cadiou/Megatron/DATA/bpass_v2.2.1_imf_chab300")) for base in candidates: test_file = base / f"reduced_spectra-bin-imf_chab300.{METAL_NAMES[0]}.dat.npy" diff --git a/yt_derived_fields/spectral_utils/pop3_stellar_spectra.py b/yt_derived_fields/spectral_utils/pop3_stellar_spectra.py index f006ab3..2eab465 100644 --- a/yt_derived_fields/spectral_utils/pop3_stellar_spectra.py +++ b/yt_derived_fields/spectral_utils/pop3_stellar_spectra.py @@ -25,7 +25,7 @@ def _resolve_data_paths(data_dir: Optional[str]) -> tuple[Path, Path]: candidates.append(Path(data_dir)) # Fallback onto known paths (glamdring, infinity) candidates.append(Path("/mnt/glacier/DATA/Pop_III_spectra")) - candidates.append(Path("/data100/cadiou/Megatron/DATA/Pop_III_spectra")) + candidates.append(Path("/data122/cadiou/Megatron/DATA/Pop_III_spectra")) for base in candidates: spec = base / "reduced_popiii_spec.npy" From d1cf73f7f9931f05eea879a07b94cef4efd38d7d Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 12 May 2026 18:02:51 +0200 Subject: [PATCH 9/9] BUG: replace deprecated np.trapz --- yt_derived_fields/spectral_utils/nebular_continuum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yt_derived_fields/spectral_utils/nebular_continuum.py b/yt_derived_fields/spectral_utils/nebular_continuum.py index 1b5d892..0585e70 100644 --- a/yt_derived_fields/spectral_utils/nebular_continuum.py +++ b/yt_derived_fields/spectral_utils/nebular_continuum.py @@ -279,7 +279,7 @@ def get_nebular_continuum_two_photon( nebc = pn.Continuum() wvls = wavelength_space(lmin=lmin, lmax=lmax, downsample=downsample, ds_nwv=ds_nwv) two_photon_generic = nebc.two_photon(1e4, 1, wvls) - two_phot_erg_s = np.trapz(two_photon_generic, wvls) + two_phot_erg_s = np.trapezoid(two_photon_generic, wvls) # Get the CIE HII fraction --> needed below to make sure cooling isn't too strong CIE_HII = coll_ion_H(temperatures) / (coll_ion_H(temperatures) + recomb_ion_H(temperatures))