unxt 0.1.dev463__py3-none-any.whl

unxt/__init__.py

@@ -0,0 +1,79 @@
+"""``unxt``: Quantities in JAX.
+
+``unxt`` is a library for working with physical quantities in JAX, supporting
+JAX's autodiff and JIT compilation, and easy integration with existing codes. If
+you're seeing this then you're in the main module of the ``unxt`` package, where
+we provide exports for the main functionality of the library. Sub-modules are
+available for more specialized functionality, such as unit systems and
+experimental features.
+
+Note that `unxt` uses multiple-dispatch to provide a flexible and extensible
+interface. In the docs you'll see the function signatures without type
+annotation and then subsections for specific function dispatches based on the
+type annotations. However dispatches registered from other modules may not be
+included in the rendered docs. To see all the dispatches execute ``<func or
+class>.methods`` in an interactive Python session. For more information on
+multiple-dispatch see the [`plum`](https://beartype.github.io/plum/intro.html)
+documentation.
+
+-----
+
+"""
+
+__all__ = (
+    "__version__",
+    # dimensions
+    "dims",  # module
+    "AbstractDimension",  # dimension type
+    "dimension",  # convenience constructor
+    "dimension_of",  # get the dimensions
+    # units
+    "units",  # module
+    "AbstractUnit",  # unit type
+    "unit",
+    "unit_of",
+    # units systems
+    "unitsystems",  # module
+    "AbstractUnitSystem",  # base class
+    "unitsystem",  # convenience constructor
+    "unitsystem_of",  # get the unit system
+    # quantities
+    "quantity",  # module
+    "Quantity",  # main class
+    "Q",  # convenience alias
+    "StaticQuantity",  # static value quantity
+    "AbstractQuantity",  # base class
+    "Angle",  # angular quantity
+    "uconvert_value",  # convert values
+    "uconvert",  # convert units
+    "ustrip",  # strip units
+    "is_unit_convertible",  # check if units can be converted
+)
+
+from .setup_package import install_import_hook
+
+with install_import_hook("unxt"):
+    from . import dims, quantity, units, unitsystems
+    from ._version import version as __version__
+    from .dims import AbstractDimension, dimension, dimension_of
+    from .quantity import (
+        AbstractQuantity,
+        Angle,
+        Q,
+        Quantity,
+        StaticQuantity,
+        is_unit_convertible,
+        uconvert,
+        uconvert_value,
+        ustrip,
+    )
+    from .units import AbstractUnit, unit, unit_of
+    from .unitsystems import AbstractUnitSystem, unitsystem, unitsystem_of
+
+from ._src import experimental  # noqa: F401
+
+# isort: split
+from . import _interop  # noqa: F401  # register interop
+
+# Clean up the namespace
+del install_import_hook

unxt/_interop/__init__.py

@@ -0,0 +1,18 @@
+"""`unxt` interoperability with other libraries."""
+# ruff:noqa: F401
+
+__all__: tuple[str, ...] = ()
+
+from .optional_deps import OptDeps
+
+# Register interoperability
+if OptDeps.ASTROPY.installed:
+    from . import unxt_interop_astropy
+
+if OptDeps.GALA.installed:
+    from . import unxt_interop_gala
+
+if OptDeps.MATPLOTLIB.installed:
+    from . import unxt_interop_mpl as interop_mpl
+
+    interop_mpl.setup_matplotlib_support_for_unxt(enable=True)

unxt/_interop/optional_deps.py

@@ -0,0 +1,13 @@
+"""Optional dependencies. Internal use only."""
+
+__all__ = ("OptDeps",)
+
+from optional_dependencies import OptionalDependencyEnum, auto
+
+
+class OptDeps(OptionalDependencyEnum):  # type: ignore[misc]  # pylint: disable=invalid-enum-extension
+    """Optional dependencies for ``unxt``."""
+
+    ASTROPY = auto()
+    GALA = auto()
+    MATPLOTLIB = auto()

unxt/_interop/unxt_interop_astropy/__init__.py

@@ -0,0 +1,5 @@
+"""Tools for representing systems of units using ``astropy.units``."""
+
+__all__: tuple[str, ...] = ()
+
+from . import dimensions, quantity, units  # noqa: F401

unxt/_interop/unxt_interop_astropy/custom_types.py

@@ -0,0 +1,12 @@
+"""Custom type annotations for astropy."""
+
+__all__ = ("APYUnits",)
+
+
+from typing import TypeAlias
+
+import astropy.units as apyu
+
+APYUnits: TypeAlias = (
+    apyu.UnitBase | apyu.Unit | apyu.FunctionUnitBase | apyu.StructuredUnit
+)

unxt/_interop/unxt_interop_astropy/dimensions.py

@@ -0,0 +1,53 @@
+"""Dimensions compatibility."""
+
+__all__: tuple[str, ...] = ()
+
+
+from dataclasses import MISSING, Field
+
+import astropy.units as apyu
+import plum
+
+from unxt.dims import AbstractDimension
+
+
+@plum.dispatch
+def dimension_of(obj: apyu.Quantity, /) -> AbstractDimension:
+    """Return the dimension of a quantity.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> import astropy.units as apyu
+
+    >>> q = apyu.Quantity(1, "m")
+    >>> u.dimension_of(q)
+    PhysicalType('length')
+
+    """
+    return dimension_of(obj.unit)
+
+
+# ===================================================
+# `Dataclassish` support
+
+
+@plum.dispatch
+def fields(obj: AbstractDimension, /) -> tuple[Field, ...]:
+    """Return the fields of a dimension.
+
+    Examples
+    --------
+    >>> import dataclassish as dc
+    >>> import astropy.units as apyu
+
+    >>> dim = apyu.get_physical_type("length")
+    >>> dc.fields(dim)
+    (Field(name='_unit',...), Field(name='_physical_type',...))
+
+    """
+    unit_field = Field(MISSING, MISSING, True, False, False, False, {}, False)  # noqa: FBT003
+    unit_field.name = "_unit"
+    physical_type_field = Field(MISSING, MISSING, True, False, False, False, {}, False)  # noqa: FBT003
+    physical_type_field.name = "_physical_type"
+    return (unit_field, physical_type_field)

unxt/_interop/unxt_interop_astropy/quantity.py

@@ -0,0 +1,308 @@
+"""Quantity compatibility."""
+
+__all__: tuple[str, ...] = ()
+
+from typing import Any, NoReturn
+
+import plum
+from astropy.coordinates import Angle as AstropyAngle, Distance as AstropyDistance
+from astropy.units import Quantity as AstropyQuantity
+from jaxtyping import ArrayLike
+
+import dataclassish as dc
+import quaxed.numpy as jnp
+
+import unxt_api as uapi
+from .custom_types import APYUnits
+from unxt.quantity import AbstractQuantity, AllowValue, BareQuantity, Quantity, ustrip
+
+# ============================================================================
+# Value Converter
+
+
+@plum.dispatch
+def convert_to_quantity_value(obj: AstropyQuantity, /) -> NoReturn:
+    """Disallow conversion of `AstropyQuantity` to a value.
+
+    >>> import astropy.units as apyu
+    >>> from unxt.quantity import convert_to_quantity_value
+
+    >>> try:
+    ...     convert_to_quantity_value(apyu.Quantity(1, "m"))
+    ... except TypeError as e:
+    ...     print(e)
+    Cannot convert 'Quantity' to a value.
+    For a Quantity, use the `.from_` constructor instead.
+
+    """
+    msg = (
+        f"Cannot convert {type(obj).__name__!r} to a value. "
+        "For a Quantity, use the `.from_` constructor instead."
+    )
+    raise TypeError(msg)
+
+
+# ============================================================================
+# AbstractQuantity
+
+
+# -----------------
+# Constructor
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity], value: AstropyQuantity, /, **kwargs: Any
+) -> AbstractQuantity:
+    """Construct a `Quantity` from another `Quantity`.
+
+    The `value` is converted to the new `unit`.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> import astropy.units as apyu
+
+    >>> u.Quantity.from_(apyu.Quantity(1, "m"))
+    Quantity(Array(1., dtype=float32), unit='m')
+
+    """
+    u = uapi.unit_of(value)
+    value = jnp.asarray(uapi.ustrip(u, value), **kwargs)
+    return cls(value, u)
+
+
+@AbstractQuantity.from_.dispatch
+def from_(
+    cls: type[AbstractQuantity], value: AstropyQuantity, u: Any, /, **kwargs: Any
+) -> AbstractQuantity:
+    """Construct a `Quantity` from another `Quantity`.
+
+    The `value` is converted to the new `unit`.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> import astropy.units as apyu
+
+    >>> u.Quantity.from_(apyu.Quantity(1, "m"), "cm")
+    Quantity(Array(100., dtype=float32), unit='cm')
+
+    """
+    u = uapi.unit(u)
+    value = jnp.asarray(uapi.ustrip(u, value), **kwargs)
+    return cls(value, u)
+
+
+# -----------------
+# Conversion Methods
+
+
+@plum.conversion_method(type_from=AbstractQuantity, type_to=AstropyQuantity)  # type: ignore[arg-type]
+def convert_unxt_quantity_to_astropy_quantity(
+    q: AbstractQuantity, /
+) -> AstropyQuantity:
+    """Convert a `unxt.AbstractQuantity` to a `astropy.units.Quantity`.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> from astropy.units import Quantity as AstropyQuantity
+    >>> from plum import convert
+
+    >>> convert(u.Q(1.0, "cm"), AstropyQuantity)
+    <Quantity 1. cm>
+
+    """
+    u = uapi.unit_of(q)
+    return AstropyQuantity(uapi.ustrip(u, q), u)
+
+
+@plum.conversion_method(type_from=AbstractQuantity, type_to=AstropyDistance)  # type: ignore[arg-type]
+def convert_unxt_quantity_to_astropy_distance(
+    q: AbstractQuantity, /
+) -> AstropyDistance:
+    """Convert a `unxt.AbstractQuantity` to a `astropy.coordinates.Distance`.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> from astropy.coordinates import Distance as AstropyDistance
+    >>> from plum import convert
+
+    >>> convert(u.Q(1.0, "cm"), AstropyDistance)
+    <Distance 1. cm>
+
+    """
+    u = uapi.unit_of(q)
+    return AstropyDistance(uapi.ustrip(u, q), u)
+
+
+@plum.conversion_method(type_from=AbstractQuantity, type_to=AstropyAngle)  # type: ignore[arg-type]
+def convert_unxt_quantity_to_astropy_angle(q: AbstractQuantity, /) -> AstropyAngle:
+    """Convert a `unxt.quantity.AbstractQuantity` to a `astropy.coordinates.Angle`.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> from astropy.coordinates import Angle as AstropyAngle
+    >>> from plum import convert
+
+    >>> convert(u.Q(1.0, "rad"), AstropyAngle)
+    <Angle 1. rad>
+
+    """
+    u = uapi.unit_of(q)
+    return AstropyAngle(uapi.ustrip(u, q), u)
+
+
+# ============================================================================
+# Quantity
+
+
+@plum.conversion_method(type_from=AstropyQuantity, type_to=Quantity)
+def convert_astropy_quantity_to_unxt_quantity(q: AstropyQuantity, /) -> Quantity:
+    """Convert a `astropy.units.Quantity` to a `unxt.Quantity`.
+
+    Examples
+    --------
+    >>> from astropy.units import Quantity as AstropyQuantity
+    >>> from plum import convert
+    >>> from unxt import Quantity
+
+    >>> convert(AstropyQuantity(1.0, "cm"), Quantity)
+    Quantity(Array(1., dtype=float32), unit='cm')
+
+    """
+    u = uapi.unit_of(q)
+    return Quantity(uapi.ustrip(u, q), u)
+
+
+# ============================================================================
+# BareQuantity
+
+
+@plum.conversion_method(type_from=AstropyQuantity, type_to=BareQuantity)  # type: ignore[arg-type]
+def convert_astropy_quantity_to_unxt_barequantity(
+    q: AstropyQuantity, /
+) -> BareQuantity:
+    """Convert a `astropy.units.Quantity` to a `unxt.BareQuantity`.
+
+    Examples
+    --------
+    >>> from astropy.units import Quantity as AstropyQuantity
+    >>> from plum import convert
+    >>> from unxt.quantity import BareQuantity
+
+    >>> convert(AstropyQuantity(1.0, "cm"), BareQuantity)
+    BareQuantity(Array(1., dtype=float32), unit='cm')
+
+    """
+    u = uapi.unit_of(q)
+    return BareQuantity(uapi.ustrip(u, q), u)
+
+
+###############################################################################
+
+
+@plum.dispatch
+def uconvert_value(uto: APYUnits, ufrom: APYUnits, x: ArrayLike, /) -> ArrayLike:
+    """Convert the value to the specified units.
+
+    Examples
+    --------
+    >>> import astropy.units as apyu
+    >>> import unxt as u
+
+    >>> u.uconvert_value(apyu.Unit("km"), apyu.Unit("m"), 1000)
+    1.0
+
+    """
+    # Hot-path: if no unit conversion is necessary
+    if ufrom == uto:
+        return x
+
+    # Compute the new value.
+    return ufrom.to(uto, x)
+
+
+@plum.dispatch
+def uconvert(u: APYUnits, x: AbstractQuantity, /) -> AbstractQuantity:
+    """Convert the quantity to the specified units.
+
+    Examples
+    --------
+    >>> import astropy.units as apyu
+    >>> import unxt as u
+
+    >>> x = u.Quantity(1000, "m")
+    >>> u.uconvert(u.unit("km"), x)
+    Quantity(Array(1., dtype=float32, ...), unit='km')
+
+    >>> x = u.Quantity([1, 2, 3], "Kelvin")
+    >>> with apyu.add_enabled_equivalencies(apyu.temperature()):
+    ...     y = x.uconvert("deg_C")
+    >>> y
+    Quantity(Array([-272.15, -271.15, -270.15], dtype=float32, ...), unit='deg_C')
+
+    >>> x = Quantity([1, 2, 3], "radian")
+    >>> with apyu.add_enabled_equivalencies(apyu.dimensionless_angles()):
+    ...     y = x.uconvert("")
+    >>> y
+    Quantity(Array([1., 2., 3.], dtype=float32, ...), unit='')
+
+    """
+    # Hot-path: if no unit conversion is necessary
+    if x.unit == u:
+        return x
+
+    value = x.unit.to(u, uapi.ustrip(x))
+
+    # If the dimensions are the same, we can just replace the value and unit.
+    if uapi.dimension_of(x.unit) == uapi.dimension_of(u):
+        return dc.replace(x, value=value, unit=u)
+
+    # If the dimensions are different, we need to create a new quantity since
+    # the dimensions can be part of the type. This won't work if the Quantity
+    # type itself needs to be changed, e.g. `unxt.Angle` -> `unxt.Quantity`.
+    # These cases are handled separately, in other dispatches.
+    fs = dict(dc.field_items(x))  # pylint: disable=unreachable
+    fs["value"] = value
+    fs["unit"] = u
+
+    return plum.type_unparametrized(x)(**fs)
+
+
+# ============================================================================
+
+
+@plum.dispatch
+def ustrip(u: Any, x: AstropyQuantity) -> Any:
+    """Strip the units from the quantity.
+
+    Examples
+    --------
+    >>> import astropy.units as apyu
+    >>> import unxt as u
+
+    >>> x = apyu.Quantity(1000, "m")
+    >>> float(u.ustrip(u.unit("m"), x))
+    1000.0
+
+    """
+    return x.to_value(u)
+
+
+@plum.dispatch  # TODO: type annotate by value
+def ustrip(flag: type[AllowValue], u: Any, x: AstropyQuantity, /) -> Any:
+    """Strip the units from a quantity.
+
+    Examples
+    --------
+    >>> import unxt as u
+    >>> q = u.Quantity(1000, "m")
+    >>> u.ustrip(AllowValue, "km", q)
+    Array(1., dtype=float32, ...)
+
+    """
+    return uapi.ustrip(u, x)

unxt/_interop/unxt_interop_astropy/units.py

@@ -0,0 +1,111 @@
+"""Astropy units compatibility."""
+
+__all__: tuple[str, ...] = ()
+
+
+import dataclasses
+
+import astropy.units as apyu
+import plum
+
+from .custom_types import APYUnits
+
+# ===================================================================
+# Register dispatches
+
+
+@plum.dispatch
+def unit(obj: apyu.UnitBase | apyu.Unit, /) -> APYUnits:
+    """Construct the units from an Astropy unit.
+
+    Examples
+    --------
+    >>> import astropy.units as apyu
+    >>> import unxt as u
+    >>> u.unit(apyu.km)
+    Unit("km")
+
+    """
+    return apyu.Unit(obj)
+
+
+@plum.dispatch
+def unit(obj: apyu.Quantity, /) -> APYUnits:
+    """Construct the units from an Astropy quantity.
+
+    Examples
+    --------
+    >>> import astropy.units as apyu
+    >>> import unxt as u
+    >>> u.unit(apyu.Quantity(2, "km"))
+    Unit("2 km")
+
+    """
+    return apyu.Unit(obj)
+
+
+# -------------------------------------------------------------------
+
+
+@plum.dispatch
+def unit_of(obj: apyu.UnitBase | apyu.Unit, /) -> APYUnits:
+    """Return the units of an object.
+
+    Examples
+    --------
+    >>> import astropy.units as apyu
+    >>> import unxt as u
+
+    >>> u.unit_of(apyu.km)
+    Unit("km")
+
+    """
+    return obj
+
+
+@plum.dispatch
+def unit_of(obj: apyu.Quantity, /) -> APYUnits:
+    """Return the units of an Astropy quantity.
+
+    Examples
+    --------
+    >>> import astropy.units as apyu
+    >>> import unxt as u
+
+    >>> u.unit_of(apyu.Quantity(1, "km"))
+    Unit("km")
+
+    """
+    return unit_of(obj.unit)
+
+
+# ===================================================
+# `Dataclassish` support
+
+
+@plum.dispatch
+def fields(obj: APYUnits, /) -> tuple[dataclasses.Field, ...]:
+    """Return the fields of a dimension.
+
+    Examples
+    --------
+    >>> import dataclassish as dc
+    >>> import astropy.units as apyu
+
+    >>> dim = apyu.Unit("m")
+    >>> dc.fields(dim)
+    (Field(name='_names',...)
+
+    """
+    st_field = dataclasses.Field(
+        dataclasses.MISSING,
+        dataclasses.MISSING,
+        True,  # noqa: FBT003
+        True,  # noqa: FBT003
+        True,  # noqa: FBT003
+        True,  # noqa: FBT003
+        {},
+        False,  # noqa: FBT003
+    )
+    st_field.name = "_names"
+    return (st_field,)

unxt/_interop/unxt_interop_gala/__init__.py

@@ -0,0 +1,5 @@
+"""Tools for representing systems of units using ``astropy.units``."""
+
+__all__: tuple[str, ...] = ()
+
+from . import unitsystems  # noqa: F401