__all__ = [
"UnitSystem",
"DimensionlessUnitSystem",
"galactic",
"dimensionless",
"solarsystem",
]
import astropy.constants as const
# Third-party
import astropy.units as u
from astropy.units.physical import _physical_unit_mapping
_greek_letters = [
"alpha",
"beta",
"gamma",
"delta",
"epsilon",
"zeta",
"eta",
"theta",
"iota",
"kappa",
"lambda",
"mu",
"nu",
"xi",
"pi",
"o",
"rho",
"sigma",
"tau",
"upsilon",
"phi",
"chi",
"psi",
"omega",
]
[docs]class UnitSystem:
_required_physical_types = [
u.get_physical_type("length"),
u.get_physical_type("time"),
u.get_physical_type("mass"),
u.get_physical_type("angle"),
]
def __init__(self, units, *args):
"""
Represents a system of units.
At minimum, this consists of a set of length, time, mass, and angle
units, but may also contain preferred representations for composite
units. For example, the base unit system could be ``{kpc, Myr, Msun,
radian}``, but you can also specify a preferred velocity unit, such as
``km/s``.
This class behaves like a dictionary with keys set by physical types. If
a unit for a particular physical type is not specified on creation, a
composite unit will be created with the base units. See the examples
below for some demonstrations.
Parameters
----------
*units
The units that define the unit system. At minimum, this must
contain length, time, mass, and angle units.
Examples
--------
If only base units are specified, any physical type specified as a key
to this object will be composed out of the base units::
>>> usys = UnitSystem(u.m, u.s, u.kg, u.radian)
>>> usys['energy'] # doctest: +SKIP
Unit("kg m2 / s2")
However, custom representations for composite units can also be
specified when initializing::
>>> usys = UnitSystem(u.m, u.s, u.kg, u.radian, u.erg)
>>> usys['energy']
Unit("erg")
This is useful for Galactic dynamics where lengths and times are usually
given in terms of ``kpc`` and ``Myr``, but velocities are given in
``km/s``::
>>> usys = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, u.km/u.s)
>>> usys['velocity']
Unit("km / s")
"""
self._core_units = []
if isinstance(units, UnitSystem):
self._registry = units._registry.copy()
self._core_units = units._core_units
return
if len(args) > 0:
units = (units,) + tuple(args)
self._registry = dict()
for unit in units:
if not isinstance(unit, u.UnitBase): # hopefully a quantity
q = unit
new_unit = u.def_unit(f"{q!s}", q)
unit = new_unit
typ = unit.physical_type
if typ in self._registry:
raise ValueError(f"Multiple units passed in with type '{typ}'")
self._registry[typ] = unit
for phys_type in self._required_physical_types:
if phys_type not in self._registry:
raise ValueError(
"You must specify a unit for the physical type" f"'{phys_type}'"
)
self._core_units.append(self._registry[phys_type])
def __getitem__(self, key):
key = u.get_physical_type(key)
if key in self._registry:
return self._registry[key]
else:
unit = None
for k, v in _physical_unit_mapping.items():
if v == key:
unit = u.Unit(" ".join([f"{x}**{y}" for x, y in k]))
break
if unit is None:
raise ValueError(
f"Physical type '{key}' doesn't exist in unit " "registry."
)
unit = unit.decompose(self._core_units)
unit._scale = 1.0
return unit
def __len__(self):
return len(self._core_units)
def __iter__(self):
for uu in self._core_units:
yield uu
def __str__(self):
core_units = ", ".join([str(uu) for uu in self._core_units])
return f"UnitSystem ({core_units})"
def __repr__(self):
return f"<{self.__str__()}>"
def __eq__(self, other):
for k in self._registry:
if not self[k] == other[k]:
return False
for k in other._registry:
if not self[k] == other[k]:
return False
return True
def __ne__(self, other):
return not self.__eq__(other)
[docs] def to_dict(self):
"""
Return a dictionary representation of the unit system with keys
set by the physical types and values set by the unit objects.
"""
return self._registry.copy()
[docs] def decompose(self, q):
"""
A thin wrapper around :meth:`astropy.units.Quantity.decompose` that
knows how to handle Quantities with physical types with non-default
representations.
Parameters
----------
q : :class:`~astropy.units.Quantity`
An instance of an astropy Quantity object.
Returns
-------
q : :class:`~astropy.units.Quantity`
A new quantity, decomposed to represented in this unit system.
"""
try:
ptype = q.unit.physical_type
except AttributeError:
raise TypeError(
"Object must be an astropy.units.Quantity, not "
f"a '{q.__class__.__name__}'."
)
if ptype in self._registry:
return q.to(self._registry[ptype])
else:
return q.decompose(self)
[docs] def get_constant(self, name):
"""
Retrieve a constant with specified name in this unit system.
Parameters
----------
name : str
The name of the constant, e.g., G.
Returns
-------
const : float
The value of the constant represented in this unit system.
Examples
--------
>>> usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun)
>>> usys.get_constant('c') # doctest: +SKIP
306.6013937855506
"""
try:
c = getattr(const, name)
except AttributeError:
raise ValueError(
f"Constant name '{name}' doesn't exist in " "astropy.constants"
)
return c.decompose(self._core_units).value
[docs]class DimensionlessUnitSystem(UnitSystem):
_required_physical_types = []
def __init__(self):
self._core_units = [u.one]
self._registry = {"dimensionless": u.one}
def __getitem__(self, key):
return u.one
def __str__(self):
return "UnitSystem (dimensionless)"
[docs] def to_dict(self):
raise ValueError("Cannot represent dimensionless unit system as dict!")
[docs] def get_constant(self, name):
raise ValueError("Cannot get constant in dimensionless units!")
# define galactic unit system
galactic = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, u.km / u.s)
# solar system units
solarsystem = UnitSystem(u.au, u.M_sun, u.yr, u.radian)
# dimensionless
dimensionless = DimensionlessUnitSystem()