unxt 0.1.dev463__py3-none-any.whl

unxt/_src/experimental.py

@@ -0,0 +1,234 @@
+r"""Experimental features.
+
+.. warning::
+
+    These features may be removed or changed in the future without notice.
+
+On some occasions JAX's automatic differentiation functions do not work well
+with quantities. This is checked by enabling runtime type-checking (see the
+docs), which will raise an error if a quantity's units do not match the expected
+input / output units of a function. In these cases, you can use the functions in
+this module to provide the units to the automatic differentiation functions.
+Instead of directly propagating the units through the automatic differentiation
+functions, the units are stripped and re-applied, while also being provided
+within the function being AD'd.
+
+To import this experimental module
+
+>>> from unxt import experimental
+
+"""
+# pylint: disable=import-error
+
+__all__ = ("grad", "hessian", "jacfwd")
+
+import functools as ft
+from collections.abc import Callable
+from typing import Any, TypeVar, TypeVarTuple
+from typing_extensions import Unpack
+
+import equinox as eqx
+import jax
+from jaxtyping import ArrayLike
+from plum import type_unparametrized
+
+from .quantity import AbstractQuantity, BareQuantity
+from .units import AbstractUnit
+from unxt_api import unit, unit_of, ustrip
+
+Args = TypeVarTuple("Args")
+R = TypeVar("R", bound=AbstractQuantity)
+
+
+def unit_or_none(obj: Any) -> AbstractUnit | None:
+    return obj if obj is None else unit(obj)
+
+
+def grad(
+    fun: Callable[[Unpack[Args]], R],
+    argnums: int = 0,
+    *,
+    units: tuple[AbstractUnit | str, ...],
+) -> Callable[[Unpack[Args]], R]:
+    """Gradient of a function with units.
+
+    In general, if you can use ``quax.quaxify(jax.grad(func))`` (or the
+    syntactic sugar ``quaxed.grad(func)``), that's the better option! The
+    difference from those functions is how this units are supported. ``quaxify``
+    will directly propagate the units through the automatic differentiation
+    functions. But sometimes that doesn't work and we need to strip the units
+    and re-apply them. This function does that, using the "units" kwarg.
+
+    See Also
+    --------
+    jax.grad : The original JAX gradient function.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+
+    >>> def cube_volume(x: u.Quantity["length"]) -> u.Quantity["volume"]:
+    ...     return x**3
+
+    >>> grad_cube_volume = u.experimental.grad(cube_volume, units=("m",))
+    >>> grad_cube_volume(u.Q(2.0, "m"))
+    Quantity(Array(12., dtype=float32, weak_type=True), unit='m2')
+
+    """
+    theunits: tuple[AbstractUnit | None, ...] = tuple(map(unit_or_none, units))
+
+    # Gradient of function, stripping and adding units
+    @ft.partial(jax.grad, argnums=argnums)
+    def gradfun_mag(*args: Any) -> ArrayLike:
+        args_ = (
+            (a if unit is None else BareQuantity(a, unit))
+            for a, unit in zip(args, theunits, strict=True)
+        )
+        return ustrip(fun(*args_))  # type: ignore[arg-type]
+
+    def gradfun(*args: *Args) -> R:
+        # Get the value of the args. They are turned back into Quantity
+        # inside the function we are taking the grad of.
+        args_ = tuple(  # type: ignore[var-annotated]
+            (a if unit is None else ustrip(unit, a))
+            for a, unit in zip(args, theunits, strict=True)  # type: ignore[arg-type]
+        )
+        # Call the grad, returning a Quantity
+        value = fun(*args)
+        grad_value = gradfun_mag(*args_)
+        # Adjust the Quantity by the units of the derivative
+        # TODO: get Quantity[unit] / unit2 -> Quantity[unit/unit2] working
+        return type_unparametrized(value)(
+            grad_value, unit_of(value) / theunits[argnums]
+        )
+
+    return gradfun
+
+
+def jacfwd(
+    fun: Callable[[Unpack[Args]], R],
+    argnums: int = 0,
+    *,
+    units: tuple[AbstractUnit | str, ...],
+) -> Callable[[Unpack[Args]], R]:
+    """Jacobian of ``fun`` evaluated column-by-column using forward-mode AD.
+
+    In general, if you can use `quax.quaxify(jax.jacfwd(func))` (or the
+    syntactic sugar `quax.jacfwd(func)`), that's the better option! The
+    difference from those functions is how this units are supported. `quaxify`
+    will directly propagate the units through the automatic differentiation
+    functions. But sometimes that doesn't work and we need to strip the units
+    and re-apply them. This function does that, using the ``units`` kwarg.
+
+    See Also
+    --------
+    jax.jacfwd : The original JAX jacfwd function.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+
+    >>> def cubbe_volume(x: u.Quantity["length"]) -> u.Quantity["volume"]:
+    ...     return x**3
+
+    >>> jacfwd_cubbe_volume = u.experimental.jacfwd(cubbe_volume, units=("m",))
+    >>> jacfwd_cubbe_volume(u.Q(2.0, "m"))
+    BareQuantity(Array(12., dtype=float32, weak_type=True), unit='m2')
+
+    """
+    argnums = eqx.error_if(
+        argnums,
+        not isinstance(argnums, int),
+        "only int argnums are currently supported",
+    )
+
+    theunits: tuple[AbstractUnit | None, ...] = tuple(map(unit_or_none, units))
+
+    @ft.partial(jax.jacfwd, argnums=argnums)
+    def jacfun_mag(*args: Any) -> R:
+        args_ = tuple(
+            (a if unit is None else BareQuantity(a, unit))
+            for a, unit in zip(args, theunits, strict=True)
+        )
+        return fun(*args_)  # type: ignore[arg-type]
+
+    def jacfun(*args: *Args) -> R:
+        # Get the value of the args. They are turned back into Quantity
+        # inside the function we are taking the Jacobian of.
+        args_ = tuple(  # type: ignore[var-annotated]
+            (a if unit is None else ustrip(unit, a))
+            for a, unit in zip(args, theunits, strict=True)  # type: ignore[arg-type]
+        )
+        # Call the Jacobian, returning a Quantity
+        value = jacfun_mag(*args_)
+        # Adjust the Quantity by the units of the derivative
+        # TODO: check the unit correction
+        # TODO: get Quantity[unit] / unit2 -> Quantity[unit/unit2] working
+        return type_unparametrized(value)(
+            ustrip(value), unit_of(value) / theunits[argnums]
+        )
+
+    return jacfun
+
+
+def hessian(
+    fun: Callable[[Unpack[Args]], R],
+    argnums: int = 0,
+    *,
+    units: tuple[AbstractUnit | str, ...],
+) -> Callable[[Unpack[Args]], R]:
+    """Hessian.
+
+    In general, if you can use `quax.quaxify(jax.hessian(func))` (or the
+    syntactic sugar `quax.hessian(func)`), that's the better option! The
+    difference from those functions is how this units are supported. `quaxify`
+    will directly propagate the units through the automatic differentiation
+    functions. But sometimes that doesn't work and we need to strip the units
+    and re-apply them. This function does that, using the ``units`` kwarg.
+
+    See Also
+    --------
+    jax.hessian : The original JAX hessian function.
+
+    Examples
+    --------
+    >>> import jax.numpy as jnp
+    >>> import unxt as u
+
+    >>> def cubbe_volume(x: u.Quantity["length"]) -> u.Quantity["volume"]:
+    ...     return x**3
+
+    >>> hessian_cubbe_volume = u.experimental.hessian(cubbe_volume, units=("m",))
+    >>> hessian_cubbe_volume(u.Q(2.0, "m"))
+    BareQuantity(Array(12., dtype=float32, weak_type=True), unit='m')
+
+    """
+    theunits: tuple[AbstractUnit, ...] = tuple(map(unit_or_none, units))
+
+    @ft.partial(jax.hessian)
+    def hessfun_mag(*args: Any) -> R:
+        args_ = tuple(
+            (a if unit is None else BareQuantity(a, unit))
+            for a, unit in zip(args, theunits, strict=True)
+        )
+        return fun(*args_)  # type: ignore[arg-type]
+
+    def hessfun(*args: *Args) -> R:
+        # Get the value of the args. They are turned back into Quantity
+        # inside the function we are taking the hessian of.
+        args_ = tuple(  # type: ignore[var-annotated]
+            (a if unit is None else ustrip(unit, a))
+            for a, unit in zip(args, units, strict=True)  # type: ignore[arg-type]
+        )
+        # Call the hessian, returning a Quantity
+        value = hessfun_mag(*args_)
+        # Adjust the Quantity by the units of the derivative
+        # TODO: check the unit correction
+        # TODO: get Quantity[unit] / unit2 -> Quantity[unit/unit2] working
+        return type_unparametrized(value)(
+            ustrip(value), unit_of(value) / theunits[argnums] ** 2
+        )
+
+    return hessfun

unxt/_src/quantity/__init__.py

@@ -0,0 +1,24 @@
+"""Quantities in JAX."""
+
+__all__ = (
+    "AbstractParametricQuantity",
+    "AbstractQuantity",
+    "Quantity",
+    "StaticQuantity",
+    "StaticValue",
+    "BareQuantity",
+    "is_any_quantity",
+    "convert_to_quantity_value",
+    "AbstractAngle",
+    "Angle",
+)
+
+
+from .angle import Angle
+from .base import AbstractQuantity, is_any_quantity
+from .base_angle import AbstractAngle
+from .base_parametric import AbstractParametricQuantity
+from .quantity import Quantity
+from .static_quantity import StaticQuantity
+from .unchecked import BareQuantity
+from .value import StaticValue, convert_to_quantity_value

unxt/_src/quantity/angle.py

@@ -0,0 +1,58 @@
+"""Angular quantities."""
+
+__all__ = ("Angle",)
+
+from typing import final
+
+import equinox as eqx
+from jaxtyping import Array, Real
+
+import unxt_api as uapi
+from .base_angle import AbstractAngle
+from .value import StaticValue, convert_to_quantity_value
+from unxt.units import AbstractUnit
+
+
+@final
+class Angle(AbstractAngle):
+    """Angular quantity.
+
+    Examples
+    --------
+    >>> import unxt as u
+
+    Create an Angle:
+
+    >>> q = u.Angle(1, "rad")
+    >>> q
+    Angle(Array(1, dtype=int32, ...), unit='rad')
+
+    Wrap an Angle to a range:
+
+    >>> q = u.Angle(370, "deg")
+    >>> q.wrap_to(u.Q(0, "deg"), u.Q(360, "deg"))
+    Angle(Array(10, dtype=int32, ...), unit='deg')
+
+    Create an Angle array:
+
+    >>> q = u.Angle([1, 2, 3], "deg")
+    >>> q
+    Angle(Array([1, 2, 3], dtype=int32), unit='deg')
+
+    Do math on an Angle:
+
+    >>> 2 * q
+    Angle(Array([2, 4, 6], dtype=int32), unit='deg')
+
+    >>> q % u.Q(4, "deg")
+    Angle(Array([1, 2, 3], dtype=int32), unit='deg')
+
+    """
+
+    value: Real[Array | StaticValue, "*shape"] = eqx.field(
+        converter=convert_to_quantity_value
+    )
+    """The value of the `unxt.AbstractQuantity`."""
+
+    unit: AbstractUnit = eqx.field(static=True, converter=uapi.unit)
+    """The unit associated with this value."""