feat: move units to unxt (GalacticDynamics#224)

* feat: move units to unxt
* fix: rebase and smoke
* fix: digits

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

pyproject.toml

@@ -25,7 +25,7 @@
     "Typing :: Typed",
   ]
   dependencies = [
-    "astropy >= 6",
+    "astropy >= 6.0",
     "beartype",
     "coordinax @ git+https://github.com/GalacticDynamics/coordinax.git",
     "diffrax >= 0.5",

src/galax/__init__.py

@@ -6,14 +6,13 @@
     "coordinates",
     "potential",
     "dynamics",
-    "units",
     "utils",
     "typing",
 ]
 
 from jax import config
 
-from . import coordinates, dynamics, potential, typing, units, utils
+from . import coordinates, dynamics, potential, typing, utils
 from ._version import __version__
 
 config.update("jax_enable_x64", True)  # noqa: FBT003

src/galax/__init__.pyi

@@ -5,15 +5,13 @@ __all__ = [
     "dynamics",
     "potential",
     "typing",
-    "units",
     "utils",
 ]
 
 from . import (
     dynamics as dynamics,
     potential as potential,
     typing as typing,
-    units as units,
     utils as utils,
 )
 from ._version import (  # type: ignore[attr-defined]

src/galax/coordinates/_psp/base.py

@@ -20,7 +20,7 @@
     Cartesian3DVector,
     represent_as as vector_represent_as,
 )
-from unxt import Quantity
+from unxt import Quantity, unitsystem
 
 from .utils import getitem_broadscalartime_index
 from galax.typing import (
@@ -30,7 +30,6 @@
     BatchVec7,
     BroadBatchFloatQScalar,
 )
-from galax.units import unitsystem
 
 if TYPE_CHECKING:
     from typing import Self
@@ -188,8 +187,8 @@ def w(self, *, units: Any) -> BatchVec6:
 
         Parameters
         ----------
-        units : `galax.units.UnitSystem`, optional keyword-only
-            The unit system. :func:`~galax.units.unitsystem` is used to
+        units : `unxt.UnitSystem`, optional keyword-only
+            The unit system. :func:`~unxt.unitsystem` is used to
             convert the input to a unit system.
 
         Returns
@@ -228,8 +227,8 @@ def wt(self, *, units: Any) -> BatchVec7:
 
         Parameters
         ----------
-        units : `galax.units.UnitSystem`, keyword-only
-            The unit system. :func:`~galax.units.unitsystem` is used to
+        units : `unxt.UnitSystem`, keyword-only
+            The unit system. :func:`~unxt.unitsystem` is used to
             convert the input to a unit system.
 
         Returns

src/galax/coordinates/_psp/operator_compat.py

@@ -176,7 +176,7 @@ def call(
     Quantity['length'](Array(2., dtype=float64), unit='kpc')
 
     >>> newpsp.t.to("Myr")
-    Quantity['time'](Array(6.52312755, dtype=float64), unit='Myr')
+    Quantity['time'](Array(6.52312732, dtype=float64), unit='Myr')
 
     This spatial translation is time independent.
 

src/galax/coordinates/_psp/psp.py

@@ -8,7 +8,7 @@
 import jax.numpy as jnp
 
 from coordinax import Abstract3DVector, Abstract3DVectorDifferential
-from unxt import Quantity
+from unxt import Quantity, UnitSystem
 
 from .base import AbstractPhaseSpacePosition
 from .utils import _p_converter, _q_converter
@@ -20,7 +20,6 @@
     QVec1,
     VecTime,
 )
-from galax.units import UnitSystem
 from galax.utils._shape import batched_shape, expand_batch_dims, vector_batched_shape
 
 

src/galax/dynamics/_dynamics/integrate/_api.py

@@ -2,9 +2,10 @@
 
 from typing import Any, Protocol, runtime_checkable
 
+from unxt import UnitSystem
+
 import galax.typing as gt
 from galax.coordinates import AbstractPhaseSpacePosition, PhaseSpacePosition
-from galax.units import UnitSystem
 from galax.utils.dataclasses import _DataclassInstance
 
 

src/galax/dynamics/_dynamics/integrate/_base.py

@@ -4,10 +4,11 @@
 
 import equinox as eqx
 
+from unxt import UnitSystem
+
 from ._api import FCallable
 from galax.coordinates import AbstractPhaseSpacePosition, PhaseSpacePosition
 from galax.typing import BatchQVecTime, BatchVec6, BatchVecTime, QVecTime, VecTime
-from galax.units import UnitSystem
 
 
 class AbstractIntegrator(eqx.Module, strict=True):  # type: ignore[call-arg, misc]

src/galax/dynamics/_dynamics/integrate/_builtin.py

@@ -10,13 +10,12 @@
 import jax
 
 import quaxed.array_api as xp
-from unxt import Quantity
+from unxt import Quantity, UnitSystem
 
 import galax.typing as gt
 from ._api import FCallable
 from ._base import AbstractIntegrator
 from galax.coordinates import AbstractPhaseSpacePosition, PhaseSpacePosition
-from galax.units import UnitSystem
 from galax.utils import ImmutableDict
 from galax.utils._jax import vectorize_method
 

src/galax/dynamics/_dynamics/integrate/_funcs.py

@@ -114,7 +114,7 @@ def evaluate_orbit(
     >>> import quaxed.array_api as xp  # preferred over `jax.numpy`
     >>> import galax.coordinates as gc
     >>> import galax.potential as gp
-    >>> from galax.units import galactic
+    >>> from unxt.unitsystems import galactic
 
     We can then create the point-mass' potential, with galactic units:
 

src/galax/dynamics/_dynamics/mockstream/df/fardal.py

@@ -152,8 +152,7 @@ def d2phidr2(
     --------
     >>> from unxt import Quantity
     >>> from galax.potential import NFWPotential
-    >>> pot = NFWPotential(m=Quantity(1e12, "Msun"), r_s=Quantity(20.0, "kpc"),
-    ...                    units="galactic")
+    >>> pot = NFWPotential(m=1e12, r_s=20.0, units="galactic")
     >>> d2phidr2(pot, xp.asarray([8.0, 0.0, 0.0]), t=0)
     Array(-0.00259193, dtype=float64)
     """
@@ -250,8 +249,7 @@ def tidal_radius(
     Examples
     --------
     >>> from galax.potential import NFWPotential
-    >>> from galax.units import galactic
-    >>> pot = NFWPotential(m=1e12, r_s=20.0, units=galactic)
+    >>> pot = NFWPotential(m=1e12, r_s=20.0, units="galactic")
     >>> x=xp.asarray([8.0, 0.0, 0.0])
     >>> v=xp.asarray([8.0, 0.0, 0.0])
     >>> tidal_radius(pot, x, v, prog_mass=1e4, t=0)

src/galax/dynamics/_dynamics/mockstream/mockstream_generator.py

@@ -13,7 +13,7 @@
 from jax.lib.xla_bridge import get_backend
 
 import quaxed.array_api as xp
-from unxt import Quantity
+from unxt import Quantity, UnitSystem
 
 from .core import MockStream
 from .df import AbstractStreamDF
@@ -24,7 +24,6 @@
 from galax.dynamics._dynamics.integrate._funcs import evaluate_orbit
 from galax.potential._potential.base import AbstractPotentialBase
 from galax.typing import BatchVec6, FloatScalar, IntScalar, QVecTime, Vec6, VecN
-from galax.units import UnitSystem
 
 Carry: TypeAlias = tuple[IntScalar, VecN, VecN]
 

src/galax/potential/_potential/base.py

@@ -20,14 +20,13 @@
 import quaxed.numpy as qnp
 import unxt
 from coordinax import Abstract3DVector, FourVector
-from unxt import Quantity
+from unxt import Quantity, UnitSystem
 
 import galax.typing as gt
 from .utils import _convert_from_3dvec, parse_to_quantity
 from galax.coordinates import AbstractPhaseSpacePosition, PhaseSpacePosition
 from galax.potential._potential.param.attr import ParametersAttribute
 from galax.potential._potential.param.utils import all_parameters
-from galax.units import UnitSystem
 from galax.utils._collections import ImmutableDict
 from galax.utils._jax import vectorize_method
 from galax.utils._shape import batched_shape, expand_arr_dims, expand_batch_dims

src/galax/potential/_potential/builtin.py

@@ -20,13 +20,12 @@
 from quax import quaxify
 
 import quaxed.array_api as xp
-from unxt import Quantity
+from unxt import Quantity, UnitSystem, unitsystem
 
 import galax.typing as gt
 from galax.potential._potential.base import default_constants
 from galax.potential._potential.core import AbstractPotential
 from galax.potential._potential.param import AbstractParameter, ParameterField
-from galax.units import UnitSystem, unitsystem
 from galax.utils import ImmutableDict
 from galax.utils._jax import vectorize_method
 
@@ -266,10 +265,10 @@ class TriaxialHernquistPotential(AbstractPotential):
         or constant, like a Quantity. See
         :class:`~galax.potential.ParameterField` for details.
 
-    units : :class:`~galax.units.UnitSystem`, keyword-only
+    units : :class:`~unxt.UnitSystem`, keyword-only
         The unit system to use for the potential.  This parameter accepts a
-        :class:`~galax.units.UnitSystem` or anything that can be converted to a
-        :class:`~galax.units.UnitSystem` using :func:`~galax.units.unitsystem`.
+        :class:`~unxt.UnitSystem` or anything that can be converted to a
+        :class:`~unxt.UnitSystem` using :func:`~unxt.unitsystem`.
 
     Examples
     --------

src/galax/potential/_potential/composite.py

@@ -10,11 +10,10 @@
 import jax
 
 import quaxed.array_api as xp
-from unxt import Quantity
+from unxt import Quantity, UnitSystem, unitsystem
 
 from .base import AbstractPotentialBase, default_constants
 from galax.typing import BatchableRealQScalar, BatchFloatQScalar, BatchQVec3
-from galax.units import UnitSystem, unitsystem
 from galax.utils import ImmutableDict
 from galax.utils._misc import first
 

src/galax/potential/_potential/core.py

@@ -7,12 +7,11 @@
 
 import equinox as eqx
 
-from unxt import Quantity
+from unxt import Quantity, UnitSystem, unitsystem
 
 from .base import AbstractPotentialBase, default_constants
 from .composite import CompositePotential
 from galax.typing import FloatQScalar, QVec3, RealScalar
-from galax.units import UnitSystem, unitsystem
 from galax.utils import ImmutableDict
 
 

src/galax/potential/_potential/frame.py

@@ -9,11 +9,10 @@
 import equinox as eqx
 
 from coordinax.operators import OperatorSequence, simplify_op
-from unxt import Quantity
+from unxt import Quantity, UnitSystem
 
 import galax.typing as gt
 from galax.potential._potential.base import AbstractPotentialBase
-from galax.units import UnitSystem
 from galax.utils import ImmutableDict
 
 
@@ -90,7 +89,7 @@ class PotentialFrame(AbstractPotentialBase):
 
     >>> op2 = cxo.GalileanTranslationOperator(Quantity([1_000, 0, 0, 0], "kpc"))
     >>> op2.translation.t.to("Myr")
-    Quantity['time'](Array(3.26156378, dtype=float64), unit='Myr')
+    Quantity['time'](Array(3.26156366, dtype=float64), unit='Myr')
 
     >>> framedpot2 = gp.PotentialFrame(potential=pot, operator=op2)
 

src/galax/potential/_potential/special.py

@@ -11,11 +11,11 @@
 import equinox as eqx
 
 from unxt import Quantity
+from unxt.unitsystems import UnitSystem, dimensionless, galactic, unitsystem
 
 from .base import AbstractPotentialBase, default_constants
 from .builtin import HernquistPotential, MiyamotoNagaiPotential, NFWPotential
 from .composite import AbstractCompositePotential
-from galax.units import UnitSystem, dimensionless, galactic, unitsystem
 from galax.utils import ImmutableDict
 
 T = TypeVar("T", bound=AbstractPotentialBase)
@@ -52,7 +52,7 @@ class MilkyWayPotential(AbstractCompositePotential):
 
     Parameters
     ----------
-    units : `~galax.units.UnitSystem` (optional)
+    units : `~unxt.UnitSystem` (optional)
         Set of non-reducable units that specify (at minimum) the
         length, mass, time, and angle units.
     disk : dict (optional)

src/galax/potential/_potential/utils.py

@@ -16,10 +16,10 @@
 
 import coordinax as cx
 from unxt import Quantity
+from unxt.unitsystems import DimensionlessUnitSystem, UnitSystem, dimensionless
 
 from galax.coordinates import AbstractPhaseSpacePosition
 from galax.typing import Unit
-from galax.units import DimensionlessUnitSystem, UnitSystem, dimensionless
 
 # --------------------------------------------------------------