unxt 0.1.dev463__py3-none-any.whl

unxt/_src/quantity/base.py

@@ -0,0 +1,1160 @@
+# pylint: disable=import-error, no-member, unsubscriptable-object
+#    b/c it doesn't understand dataclass fields
+
+__all__ = ("AbstractQuantity", "is_any_quantity")
+
+import functools as ft
+import warnings
+from collections.abc import Mapping
+from types import ModuleType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Literal,
+    NoReturn,
+    TypeAlias,
+    TypeGuard,
+    cast,
+)
+from typing_extensions import override
+
+import equinox as eqx
+import jax
+import jax.core
+import numpy as np
+import quax_blocks
+import wadler_lindig as wl
+from jax._src.numpy.array_methods import _IndexUpdateHelper, _IndexUpdateRef
+from jaxtyping import Array, ArrayLike, Bool, ScalarLike, Shaped
+from plum import add_promotion_rule, dispatch, type_nonparametric
+from quax import ArrayValue
+
+import quaxed.numpy as jnp
+from dataclassish import field_items, replace
+
+import unxt_api as uapi
+from .mixins import AstropyQuantityCompatMixin, IPythonReprMixin, NumPyCompatMixin
+from .value import StaticValue
+from unxt.units import AbstractUnit
+
+if TYPE_CHECKING:
+    from typing import Self
+
+
+ArrayLikeSequence: TypeAlias = list[ScalarLike] | tuple[ScalarLike, ...]
+
+
+class AbstractQuantity(
+    AstropyQuantityCompatMixin,
+    NumPyCompatMixin,
+    IPythonReprMixin,
+    ArrayValue,
+    quax_blocks.NumpyBinaryOpsMixin[Any, "AbstractQuantity"],
+    quax_blocks.NumpyComparisonMixin[Any, Bool[Array, "*shape"]],  # TODO: shape hint
+    quax_blocks.NumpyUnaryMixin["AbstractQuantity"],
+    quax_blocks.NumpyRoundMixin["AbstractQuantity"],
+    quax_blocks.NumpyTruncMixin["AbstractQuantity"],
+    quax_blocks.NumpyFloorMixin["AbstractQuantity"],
+    quax_blocks.NumpyCeilMixin["AbstractQuantity"],
+    quax_blocks.LaxLenMixin,
+    quax_blocks.LaxLengthHintMixin,
+):
+    """Represents a Quantity with a unit.
+
+    Attributes
+    ----------
+    short_name : str | None
+        Optional short name for the class used in wadler-lindig printing when
+        ``use_short_name=True``. Defaults to ``None``.
+
+    Examples
+    --------
+    >>> import unxt as u
+
+    From an integer:
+
+    >>> u.Quantity(1, "m")
+    Quantity(Array(1, dtype=int32, ...), unit='m')
+
+    From a float:
+
+    >>> u.Quantity(1.0, "m")
+    Quantity(Array(1., dtype=float32, ...), unit='m')
+
+    From a list:
+
+    >>> u.Quantity([1, 2, 3], "m")
+    Quantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+    From a tuple:
+
+    >>> u.Quantity((1, 2, 3), "m")
+    Quantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+    From a `numpy.ndarray`:
+
+    >>> import numpy as np
+    >>> u.Quantity(np.array([1, 2, 3]), "m")
+    Quantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+    From a `jax.Array`:
+
+    >>> import jax.numpy as jnp
+    >>> u.Quantity(jnp.array([1, 2, 3]), "m")
+    Quantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+    The unit can also be given as a `astropy.units.Unit`:
+
+    >>> import astropy.units as apyu
+    >>> u.Quantity(1, apyu.m)
+    Quantity(Array(1, dtype=int32, ...), unit='m')
+
+    """
+
+    value: eqx.AbstractVar[Shaped[Array | StaticValue, "*shape"]]
+    """The value of the `AbstractQuantity`."""
+
+    unit: eqx.AbstractVar[AbstractUnit]
+    """The unit associated with this value."""
+
+    # ---------------------------------------------------------------
+    # Constructors
+
+    @classmethod
+    @dispatch.abstract
+    def from_(
+        cls: "type[AbstractQuantity]",
+        *args: Any,
+        **kwargs: Any,
+    ) -> "AbstractQuantity":
+        raise NotImplementedError  # pragma: no cover
+
+    # See below for additional constructors
+
+    # ===============================================================
+    # Quantity API
+
+    def uconvert(self, u: Any, /) -> "AbstractQuantity":
+        """Convert the quantity to the given units.
+
+        See Also
+        --------
+        `unxt.uconvert` : convert a quantity to a new unit.
+
+        Examples
+        --------
+        >>> import unxt as u
+
+        >>> q = u.Quantity(1, "m")
+        >>> q.uconvert("cm")
+        Quantity(Array(100., dtype=float32, ...), unit='cm')
+
+        """
+        return uapi.uconvert(u, self)
+
+    def ustrip(self, u: Any, /) -> Array:
+        """Return the value in the given units.
+
+        See Also
+        --------
+        `unxt.ustrip` : strip the units from a quantity.
+
+        Examples
+        --------
+        >>> import unxt as u
+
+        >>> q = u.Quantity(1, "m")
+        >>> q.ustrip("cm")
+        Array(100., dtype=float32, weak_type=True)
+
+        """
+        return uapi.ustrip(u, self)
+
+    # ===============================================================
+    # Quax API
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Shape of the array."""
+        return self.value.shape
+
+    def materialise(self) -> NoReturn:
+        msg = f"Refusing to materialise `{type(self).__name__}`."
+        raise RuntimeError(msg)
+
+    def aval(self) -> jax.core.ShapedArray:
+        value = self.value
+        if isinstance(value, StaticValue):
+            value = jnp.asarray(value.array)
+        return jax.typeof(value)
+
+    def enable_materialise(self, _: bool = True) -> "Self":  # noqa: FBT001, FBT002
+        return replace(self, value=self.value, unit=self.unit)
+
+    # ===============================================================
+    # Plum API
+
+    #: This tells `plum` that this type can be efficiently cached.
+    __faithful__: ClassVar[bool] = True
+
+    # ===============================================================
+    # Array API
+
+    def __array_namespace__(self, *, api_version: Any = None) -> ModuleType:
+        """Return the namespace for the array API.
+
+        Here we return the `quaxed.numpy` module, which is a drop-in replacement
+        for `jax.numpy`, but allows for array-ish objects to be used in place of
+        `jax` arrays. See `quax` for more information.
+
+        """
+        return jnp  # quaxed.numpy
+
+    # ---------------------------------------------------------------
+    # attributes
+
+    @property
+    def dtype(self) -> jax.numpy.dtype:
+        """Data type of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> u.Quantity(1, "m").dtype
+        dtype('int32')
+
+        """
+        return self.value.dtype
+
+    @property
+    def device(self) -> jax.Device:
+        """Device where the array is located.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> u.Quantity(1, "m").device
+        CpuDevice(id=0)
+
+        """
+        return self.value.devices().pop()
+
+    @property
+    def mT(self) -> "AbstractQuantity":  # noqa: N802
+        """Matrix transpose of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([[0, 1], [1, 2]], "m")
+        >>> q.mT
+        Quantity(Array([[0, 1],
+                                  [1, 2]], dtype=int32), unit='m')
+
+        """
+        return replace(self, value=jnp.matrix_transpose(self.value))
+
+    @property
+    def ndim(self) -> int:
+        """Number of dimensions.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([[1]], "m")
+        >>> q.ndim
+        2
+
+        """
+        return self.value.ndim
+
+    @property
+    def size(self) -> int:
+        """Total number of elements.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.size
+        3
+
+        """
+        return self.value.size
+
+    @property
+    def T(self) -> "AbstractQuantity":  # noqa: N802
+        """Transpose of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([[0, 1], [1, 2]], "m")
+        >>> q.T
+        Quantity(Array([[0, 1],
+                                  [1, 2]], dtype=int32), unit='m')
+
+        """
+        return replace(self, value=self.value.T)
+
+    # required to override mixin methods
+    __eq__ = quax_blocks.NumpyEqMixin.__eq__  # type: ignore[assignment,unused-ignore]
+
+    # ---------------------------------------------------------------
+    # methods
+
+    def __bool__(self) -> bool:
+        """Convert a zero-dimensional array to a Python bool object.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([0, 1], "m")
+
+        >>> bool(q[0])
+        False
+
+        >>> bool(q[1])
+        True
+
+        """
+        return bool(self.value)
+
+    def __complex__(self) -> complex:
+        """Convert the array to a Python complex object.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "")
+        >>> complex(q)
+        (1+0j)
+
+        """
+        return complex(uapi.ustrip("", self))
+
+    def __dlpack__(self, *args: Any, **kwargs: Any) -> Any:
+        """Export the array for consumption as a DLPack capsule."""
+        raise NotImplementedError
+
+    def __dlpack_device__(self, *args: Any, **kwargs: Any) -> Any:
+        """Return device type and device ID in DLPack format."""
+        raise NotImplementedError
+
+    def __float__(self) -> float:
+        """Convert the array to a Python float object.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "")
+        >>> float(q)
+        1.0
+
+        """
+        return float(uapi.ustrip("", self))
+
+    def __getitem__(self, key: Any) -> "AbstractQuantity":
+        """Get an item from the array.
+
+        This is a simple wrapper around the `__getitem__` method of the array,
+        calling `replace` to only update the value.
+
+        """
+        return replace(self, value=self.value[key])
+
+    def __index__(self) -> int:
+        """Convert a zero-dimensional integer array to a Python int object.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "")
+        >>> q.__index__()
+        1
+
+        """
+        return uapi.ustrip("", self).__index__()
+
+    def __int__(self) -> int:
+        """Convert a zero-dimensional array to a Python int object.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "")
+        >>> int(q)
+        1
+
+        """
+        return int(uapi.ustrip("", self))
+
+    def __setitem__(self, key: Any, value: Any) -> None:
+        """Set an item in the array.
+
+        This is a simple wrapper around the `__setitem__` method of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> try:
+        ...     q[0] = 2
+        ... except Exception as e:
+        ...     print("jax arrays do not support in-place item assignment")
+        jax arrays do not support in-place item assignment
+
+        """
+        self.value[key] = value
+
+    def to_device(self, device: None | jax.Device = None) -> "AbstractQuantity":
+        """Move the array to a new device.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "m")
+        >>> q.to_device(None)
+        Quantity(Array(1, dtype=int32, weak_type=True), unit='m')
+
+        """
+        return replace(self, value=self.value.to_device(device))
+
+    # ===============================================================
+    # JAX API
+
+    def __iter__(self) -> Any:
+        """Iterate over the Quantity's value.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> [x for x in q]
+        [Quantity(Array(1, dtype=int32), unit='m'),
+         Quantity(Array(2, dtype=int32), unit='m'),
+         Quantity(Array(3, dtype=int32), unit='m')]
+
+        """
+        yield from (self[i] for i in range(len(self.value)))
+
+    def argmax(self, *args: Any, **kwargs: Any) -> Array:
+        """Return the indices of the maximum value.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.argmax()
+        Array(2, dtype=int32)
+
+        """
+        return self.value.argmax(*args, **kwargs)
+
+    def argmin(self, *args: Any, **kwargs: Any) -> Array:
+        """Return the indices of the minimum value.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.argmin()
+        Array(0, dtype=int32)
+
+        """
+        return self.value.argmin(*args, **kwargs)
+
+    def astype(self, *args: Any, **kwargs: Any) -> "AbstractQuantity":
+        """Copy the array and cast to a specified dtype.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.dtype
+        dtype('int32')
+
+        >>> q.astype(float)
+        Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+        """
+        return replace(self, value=self.value.astype(*args, **kwargs))
+
+    @ft.partial(property, doc=jax.Array.at.__doc__)
+    def at(self) -> "_QuantityIndexUpdateHelper":
+        return _QuantityIndexUpdateHelper(self)  # type: ignore[arg-type]
+
+    def block_until_ready(self) -> "AbstractQuantity":
+        """Block until the array is ready.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "m")
+        >>> q.block_until_ready() is q
+        True
+
+        """
+        _ = self.value.block_until_ready()
+        return self
+
+    def devices(self) -> set[jax.Device]:
+        """Return the devices where the array is located.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "m")
+        >>> q.devices()
+        {CpuDevice(id=0)}
+
+        """
+        return self.value.devices()
+
+    def flatten(self) -> "AbstractQuantity":
+        """Return a flattened version of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([[1, 2], [3, 4]], "m")
+        >>> q.flatten()
+        Quantity(Array([1, 2, 3, 4], dtype=int32), unit='m')
+
+        """
+        return replace(self, value=self.value.flatten())
+
+    def max(self, *args: Any, **kwargs: Any) -> "AbstractQuantity":
+        """Return the maximum value.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.max()
+        Quantity(Array(3, dtype=int32), unit='m')
+
+        """
+        return replace(self, value=self.value.max(*args, **kwargs))
+
+    def mean(self, *args: Any, **kwargs: Any) -> "AbstractQuantity":
+        """Return the mean value.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.mean()
+        Quantity(Array(2., dtype=float32), unit='m')
+
+        """
+        return replace(self, value=self.value.mean(*args, **kwargs))
+
+    def min(self, *args: Any, **kwargs: Any) -> "AbstractQuantity":
+        """Return the minimum value.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.min()
+        Quantity(Array(1, dtype=int32), unit='m')
+
+        """
+        return replace(self, value=self.value.min(*args, **kwargs))
+
+    def ravel(self) -> "AbstractQuantity":
+        """Return a flattened version of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([[1, 2], [3, 4]], "m")
+        >>> q.ravel()
+        Quantity(Array([1, 2, 3, 4], dtype=int32), unit='m')
+
+        """
+        return replace(self, value=self.value.ravel())
+
+    def reshape(self, *args: Any, order: str = "C") -> "AbstractQuantity":
+        """Return a reshaped version of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3, 4], "m")
+        >>> q.reshape(2, 2)
+        Quantity(Array([[1, 2],
+                                  [3, 4]], dtype=int32), unit='m')
+
+        """
+        __tracebackhide__ = True  # pylint: disable=unused-variable
+        return replace(self, value=self.value.reshape(*args, order=order))
+
+    def round(self, *args: Any, **kwargs: Any) -> "AbstractQuantity":
+        """Round the array to the given number of decimals.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1.1, 2.2, 3.3], "m")
+        >>> q.round(0)
+        Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+        """
+        return replace(self, value=self.value.round(*args, **kwargs))
+
+    @property
+    def sharding(self) -> Any:
+        """Return the sharding configuration of the array.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3], "m")
+        >>> q.sharding
+        SingleDeviceSharding(device=..., memory_kind=...)
+
+        """
+        return self.value.sharding
+
+    def squeeze(self, *args: Any, **kwargs: Any) -> "AbstractQuantity":
+        """Return the array with all single-dimensional entries removed.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([[[1], [2], [3]]], "m")
+        >>> q.squeeze()
+        Quantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+        """
+        return replace(self, value=self.value.squeeze(*args, **kwargs))
+
+    # ===============================================================
+    # Python stuff
+
+    def __hash__(self) -> int:
+        """Return the hash of the quantity.
+
+        This is the hash of the value and the unit, however since the value is
+        generally not hashable this will generally raise an exception.
+        Defining the `__hash__` method is required for the `AbstractQuantity` to
+        be considered immutable, e.g. by `dataclasses.dataclass`.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity(1, "m")
+        >>> try:
+        ...     hash(q)
+        ... except TypeError as e:
+        ...     print(e)
+        unhashable type: 'jaxlib...ArrayImpl'
+
+        """
+        return hash((self.value, self.unit))
+
+    def __pdoc__(
+        self,
+        *,
+        named_unit: bool = True,
+        short_arrays: bool | Literal["compact"] = True,
+        use_short_name: bool = False,
+        **kwargs: Any,
+    ) -> wl.AbstractDoc:
+        """Return the Wadler-Lindig representation of this class.
+
+        This is used for the `__repr__` and `__str__` methods or when using the
+        `wadler_lindig` library.
+
+        Parameters
+        ----------
+        named_unit
+            If `True`, the unit is included in the representation as a named
+            argument. If `False`, the unit is included as a positional argument.
+            For example, ``Quantity(<array>, unit='m')`` versus
+            ``Quantity(<array>, 'm')``.
+        short_arrays
+            If 'compact', the array is printed in a compact form,
+            ``Quantity(Array([1., 2.], dtype=float32), unit='')`` to
+            ``Quantity([1., 2.], unit='')``. If `True` or `False` it uses the
+            default `wadler_lindig` behavior.
+        use_short_name
+            If `True` and the class has a ``short_name`` class variable, use the
+            short name instead of the full class name. For example, ``Q(...)``
+            instead of ``Quantity(...)``.
+        kwargs
+            Additional keyword arguments ``wadler_lindig.pdoc`` method for
+            formatting the value, stringified unit, and other fields.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> import wadler_lindig as wl
+
+        >>> q = u.quantity.BareQuantity([1, 2, 3], "m")
+
+        The default pretty printing:
+
+        >>> wl.pprint(q)
+        BareQuantity(i32[3], unit='m')
+
+        The `str` method uses this as well:
+
+        >>> print(q)
+        BareQuantity(i32[3], unit='m')
+
+        Arrays can be printed in full:
+
+        >>> wl.pprint(q, short_arrays=False)
+        BareQuantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+        The `repr` method uses this setting:
+
+        >>> print(repr(q))
+        BareQuantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+        The units can be turned from a named argument to a positional argument
+        by setting `named_unit=False`:
+
+        >>> wl.pprint(q, named_unit=False)
+        BareQuantity(i32[3], 'm')
+
+        The arrays can be printed in a compact but informative form by setting
+        ``short_arrays="compact"``:
+
+        >>> wl.pprint(q, short_arrays="compact")
+        BareQuantity([1, 2, 3], unit='m')
+
+        The class short name can be used if available:
+
+        >>> wl.pprint(q, use_short_name=True)
+        BareQuantity(i32[3], unit='m')
+
+        """
+        # Class Name
+        if use_short_name and getattr(self, "short_name", None) is not None:
+            cls_name = wl.TextDoc(self.short_name)
+        else:
+            cls_name = wl.TextDoc(type_nonparametric(self).__name__)
+
+        # Object fields
+        fs = dict(field_items(self))
+        del fs["value"]
+        del fs["unit"]
+
+        # Customize value representation
+        # (backward compatibility for `compact_arrays` argument)
+        # TODO: remove in v2.0
+        if "compact_arrays" in kwargs:
+            short_arrays = kwargs.pop("compact_arrays")
+            warnings.warn(
+                "`compact_arrays` argument is deprecated; use "
+                "`short_arrays='compact'` instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        match short_arrays:
+            case "compact":
+                kwargs["custom"] = custom_pdoc_noarray
+            case True:
+                kwargs["custom"] = custom_pdoc_no_kind
+                kwargs["short_arrays"] = True
+            case False:
+                kwargs["short_arrays"] = False
+
+        # Make the pdocs for the base and extra fields.
+        base_fields = [
+            wl.pdoc(self.value, **kwargs),
+            wl.TextDoc("unit=" if named_unit else "")
+            + wl.pdoc(self.unit.to_string(), **kwargs),
+        ]
+        extra_fields = wl.named_objs(tuple(fs.items()), **kwargs)
+
+        # Construct and return the Wadler-Lindig document.
+        return (
+            cls_name
+            + wl.TextDoc("(")
+            + wl.join(wl.comma, base_fields + extra_fields)
+            .group()
+            .nest(kwargs.get("indent", 4))
+            + wl.TextDoc(")")
+        )
+
+    def __repr__(self) -> str:
+        # TODO: make named_unit False?
+        return wl.pformat(self, short_arrays=False, named_unit=True, indent=4)
+
+    def __str__(self) -> str:
+        # TODO: make named_unit False?
+        return wl.pformat(self, short_arrays=True, named_unit=True, indent=4)
+
+
+# -----------------------------------------------
+# Register additional constructors
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity],
+    value: ArrayLike | ArrayLikeSequence,
+    unit: Any,
+    /,
+    *,
+    dtype: Any = None,
+) -> AbstractQuantity:
+    """Construct a `unxt.Quantity` from an array-like value and a unit.
+
+    :param value: The array-like value.
+    :param unit: The unit of the value.
+    :param dtype: The data type of the array (keyword-only).
+
+    Examples
+    --------
+    For this example we'll use the `Quantity` class. The same applies to
+    any subclass of `AbstractQuantity`.
+
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+
+    >>> x = jnp.array([1.0, 2, 3])
+    >>> u.Quantity.from_(x, "m")
+    Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+    >>> u.Quantity.from_([1.0, 2, 3], "m")
+    Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+    >>> u.Quantity.from_((1.0, 2, 3), "m")
+    Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+    """
+    # Dispatch on both arguments.
+    # Construct using the standard `__init__` method.
+    return cls(jnp.asarray(value, dtype=dtype), unit)
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity],
+    value: ArrayLike | ArrayLikeSequence,
+    /,
+    *,
+    unit: Any,
+    dtype: Any = None,
+) -> AbstractQuantity:
+    """Make a `unxt.AbstractQuantity` from an array-like value and a unit kwarg.
+
+    Examples
+    --------
+    For this example we'll use the `unxt.Quantity` class. The same applies
+    to any subclass of `unxt.AbstractQuantity`.
+
+    >>> import unxt as u
+    >>> u.Quantity.from_([1.0, 2, 3], unit="m")
+    Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+    """
+    # Dispatch on the `value` only. Dispatch to the full constructor.
+    return cls.from_(value, unit, dtype=dtype)
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity], *, value: Any, unit: Any, dtype: Any = None
+) -> AbstractQuantity:
+    """Construct a `AbstractQuantity` from value and unit kwargs.
+
+    Examples
+    --------
+    For this example we'll use the `Quantity` class. The same applies to
+    any subclass of `AbstractQuantity`.
+
+    >>> import unxt as u
+    >>> u.Quantity.from_(value=[1.0, 2, 3], unit="m")
+    Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+    """
+    # Dispatched on no argument. Dispatch to the full constructor.
+    return cls.from_(value, unit, dtype=dtype)
+
+
+@AbstractQuantity.from_.dispatch
+def from_(cls: type[AbstractQuantity], mapping: Mapping[str, Any]) -> AbstractQuantity:
+    """Construct a `Quantity` from a Mapping.
+
+    Examples
+    --------
+    For this example we'll use the `Quantity` class. The same applies to
+    any subclass of `AbstractQuantity`.
+
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+
+    >>> x = jnp.array([1.0, 2, 3])
+    >>> q = u.Quantity.from_({"value": x, "unit": "m"})
+    >>> q
+    Quantity(Array([1., 2., 3.], dtype=float32), unit='m')
+
+    >>> u.Quantity.from_({"value": q, "unit": "km"})
+    Quantity(Array([0.001, 0.002, 0.003], dtype=float32), unit='km')
+
+    """
+    # Dispatch on both arguments.
+    # Construct using the standard `__init__` method.
+    return cls.from_(**mapping)
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity],
+    value: AbstractQuantity,
+    unit: Any,
+    /,
+    *,
+    dtype: Any = None,
+) -> AbstractQuantity:
+    """Construct a `Quantity` from another `Quantity`.
+
+    The `value` is converted to the new `unit`.
+
+    Examples
+    --------
+    >>> import unxt as u
+
+    >>> q = u.Quantity(1, "m")
+    >>> u.Quantity.from_(q, "cm")
+    Quantity(Array(100., dtype=float32, ...), unit='cm')
+
+    """
+    value = jnp.asarray(uapi.uconvert(unit, value), dtype=dtype)
+    return cls(uapi.ustrip(unit, value), unit)
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity],
+    value: AbstractQuantity,
+    unit: None,
+    /,
+    *,
+    dtype: Any = None,
+) -> AbstractQuantity:
+    """Construct a `Quantity` from another `Quantity`.
+
+    The `value` is converted to the new `unit`.
+
+    Examples
+    --------
+    >>> import unxt as u
+
+    >>> q = u.Quantity(1, "m")
+    >>> u.Quantity.from_(q, None)
+    Quantity(Array(1, dtype=int32, ...), unit='m')
+
+    """
+    value = jnp.asarray(value, dtype=dtype)
+    unit = uapi.unit_of(value)
+    return cls(uapi.ustrip(unit, value), unit)
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity],
+    value: AbstractQuantity,
+    /,
+    *,
+    unit: Any | None = None,
+    dtype: Any = None,
+) -> AbstractQuantity:
+    """Construct a `Quantity` from another `Quantity`, with no unit change."""
+    unit = value.unit if unit is None else unit
+    value = jnp.asarray(uapi.uconvert(unit, value), dtype=dtype)
+    return cls(uapi.ustrip(unit, value), unit)
+
+
+# -----------------------------------------------
+# Promotion rules
+
+add_promotion_rule(AbstractQuantity, AbstractQuantity, AbstractQuantity)
+
+
+# ===============================================================
+# Support for ``at``.
+
+
+# `_QuantityIndexUpdateHelper` is defined up here because it is used in the
+# runtime-checkable type annotation in `AbstractQuantity.at`.
+# `_QuantityIndexUpdateRef` is defined after `AbstractQuantity` because it
+# references `AbstractQuantity` in its runtime-checkable type annotations.
+class _QuantityIndexUpdateHelper(_IndexUpdateHelper):
+    def __getitem__(self, index: Any) -> "_IndexUpdateRef":
+        return _QuantityIndexUpdateRef(self.array, index)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the object.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> q = u.Quantity([1, 2, 3, 4], "m")
+        >>> q.at
+        _QuantityIndexUpdateHelper(Quantity(Array([1, 2, 3, 4], dtype=int32), unit='m'))
+
+        """
+        return f"_QuantityIndexUpdateHelper({self.array!r})"
+
+
+class _QuantityIndexUpdateRef(_IndexUpdateRef):
+    # This is a subclass of `_IndexUpdateRef` that is used to implement the `at`
+    # attribute of `AbstractQuantity`. See also `_QuantityIndexUpdateHelper`.
+
+    array: AbstractQuantity  # type: ignore[assignment]
+
+    def __repr__(self) -> str:
+        return super().__repr__().replace("_IndexUpdateRef", "_QuantityIndexUpdateRef")
+
+    @override
+    def get(  # type: ignore[override]
+        self, *, fill_value: AbstractQuantity | None = None, **kw: Any
+    ) -> AbstractQuantity:
+        # TODO: by quaxified super
+        value = self.array.value.at[self.index].get(
+            fill_value=(
+                fill_value
+                if fill_value is None
+                else uapi.ustrip(self.array.unit, fill_value)
+            ),
+            **kw,
+        )
+        return replace(self.array, value=value)
+
+    @override
+    def set(self, values: AbstractQuantity, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        # TODO: by quaxified super
+        value = self.array.value.at[self.index].set(
+            uapi.ustrip(self.array.unit, values), **kw
+        )
+        return replace(self.array, value=value)
+
+    @override
+    def apply(self, func: Any, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        raise NotImplementedError  # TODO: by quaxified super
+
+    @override
+    def add(self, values: AbstractQuantity, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        # TODO: by quaxified super
+        value = self.array.value.at[self.index].add(
+            uapi.ustrip(self.array.unit, values), **kw
+        )
+        return replace(self.array, value=value)
+
+    @override
+    def multiply(self, values: ArrayLike, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        values = eqx.error_if(
+            values, isinstance(values, AbstractQuantity), "values cannot be a Quantity"
+        )  # TODO: can permit dimensionless quantities.
+
+        # TODO: by quaxified super
+        value = self.array.value.at[self.index].multiply(values, **kw)
+        return replace(self.array, value=value)
+
+    mul = multiply  # type: ignore[assignment]
+
+    @override
+    def divide(self, values: ArrayLike, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        values = eqx.error_if(
+            values, isinstance(values, AbstractQuantity), "values cannot be a Quantity"
+        )  # TODO: can permit dimensionless quantities.
+
+        # TODO: by quaxified super
+        value = self.array.value.at[self.index].divide(values, **kw)
+        return replace(self.array, value=value)
+
+    @override
+    def power(self, values: ArrayLike, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        raise NotImplementedError
+
+    @override
+    def min(self, values: AbstractQuantity, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        # TODO: by quaxified super
+        value = self.array.value.at[self.index].min(
+            uapi.ustrip(self.array.unit, values), **kw
+        )
+        return replace(self.array, value=value)
+
+    @override
+    def max(self, values: AbstractQuantity, **kw: Any) -> AbstractQuantity:  # type: ignore[override]
+        # TODO: by quaxified super
+        value = self.array.value.at[self.index].max(
+            uapi.ustrip(self.array.unit, values), **kw
+        )
+        return replace(self.array, value=value)
+
+
+# This is public!
+def is_any_quantity(obj: Any, /) -> TypeGuard[AbstractQuantity]:
+    """Check if an object is an instance of `unxt.quantity.AbstractQuantity`.
+
+    Examples
+    --------
+    >>> import unxt as u
+
+    >>> q = u.Quantity(1, "m")
+    >>> is_any_quantity(q)
+    True
+
+    """
+    return isinstance(obj, AbstractQuantity)
+
+
+# TODO: replace with `equinox.internal.TreeWLCustom` when available.
+def custom_pdoc_no_kind(obj: Any) -> wl.AbstractDoc | None:
+    """Return custom pdoc for ``AbstractQuantity`` objects."""
+    if isinstance(obj, jax.Array):
+        dtype = obj.dtype.name
+        # Added in JAX 0.4.32 to `ShapeDtypeStruct`
+        if getattr(obj, "weak_type", False):
+            dtype = f"weak_{dtype}"
+        return wl.array_summary(obj.shape, dtype, kind=None)
+    return None
+
+
+def custom_pdoc_noarray(obj: Any) -> wl.AbstractDoc | None:
+    """Return custom pdoc for ``AbstractQuantity`` objects."""
+    if isinstance(obj, jax.Array):
+        return wl.TextDoc(np.array2string(cast("np.ndarray", obj), separator=", "))
+    return None
+
+
+@StaticValue.from_.dispatch
+def from_(cls: StaticValue, value: AbstractQuantity, /) -> StaticValue:
+    """Disallow conversion of `AbstractQuantity` to a value."""
+    msg = (
+        f"Cannot convert '{type(value).__name__}' to a value. "
+        "For a Quantity, use the `.from_` constructor instead."
+    )
+    raise TypeError(msg)
+
+
+@dispatch
+def convert_to_quantity_value(obj: AbstractQuantity, /) -> NoReturn:
+    """Disallow conversion of `AbstractQuantity` to a value.
+
+    >>> import unxt as u
+    >>> from unxt.quantity import convert_to_quantity_value
+
+    >>> try:
+    ...     convert_to_quantity_value(u.Quantity(1, "m"))
+    ... except TypeError as e:
+    ...     print(e)
+    Cannot convert 'Quantity[PhysicalType('length')]' to a value.
+    For a Quantity, use the `.from_` constructor instead.
+
+    """
+    msg = (
+        f"Cannot convert '{type(obj).__name__}' to a value. "
+        "For a Quantity, use the `.from_` constructor instead."
+    )
+    raise TypeError(msg)