Add a simulation unit system class (#392)

* add a simulation unit system class

* hah, oops for removing the greek letters that I thought were vestigial! I'll move them in a different pr

* Added a changelog entry

* Added examples to the docstring

* fix error with length and velocity unit

* add sim units test and reformat

CHANGES.rst

@@ -4,6 +4,9 @@
 New Features
 ------------
 
+- Added a new ``SimulationUnitSystem`` class for handling unit systems in
+  simulations, especially for N-body simulations.
+
 Bug fixes
 ---------
 

gala/tests/test_units.py

@@ -2,17 +2,15 @@
     Test the unit system.
 """
 
-# Standard library
+import itertools
 import pickle
 
-# Third party
 import astropy.units as u
-from astropy.constants import G, c
 import numpy as np
 import pytest
+from astropy.constants import G, c
 
-# This package
-from ..units import UnitSystem, DimensionlessUnitSystem
+from ..units import DimensionlessUnitSystem, SimulationUnitSystem, UnitSystem
 
 
 def test_create():
@@ -37,35 +35,58 @@ def test_create():
 
 def test_constants():
     usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun)
-    assert np.allclose(usys.get_constant('G'), G.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value)
-    assert np.allclose(usys.get_constant('c'), c.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value)
+    assert np.allclose(
+        usys.get_constant("G"), G.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value
+    )
+    assert np.allclose(
+        usys.get_constant("c"), c.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value
+    )
 
 
 def test_decompose():
-    usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.km/u.s)
-    q = 15.*u.km/u.s
-    assert q.decompose(usys).unit == u.kpc/u.Myr  # uses the core units
-    assert usys.decompose(q).unit == u.km/u.s
+    usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.km / u.s)
+    q = 15.0 * u.km / u.s
+    assert q.decompose(usys).unit == u.kpc / u.Myr  # uses the core units
+    assert usys.decompose(q).unit == u.km / u.s
 
 
 def test_dimensionless():
     usys = DimensionlessUnitSystem()
-    assert usys['dimensionless'] == u.one
-    assert usys['length'] == u.one
+    assert usys["dimensionless"] == u.one
+    assert usys["length"] == u.one
 
     with pytest.raises(ValueError):
-        (15*u.kpc).decompose(usys)
+        (15 * u.kpc).decompose(usys)
 
     with pytest.raises(ValueError):
-        usys.decompose(15*u.kpc)
+        usys.decompose(15 * u.kpc)
+
+
+@pytest.mark.parametrize(
+    "nu1, nu2",
+    itertools.combinations(
+        {
+            "length": 15 * u.kpc,
+            "mass": 1e6 * u.Msun,
+            "time": 5e2 * u.Myr,
+            "velocity": 150 * u.km / u.s,
+        }.items(),
+        2,
+    ),
+)
+def test_simulation(nu1, nu2):
+    print(nu1, nu2)
+    name1, unit1 = nu1
+    name2, unit2 = nu2
+    SimulationUnitSystem(**{name1: unit1, name2: unit2})
 
 
 def test_compare():
-    usys1 = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.mas/u.yr)
-    usys1_clone = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.mas/u.yr)
+    usys1 = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.mas / u.yr)
+    usys1_clone = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.mas / u.yr)
 
-    usys2 = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.kiloarcsecond/u.yr)
-    usys3 = UnitSystem(u.kpc, u.Myr, u.radian, u.kg, u.mas/u.yr)
+    usys2 = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.kiloarcsecond / u.yr)
+    usys3 = UnitSystem(u.kpc, u.Myr, u.radian, u.kg, u.mas / u.yr)
 
     assert usys1 == usys1_clone
     assert usys1_clone == usys1
@@ -80,19 +101,15 @@ def test_compare():
 def test_pickle(tmpdir):
     usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun)
 
-    with open(tmpdir / 'test.pkl', 'wb') as f:
+    with open(tmpdir / "test.pkl", "wb") as f:
         pickle.dump(usys, f)
 
-    with open(tmpdir / 'test.pkl', 'rb') as f:
+    with open(tmpdir / "test.pkl", "rb") as f:
         usys2 = pickle.load(f)
 
 
 def test_quantity_units():
-    usys = UnitSystem(
-        5 * u.kpc,
-        50 * u.Myr,
-        1e5 * u.Msun,
-        u.rad)
-
-    assert np.isclose((8*u.Myr).decompose(usys).value, 8/50)
-    usys.get_constant('G')
+    usys = UnitSystem(5 * u.kpc, 50 * u.Myr, 1e5 * u.Msun, u.rad)
+
+    assert np.isclose((8 * u.Myr).decompose(usys).value, 8 / 50)
+    usys.get_constant("G")

gala/units.py

@@ -1,15 +1,15 @@
 __all__ = [
     "UnitSystem",
     "DimensionlessUnitSystem",
+    "SimulationUnitSystem",
     "galactic",
     "dimensionless",
     "solarsystem",
 ]
 
 import astropy.constants as const
-
-# Third-party
 import astropy.units as u
+import numpy as np
 from astropy.units.physical import _physical_unit_mapping
 
 _greek_letters = [
@@ -112,7 +112,7 @@ def __init__(self, units, *args):
                 new_unit = u.def_unit(f"{q!s}", q)
                 unit = new_unit
 
-            typ = unit.physical_type
+            typ = unit.decompose().physical_type
             if typ in self._registry:
                 raise ValueError(f"Multiple units passed in with type '{typ}'")
             self._registry[typ] = unit
@@ -262,6 +262,73 @@ def get_constant(self, name):
         raise ValueError("Cannot get constant in dimensionless units!")
 
 
+l_pt = u.get_physical_type("length")
+m_pt = u.get_physical_type("mass")
+t_pt = u.get_physical_type("time")
+v_pt = u.get_physical_type("velocity")
+a_pt = u.get_physical_type("angle")
+
+
+class SimulationUnitSystem(UnitSystem):
+    def __init__(
+        self,
+        length: u.Unit | u.Quantity[l_pt] = None,
+        mass: u.Unit | u.Quantity[m_pt] = None,
+        time: u.Unit | u.Quantity[t_pt] = None,
+        velocity: u.Unit | u.Quantity[v_pt] = None,
+        G: float | u.Quantity = 1.0,
+        angle: u.Unit | u.Quantity[a_pt] = u.radian,
+    ):
+        """
+        Represents a system of units for a (dynamical) simulation.
+
+        A common assumption is that G=1. If this is the case, then you only have to
+        specify two of the three fundamental unit types (length, mass, time) and the
+        rest will be derived from these. You may also optionally specify a velocity with
+        one of the base unit types (length, mass, time).
+
+        Examples
+        --------
+        To convert simulation positions and velocities to physical units, you can
+        use this unit system::
+
+            usys = SimulationUnitSystem(length=10 * u.kpc, time=50 * u.Myr)
+            (sim_pos * usys["length"]).to(u.kpc)
+            (sim_vel * usys["velocity"]).to(u.km/u.s)
+
+        Or, to convert positions and velocities from physical units to simulation
+        units::
+
+            (100 * u.kpc).to(usys["length"])
+
+        """
+        G = G * const.G.unit
+
+        if length is not None and mass is not None:
+            time = 1 / np.sqrt(G * mass / length**3)
+        elif length is not None and time is not None:
+            mass = 1 / G * length**3 / time**2
+        elif length is not None and velocity is not None:
+            time = length / velocity
+            mass = velocity**2 / G * length
+        elif mass is not None and time is not None:
+            length = np.cbrt(G * mass * time**2)
+        elif mass is not None and velocity is not None:
+            length = G * mass / velocity**2
+            time = length / velocity
+        elif time is not None and velocity is not None:
+            mass = 1 / G * velocity**3 * time
+            length = G * mass / velocity**2
+        else:
+            msg = (
+                "You must specify at least two of the three fundamental unit types "
+                "(length, mass, time) or a velocity unit."
+            )
+            raise ValueError(msg)
+
+        super().__init__(length, mass, time, angle)
+
+
 # define galactic unit system
 galactic = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, u.km / u.s)