From 65ed224db7264d98047813f45d80b9660260a40d Mon Sep 17 00:00:00 2001 From: Lenoble Romain Date: Fri, 28 Apr 2023 14:09:36 +0200 Subject: [PATCH] ENH: support for sinks in csv format for RAMSES frontend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Corentin Cadiou Co-authored-by: Clément Robert --- tests/pytest_mpl_baseline | 2 +- yt/frontends/ramses/fields.py | 17 +++ yt/frontends/ramses/io.py | 142 ++++++++++++++++++++--- yt/frontends/ramses/particle_handlers.py | 46 +++++++- yt/utilities/on_demand_imports.py | 6 + 5 files changed, 192 insertions(+), 21 deletions(-) diff --git a/tests/pytest_mpl_baseline b/tests/pytest_mpl_baseline index b02c4ed360..301246ebb6 160000 --- a/tests/pytest_mpl_baseline +++ b/tests/pytest_mpl_baseline @@ -1 +1 @@ -Subproject commit b02c4ed3603969308f92b5022c2ee075d7f68db0 +Subproject commit 301246ebb68f0ad15c7481db2ecb08987c57fdd2 diff --git a/yt/frontends/ramses/fields.py b/yt/frontends/ramses/fields.py index d81745e75e..39c0321fef 100644 --- a/yt/frontends/ramses/fields.py +++ b/yt/frontends/ramses/fields.py @@ -131,6 +131,23 @@ class RAMSESFieldInfo(FieldInfoContainer): ("particle_metallicity", ("", [], None)), ("particle_family", ("", [], None)), ("particle_tag", ("", [], None)), + # sink field parameters + ("particle_mass", ("code_mass", [], None)), + ("particle_angular_momentum_x", (ang_mom_units, [], None)), + ("particle_angular_momentum_y", (ang_mom_units, [], None)), + ("particle_angular_momentum_z", (ang_mom_units, [], None)), + ("particle_formation_time", ("code_time", [], None)), + ("particle_accretion_Rate", ("code_mass/code_time", [], None)), + ("particle_delta_mass", ("code_mass", [], None)), + ("particle_rho_gas", (rho_units, [], None)), + ("particle_cs**2", (vel_units, [], None)), + ("particle_etherm", (ener_units, [], None)), + ("particle_velocity_x_gas", (vel_units, [], None)), + ("particle_velocity_y_gas", (vel_units, [], None)), + ("particle_velocity_z_gas", (vel_units, [], None)), + ("particle_mass_bh", ("code_mass", [], None)), + ("particle_level", ("", [], None)), + ("particle_radius_star", ("code_length", [], None)), ) known_sink_fields: KnownFieldsT = ( diff --git a/yt/frontends/ramses/io.py b/yt/frontends/ramses/io.py index b571a37594..d217ae337b 100644 --- a/yt/frontends/ramses/io.py +++ b/yt/frontends/ramses/io.py @@ -72,17 +72,13 @@ def convert_ramses_conformal_time_to_physical_age( return (tsim - t) * t_scale -def _ramses_particle_file_handler(fname, foffsets, data_types, subset, fields, count): - """General file handler, called by _read_particle_subset +def _ramses_particle_binary_file_handler(particle, subset, fields, count): + """General file handler for binary file, called by _read_particle_subset Parameters ---------- - fname : string - filename to read from - foffsets: dict - Offsets in file of the fields - data_types: dict - Data type of the fields + particle : ``ParticleFileHandler`` + the particle class we want to read subset: ``RAMSESDomainSubset`` A RAMSES domain subset object fields: list of tuple @@ -93,6 +89,9 @@ def _ramses_particle_file_handler(fname, foffsets, data_types, subset, fields, c 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 with FortranFile(fname) as fd: # We do *all* conversion into boxlen here. # This means that no other conversions need to be applied to convert @@ -118,6 +117,50 @@ def _ramses_particle_file_handler(fname, foffsets, data_types, subset, fields, c return tr +def _ramses_particle_csv_file_handler(particle, subset, fields, count): + """General file handler for csv file, called by _read_particle_subset + + Parameters + ---------- + particle : ``ParticleFileHandler`` + the particle class we want to read + subset: ``RAMSESDomainSubset`` + A RAMSES domain subset object + fields: list of tuple + The fields to read + count: integer + The number of elements to count + """ + tr = {} + ds = subset.domain.ds + current_time = ds.current_time.in_units("code_time").v + foffsets = particle.field_offsets + fname = particle.fname + + from yt.utilities.on_demand_imports import _pandas as pd + + for ind, field in enumerate(sorted(fields, key=lambda a: foffsets[a])): + ind = foffsets[field] + dat = pd.read_csv( + fname, delimiter=",", usecols=[ind], skiprows=2, header=None + ) # read only selected fields + tr[field] = np.array(dat[ind].to_list()) + + 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 + + return tr + + class IOHandlerRAMSES(BaseIOHandler): _dataset_type = "ramses" @@ -189,12 +232,29 @@ def _read_particle_fields(self, chunks, ptf, selector): fields = [ (ptype, fname) for ptype, field_list in ptf.items() for fname in field_list ] + for ptype, field_list in sorted(ptf.items()): for ax in "xyz": if pn % ax not in field_list: fields.append((ptype, pn % ax)) - for chunk in chunks: - for subset in chunk.objs: + + if ptype != "sink_csv": + for chunk in chunks: + for subset in chunk.objs: + rv = self._read_particle_subset(subset, fields) + for ptype, field_list in sorted(ptf.items()): + x, y, z = ( + np.asarray(rv[ptype, pn % ax], "=f8") for ax in "xyz" + ) + mask = selector.select_points(x, y, z, 0.0) + if mask is None: + mask = [] + for field in field_list: + data = np.asarray(rv.pop((ptype, field))[mask], "=f8") + yield (ptype, field), data + + else: + subset = chunks[0].objs[0] rv = self._read_particle_subset(subset, fields) for ptype, field_list in sorted(ptf.items()): x, y, z = (np.asarray(rv[ptype, pn % ax], "=f8") for ax in "xyz") @@ -217,7 +277,6 @@ def _read_particle_subset(self, subset, fields): ok = False for ph in subset.domain.particle_handlers: if ph.ptype == ptype: - fname = ph.fname foffsets = ph.field_offsets data_types = ph.field_types ok = True @@ -235,16 +294,12 @@ def _read_particle_subset(self, subset, fields): ptype, "particle_birth_time" ] - tr.update( - _ramses_particle_file_handler( - fname, foffsets, data_types, subset, subs_fields, count=count - ) - ) + tr.update(ph.reader(subset, subs_fields, count)) return tr -def _read_part_file_descriptor(fname: Union[str, "os.PathLike[str]"]): +def _read_part_binary_file_descriptor(fname: Union[str, "os.PathLike[str]"]): """ Read a file descriptor and returns the array of the fields found. """ @@ -300,6 +355,59 @@ def _read_part_file_descriptor(fname: Union[str, "os.PathLike[str]"]): return fields +def _read_part_csv_file_descriptor(fname: Union[str, "os.PathLike[str]"]): + """ + Read the file from the csv sink particles output. + """ + from yt.utilities.on_demand_imports import _pandas as pd + + # Fields name from the default csv RAMSES sink algorithm in the yt default convention + mapping_list = [ + (" # id", "particle_identifier"), + ("msink", "particle_mass"), + ("x", "particle_position_x"), + ("y", "particle_position_y"), + ("z", "particle_position_z"), + ("vx", "particle_velocity_x"), + ("vy", "particle_velocity_y"), + ("vz", "particle_velocity_z"), + ("lx", "particle_angular_momentum_x"), + ("ly", "particle_angular_momentum_y"), + ("lz", "particle_angular_momentum_z"), + ("tform", "particle_formation_time"), + ("acc_Rate", "particle_accretion_Rate"), + ("del_mass", "particle_delta_mass"), + ("rho_gas", "particle_rho_gas"), + ("cs**2", "particle_sound_speed"), + ("etherm", "particle_etherm"), + ("vx_gas", "particle_velocity_x_gas"), + ("vy_gas", "particle_velocity_y_gas"), + ("vz_gas", "particle_velocity_z_gas"), + ("mbh", "particle_mass_bh"), + ("level", "particle_level"), + ("rsink_star", "particle_radius_star"), + ] + + # Convert to dictionary + mapping = {k: v for k, v in mapping_list} + + dat = pd.read_csv( + fname, delimiter="," + ) # read the all file to get the number of particle + fields = [] + local_particle_count = len(dat) + + for varname in dat.columns: + if varname in mapping: + varname = mapping[varname] + else: + varname = f"particle_{varname}" + + fields.append(varname) + + return fields, local_particle_count + + def _read_fluid_file_descriptor(fname: Union[str, "os.PathLike[str]"], *, prefix: str): """ Read a file descriptor and returns the array of the fields found. diff --git a/yt/frontends/ramses/particle_handlers.py b/yt/frontends/ramses/particle_handlers.py index f6215ecc6d..566bfc2769 100644 --- a/yt/frontends/ramses/particle_handlers.py +++ b/yt/frontends/ramses/particle_handlers.py @@ -8,7 +8,12 @@ from yt.utilities.cython_fortran_utils import FortranFile from .field_handlers import HandlerMixin -from .io import _read_part_file_descriptor +from .io import ( + _ramses_particle_binary_file_handler, + _ramses_particle_csv_file_handler, + _read_part_binary_file_descriptor, + _read_part_csv_file_descriptor, +) PARTICLE_HANDLERS: Set[Type["ParticleFileHandler"]] = set() @@ -101,6 +106,7 @@ class DefaultParticleFileHandler(ParticleFileHandler): fname = "part_{iout:05d}.out{icpu:05d}" file_descriptor = "part_file_descriptor.txt" config_field = "ramses-particles" + reader = _ramses_particle_binary_file_handler attrs = ( ("ncpu", 1, "i"), @@ -144,7 +150,7 @@ def read_header(self): extra_particle_fields = self.ds._extra_particle_fields if self.has_descriptor: - particle_fields = _read_part_file_descriptor(self.file_descriptor) + particle_fields = _read_part_binary_file_descriptor(self.file_descriptor) else: particle_fields = list(self.known_fields) @@ -202,6 +208,7 @@ class SinkParticleFileHandler(ParticleFileHandler): fname = "sink_{iout:05d}.out{icpu:05d}" file_descriptor = "sink_file_descriptor.txt" config_field = "ramses-sink-particles" + reader = _ramses_particle_binary_file_handler attrs = (("nsink", 1, "i"), ("nindsink", 1, "i")) @@ -258,7 +265,7 @@ def read_header(self): # Read the fields + add the sink properties if self.has_descriptor: - fields = _read_part_file_descriptor(self.file_descriptor) + fields = _read_part_binary_file_descriptor(self.file_descriptor) else: fields = list(self.known_fields) @@ -282,3 +289,36 @@ def read_header(self): self.field_offsets = field_offsets self.field_types = _pfields fd.close() + + +class SinkParticleFileHandlerCsv(ParticleFileHandler): + """Handle sink files from a csv file, the format from the sink particule in ramses""" + + ptype = "sink_csv" + fname = "sink_{iout:05d}.csv" + file_descriptor = None + config_field = "ramses-sink-particles" + reader = _ramses_particle_csv_file_handler + attrs = (("nsink", 1, "i"), ("nindsink", 1, "i")) + second_read = False + + def read_header(self): + if not self.exists: + self.field_offsets = {} + self.field_types = {} + self.local_particle_count = 0 + return + field_offsets = {} + _pfields = {} + + fields, self.local_particle_count = _read_part_csv_file_descriptor(self.fname) + + for ind, field in enumerate(fields): + field_offsets[self.ptype, field] = ind + _pfields[self.ptype, field] = "d" + + self.field_offsets = field_offsets + self.field_types = _pfields + + def sec_read(self): + pass diff --git a/yt/utilities/on_demand_imports.py b/yt/utilities/on_demand_imports.py index 07021988cb..9a17ba6f13 100644 --- a/yt/utilities/on_demand_imports.py +++ b/yt/utilities/on_demand_imports.py @@ -485,6 +485,12 @@ def concat(self): from pandas import concat return concat + + @safe_import + def read_csv(self): + from pandas import read_csv + + return read_csv _pandas = pandas_imports()