diff --git a/csiborgtools/units/__init__.py b/csiborgtools/units/__init__.py index 5e00d8f..c45fa82 100644 --- a/csiborgtools/units/__init__.py +++ b/csiborgtools/units/__init__.py @@ -15,4 +15,5 @@ from .transforms import (cartesian_to_radec, convert_mass_cols, # noqa convert_position_cols) # noqa -from .box_units import BoxUnits # noqa +from .box_units import (BoxUnits, convert_from_boxunits) # noqa + diff --git a/csiborgtools/units/box_units.py b/csiborgtools/units/box_units.py index de9432d..8904bd2 100644 --- a/csiborgtools/units/box_units.py +++ b/csiborgtools/units/box_units.py @@ -21,6 +21,14 @@ from astropy.cosmology import LambdaCDM from astropy import (constants, units) from ..io import read_info +# Map of unit conversions +CONV_NAME = { + "length": ["peak_x", "peak_y", "peak_z", "Rs", "rmin", "rmax", "r200", + "r500"], + "mass": ["mass_cl", "totpartmass", "m200", "m500", "mass_mmain"], + "density": ["rho0"] + } + class BoxUnits: r""" @@ -173,6 +181,23 @@ class BoxUnits: """ return length / (self._unit_l / units.kpc.to(units.cm) / self._aexp) + def box2mpc(self, length): + r""" + Convert length from box units to :math:`\mathrm{cMpc}` (with + :math:`h=0.705`). + + Parameters + ---------- + length : float + Length in box units. + + Returns + ------- + length : foat + Length in :math:`\mathrm{ckpc}` + """ + return self.box2kpc(length) * 1e-3 + def solarmass2box(self, mass): r""" Convert mass from :math:`M_\odot` (with :math:`h=0.705`) to box units. @@ -242,3 +267,60 @@ class BoxUnits: """ return (density / self._unit_d * self._Msuncgs / (units.pc.to(units.cm))**3) + + +def convert_from_boxunits(data, names, boxunits): + r""" + Convert columns named `names` in array `data` from box units to physical + units, such that + - length -> Mpc, + - mass -> solar mass, + - density -> solar mass per cubed kpc. + Any other conversions are currently not implemented. Note that the array + is passed by reference and directly modified, even though it is also + explicitly returned. + + Parameters + ---------- + data : structured array + Input array. + names : list of str + Columns to be converted. + boxunits : `BoxUnits` + Box units class of the simulation and snapshot. + + Returns + ------- + data : structured array + Input array with converted columns. + """ + if not isinstance(boxunits, BoxUnits): + raise TypeError("`boxunits` must be of type `{}`. Currently `{}`." + .format(BoxUnits, type(boxunits))) + names = [names] if isinstance(names, str) else names + + # Shortcut for the transform functions + transforms = { + "length": boxunits.box2mpc, + "mass": boxunits.box2solarmass, + "density": boxunits.box2dens + } + + for name in names: + # Check that the name is even in the array + if name not in data.dtype.names: + raise ValueError("Name `{}` is not in `data` array.".format(name)) + + # Convert + found = False + for unittype, suppnames in CONV_NAME.items(): + if name in suppnames: + data[name] = transforms[unittype](data[name]) + found = True + continue + # If nothing found + if not found: + raise NotImplementedError( + "Conversion of `{}` is not defined.".format(name)) + + return data