unxt 0.1.dev463__py3-none-any.whl

unxt/_src/quantity/static_quantity.py

@@ -0,0 +1,90 @@
+"""Static Quantity Class."""
+
+# pylint: disable=import-error, no-member, unsubscriptable-object
+#    b/c it doesn't understand dataclass fields
+
+__all__ = ("StaticQuantity",)
+
+from typing import Any, final
+
+import equinox as eqx
+import wadler_lindig as wl
+from numpy.typing import NDArray
+from plum import add_promotion_rule, parametric
+
+from .base_parametric import AbstractParametricQuantity
+from .quantity import Quantity
+from .unchecked import BareQuantity
+from .value import StaticValue
+from unxt.units import AbstractUnit, unit as parse_unit
+
+
+@final
+@parametric
+class StaticQuantity(AbstractParametricQuantity):
+    """Static quantities with associated units.
+
+    This class is parametrized by the dimensions of the units, just like
+    {class}`~unxt.quantity.Quantity`, but its value is always stored as a static
+    NumPy array. It accepts Python scalars and array-like inputs that can be
+    converted to NumPy arrays, and it rejects JAX arrays.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import unxt as u
+
+    Basic construction:
+
+    >>> q = u.StaticQuantity(np.array([1.0, 2.0]), "m")
+    >>> q
+    StaticQuantity(array([1., 2.]), unit='m')
+
+    Values are static and hashable:
+
+    >>> isinstance(hash(q), int)
+    True
+
+    JAX arrays are rejected:
+
+    >>> import jax.numpy as jnp
+    >>> try:
+    ...     u.StaticQuantity(jnp.array([1.0, 2.0]), "m")
+    ... except TypeError as e:
+    ...     print(e)
+    StaticQuantity does not accept JAX arrays. Use Quantity for traced values.
+
+    The Wadler-Lindig representation hides the internal static wrapper:
+
+    >>> import wadler_lindig as wl
+    >>> wl.pprint(q, short_arrays=False)
+    StaticQuantity(array([1., 2.]), unit='m')
+
+    """
+
+    value: StaticValue = eqx.field(  # type: ignore[assignment]
+        static=True, converter=StaticValue.from_
+    )
+    """The static value of the `AbstractQuantity`."""
+
+    unit: AbstractUnit = eqx.field(static=True, converter=parse_unit)
+    """The unit associated with this value."""
+
+    def __hash__(self) -> int:
+        """Return the hash of the quantity."""
+        return hash((self.value, self.unit))
+
+    def __eq__(self, other: Any, /) -> bool | NDArray[bool]:  # type: ignore[override]
+        """Return structural equality for static quantities."""
+        if isinstance(other, StaticQuantity):
+            return self.unit == other.unit and self.value == other.value
+        return super().__eq__(other)
+
+    def __pdoc__(self, *, show_wrapper: bool = False, **kwargs: Any) -> wl.AbstractDoc:
+        """Return the Wadler-Lindig representation of this class."""
+        return super().__pdoc__(show_wrapper=False, **kwargs)
+
+
+add_promotion_rule(StaticQuantity, StaticQuantity, StaticQuantity)
+add_promotion_rule(StaticQuantity, Quantity, Quantity)
+add_promotion_rule(StaticQuantity, BareQuantity, BareQuantity)

unxt/_src/quantity/unchecked.py

@@ -0,0 +1,45 @@
+# pylint: disable=import-error, no-member, unsubscriptable-object
+#    b/c it doesn't understand dataclass fields
+
+__all__ = ("BareQuantity",)
+
+from typing import Any
+
+import equinox as eqx
+from jaxtyping import Array, Shaped
+from plum import add_promotion_rule
+
+from .base import AbstractQuantity
+from .quantity import Quantity
+from .value import StaticValue, convert_to_quantity_value
+from unxt.units import AbstractUnit, unit as parse_unit
+
+
+class BareQuantity(AbstractQuantity):
+    """A fast implementation of the Quantity class.
+
+    This class is not parametrized by its dimensionality.
+    """
+
+    value: Shaped[Array | StaticValue, "*shape"] = eqx.field(
+        converter=convert_to_quantity_value
+    )
+    """The value of the `AbstractQuantity`."""
+
+    unit: AbstractUnit = eqx.field(static=True, converter=parse_unit)
+    """The unit associated with this value."""
+
+    def __class_getitem__(cls: "type[BareQuantity]", item: Any) -> "type[BareQuantity]":
+        """No-op support for `BareQuantity[...]` syntax.
+
+        This method is called when the class is subscripted, e.g.:
+
+        >>> from unxt.quantity import BareQuantity
+        >>> BareQuantity["length"]
+        <class 'unxt...quantity...BareQuantity'>
+
+        """
+        return cls
+
+
+add_promotion_rule(BareQuantity, Quantity, Quantity)

unxt/_src/quantity/value.py

@@ -0,0 +1,289 @@
+"""Conversion to quantity value types."""
+
+__all__ = ("StaticValue", "convert_to_quantity_value")
+
+import operator
+import warnings
+from typing import Any, final
+from typing_extensions import override
+
+import jax
+import numpy as np
+import quax
+import wadler_lindig as wl
+from jaxtyping import Array, ArrayLike
+from numpy.typing import NDArray
+from plum import dispatch
+
+import quaxed.numpy as jnp
+
+
+@final
+class StaticValue:
+    """Immutable static value wrapper for `StaticQuantity`.
+
+    This stores a read-only NumPy array and is used to keep the value static
+    while avoiding Equinox's static-array warnings. Arithmetic operations
+    degrade to the wrapped array unless both operands are `StaticValue`, in
+    which case a `StaticValue` is returned.
+
+    Note that since the array is immutable, hashing is supported.
+    The hash is computed from the array's dtype, shape, and bytes.
+
+    """
+
+    __slots__ = ("_array",)
+
+    def __init__(self, array: np.ndarray, /) -> None:
+        value = np.asarray(array).copy()
+        value.setflags(write=False)
+        self._array = value
+
+    @property
+    def array(self) -> np.ndarray:
+        """Return the contained NumPy array."""
+        return self._array
+
+    @property
+    def _jnparray(self) -> Array:
+        """Return the contained array as a JAX array."""
+        return jnp.asarray(self._array)
+
+    # Constructor
+    @classmethod
+    @dispatch.abstract
+    def from_(cls: type["StaticValue"], *args: Any, **kwargs: Any) -> "StaticValue":
+        """Create a `StaticValue` from given arguments."""
+        raise NotImplementedError  # pragma: no cover
+
+    # ==================================================
+    # NumPy API
+
+    def __array__(self, dtype: Any = None) -> np.ndarray:
+        if dtype is None:
+            return np.asarray(self._array)
+        return np.asarray(self._array, dtype=dtype)
+
+    # ==================================================
+    # Wadler-Lindig API
+
+    def __pdoc__(self, *, show_wrapper: bool = True, **kwargs: Any) -> wl.AbstractDoc:
+        """Return the Wadler-Lindig representation of this class."""
+        if not show_wrapper:
+            return wl.pdoc(self._array, **kwargs)
+        return (
+            wl.TextDoc("StaticValue(")
+            + wl.pdoc(self._array, **kwargs)
+            + wl.TextDoc(")")
+        )
+
+    # ==================================================
+    # Python container API
+
+    def __len__(self) -> int:
+        return len(self._array)
+
+    def __iter__(self) -> Any:
+        return iter(self._array)
+
+    def __getitem__(self, key: Any) -> Any:
+        return self._array[key]
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(jnp.asarray(self._array), name)
+
+    def __repr__(self) -> str:
+        return wl.pformat(self, short_arrays=False, max_width=80)
+
+    @override
+    def __eq__(self, other: object, /) -> bool | NDArray[np.bool_]:  # type: ignore[override]
+        if isinstance(other, StaticValue):
+            return bool(np.array_equal(self._array, other._array))
+        if isinstance(other, (np.ndarray, Array, list, tuple)):
+            return self._array == other
+        return NotImplemented
+
+    def __ne__(self, other: object, /) -> bool | NDArray[np.bool_]:  # type: ignore[override]
+        if isinstance(other, StaticValue):
+            return not bool(np.array_equal(self._array, other._array))
+        if isinstance(other, (np.ndarray, Array, list, tuple)):
+            return self._array != other
+        return NotImplemented
+
+    def __lt__(self, other: Any, /) -> NDArray[np.bool_]:
+        if isinstance(other, StaticValue):
+            other = other._array
+        return self._array < other
+
+    def __le__(self, other: Any, /) -> NDArray[np.bool_]:
+        if isinstance(other, StaticValue):
+            other = other._array
+        return self._array <= other
+
+    def __gt__(self, other: Any, /) -> NDArray[np.bool_]:
+        if isinstance(other, StaticValue):
+            other = other._array
+        return self._array > other
+
+    def __ge__(self, other: Any, /) -> NDArray[np.bool_]:
+        if isinstance(other, StaticValue):
+            other = other._array
+        return self._array >= other
+
+    def __hash__(self) -> int:
+        return hash((self._array.dtype.str, self._array.shape, self._array.tobytes()))
+
+    def _binary_op(self, other: Any, op: Any) -> Any:
+        if isinstance(other, StaticValue):
+            result = op(self._jnparray, jnp.asarray(other.array))
+            return StaticValue(np.asarray(result))
+        return op(self._jnparray, other)
+
+    def _rbinary_op(self, other: Any, op: Any) -> Any:
+        if isinstance(other, StaticValue):
+            result = op(jnp.asarray(other.array), self._jnparray)
+            return StaticValue(np.asarray(result))
+        return op(other, self._jnparray)
+
+    def __add__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.add)
+
+    def __radd__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.add)
+
+    def __sub__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.sub)
+
+    def __rsub__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.sub)
+
+    def __mul__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.mul)
+
+    def __rmul__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.mul)
+
+    def __truediv__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.truediv)
+
+    def __rtruediv__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.truediv)
+
+    def __floordiv__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.floordiv)
+
+    def __rfloordiv__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.floordiv)
+
+    def __pow__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.pow)
+
+    def __rpow__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.pow)
+
+    def __mod__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.mod)
+
+    def __rmod__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.mod)
+
+    def __matmul__(self, other: Any) -> Any:
+        return self._binary_op(other, operator.matmul)
+
+    def __rmatmul__(self, other: Any) -> Any:
+        return self._rbinary_op(other, operator.matmul)
+
+    def __neg__(self) -> Any:
+        return -self._jnparray
+
+    def __pos__(self) -> Any:
+        return +self._jnparray
+
+    def __abs__(self) -> Any:
+        return abs(self._jnparray)
+
+
+@StaticValue.from_.dispatch
+def from_(cls: type[StaticValue], value: object, /) -> StaticValue:
+    """Convert a value for `StaticQuantity`."""
+    return cls(np.asarray(value))
+
+
+@StaticValue.from_.dispatch
+def from_(cls: type[StaticValue], value: StaticValue, /) -> StaticValue:
+    """Convert a value for `StaticQuantity`."""
+    return value
+
+
+@StaticValue.from_.dispatch
+def from_(cls: type[StaticValue], value: jax.Array | jax.core.Tracer, /) -> StaticValue:
+    """Reject JAX arrays for `StaticQuantity`."""
+    msg = "StaticQuantity does not accept JAX arrays. Use Quantity for traced values."
+    raise TypeError(msg)
+
+
+# ==================================================================
+# Converters for quantity value field
+
+
+@dispatch.abstract
+def convert_to_quantity_value(obj: Any, /) -> Any:
+    """Convert for the value field of an `AbstractQuantity` subclass."""
+    raise NotImplementedError  # pragma: no cover
+
+
+@dispatch
+def convert_to_quantity_value(obj: StaticValue, /) -> StaticValue:
+    """Allow static values in `Quantity`-like classes."""
+    return obj
+
+
+@dispatch
+def convert_to_quantity_value(obj: quax.ArrayValue, /) -> Any:
+    """Convert a `quax.ArrayValue` for the value field.
+
+    >>> import warnings
+    >>> import jax
+    >>> import jax.numpy as jnp
+    >>> from jaxtyping import Array
+    >>> from quax import ArrayValue
+
+    >>> class MyArray(ArrayValue):
+    ...     value: Array
+    ...
+    ...     def aval(self):
+    ...         return jax.core.ShapedArray(self.value.shape, self.value.dtype)
+    ...
+    ...     def materialise(self):
+    ...         return self.value
+
+    >>> x = MyArray(jnp.array([1, 2, 3]))
+    >>> with warnings.catch_warnings(record=True, action="always") as w:
+    ...     y = convert_to_quantity_value(x)
+    >>> y
+    MyArray(value=i32[3])
+    >>> print(f"Warning caught: {w[-1].message}")
+    Warning caught: 'quax.ArrayValue' subclass 'MyArray' ...
+
+    """
+    warnings.warn(
+        f"'quax.ArrayValue' subclass {type(obj).__name__!r} does not have a registered "
+        "converter. Returning the object as is.",
+        category=UserWarning,
+        stacklevel=2,
+    )
+    return obj
+
+
+@dispatch
+def convert_to_quantity_value(obj: ArrayLike | list[Any] | tuple[Any, ...], /) -> Array:
+    """Convert an array-like object to a `jax.numpy.ndarray`.
+
+    >>> import jax.numpy as jnp
+    >>> from unxt.quantity import convert_to_quantity_value
+
+    >>> convert_to_quantity_value([1, 2, 3])
+    Array([1, 2, 3], dtype=int32)
+
+    """
+    return jnp.asarray(obj)

unxt/_src/units.py

@@ -0,0 +1,103 @@
+"""Units objects in JAX.
+
+Copyright (c) 2023 Galactic Dynamics. All rights reserved.
+"""
+
+__all__ = ("unit", "unit_of", "AbstractUnit")
+
+from typing import Any, TypeAlias
+
+import astropy.units as apyu
+from plum import dispatch
+
+import unxt_api as uapi
+from unxt.dims import AbstractDimension
+
+AbstractUnit: TypeAlias = apyu.Unit | apyu.UnitBase | apyu.CompositeUnit
+
+
+# ===================================================================
+# Construct units
+
+
+@dispatch
+def unit(obj: AbstractUnit, /) -> AbstractUnit:
+    """Construct the units from a units object.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> m = u.unit("m")
+
+    >>> u.unit(m) is m
+    True
+
+    """
+    return obj
+
+
+@dispatch
+def unit(obj: str, /) -> AbstractUnit:
+    """Construct units from a string.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> m = u.unit("m")
+    >>> m
+    Unit("m")
+
+    """
+    return apyu.Unit(obj)
+
+
+# ===================================================================
+# Get units
+
+
+@dispatch
+def unit_of(obj: Any, /) -> None:
+    """Return the units of an object.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> print(u.unit_of(1))
+    None
+
+    """
+    return None  # noqa: RET501
+
+
+@dispatch
+def unit_of(obj: AbstractUnit, /) -> AbstractUnit:
+    """Return the units of an unit.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> m = u.unit("m")
+
+    >>> u.unit_of(m)
+    Unit("m")
+
+    """
+    return obj
+
+
+# ===================================================================
+# Get dimensions
+
+
+@dispatch
+def dimension_of(obj: AbstractUnit, /) -> AbstractDimension:
+    """Return the dimensions of the given units.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> u.dimension_of(u.unit("km"))
+    PhysicalType('length')
+
+    """
+    return uapi.dimension(obj.physical_type)

unxt/_src/unitsystems/__init__.py

@@ -0,0 +1,40 @@
+"""Tools for representing systems of units using ``astropy.units``."""
+# pylint: disable=duplicate-code
+
+__all__ = (
+    "unitsystem",
+    "unitsystem_of",
+    # classes
+    "AbstractUnitSystem",
+    "UNITSYSTEMS_REGISTRY",
+    "DimensionlessUnitSystem",
+    "LTMAUnitSystem",
+    # unit system instance
+    "galactic",
+    "dimensionless",
+    "solarsystem",
+    "si",
+    "cgs",
+    # unit system alias
+    "NAMED_UNIT_SYSTEMS",
+    # unit system flags
+    "AbstractUSysFlag",
+    "StandardUSysFlag",
+    "DynamicalSimUSysFlag",
+    # functions
+    "equivalent",
+)
+
+from .base import UNITSYSTEMS_REGISTRY, AbstractUnitSystem
+from .builtin import DimensionlessUnitSystem, LTMAUnitSystem
+from .compare import equivalent
+from .core import unitsystem, unitsystem_of
+from .flags import AbstractUSysFlag, DynamicalSimUSysFlag, StandardUSysFlag
+from .realizations import (
+    NAMED_UNIT_SYSTEMS,
+    cgs,
+    dimensionless,
+    galactic,
+    si,
+    solarsystem,
+)