unxt 0.1.dev463__py3-none-any.whl

unxt/_src/unitsystems/core.py

@@ -0,0 +1,307 @@
+"""Tools for representing systems of units using ``astropy.units``."""
+
+__all__ = ("unitsystem", "unitsystem_of")
+
+
+from collections.abc import Sequence
+from dataclasses import field, make_dataclass
+from typing import Annotated, Any
+
+import equinox as eqx
+import jax.tree_util as jtu
+import numpy as np
+from astropy.constants import G as const_G  # noqa: N811, pylint: disable=E0611
+from astropy.units import UnitBase as AstropyUnitBase
+from plum import dispatch
+
+from . import builtin_dimensions as ud
+from .base import UNITSYSTEMS_REGISTRY, AbstractUnitSystem
+from .builtin import DimensionlessUnitSystem
+from .flags import AbstractUSysFlag, DynamicalSimUSysFlag, StandardUSysFlag
+from .realizations import NAMED_UNIT_SYSTEMS, dimensionless
+from .utils import parse_dimlike_name
+from unxt.dims import dimension_of
+from unxt.units import unit
+
+# ===================================================================
+# `unitsystem` function
+
+
+@dispatch
+def unitsystem(usys: AbstractUnitSystem, /) -> AbstractUnitSystem:
+    """Convert a UnitSystem to a UnitSystem.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem
+    >>> usys = unitsystem("kpc", "Myr", "Msun", "radian")
+    >>> usys
+    unitsystem(kpc, Myr, solMass, rad)
+
+    >>> unitsystem(usys) is usys
+    True
+
+    """
+    return usys
+
+
+@dispatch
+def unitsystem(seq: Sequence[Any], /) -> AbstractUnitSystem:
+    """Convert a UnitSystem or tuple of arguments to a UnitSystem.
+
+    Examples
+    --------
+    >>> import unxt as u
+
+    >>> u.unitsystem(())
+    DimensionlessUnitSystem()
+
+    >>> u.unitsystem(("kpc", "Myr", "Msun", "radian"))
+    unitsystem(kpc, Myr, solMass, rad)
+
+    >>> u.unitsystem(["kpc", "Myr", "Msun", "radian"])
+    unitsystem(kpc, Myr, solMass, rad)
+
+    """
+    return unitsystem(*seq) if len(seq) > 0 else dimensionless
+
+
+@dispatch
+def unitsystem(_: None, /) -> DimensionlessUnitSystem:
+    """Dimensionless unit system from None.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem
+    >>> unitsystem(None)
+    DimensionlessUnitSystem()
+
+    """
+    return dimensionless
+
+
+@dispatch
+def unitsystem(*args: Any) -> AbstractUnitSystem:
+    """Convert a set of arguments to a UnitSystem.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem
+
+    >>> unitsystem("kpc", "Myr", "Msun", "radian")
+    unitsystem(kpc, Myr, solMass, rad)
+
+    """
+    # Convert everything to a unit
+    args = tuple(map(unit, args))
+
+    # Check that the units all have different dimensions
+    dims = tuple(map(dimension_of, args))
+    dims = eqx.error_if(
+        dims,
+        len(set(dims)) < len(dims),
+        "some dimensions are repeated",
+    )
+
+    # Return if the unit system is already registered
+    if dims in UNITSYSTEMS_REGISTRY:
+        return UNITSYSTEMS_REGISTRY[dims](*args)
+
+    # Otherwise, create a new unit system
+    # dimension names of all the units
+    du = {parse_dimlike_name(x).replace(" ", "_"): dimension_of(x) for x in args}
+    # name: physical types
+    cls_name = "".join(k.title().replace("_", "") for k in du) + "UnitSystem"  # pylint: disable=unreachable
+    # fields: name, unit
+    fields = [
+        (
+            k,
+            Annotated[AstropyUnitBase, v],
+            field(init=True, repr=True, hash=True, compare=True),  # pylint: disable=invalid-field-call
+        )
+        for k, v in du.items()
+    ]
+
+    def _reduce_(self: AbstractUnitSystem) -> tuple:
+        return (_call_unitsystem, self.base_units, None, None, None, None)
+
+    # Make and register the dataclass class
+    unitsystem_cls: type[AbstractUnitSystem] = make_dataclass(
+        cls_name,
+        fields,
+        bases=(AbstractUnitSystem,),
+        namespace={"__reduce__": _reduce_},
+        frozen=True,
+        slots=True,
+        eq=True,
+        repr=True,
+        init=True,
+    )
+
+    # Register the dynamically created unit system as a static PyTree node
+    jtu.register_static(unitsystem_cls)
+
+    # Make the dataclass instance
+    return unitsystem_cls(*args)
+
+
+@dispatch
+def unitsystem(name: str, /) -> AbstractUnitSystem:
+    """Return unit system from name.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem
+    >>> unitsystem("galactic")
+    unitsystem(kpc, Myr, solMass, rad)
+
+    >>> unitsystem("solarsystem")
+    unitsystem(AU, yr, solMass, rad)
+
+    >>> unitsystem("dimensionless")
+    DimensionlessUnitSystem()
+
+    """
+    return NAMED_UNIT_SYSTEMS[name]
+
+
+@dispatch
+def unitsystem(usys: AbstractUnitSystem, *args: Any) -> AbstractUnitSystem:
+    """Create a unit system from an existing unit system and additional units.
+
+    Examples
+    --------
+    We can add a new unit definition to an existing unit system:
+
+    >>> from unxt.unitsystems import unitsystem
+    >>> usys = unitsystem("galactic")
+    >>> unitsystem(usys, "km/s")
+    LengthTimeMassAngleSpeedUnitSystem(length=Unit("kpc"), time=Unit("Myr"), mass=Unit("solMass"), angle=Unit("rad"), speed=Unit("km / s"))
+
+    We can also override the base unit of an existing unit system:
+
+    >>> new_usys = unitsystem(usys, "pc")
+    >>> new_usys
+    TimeMassAngleLengthUnitSystem(time=Unit("Myr"), mass=Unit("solMass"), angle=Unit("rad"), length=Unit("pc"))
+
+    """  # noqa: E501
+    # TODO: not need this hack for single-string inputs
+    # TODO: process new units without making a whole unit system
+    if len(args) == 1 and isinstance(args[0], str):
+        new_usys = unitsystem(unit(args[0]))
+    else:
+        new_usys = unitsystem(*args)
+
+    current_units = [
+        u for u in usys.base_units if dimension_of(u) not in new_usys.base_dimensions
+    ]
+    return unitsystem(*current_units, *args)  # pylint: disable=unreachable
+
+
+@dispatch
+def unitsystem(flag: type[AbstractUSysFlag], *_: Any) -> AbstractUnitSystem:
+    """Raise an exception since the flag is abstract."""
+    msg = "Do not use the AbstractUSysFlag directly, only use subclasses."
+    raise TypeError(msg)
+
+
+@dispatch
+def unitsystem(flag: type[StandardUSysFlag], *args: Any) -> AbstractUnitSystem:
+    """Create a standard unit system using the inputted units.
+
+    Examples
+    --------
+    >>> from unxt import unitsystem, unitsystems
+    >>> unitsystem(unitsystems.StandardUSysFlag, "kpc", "Myr", "Msun")
+    LengthTimeMassUnitSystem(length=Unit("kpc"), time=Unit("Myr"), mass=Unit("solMass"))
+
+    """
+    return unitsystem(*args)
+
+
+@dispatch
+def unitsystem(
+    flag: type[DynamicalSimUSysFlag],
+    *args: Any,
+    G: float | int = 1.0,  # noqa: N803
+) -> AbstractUnitSystem:
+    """Make a dynamical unit system.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem, DynamicalSimUSysFlag
+
+    >>> unitsystem(DynamicalSimUSysFlag, "m", "kg")
+    LengthMassTimeUnitSystem(length=Unit("m"), mass=Unit("kg"), time=Unit("122404 s"))
+
+    """
+    tmp = unitsystem(*args)
+
+    # Use G for computing the missing units below:
+    G = G * const_G
+
+    added = ()
+    if ud.length in tmp.base_dimensions and ud.mass in tmp.base_dimensions:
+        time = 1 / np.sqrt(G * tmp["mass"] / tmp["length"] ** 3)
+        added = (time,)
+    elif ud.length in tmp.base_dimensions and ud.time in tmp.base_dimensions:
+        mass = 1 / G * tmp["length"] ** 3 / tmp["time"] ** 2
+        added = (mass,)
+    elif ud.length in tmp.base_dimensions and ud.speed in tmp.base_dimensions:
+        time = tmp["length"] / tmp["velocity"]
+        mass = tmp["velocity"] ** 2 / G * tmp["length"]
+        added = (time, mass)
+    elif ud.mass in tmp.base_dimensions and ud.time in tmp.base_dimensions:
+        length = np.cbrt(G * tmp["mass"] * tmp["time"] ** 2)
+        added = (length,)
+    elif ud.mass in tmp.base_dimensions and ud.speed in tmp.base_dimensions:
+        length = G * tmp["mass"] / tmp["velocity"] ** 2
+        time = length / tmp["velocity"]
+        added = (length, time)
+    elif ud.time in tmp.base_dimensions and ud.speed in tmp.base_dimensions:
+        mass = 1 / G * tmp["velocity"] ** 3 * tmp["time"]
+        length = G * mass / tmp["velocity"] ** 2
+        added = (mass, length)
+
+    return unitsystem(*tmp, *added)
+
+
+# ----
+
+
+def _call_unitsystem(*args: Any) -> AbstractUnitSystem:
+    return unitsystem(*args)
+
+
+# ===================================================================
+# `unitsystem_of` function
+
+
+@dispatch
+def unitsystem_of(obj: Any, /) -> DimensionlessUnitSystem:
+    """Return the unit system of the object.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem_of
+
+    >>> unitsystem_of(1)
+    DimensionlessUnitSystem()
+
+    """
+    return dimensionless
+
+
+@dispatch
+def unitsystem_of(obj: AbstractUnitSystem, /) -> AbstractUnitSystem:
+    """Return the unit system from the unit system.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import galactic, unitsystem_of
+
+    >>> unitsystem_of(galactic) is galactic
+    True
+
+    """
+    return obj

unxt/_src/unitsystems/flags.py

@@ -0,0 +1,80 @@
+"""Flags for unit system definitions."""
+
+__all__ = ("AbstractUSysFlag", "DynamicalSimUSysFlag", "StandardUSysFlag")
+
+from typing import Any
+
+
+class AbstractUSysFlag:
+    """Abstract class for unit system flags to provide dispatch control.
+
+    Unit system flags are used to indicate the type of unit system being
+    defined. They are not intended to be instantiated and are used to provide
+    dispatch control for defining unit systems.
+
+    Raises
+    ------
+    ValueError
+        If an attempt is made to instantiate a unit system flag class.
+
+    See Also
+    --------
+    `unxt.unitsystems.unitsystem` :
+        Function to define a unit system. The `AbstractUSysFlag` can be provided
+        as the first argument to indicate the type of unit system being defined.
+        This is useful for multiple-dispatching to the correct constructor.
+
+    Examples
+    --------
+    For this example we will use the `unxt.unitsystems.StandardUSysFlag`, which
+    indicates the standard construction of a unit system directly from the set
+    of base units. Flags like `unxt.unitsystems.DynamicalSimUSysFlag` can be
+    used to create a unit system where the gravitational constant ``G`` is
+    unit-less and units are defined accordingly.
+
+    >>> import unxt
+
+    Define a unit system with the standard flag:
+
+    >>> unxt.unitsystem(unxt.unitsystems.StandardUSysFlag, "m", "kg", "s")
+    LengthMassTimeUnitSystem(length=Unit("m"), mass=Unit("kg"), time=Unit("s"))
+
+    """
+
+    def __new__(cls, *_: Any, **__: Any) -> None:  # type: ignore[misc]
+        msg = "unit system flag classes cannot be instantiated."
+        raise ValueError(msg)
+
+
+class StandardUSysFlag(AbstractUSysFlag):
+    """Flag to indicate a standard unit system with no additional arguments.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem, StandardUSysFlag
+
+    Define a unit system with the standard flag:
+
+    >>> unitsystem(StandardUSysFlag, "m", "kg", "s")
+    LengthMassTimeUnitSystem(length=Unit("m"), mass=Unit("kg"), time=Unit("s"))
+
+    Further examples may be found in the ``unitsystem`` docs.
+
+    """
+
+
+class DynamicalSimUSysFlag(AbstractUSysFlag):
+    """Flag to indicate a unit system with optional definition of G.
+
+    Examples
+    --------
+    >>> from unxt.unitsystems import unitsystem, DynamicalSimUSysFlag
+
+    Define a unit system with the dynamical simulation flag:
+
+    >>> unitsystem(DynamicalSimUSysFlag, "m", "kg")
+    LengthMassTimeUnitSystem(length=Unit("m"), mass=Unit("kg"), time=Unit("122404 s"))
+
+    Further examples may be found in the ``unitsystem`` docs.
+
+    """

unxt/_src/unitsystems/realizations.py

@@ -0,0 +1,73 @@
+"""Realizations of unit systems."""
+# pylint: disable=no-member
+
+__all__ = (
+    # unit system instance
+    "galactic",
+    "dimensionless",
+    "solarsystem",
+    "si",
+    "cgs",
+    # unit system alias
+    "NAMED_UNIT_SYSTEMS",
+)
+
+from .base import AbstractUnitSystem
+from .builtin import (
+    CGSUnitSystem,
+    DimensionlessUnitSystem,
+    LTMAUnitSystem,
+    SIUnitSystem,
+)
+from unxt.units import unit
+
+# Dimensionless. This is a singleton.
+dimensionless = DimensionlessUnitSystem()
+
+# Galactic unit system
+galactic = LTMAUnitSystem(
+    length=unit("kpc"),
+    time=unit("Myr"),
+    mass=unit("Msun"),
+    angle=unit("radian"),
+)
+
+# Solar system units
+solarsystem = LTMAUnitSystem(
+    length=unit("au"), time=unit("yr"), mass=unit("Msun"), angle=unit("radian")
+)
+
+#: International System of Units
+si = SIUnitSystem(
+    length=unit("meter"),
+    time=unit("second"),
+    mass=unit("kilogram"),
+    electric_current=unit("ampere"),
+    temperature=unit("Kelvin"),
+    amount=unit("mole"),
+    luminous_intensity=unit("candela"),
+    angle=unit("radian"),
+)
+
+#: Centimeter, gram, second unit system
+cgs = CGSUnitSystem(
+    length=unit("centimeter"),
+    time=unit("second"),
+    mass=unit("gram"),
+    angle=unit("radian"),
+    force=unit("dyne"),
+    energy=unit("erg"),
+    pressure=unit("barye"),
+    dynamic_viscosity=unit("poise"),
+    kinematic_viscosity=unit("stokes"),
+)
+
+
+#: Named unit systems
+NAMED_UNIT_SYSTEMS: dict[str, AbstractUnitSystem] = {
+    "galactic": galactic,
+    "solarsystem": solarsystem,
+    "dimensionless": dimensionless,
+    "si": si,
+    "cgs": cgs,
+}

unxt/_src/unitsystems/utils.py

@@ -0,0 +1,95 @@
+"""Unit system utils."""
+
+__all__: tuple[str, ...] = ()
+
+import re
+from typing import Any
+
+from plum import dispatch
+
+from zeroth import zeroth
+
+from . import builtin_dimensions as bdims
+from unxt.dims import AbstractDimension, dimension_of
+from unxt.units import AbstractUnit
+
+# ------------------------------------
+
+
+@dispatch.abstract
+def parse_dimlike_name(pt: Any, /) -> str:
+    """Get the dimension name from the object."""
+    raise NotImplementedError  # pragma: no cover
+
+
+@dispatch
+def parse_dimlike_name(pt: str, /) -> str:
+    """Return the dimension name.
+
+    Note that this does not check for the existence of that dimension.
+
+    Examples
+    --------
+    >>> from unxt._src.unitsystems.utils import parse_dimlike_name
+
+    >>> parse_dimlike_name("length")
+    'length'
+
+    >>> parse_dimlike_name("not real")
+    'not_real'
+
+    >>> try:
+    ...     parse_dimlike_name("*62")
+    ... except ValueError as e:
+    ...     print(e)
+    Input contains non-letter characters
+
+    """
+    # A regex search to match anything that's not a letter or a whitespace.
+    if re.search(r"[^a-zA-Z_ ]", pt):
+        msg = "Input contains non-letter characters"
+        raise ValueError(msg)
+
+    return pt.replace(" ", "_")
+
+
+@dispatch
+def parse_dimlike_name(pt: AbstractDimension, /) -> str:
+    """Return the dimension name from a dimension.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> from unxt._src.unitsystems.utils import parse_dimlike_name
+
+    >>> parse_dimlike_name(u.dimension("length"))
+    'length'
+
+    >>> parse_dimlike_name(u.dimension("speed"))
+    'speed'
+
+    """
+    # Note: this is not deterministic b/c ``_physical_type`` is a set
+    #       that's why the `if` statement is needed.
+    match pt:
+        case bdims.speed:
+            out = "speed"
+        case _:
+            out = parse_dimlike_name(zeroth(pt._physical_type))  # noqa: SLF001
+    return out
+
+
+@dispatch
+def parse_dimlike_name(pt: AbstractUnit, /) -> str:
+    """Return the dimension name from a unit.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> from unxt._src.unitsystems.utils import parse_dimlike_name
+
+    >>> parse_dimlike_name(u.unit("km"))
+    'length'
+
+    """
+    return parse_dimlike_name(dimension_of(pt))

unxt/_src/utils.py

@@ -0,0 +1,112 @@
+"""Utilities.
+
+Copyright (c) 2023 Galactic Dynamics. All rights reserved.
+"""
+
+__all__: tuple[str, ...] = ()
+
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast, runtime_checkable
+
+from jax import dtypes
+from jax.numpy import dtype as DType  # noqa: N812
+
+import quaxed.lax as qlax
+
+if TYPE_CHECKING:
+    from typing import Self
+
+
+_singleton_insts: dict[type, object] = {}
+
+
+class SingletonMixin:
+    """Singleton class.
+
+    This class is a mixin that can be used to create singletons.
+
+    Examples
+    --------
+    >>> class MySingleton(SingletonMixin):
+    ...     pass
+
+    >>> a = MySingleton()
+    >>> b = MySingleton()
+    >>> a is b
+    True
+
+    """
+
+    def __new__(cls, /, *_: Any, **__: Any) -> "Self":
+        # Check if instance already exists
+        if cls in _singleton_insts:
+            return cast("Self", _singleton_insts[cls])
+        # Create new instance and cache it
+        self = object.__new__(cls)
+        _singleton_insts[cls] = self
+        return self
+
+
+@runtime_checkable
+class HasDType(Protocol):
+    """Protocol for objects that have a dtype attribute."""
+
+    @property
+    def dtype(self) -> DType:
+        """The dtype of the object."""
+
+
+HasDTypeT = TypeVar("HasDTypeT", bound=HasDType)
+
+
+def promote_dtypes(*arrays: HasDTypeT) -> tuple[HasDTypeT, ...]:
+    """Promotes all input arrays to a common dtype.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+
+    >>> x1 = jnp.array([1, 2, 3], dtype=jnp.int32)
+    >>> x2 = jnp.array([4, 5, 6], dtype=jnp.float32)
+
+    >>> x1, x2 = promote_dtypes(x1, x2)
+    >>> x1.dtype, x2.dtype
+    (dtype('float32'), dtype('float32'))
+
+    >>> q1 = u.Quantity.from_([1, 2, 3], "m", dtype=int)
+    >>> q2 = u.Q([4.0, 5, 6], unit="km")
+    >>> q1, q2 = promote_dtypes(q1, q2)
+    >>> q1.dtype, q2.dtype
+    (dtype('float32'), dtype('float32'))
+
+    """
+    common_dtype = dtypes.result_type(*arrays)
+    # TODO: check if this copies.
+    return tuple([qlax.convert_element_type(arr, common_dtype) for arr in arrays])  # type: ignore[arg-type,misc]  # pylint: disable=R1728
+
+
+def promote_dtypes_if_needed(
+    original_dtypes: tuple[DType, ...], /, *args: HasDTypeT
+) -> tuple[HasDTypeT, ...]:
+    """Promotes dtypes of args if they differ from original_dtypes.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+    >>> from unxt._src.utils import promote_dtypes_if_needed
+    >>> x1 = jnp.array([1, 2, 3], dtype=jnp.int32)
+    >>> x2 = jnp.array([4, 5, 6], dtype=jnp.float32)
+    >>> original_dtypes = (x1.dtype, x1.dtype)
+    >>> x1_new, x2_new = promote_dtypes_if_needed(original_dtypes, x1, x2)
+    >>> x1_new.dtype, x2_new.dtype
+    (dtype('float32'), dtype('float32'))
+
+    """
+    # Compare equality of all `original_dtypes` and all `args` dtypes.
+    all_orig_dtypes_eq = all(dt == original_dtypes[0] for dt in original_dtypes)
+    all_new_dtypes_eq = all(arg.dtype == args[0].dtype for arg in args)
+    if all_orig_dtypes_eq and not all_new_dtypes_eq:
+        return promote_dtypes(*args)
+
+    return args

unxt/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.1.dev463'
+__version_tuple__ = version_tuple = (0, 1, 'dev463')
+
+__commit_id__ = commit_id = None

unxt/_version.pyi

@@ -0,0 +1,2 @@
+version: str
+version_tuple: tuple[int, int, int] | tuple[int, int, int, str, str]