Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[ENH] Simplify handling of conformal birth times in RAMSES #4706

Merged
merged 13 commits into from
Oct 11, 2024
Merged

[ENH] Simplify handling of conformal birth times in RAMSES #4706

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
1ec01a1
Simplify handling of conformal birth times in RAMSES
cphyc May 17, 2024
a4e0068
Use analytical integration rather than approximate Euler method
cphyc Nov 2, 2023
0109812
Also work for non-conformal simulations
cphyc Nov 2, 2023
1a71233
Cleanup cosmology_time
cphyc Nov 2, 2023
68cb79e
Convert to code time
May 17, 2024
68fd745
Interp to min/max time when running out of bounds
May 17, 2024
7494050
Fix linting
May 17, 2024
1e6055e
Require scipy for RAMSES only
May 17, 2024
08ff40d
Clip min/max times to 0, tsim
cphyc May 23, 2024
c37f115
Strip unit before clipping/interpolating
cphyc May 24, 2024
9b6fbf6
Rename functions for better good
cphyc May 24, 2024
2e76412
Apply suggestions from code review
cphyc Oct 3, 2024
2769152
Remove commented-out lines
cphyc Oct 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading