Skip to content

Commit

Permalink
Merge pull request #4706 from cphyc/read-part-birth-files
Browse files Browse the repository at this point in the history 
[ENH] Simplify handling of conformal birth times in RAMSES
  • Loading branch information
@chrishavlin
chrishavlin authored Oct 11, 2024
2 parents 8640daf + 2769152 commit fae33d6
Show file tree
Hide file tree
Showing 7 changed files with 238 additions and 183 deletions.
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand Down
59 changes: 20 additions & 39 deletions yt/frontends/ramses/data_structures.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -809,6 +804,7 @@ def __init__(
max_level=None,
max_level_convention=None,
default_species_fields=None,
use_conformal_time=None,
):
# Here we want to initiate a traceback, if the reader is not built.
if isinstance(fields, str):
Expand Down Expand Up @@ -879,6 +875,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

@staticmethod
Expand Down Expand Up @@ -1055,15 +1061,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"]
Expand All @@ -1074,38 +1079,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
Expand Down
40 changes: 29 additions & 11 deletions yt/frontends/ramses/fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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"],
)

Expand Down
114 changes: 46 additions & 68 deletions yt/frontends/ramses/io.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
---------
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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])
Expand All @@ -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


Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit fae33d6

Please sign in to comment.