Merge branch 'main' into RAMSES-self-shielding

.github/workflows/wheels.yaml

@@ -33,7 +33,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Build wheels for CPython
-        uses: pypa/[email protected].2
+        uses: pypa/[email protected].3
         with:
           output-dir: dist
 

pyproject.toml

@@ -121,7 +121,7 @@ open-pmd = ["yt[HDF5]"]
 owls = ["yt[HDF5]"]
 owls-subfind = ["yt[HDF5]"]
 parthenon = ["yt[HDF5]"]
-ramses = ["yt[Fortran]"]
+ramses = ["yt[Fortran]", "scipy"]
 rockstar = []
 sdf = ["requests>=2.20.0"]
 stream = []

yt/frontends/ramses/data_structures.py

@@ -16,10 +16,9 @@
 from yt.geometry.oct_container import RAMSESOctreeContainer
 from yt.geometry.oct_geometry_handler import OctreeIndex
 from yt.utilities.cython_fortran_utils import FortranFile as fpu
-from yt.utilities.lib.cosmology_time import friedman
+from yt.utilities.lib.cosmology_time import t_frw, tau_frw
 from yt.utilities.on_demand_imports import _f90nml as f90nml
 from yt.utilities.physical_constants import kb, mp
-from yt.utilities.physical_ratios import cm_per_mpc
 
 from .definitions import (
     OUTPUT_DIR_EXP,
@@ -640,10 +639,6 @@ def _detect_output_fields(self):
             dsl.update(set(ph.field_offsets.keys()))
 
         self.particle_field_list = list(dsl)
-        cosmo = self.ds.cosmological_simulation
-        for f in self.particle_field_list[:]:
-            if f[1] == "particle_birth_time" and cosmo:
-                self.particle_field_list.append((f[0], "conformal_birth_time"))
 
         # Get the detected fields
         dsl = set()
@@ -816,6 +811,7 @@ def __init__(
         max_level_convention=None,
         default_species_fields=None,
         self_shielding=None,
+        use_conformal_time=None,
     ):
         # Here we want to initiate a traceback, if the reader is not built.
         if isinstance(fields, str):
@@ -890,6 +886,16 @@ def __init__(
             if FH.any_exist(self):
                 self.fluid_types += (FH.ftype,)
 
+        if use_conformal_time is not None:
+            self.use_conformal_time = use_conformal_time
+        elif self.cosmological_simulation:
+            if "rt" in self.fluid_types:
+                self.use_conformal_time = False
+            else:
+                self.use_conformal_time = True
+        else:
+            self.use_conformal_time = False
+
         self.storage_filename = storage_filename
 
         self.self_shielding = self_shielding
@@ -1097,15 +1103,14 @@ def caster(val):
             is_cosmological = not (
                 rheader["time"] >= 0 and rheader["H0"] == 1 and rheader["aexp"] == 1
             )
-
         if not is_cosmological:
-            self.cosmological_simulation = 0
+            self.cosmological_simulation = False
             self.current_redshift = 0
             self.hubble_constant = 0
             self.omega_matter = 0
             self.omega_lambda = 0
         else:
-            self.cosmological_simulation = 1
+            self.cosmological_simulation = True
             self.current_redshift = (1.0 / rheader["aexp"]) - 1.0
             self.omega_lambda = rheader["omega_l"]
             self.omega_matter = rheader["omega_m"]
@@ -1116,38 +1121,14 @@ def caster(val):
             force_max_level += self.min_level + 1
         self.max_level = min(force_max_level, rheader["levelmax"]) - self.min_level - 1
 
-        if self.cosmological_simulation == 0:
+        if not self.cosmological_simulation:
             self.current_time = self.parameters["time"]
         else:
-            self.tau_frw, self.t_frw, self.dtau, self.n_frw, self.time_tot = friedman(
-                self.omega_matter,
-                self.omega_lambda,
-                1.0 - self.omega_matter - self.omega_lambda,
-            )
-
-            age = self.parameters["time"]
-            iage = 1 + int(10.0 * age / self.dtau)
-            iage = np.min([iage, self.n_frw // 2 + (iage - self.n_frw // 2) // 10])
-
-            try:
-                self.time_simu = self.t_frw[iage] * (age - self.tau_frw[iage - 1]) / (
-                    self.tau_frw[iage] - self.tau_frw[iage - 1]
-                ) + self.t_frw[iage - 1] * (age - self.tau_frw[iage]) / (
-                    self.tau_frw[iage - 1] - self.tau_frw[iage]
-                )
-
-                self.current_time = (
-                    (self.time_tot + self.time_simu)
-                    / (self.hubble_constant * 1e7 / cm_per_mpc)
-                    / self.parameters["unit_t"]
-                )
-            except IndexError:
-                mylog.warning(
-                    "Yt could not convert conformal time to physical time. "
-                    "Yt will assume the simulation is *not* cosmological."
-                )
-                self.cosmological_simulation = 0
-                self.current_time = self.parameters["time"]
+            aexp_grid = np.geomspace(1e-3, 1, 2_000, endpoint=False)
+            z_grid = 1 / aexp_grid - 1
+            self.tau_frw = tau_frw(self, z_grid)
+            self.t_frw = t_frw(self, z_grid)
+            self.current_time = t_frw(self, self.current_redshift).to("Gyr")
 
         if self.num_groups > 0:
             self.group_size = rheader["ncpu"] // self.num_groups

yt/frontends/ramses/fields.py

@@ -8,8 +8,9 @@
 from yt._typing import KnownFieldsT
 from yt.fields.field_detector import FieldDetector
 from yt.fields.field_info_container import FieldInfoContainer
-from yt.frontends.ramses.io import convert_ramses_conformal_time_to_physical_age
+from yt.frontends.ramses.io import convert_ramses_conformal_time_to_physical_time
 from yt.utilities.cython_fortran_utils import FortranFile
+from yt.utilities.lib.cosmology_time import t_frw
 from yt.utilities.linear_interpolators import BilinearFieldInterpolator
 from yt.utilities.logger import ytLogger as mylog
 from yt.utilities.physical_constants import (
@@ -179,21 +180,38 @@ class RAMSESFieldInfo(FieldInfoContainer):
     def setup_particle_fields(self, ptype):
         super().setup_particle_fields(ptype)
 
+        def star_age_from_conformal_cosmo(field, data):
+            conformal_age = data[ptype, "conformal_birth_time"]
+            birth_time = convert_ramses_conformal_time_to_physical_time(
+                data.ds, conformal_age
+            )
+            return data.ds.current_time - birth_time
+
+        def star_age_from_physical_cosmo(field, data):
+            H0 = float(
+                data.ds.quan(data.ds.hubble_constant * 100, "km/s/Mpc").to("1/Gyr")
+            )
+            times = data[ptype, "conformal_birth_time"].value
+            time_tot = float(t_frw(data.ds, 0) * H0)
+            birth_time = (time_tot + times) / H0
+            t_out = float(data.ds.current_time.to("Gyr"))
+            return data.apply_units(t_out - birth_time, "Gyr")
+
         def star_age(field, data):
-            if data.ds.cosmological_simulation:
-                conformal_age = data[ptype, "conformal_birth_time"]
-                physical_age = convert_ramses_conformal_time_to_physical_age(
-                    data.ds, conformal_age
-                )
-                return data.ds.arr(physical_age, "code_time")
-            else:
-                formation_time = data[ptype, "particle_birth_time"]
-                return data.ds.current_time - formation_time
+            formation_time = data[ptype, "particle_birth_time"]
+            return data.ds.current_time - formation_time
+
+        if self.ds.cosmological_simulation and self.ds.use_conformal_time:
+            fun = star_age_from_conformal_cosmo
+        elif self.ds.cosmological_simulation:
+            fun = star_age_from_physical_cosmo
+        else:
+            fun = star_age
 
         self.add_field(
             (ptype, "star_age"),
             sampling_type="particle",
-            function=star_age,
+            function=fun,
             units=self.ds.unit_system["time"],
         )
 

yt/frontends/ramses/io.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Union
 
 import numpy as np
+from unyt import unyt_array
 
 from yt._maintenance.deprecation import issue_deprecation_warning
 from yt.frontends.ramses.definitions import VAR_DESC_RE, VERSION_RE
@@ -14,7 +15,6 @@
 )
 from yt.utilities.io_handler import BaseIOHandler
 from yt.utilities.logger import ytLogger as mylog
-from yt.utilities.physical_ratios import cm_per_km, cm_per_mpc
 
 if TYPE_CHECKING:
     import os
@@ -29,14 +29,14 @@ def convert_ramses_ages(ds, conformal_ages):
         stacklevel=3,
         since="4.0.3",
     )
-    return convert_ramses_conformal_time_to_physical_age(ds, conformal_ages)
+    return convert_ramses_conformal_time_to_physical_time(ds, conformal_ages)
 
 
-def convert_ramses_conformal_time_to_physical_age(
+def convert_ramses_conformal_time_to_physical_time(
     ds, conformal_time: np.ndarray
-) -> np.ndarray:
+) -> unyt_array:
     """
-    Convert conformal times (as defined in RAMSES) to physical age.
+    Convert conformal times (as defined in RAMSES) to physical times.
 
     Arguments
     ---------
@@ -50,32 +50,30 @@ def convert_ramses_conformal_time_to_physical_age(
     physical_age : np.ndarray
         The physical age in code units
     """
-    tf = ds.t_frw
-    dtau = ds.dtau
-    tauf = ds.tau_frw
-    h100 = ds.hubble_constant
-    nOver2 = ds.n_frw / 2
-    unit_t = ds.parameters["unit_t"]
-    tsim = ds.time_simu
-
-    t_scale = 1.0 / (h100 * 100 * cm_per_km / cm_per_mpc) / unit_t
-
-    # calculate index into lookup table (n_frw elements in
-    # lookup table)
-    dage = 1 + (10 * conformal_time / dtau)
-    dage = np.minimum(dage, nOver2 + (dage - nOver2) / 10.0)
-    iage = np.array(dage, dtype=np.int32)
-
-    # linearly interpolate physical times from tf and tauf lookup
-    # tables.
-    t = tf[iage] * (conformal_time - tauf[iage - 1]) / (tauf[iage] - tauf[iage - 1])
-    t = t + (
-        tf[iage - 1] * (conformal_time - tauf[iage]) / (tauf[iage - 1] - tauf[iage])
+    h0 = ds.hubble_constant
+    tau_bins = ds.tau_frw * h0
+    t_bins = ds.t_frw
+
+    min_time = 0
+    max_time = ds.current_time.to(t_bins.units)
+
+    return ds.arr(
+        np.clip(
+            np.interp(
+                conformal_time,
+                tau_bins,
+                t_bins.value,
+                right=max_time,
+                left=min_time,
+            ),
+            min_time,
+            max_time.value,
+        ),
+        t_bins.units,
     )
-    return (tsim - t) * t_scale
 
 
-def _ramses_particle_binary_file_handler(particle, subset, fields, count):
+def _ramses_particle_binary_file_handler(particle_handler, subset, fields, count):
     """General file handler for binary file, called by _read_particle_subset
 
     Parameters
@@ -91,10 +89,9 @@ def _ramses_particle_binary_file_handler(particle, subset, fields, count):
     """
     tr = {}
     ds = subset.domain.ds
-    current_time = ds.current_time.in_units("code_time").v
-    foffsets = particle.field_offsets
-    fname = particle.fname
-    data_types = particle.field_types
+    foffsets = particle_handler.field_offsets
+    fname = particle_handler.fname
+    data_types = particle_handler.field_types
     with FortranFile(fname) as fd:
         # We do *all* conversion into boxlen here.
         # This means that no other conversions need to be applied to convert
@@ -103,24 +100,23 @@ def _ramses_particle_binary_file_handler(particle, subset, fields, count):
             if count == 0:
                 tr[field] = np.empty(0, dtype=data_types[field])
                 continue
-            fd.seek(foffsets[field])
-            dt = data_types[field]
-            tr[field] = fd.read_vector(dt)
+            # Sentinel value: -1 means we don't have this field
+            if foffsets[field] == -1:
+                tr[field] = np.empty(count, dtype=data_types[field])
+            else:
+                fd.seek(foffsets[field])
+                dt = data_types[field]
+                tr[field] = fd.read_vector(dt)
             if field[1].startswith("particle_position"):
                 np.divide(tr[field], ds["boxlen"], tr[field])
-            if ds.cosmological_simulation and field[1] == "particle_birth_time":
-                conformal_time = tr[field]
-                physical_age = convert_ramses_conformal_time_to_physical_age(
-                    ds, conformal_time
-                )
-                tr[field] = current_time - physical_age
-                # arbitrarily set particles with zero conformal_age to zero
-                # particle_age. This corresponds to DM particles.
-                tr[field][conformal_time == 0] = 0
+
+            # Hand over to field handler for special cases, like particle_birth_times
+            particle_handler.handle_field(field, tr)
+
     return tr
 
 
-def _ramses_particle_csv_file_handler(particle, subset, fields, count):
+def _ramses_particle_csv_file_handler(particle_handler, subset, fields, count):
     """General file handler for csv file, called by _read_particle_subset
 
     Parameters
@@ -138,9 +134,8 @@ def _ramses_particle_csv_file_handler(particle, subset, fields, count):
 
     tr = {}
     ds = subset.domain.ds
-    current_time = ds.current_time.in_units("code_time").v
-    foffsets = particle.field_offsets
-    fname = particle.fname
+    foffsets = particle_handler.field_offsets
+    fname = particle_handler.fname
 
     list_field_ind = [
         (field, foffsets[field]) for field in sorted(fields, key=lambda a: foffsets[a])
@@ -159,15 +154,9 @@ def _ramses_particle_csv_file_handler(particle, subset, fields, count):
         tr[field] = dat[ind].to_numpy()
         if field[1].startswith("particle_position"):
             np.divide(tr[field], ds["boxlen"], tr[field])
-        if ds.cosmological_simulation and field[1] == "particle_birth_time":
-            conformal_time = tr[field]
-            physical_age = convert_ramses_conformal_time_to_physical_age(
-                ds, conformal_time
-            )
-            tr[field] = current_time - physical_age
-            # arbitrarily set particles with zero conformal_age to zero
-            # particle_age. This corresponds to DM particles.
-            tr[field][conformal_time == 0] = 0
+
+        particle_handler.handle_field(field, tr)
+
     return tr
 
 
@@ -294,23 +283,12 @@ def _read_particle_subset(self, subset, fields):
             ok = False
             for ph in subset.domain.particle_handlers:
                 if ph.ptype == ptype:
-                    foffsets = ph.field_offsets
-                    data_types = ph.field_types
                     ok = True
                     count = ph.local_particle_count
                     break
             if not ok:
                 raise YTFieldTypeNotFound(ptype)
 
-            cosmo = self.ds.cosmological_simulation
-            if (ptype, "particle_birth_time") in foffsets and cosmo:
-                foffsets[ptype, "conformal_birth_time"] = foffsets[
-                    ptype, "particle_birth_time"
-                ]
-                data_types[ptype, "conformal_birth_time"] = data_types[
-                    ptype, "particle_birth_time"
-                ]
-
             tr.update(ph.reader(subset, subs_fields, count))
 
         return tr