unxt 0.1.dev463__py3-none-any.whl

unxt/_src/quantity/base_angle.py

@@ -0,0 +1,89 @@
+"""Base classes for angular quantities."""
+
+__all__ = ("AbstractAngle",)
+
+import equinox as eqx
+from jaxtyping import Array, Shaped
+from plum import add_promotion_rule
+
+import unxt_api as uapi
+from .base import AbstractQuantity
+from .quantity import Quantity
+from .unchecked import BareQuantity
+from unxt._src.dimensions import dimension, dimension_of
+from unxt._src.quantity.value import StaticValue
+from unxt.units import AbstractUnit
+
+angle_dimension = dimension("angle")
+
+
+class AbstractAngle(AbstractQuantity):
+    """Angular Quantity.
+
+    See Also
+    --------
+    `unxt.Angle` : a concrete implementation of this class.
+
+    Examples
+    --------
+    For this example, we will use the concrete implementation of
+    `unxt.AbstractAngle`, `unxt.Angle`.
+
+    >>> from unxt import Angle
+
+    >>> Angle(90, "deg")
+    Angle(Array(90, dtype=int32, ...), unit='deg')
+
+    Angles have to have dimensions of angle.
+
+    >>> try:
+    ...     Angle(90, "m")
+    ... except ValueError as e:
+    ...     print(e)
+    Angle must have units with angular dimensions.
+
+    """
+
+    value: eqx.AbstractVar[Shaped[Array | StaticValue, "*shape"]]
+    """The value of the `unxt.AbstractQuantity`."""
+
+    unit: eqx.AbstractVar[AbstractUnit]
+    """The unit associated with this value."""
+
+    def __check_init__(self) -> None:
+        """Check the initialization."""
+        if dimension_of(self) != angle_dimension:
+            msg = f"{type(self).__name__} must have units with angular dimensions."
+            raise ValueError(msg)
+
+    def wrap_to(
+        self, /, min: AbstractQuantity, max: AbstractQuantity
+    ) -> "AbstractAngle":
+        """Wrap the angle to the range [min, max).
+
+        Parameters
+        ----------
+        min, max
+            The minimum, maximum value of the range.
+
+        See Also
+        --------
+        `unxt.quantity.wrap_to` : functional version of this method.
+
+        Examples
+        --------
+        >>> import unxt as u
+        >>> angle = u.Angle(370, "deg")
+        >>> angle.wrap_to(min=u.Q(0, "deg"), max=u.Q(360, "deg"))
+        Angle(Array(10, dtype=int32, ...), unit='deg')
+
+        """
+        return uapi.wrap_to(self, min, max)
+
+
+# Add a rule that when a AbstractAngle interacts with a Quantity, the
+# angle degrades to a Quantity. This is necessary for many operations, e.g.
+# division of an angle by non-dimensionless quantity where the resulting units
+# are not those of an angle.
+add_promotion_rule(AbstractAngle, Quantity, Quantity)
+add_promotion_rule(AbstractAngle, BareQuantity, BareQuantity)

unxt/_src/quantity/base_parametric.py

@@ -0,0 +1,290 @@
+# pylint: disable=import-error, no-member, unsubscriptable-object
+#    b/c it doesn't understand dataclass fields
+
+__all__ = ("AbstractParametricQuantity",)
+
+from collections.abc import Callable
+from functools import partial
+from typing import Any
+
+import equinox as eqx
+import jax
+import wadler_lindig as wl
+from astropy.units import PhysicalType
+from jaxtyping import Array, Shaped
+from plum import dispatch, parametric, type_unparametrized
+
+from dataclassish import field_items, fields
+
+from .base import AbstractQuantity
+from .value import StaticValue
+from unxt._src.dimensions import name_of
+from unxt.dims import AbstractDimension, dimension, dimension_of
+from unxt.units import AbstractUnit, unit as parse_unit
+
+
+@parametric
+class AbstractParametricQuantity(AbstractQuantity):
+    """Arrays with associated units.
+
+    This class is parametrized by the dimensions of the units.
+
+    """
+
+    value: eqx.AbstractVar[Shaped[Array | StaticValue, "*shape"]]
+    """The value of the `AbstractQuantity`."""
+
+    unit: eqx.AbstractVar[AbstractUnit]
+    """The unit associated with this value."""
+
+    def __post_init__(self) -> None:
+        """Check whether the arguments are valid."""
+        self._type_parameter: AbstractDimension
+
+    def __check_init__(self) -> None:
+        """Check whether the arguments are valid."""
+        expected_dimensions = self._type_parameter
+        got_dimensions = dimension_of(self.unit)
+        if got_dimensions != expected_dimensions:  # pylint: disable=unreachable
+            msg = "Physical type mismatch."  # TODO: better error message
+            raise ValueError(msg)
+
+    # ---------------------------------------------------------------
+    # Plum stuff
+
+    @classmethod
+    @dispatch
+    def __init_type_parameter__(
+        cls, dims: AbstractDimension, /
+    ) -> tuple[AbstractDimension]:
+        """Check whether the type parameters are valid."""
+        return (dims,)
+
+    @classmethod
+    @dispatch
+    def __init_type_parameter__(cls, dims: str, /) -> tuple[AbstractDimension]:
+        """Check whether the type parameters are valid."""
+        dims_: AbstractDimension
+        try:
+            dims_ = dimension(dims)
+        except ValueError:
+            dims_ = dimension_of(parse_unit(dims))
+        return (dims_,)
+
+    @classmethod
+    @dispatch
+    def __init_type_parameter__(cls, unit: AbstractUnit, /) -> tuple[AbstractDimension]:
+        """Infer the type parameter from the arguments."""
+        dims = dimension_of(unit)
+        if dims != "unknown":  # pylint: disable=unreachable
+            return (dims,)
+        return (PhysicalType(unit, unit.to_string(fraction=False)),)
+
+    @classmethod
+    def __infer_type_parameter__(
+        cls, value: Any, unit: Any, **kwargs: Any
+    ) -> tuple[AbstractDimension]:
+        """Infer the type parameter from the arguments."""
+        return (dimension_of(parse_unit(unit)),)
+
+    @classmethod
+    @dispatch
+    def __le_type_parameter__(
+        cls, left: tuple[AbstractDimension], right: tuple[AbstractDimension]
+    ) -> bool:
+        """Define an order on type parameters.
+
+        That is, check whether `left <= right` or not.
+        """
+        (dim_left,) = left
+        (dim_right,) = right
+        return dim_left == dim_right
+
+    # ---------------------------------------------------------------
+    # misc
+
+    def __getnewargs_ex__(self) -> tuple[tuple[Any, ...], dict[str, Any]]:
+        """Return args, kwargs for ``__new__``.
+
+        In protocols 2+, this is used to determine the values (args,
+        and kwargs) passed to ``__new__``. We implement ``__getnewargs_ex__``
+        instead of ``__getnewargs__`` because the latter does not support
+        keyword-only arguments.
+
+        Examples
+        --------
+        >>> import copy as pycopy
+        >>> import unxt as u
+
+        >>> x = u.Quantity([1, 2, 3], "m")
+        >>> pycopy.copy(x)
+        Quantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+        """
+        return (), field_items(self)
+
+    # TODO: see if pickling can be accomplished without reduce.
+    # https://docs.python.org/3.12/library/pickle.html
+    def __reduce__(
+        self,
+    ) -> tuple[Callable[..., "AbstractParametricQuantity"], tuple[Any, ...]]:
+        r"""Return the reduced value: a constructor and arguments.
+
+        The ``__reduce__`` protocol has limited support for keyword-only
+        argument. The only built-in means to pass kwargs is to the
+        ``__setstate__`` method or through a callable with signature ``(obj,
+        state)``. Neither of these methods allow for the kwargs to be passed to
+        the constructor, only after the object is partially initialized. This
+        does not work well for JAX, Equinox, or particularly run-time
+        typechecking. To get around this, instead of the class type as the
+        customary first element of reduced value we will instead return a
+        `functools.partial`-wrapping of the class type with the kwargs bundled
+        into the partial object. This is allowed by the ``__reduce__`` protocol,
+        which says the first returned item must be any "callable object that
+        will be called to create the initial version of the object." In
+        conjunction with the standard args, the `functools.partial`-wrapped
+        class type will properly construct any parametric subclass.
+
+        Returns
+        -------
+        functools.partial[type]
+            The `plum.type_unparametrized` form of this class object, which is
+            the parametric class without the specialization to the specific
+            parameter. The specialized subtype cannot be pickled since it is
+            dynamically produced. The type is wrapped into a `functools.partial`
+            with the keyword-only arguments to the constructor.
+        tuple[Any, ...]
+            The arguments to this class. Note: the keyword-only arguments are
+            bundled with the type.
+
+        Examples
+        --------
+        >>> import pickle
+        >>> import unxt as u
+
+        >>> x = u.Quantity([1, 2, 3], "m")
+        >>> pickle.dumps(x)
+        b'...'
+
+        """
+        args, kwargs = [], {}
+        for f in fields(self):
+            v = getattr(self, f.name)
+            if f.kw_only:
+                kwargs[f.name] = v
+            else:
+                args.append(v)
+
+        return partial(type_unparametrized(self), **kwargs), tuple(args)
+
+    def __pdoc__(
+        self,
+        *,
+        include_params: bool = False,
+        named_unit: bool = 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
+        ----------
+        include_params
+            If `True`, the type parameter is included in the representation. If
+            `False`, the type parameter is not included in the representation.
+            For example, ``Quantity['length'][i32]`` versus ``Quantity[i32]``.
+        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')``.
+        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([1, 2, 3], "m")
+
+        The default pretty printing:
+
+        >>> wl.pprint(q)
+        Quantity(i32[3], unit='m')
+
+        The type parameter can be included in the representation:
+
+        >>> wl.pprint(q, include_params=True)
+        Quantity['length'](i32[3], unit='m')
+
+        The `str` method uses this as well:
+
+        >>> print(q)
+        Quantity['length']([1, 2, 3], unit='m')
+
+        Arrays can be printed in full:
+
+        >>> wl.pprint(q, short_arrays=False)
+        Quantity(Array([1, 2, 3], dtype=int32), unit='m')
+
+        The `repr` method uses this setting:
+
+        >>> print(repr(q))
+        Quantity(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)
+        Quantity(i32[3], 'm')
+
+        The class can be printed with its short name:
+
+        >>> wl.pprint(q, use_short_name=True)
+        Q(i32[3], unit='m')
+
+        Short name with type parameter:
+
+        >>> wl.pprint(q, use_short_name=True, include_params=True)
+        Q['length'](i32[3], unit='m')
+
+        """
+        pdoc = super().__pdoc__(
+            include_params=include_params,
+            named_unit=named_unit,
+            use_short_name=use_short_name,
+            **kwargs,
+        )
+
+        # Type Parameter
+        if not include_params:
+            param = wl.TextDoc("")
+        else:
+            param = wl.TextDoc(f"[{name_of(self._type_parameter)!r}]")
+
+        return wl.ConcatDoc(pdoc.children[0], param, *pdoc.children[1:])
+
+    def __repr__(self) -> str:
+        return wl.pformat(
+            self, include_params=False, named_unit=True, short_arrays=False, indent=4
+        )
+
+    def __str__(self) -> str:
+        return wl.pformat(
+            self,
+            include_params=True,
+            named_unit=True,
+            short_arrays="compact"
+            if not isinstance(self.value, jax.core.Tracer)
+            else True,
+            indent=4,
+        )

unxt/_src/quantity/flag.py

@@ -0,0 +1,137 @@
+"""Flags for quantity operations."""
+
+__all__ = ("AllowValue",)
+
+from typing import Any, NoReturn
+
+from plum import dispatch
+
+import unxt_api as uapi
+from .base import AbstractQuantity
+
+
+class AllowValue:
+    """A flag to allow a value to be passed through `unxt.ustrip`.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+    >>> from unxt.quantity import AllowValue
+
+    >>> x = jnp.array(1)
+    >>> y = u.ustrip(AllowValue, "km", x)
+    >>> y is x
+    True
+
+    >>> u.ustrip(AllowValue, "km", u.Q(1000, "m"))
+    Array(1., dtype=float32, ...)
+
+    This is a flag, so it cannot be instantiated.
+
+    >>> try:
+    ...     AllowValue()
+    ... except TypeError as e:
+    ...     print(e)
+    Cannot instantiate AllowValue
+
+    """
+
+    def __new__(cls) -> NoReturn:
+        msg = "Cannot instantiate AllowValue"
+        raise TypeError(msg)
+
+
+@dispatch
+def ustrip(flag: type[AllowValue], unit: Any, x: Any, /) -> Any:
+    """Strip the units from a value. This is a no-op.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+    >>> from unxt.quantity import AllowValue
+
+    >>> x = jnp.array(1)
+    >>> y = u.ustrip(AllowValue, "km", x)
+    >>> y is x
+    True
+
+    >>> x = 1_000
+    >>> y = u.ustrip(AllowValue, "km", x)
+    >>> y is x
+    True
+
+    >>> x = "hello"
+    >>> y = u.ustrip(AllowValue, "km", x)
+    >>> y is x
+    True
+
+    """
+    return x
+
+
+@dispatch
+def ustrip(flag: type[AllowValue], x: Any, /) -> Any:
+    """Strip the units from a value. This is a no-op.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+    >>> from unxt.quantity import AllowValue
+
+    >>> x = jnp.array(1)
+    >>> y = u.ustrip(AllowValue, x)
+    >>> y is x
+    True
+
+    >>> x = 1_000
+    >>> y = u.ustrip(AllowValue, x)
+    >>> y is x
+    True
+
+    >>> x = "hello"
+    >>> y = u.ustrip(AllowValue, x)
+    >>> y is x
+    True
+
+    """
+    return x
+
+
+@dispatch  # TODO: type annotate by value
+def ustrip(flag: type[AllowValue], unit: Any, x: AbstractQuantity, /) -> Any:
+    """Strip the units from a quantity.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> from unxt.quantity import AllowValue
+    >>> q = u.Q(1000, "m")
+    >>> u.ustrip(AllowValue, "km", q)
+    Array(1., dtype=float32, ...)
+
+    """
+    return uapi.ustrip(unit, x)
+
+
+@dispatch
+def ustrip(flag: type[AllowValue], x: AbstractQuantity, /) -> Any:
+    """Strip the units from a quantity.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+    >>> from unxt.quantity import AllowValue
+
+    >>> x = u.Q(1, "kpc")
+    >>> y = u.ustrip(AllowValue, x)
+    >>> not isinstance(y, u.Quantity)
+    True
+    >>> y == 1
+    Array(True, dtype=bool, ...)
+
+    """
+    return uapi.ustrip(x)