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 fb00fdd..fbcfd17 100644 --- a/yt_derived_fields/cutouts/loader.py +++ b/yt_derived_fields/cutouts/loader.py @@ -1,132 +1,362 @@ +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 import yt -from scipy.io import FortranFile +from cython_fortran_file import FortranFile from tqdm import tqdm +from yt.utilities.cosmology import Cosmology from yt_experiments.octree.converter import OctTree 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 -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): +star_headers: dict[int, list[tuple[str, Scale, str, str]]] = { + 1: [], + 2: [ + ("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"), + ("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"), + ], +} + +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, "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, "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"), + ("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"), + ], +} + + +@dataclass +class IOHandler: + 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.field_fp = FortranFile(self.field_filename) + self.particles_fp = FortranFile(self.particle_filename) + + def __delete__(self, instance): + 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: + raw_data = 10**raw_data + + return raw_data + + 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 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 = {} + # 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] = ( + self.read_particle(k), + 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.0, + omega_m=0.313899993896484, + omega_l=0.686094999313354, + omega_b=0.04916, +): """Load a Megatron cutout file as a yt dataset. Parameters @@ -135,47 +365,83 @@ def load_cutout(filename: str | Path, boxsize: float = 50, h0: float = 0.6727, v 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, 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. + 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 + ------- + yt.Dataset + The loaded yt dataset. """ original_path = path = Path(filename) - try: - import numexpr as ne - - use_numexpr = True - except ImportError: - use_numexpr = False + 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( + 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) + endpos = ff.tell() + ff.seek(0) + prog = tqdm if verbose and yt.is_root() else lambda x, *args, **kwargs: x - for name, scale, _unit in prog(header, desc="Loading cutout"): - # Read in the quantity - raw_data = ff.read_reals("float64") - if scale == Scale.LOG: - if use_numexpr: + 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) - else: + elif scale == Scale.LOG: raw_data = 10**raw_data - 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 + data[name] = raw_data + else: + 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) # Create a unyt registry boxsize_physical = boxsize * unyt.Mpc * aexp / h0 - registry = unyt.UnitRegistry() # Get xc (no need for unit conversion thus) xc = np.stack([data.pop(_) for _ in "xyz"], axis=-1) @@ -185,12 +451,6 @@ def load_cutout(filename: str | Path, boxsize: float = 50, h0: float = 0.6727, v # 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 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) @@ -225,17 +485,24 @@ def load_cutout(filename: str | Path, boxsize: float = 50, h0: float = 0.6727, v 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( + ftype="gas", + ptype="star", + leaf_order=leaf_order, + nan_mask=nan_mask, + ) - yt.mylog.debug("Reordering data according to octree leaf order") - data = {("gas", k): reorder(v) for k, v in data.items()} + 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, + "omega_baryon": omega_b, } yt.mylog.debug("Loading octree dataset") @@ -246,8 +513,19 @@ def reorder(dt): 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"), + 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") yt.mylog.debug("---------------------------------------------") 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/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/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)) diff --git a/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py b/yt_derived_fields/spectral_utils/pop2_stellar_spectra.py index 01459d4..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" @@ -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 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"