unxt 0.1.dev463__py3-none-any.whl

unxt/_src/quantity/mixins.py

@@ -0,0 +1,238 @@
+"""Orthogonal Mixin classes for Quantity classes."""
+
+__all__: tuple[str, ...] = ()
+
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any, cast
+
+import equinox as eqx
+import numpy as np
+import numpy.typing as npt
+from astropy.units import CompositeUnit
+from jax.typing import ArrayLike
+from jaxtyping import Array
+
+from dataclassish import replace
+
+import unxt_api as uapi
+from unxt.units import AbstractUnit, unit as parse_unit
+
+if TYPE_CHECKING:
+    import unxt.quantity
+
+
+class AstropyQuantityCompatMixin:
+    """Mixin for compatibility with `astropy.units.Quantity`."""
+
+    value: eqx.AbstractVar[ArrayLike]
+    unit: eqx.AbstractVar[AbstractUnit]
+    uconvert: Callable[[Any], "unxt.quantity.AbstractQuantity"]
+    ustrip: Callable[[Any], ArrayLike]
+
+    def to(self, u: Any, /) -> "unxt.quantity.AbstractQuantity":
+        """Convert the quantity to the given units.
+
+        See `unxt.quantity.AbstractQuantity.uconvert`.
+
+        Examples
+        --------
+        >>> from unxt import Quantity
+
+        >>> q = Quantity(1, "m")
+        >>> q.to("cm")
+        Quantity(Array(100., dtype=float32, ...), unit='cm')
+
+        """
+        return uapi.uconvert(u, self)  # redirect to the standard method
+
+    def to_value(self, u: Any, /) -> ArrayLike:
+        """Return the value in the given units.
+
+        See `unxt.AbstractQuantity.ustrip`.
+
+        Examples
+        --------
+        >>> from unxt import Quantity
+
+        >>> q = Quantity(1, "m")
+        >>> q.to_value("cm")
+        Array(100., dtype=float32, weak_type=True)
+
+        """
+        return uapi.ustrip(u, self)  # redirect to the standard method
+
+    def decompose(
+        self, bases: Sequence[AbstractUnit | str], /
+    ) -> "unxt.quantity.AbstractQuantity":
+        """Decompose the quantity into the given bases.
+
+        Examples
+        --------
+        >>> from unxt import Quantity
+
+        >>> q = Quantity(1, "m")
+        >>> q.decompose(["cm", "s"])
+        Quantity(Array(100., dtype=float32, ...), unit='cm')
+
+        """
+        bases_ = [parse_unit(b) for b in bases]
+        du = self.unit.decompose(bases_)  # decomposed units
+        base_units = CompositeUnit(scale=1, bases=du.bases, powers=du.powers)
+        return replace(self, value=self.value * du.scale, unit=base_units)
+
+
+#####################################################################
+
+
+SUPPORTED_IPYTHON_REPR_FORMATS: dict[str, str] = {
+    "text/plain": "__repr__",
+    "text/html": "_repr_html_",
+    "text/latex": "_repr_latex_",
+}
+
+
+class IPythonReprMixin:
+    """Mixin class for IPython representation of a Quantity."""
+
+    value: Array
+    unit: AbstractUnit
+
+    def _repr_mimebundle_(
+        self,
+        *,
+        include: Sequence[str] | None = None,
+        exclude: Sequence[str] | None = None,
+    ) -> dict[str, str]:
+        r"""Return a MIME bundle representation of the Quantity.
+
+        :param include: The set of keys to include in the MIME bundle. If not
+            provided, all supported formats are included.
+        :param exclude: The set of keys to exclude in the MIME bundle. If not
+            provided, all supported formats are included. 'include' has
+            precedence over 'exclude'.
+
+        Examples
+        --------
+        >>> from unxt import Quantity
+
+        >>> q = Quantity([1.0, 2, 3, 4], "m")
+        >>> q._repr_mimebundle_()
+        {'text/plain': "Quantity(Array([1., 2., 3., 4.], dtype=float32), unit='m')",
+         'text/html': '<span>[1., 2., 3., 4.]</span> * <span>Unit("m")</span>',
+         'text/latex': '$[1.,~2.,~3.,~4.] \\; \\mathrm{m}$'}
+
+        >>> q._repr_mimebundle_(include=["text/plain"])
+        {'text/plain': "Quantity(Array([1., 2., 3., 4.], dtype=float32), unit='m')"}
+
+        >>> q._repr_mimebundle_(exclude=["text/html", "text/latex"])
+        {'text/plain': "Quantity(Array([1., 2., 3., 4.], dtype=float32), unit='m')"}
+
+        """
+        # Determine the set of keys to include in the MIME bundle
+        keys: Sequence[str]
+        if include is None and exclude is None:
+            keys = tuple(SUPPORTED_IPYTHON_REPR_FORMATS)
+        elif include is not None:
+            keys = [key for key in include if key in SUPPORTED_IPYTHON_REPR_FORMATS]
+        else:
+            keys = [
+                k
+                for k in SUPPORTED_IPYTHON_REPR_FORMATS
+                if k not in cast("str", exclude)
+            ]
+
+        # Create the MIME bundle
+        return {
+            key: getattr(self, SUPPORTED_IPYTHON_REPR_FORMATS[key])() for key in keys
+        }
+
+    def _repr_html_(self) -> str:
+        """Return an HTML representation of the Quantity.
+
+        Examples
+        --------
+        >>> from unxt import Quantity
+
+        >>> q = Quantity([1.0, 2, 3, 4], "m")
+        >>> q._repr_html_()
+        '<span>[1., 2., 3., 4.]</span> * <span>Unit("m")</span>'
+
+        """
+        unit_repr = getattr(self.unit, "_repr_html_", self.unit.__repr__)()
+        value_repr = np.array2string(self.value, separator=", ")  # type: ignore[call-overload]
+
+        return f"<span>{value_repr}</span> * <span>{unit_repr}</span>"
+
+    def _repr_latex_(self) -> str:
+        r"""Return a LaTeX representation of the Quantity.
+
+        Examples
+        --------
+        >>> from unxt import Quantity
+
+        >>> q = Quantity([1.0, 2, 3, 4], "m")
+        >>> q._repr_latex_()
+        '$[1.,~2.,~3.,~4.] \\; \\mathrm{m}$'
+
+        """
+        unit_repr = getattr(self.unit, "_repr_latex_", self.unit.__repr__)()
+        value_repr = np.array2string(self.value, separator=",~")  # type: ignore[call-overload]
+        return f"${value_repr} \\; {unit_repr[1:-1]}$"
+
+    # TODO: implement:
+    # - _repr_markdown_
+    # - _repr_json_
+
+
+#####################################################################
+
+
+class NumPyCompatMixin:
+    """Mixin for compatibility with numpy arrays."""
+
+    unit: AbstractUnit
+
+    __array_namespace__: Callable[[], Any]
+
+    def __array__(self, *args: object, **kw: object) -> npt.NDArray[Any]:
+        """Return the array as a numpy array, stripping the units.
+
+        Examples
+        --------
+        >>> from unxt import Quantity
+        >>> import numpy as np
+
+        >>> q = Quantity(1.01, "m")
+        >>> np.array(q)
+        array(1.01, dtype=float32)
+
+        """
+        return np.asarray(uapi.ustrip(self.unit, self), *args, **kw)
+
+    # TODO: why doesn't `__array_namespace__` supersede this?
+    def __array_function__(
+        self,
+        func: Callable[..., Any],
+        types: Any,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        """Dispatch to the corresponding jax.numpy function.
+
+        Examples
+        --------
+        >>> import numpy as np
+        >>> import unxt as u
+
+        >>> q = u.Quantity([1.0, 2, 3, 4], "m")
+        >>> np.sum(q)
+        Quantity(Array(10., dtype=float32), unit='m')
+
+        >>> np.stack([q, q])
+        Quantity(Array([[1., 2., 3., 4.],
+                                  [1., 2., 3., 4.]], dtype=float32), unit='m')
+
+        """
+        xp = self.__array_namespace__()
+        xfunc = getattr(xp, func.__name__)
+        return xfunc(*args, **kwargs)

unxt/_src/quantity/quantity.py

@@ -0,0 +1,123 @@
+# pylint: disable=import-error, no-member, unsubscriptable-object
+#    b/c it doesn't understand dataclass fields
+
+__all__ = ("Quantity", "Q")
+
+from typing import ClassVar, final
+
+import equinox as eqx
+from jaxtyping import Array, Shaped
+from plum import parametric
+
+from .base_parametric import AbstractParametricQuantity
+from .value import StaticValue, convert_to_quantity_value
+from unxt.units import AbstractUnit, unit as parse_unit
+
+
+@final
+@parametric
+class Quantity(AbstractParametricQuantity):
+    """Arrays with associated units.
+
+    This class is parametrized by the dimensions of the units.
+
+    Attributes
+    ----------
+    short_name : str
+        Short name 'Q' used for compact wadler-lindig printing.
+
+    Examples
+    --------
+    >>> import unxt as u
+
+    From an integer:
+
+    >>> u.Quantity(1, "m")
+    Quantity(Array(1, dtype=int32, ...), unit='m')
+
+    From a float:
+
+    >>> u.Q(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 units object:
+
+    >>> u.Quantity(1, u.unit("m"))
+    Quantity(Array(1, dtype=int32, ...), unit='m')
+
+    In the previous examples, the dimension parameter was inferred from the
+    values. It can also be given explicitly:
+
+    >>> u.Quantity["length"](1, "m")
+    Quantity(Array(1, dtype=int32, ...), unit='m')
+
+    This can be used for runtime checking of the input dimension!
+
+    >>> try:
+    ...     u.Quantity["length"](1, "s")
+    ... except Exception as e:
+    ...     print(e)
+    Physical type mismatch.
+
+    The dimension can also be given as a dimension object:
+
+    >>> dims = u.dimension("length")
+    >>> dims
+    PhysicalType('length')
+    >>> u.Quantity[dims](1.0, "m")
+    Quantity(Array(1., dtype=float32, ...), unit='m')
+
+    Or as a unit:
+
+    >>> u.Quantity[u.unit("m")](1.0, "m")
+    Quantity(Array(1., dtype=float32, ...), unit='m')
+
+    Some tricky cases are when the physical type is unknown:
+
+    >>> unit = u.unit("m2 / (kg s2)")
+    >>> u.dimension_of(unit)
+    PhysicalType('unknown')
+
+    The dimension can be given as a string in all cases, but is necessary when
+    the physical type is unknown:
+
+    >>> print(u.Quantity["m2 kg-1 s-2"](1.0, unit))  # to show the [dim]
+    Quantity['m2 kg-1 s-2'](1., unit='m2 / (kg s2)')
+
+    """
+
+    short_name: ClassVar[str] = "Q"
+    """Short name for compact printing."""
+
+    value: Shaped[Array, "*shape"] | StaticValue = 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."""
+
+
+Q = Quantity  # convenience alias