Source code for gravhopper.unitconverter

"""Converts between Astropy-style and Pynbody-style units, which are not
directly compatible."""

from astropy import units as u, constants as const
from fractions import Fraction
import pynbody.units as pynu

__all__=['pynbody_to_astropy','astropy_to_pynbody']

# Direct mapping of astropy units for all pynbody units
unitmapper_pyn2ap = {
    'm':u.m,
    's':u.s,
    'kg':u.kg,
    'K':u.K,
    'rad':u.radian,
    'yr':u.yr,
    'kyr':u.kyr,
    'Myr':u.Myr,
    'Gyr':u.Gyr,
    'Hz':u.Hz,
    'kHz':u.kHz,
    'MHz':u.MHz,
    'GHz':u.GHz,
    'THz':u.THz,
    'angst':u.AA,
    'cm':u.cm,
    'mm':u.mm,
    'nm':u.nm,
    'km':u.km,
    'au':u.au,
    'pc':u.pc,
    'kpc':u.kpc,
    'Mpc':u.Mpc,
    'Gpc':u.Gpc,
    'sr':u.sr,
    'deg':u.deg,
    'arcmin':u.arcmin,
    'arcsec':u.arcsec,
    'Msol':u.Msun,
    'g':u.g,
    'm_p':const.m_p,
    'm_e':const.m_e,
    'N':u.N,
    'dyn':u.dyn,
    'J':u.J,
    'erg':u.erg,
    'eV':u.eV,
    'keV':u.keV,
    'MeV':u.MeV,
    'W':u.W,
    'Jy':u.Jy,
    'Pa':u.Pa,
    'k':const.k_B,
    'c':const.c,
    'G':const.G,
    'hP':const.h
}
# Invert mapping from astropy units to pynbody base equivalents
unitmapper_ap2pyn = {}
for k, v in unitmapper_pyn2ap.items():
    # If it's a unit, use the Astropy string repr, otherwise use the Pynbody symbol too.
    if isinstance(v, u.UnitBase):
        unitmapper_ap2pyn[str(v)] = k
    else:
        unitmapper_ap2pyn[k] = k
    

[docs] def pynbody_to_astropy(pynunit): """Return an astropy unit equivalent to the pynbody unit input. Parameters ---------- pynunit : pynbody.units.Unit Pynbody unit to convert Returns ------- apunit : astropy.unit.Unit Equivalent astropy unit. If the input unit is a composite, the function will attempt to compose it in the same way, but sometimes that is not possible; however, it will always be equivalent numerically. Raises ------ pynbody.UnitsException If it cannot successfully convert the unit """ # See also # -------- # `astropy_to_pynbody` : Converts astropy units to pynbody units. # :meth:`~gravhopper.IC.from_pyn_snap` : Converts a pynbody SimSnap to Gravhopper ICs # """ pynunit_str = str(pynunit) # Does it have a direct astropy equivalent? if pynunit_str in unitmapper_pyn2ap: return unitmapper_pyn2ap[pynunit_str] # Must be a composite. if not isinstance(pynunit, pynu.CompositeUnit): raise pynu.UnitsException(f"Pynbody unit {pynunit_str} has no astropy equivalent and is not composite") # Decompose it astropy_equivalent = pynunit._scale for base, power in zip(pynunit._bases, pynunit._powers): # Each base had better have an astropy equivalent base_str = str(base) if base_str in unitmapper_pyn2ap: astropy_base = unitmapper_pyn2ap[base_str] else: raise pynu.UnitsException(f"Pynbody unit {base_str} has no astropy equivalent") astropy_equivalent *= astropy_base**power return u.Unit(astropy_equivalent)
[docs] def astropy_to_pynbody(apunit): """Return a pynbody unit equivalent to the astropy unit input. Parameters ---------- apunit : astropy.unit.Unit Astropy unit to convert Returns ------- pynunit : pynbody.units.Unit Equivalent Pynbody unit. If the input unit is a composite, the function will attempt to compose it in the same way, but sometimes that is not possible; however, it will always be equivalent numerically. Raises ------ pynbody.UnitsException If it cannot successfully convert the unit """ # See also # -------- # `pynbody_to_astropy` : Converts astropy units to pynbody units. # :meth:`~gravhopper.Simulation.pyn_snap` : Converts a Simulation snapshot to a Pynbody SimSnap # """ # Emergency base units for each dimension. defaults = (u.kpc, u.s, u.Msun, u.K, u.radian, u.sr) # Does it have a direct pynbody equivalent? apunit_str = str(apunit) if apunit_str in unitmapper_ap2pyn: return pynu.Unit(unitmapper_ap2pyn[apunit_str]) # If it's a quantity, split into scale and unit if isinstance(apunit, u.Quantity): scale = apunit.value decomp_unit = apunit.unit elif isinstance(apunit, (u.Unit, u.CompositeUnit)): scale = apunit.scale decomp_unit = apunit else: raise pynu.UnitsException(f"Error decomposing Astropy unit {apunit}") # Construct a CompositeUnit pynbody_equivalent = pynu.CompositeUnit(scale, [], []) # Go through each decomposed piece for base, power in zip(decomp_unit.bases, decomp_unit.powers): # See if each base has a pynbody equivalent base_str = str(base) if base_str in unitmapper_ap2pyn: compose_pynbody_unit(pynbody_equivalent, 1, [base], [power]) elif isinstance(base, u.PrefixUnit): # If it's a prefix unit, decompose it and try again prefix_decomp = base.decompose() compose_pynbody_unit(pynbody_equivalent, prefix_decomp.scale, prefix_decomp.bases, prefix_decomp.powers) else: # Decompose into emergency bases default_decomp = base.decompose(defaults) compose_pynbody_unit(pynbody_equivalent, default_decomp.scale, default_decomp.bases, default_decomp.powers) return pynbody_equivalent
def compose_pynbody_unit(unit_so_far, scale, bases, powers): """Add bases and powers and scale to existing pynbody composite unit.""" unit_so_far._scale *= scale for base, power in zip(bases, powers): # See if each base has a pynbody equivalent base_str = str(base) if base_str in unitmapper_ap2pyn: pyn_base = pynu.Unit(unitmapper_ap2pyn[base_str]) unit_so_far._bases.append(pyn_base) if isinstance(power, float): fracpower = Fraction(power) else: fracpower = power unit_so_far._powers.append(fracpower) else: raise pynu.UnitsException(f"Error decomposing Astropy unit {base_str}")