unxt 0.1.dev463__py3-none-any.whl

unxt/dims.py

@@ -0,0 +1,19 @@
+"""Working with dimensions.
+
+The main features are:
+
+- ``unxt.dims.dimension``: a function to construct a dimension object.
+- ``unxt.dims.dimension_of``: a function to get the dimensions of an object.
+
+"""
+
+__all__ = ("dimension", "dimension_of", "AbstractDimension")
+
+from .setup_package import install_import_hook
+
+with install_import_hook("unxt.dims"):
+    from ._src.dimensions import AbstractDimension
+    from unxt_api import dimension, dimension_of
+
+# Clean up the namespace
+del install_import_hook

unxt/quantity.py

@@ -0,0 +1,172 @@
+"""Quantities in JAX.
+
+This module provides JAX-compatible quantity types with automatic unit handling,
+conversions, and a rich set of mathematical operations that preserve dimensional
+correctness.
+
+## Core Classes
+
+- **`Quantity`**: The main quantity class with dimension parametrization and
+  full unit checking. Aliased as `Q` for convenience.
+- **`StaticQuantity`**: Parametric quantity with static NumPy values for JAX
+  static arguments.
+- **`BareQuantity`**: Lightweight quantity without dimension parametrization.
+- **`Angle`**: Specialized quantity type for angular measurements with wrapping
+  support.
+
+## Key Functions
+
+- **`uconvert`**: Convert a quantity to different units.
+- **`ustrip`**: Extract the numerical value from a quantity.
+- **`is_unit_convertible`**: Check if two units are convertible to each other.
+- **`wrap_to`**: Wrap angular quantities to a specified range.
+
+## Examples
+
+### Creating quantities with units:
+
+>>> import unxt as u
+>>> import jax
+>>> import jax.numpy as jnp
+
+>>> distance = u.Quantity(100, "m")
+>>> distance
+Quantity(Array(100, dtype=int32, ...), unit='m')
+
+Create from arrays
+
+>>> velocities = u.Quantity([10, 20, 30], "m/s")
+>>> velocities
+Quantity(Array([10, 20, 30], dtype=int32), unit='m / s')
+
+Unit conversions and arithmetic:
+
+Convert units
+
+>>> distance_km = u.uconvert("km", distance)
+>>> distance_km
+Quantity(Array(0.1, dtype=float32, ...), unit='km')
+
+Arithmetic preserves units.
+
+>>> time = u.Q(5, "s")  # use Quantity alias
+>>> velocity = distance / time
+>>> velocity
+Quantity(Array(20., dtype=float32, ...), unit='m / s')
+
+Strip units for numerical operations
+
+>>> u.ustrip("m", distance)
+Array(100, dtype=int32, ...)
+
+
+### Working with angles:
+
+Create angle quantities
+
+>>> theta = u.Angle(180, "deg")
+>>> theta
+Angle(Array(180, dtype=int32, ...), unit='deg')
+
+Convert to radians
+
+>>> theta_rad = u.uconvert("rad", theta)
+>>> theta_rad
+Angle(Array(3.1415927, dtype=float32, ...), unit='rad')
+
+Wrap angles to a range
+
+>>> angle = u.Angle(450, "deg")
+>>> wrapped = u.quantity.wrap_to(angle, u.Angle(0, "deg"), u.Angle(360, "deg"))
+>>> wrapped
+Angle(Array(90, dtype=int32, ...), unit='deg')
+
+
+### Advanced usage with JAX transformations:
+
+Quantities work with JAX transformations
+
+>>> def kinetic_energy(mass, velocity):
+...     return 0.5 * mass * velocity**2
+
+>>> mass = u.Q(2.0, "kg")
+>>> vel = u.Q(10.0, "m/s")
+>>> energy = kinetic_energy(mass, vel)
+>>> energy
+Quantity(Array(100., dtype=float32, ...), unit='m2 kg / s2')
+
+Convert to standard energy units
+
+>>> u.uconvert("J", energy)
+Quantity(Array(100., dtype=float32, ...), unit='m2 kg / s2')
+
+JIT compilation works seamlessly
+
+>>> @jax.jit
+... def compute_force(mass, accel):
+...     return mass * accel
+>>> force = compute_force(u.Q(5.0, "kg"), u.Q(9.8, "m/s^2"))
+>>> force
+Quantity(Array(49., dtype=float32, ...), unit='kg m / s2')
+
+"""
+# pylint: disable=duplicate-code
+
+__all__ = (
+    # Core
+    "Quantity",
+    "Q",  # convenience alias
+    "StaticQuantity",
+    "StaticValue",
+    # Base
+    "AbstractQuantity",
+    # Fast
+    "BareQuantity",
+    # Angles
+    "AbstractAngle",
+    "Angle",
+    # Base Parametric
+    "AbstractParametricQuantity",
+    # Functional
+    "uconvert_value",
+    "uconvert",
+    "ustrip",
+    "is_unit_convertible",
+    "wrap_to",
+    "is_any_quantity",
+    "convert_to_quantity_value",
+    "AllowValue",
+)
+
+
+from .setup_package import install_import_hook
+
+with install_import_hook("unxt.quantity"):
+    from ._src.quantity.angle import Angle
+    from ._src.quantity.base import AbstractQuantity, is_any_quantity
+    from ._src.quantity.base_angle import AbstractAngle
+    from ._src.quantity.base_parametric import AbstractParametricQuantity
+    from ._src.quantity.flag import AllowValue
+    from ._src.quantity.quantity import Q, Quantity
+    from ._src.quantity.static_quantity import StaticQuantity
+    from ._src.quantity.unchecked import BareQuantity
+    from ._src.quantity.value import StaticValue, convert_to_quantity_value
+    from unxt_api import is_unit_convertible, uconvert, uconvert_value, ustrip, wrap_to
+
+    # isort: split
+    # Register dispatches and conversions
+    from ._src.quantity import (
+        register_api,
+        register_conversions,
+        register_dispatches,
+        register_primitives,
+    )
+
+# Clean up namespace
+del (
+    register_conversions,
+    register_api,
+    register_dispatches,
+    register_primitives,
+    install_import_hook,
+)

unxt/setup_package.py

@@ -0,0 +1,60 @@
+"""Package setup information.
+
+Note that this module is NOT public API nor are any of its contents.
+Stability is NOT guaranteed.
+This module exposes package setup information for the `unxt` package.
+
+"""
+
+__all__: tuple[str, ...] = ()
+
+import contextlib
+import os
+from collections.abc import Sequence
+from typing import Any, Final, Literal
+
+from jaxtyping import install_import_hook as _install_import_hook
+
+_RUNTIME_TYPECHECKER: str | None | Literal[False]
+match os.getenv("UNXT_ENABLE_RUNTIME_TYPECHECKING", "False"):
+    case "False":
+        _RUNTIME_TYPECHECKER = False
+    case "None":
+        _RUNTIME_TYPECHECKER = None
+    case str() as _name:
+        _RUNTIME_TYPECHECKER = _name
+
+RUNTIME_TYPECHECKER: Final[str | None | Literal[False]] = _RUNTIME_TYPECHECKER
+"""Runtime type checking variable "UNXT_ENABLE_RUNTIME_TYPECHECKING".
+
+Set to "False" to disable runtime typechecking (default).
+Set to "None" to only enable typechecking for `@jaxtyped`-decorated functions.
+Set to "beartype.beartype" to enable runtime typechecking.
+
+See https://docs.kidger.site/jaxtyping/api/runtime-type-checking for more
+information on options.
+
+"""
+
+
+def install_import_hook(
+    modules: str | Sequence[str], /
+) -> contextlib.AbstractContextManager[Any, None]:
+    """Install the jaxtyping import hook for the given modules.
+
+    Parameters
+    ----------
+    modules
+        Module name or sequence of module names to install the import hook for.
+
+    Returns
+    -------
+    contextlib.AbstractContextManager
+        Context manager that installs the import hook on entry and removes it on exit.
+
+    """
+    return (
+        _install_import_hook(modules, RUNTIME_TYPECHECKER)
+        if RUNTIME_TYPECHECKER is not False
+        else contextlib.nullcontext()
+    )

unxt/units.py

@@ -0,0 +1,19 @@
+"""Working with units.
+
+The main features are:
+
+- ``unxt.units.unit``: a function to construct units.
+- ``unxt.units.unit_of``: a function to get the units of an object.
+
+"""
+
+__all__ = ("unit", "unit_of", "AbstractUnit")
+
+from .setup_package import install_import_hook
+
+with install_import_hook("unxt.units"):
+    from ._src.units import AbstractUnit
+    from unxt_api import unit, unit_of
+
+# Clean up the namespace
+del install_import_hook

unxt/unitsystems.py

@@ -0,0 +1,67 @@
+"""Systems of units.
+
+A unit system is a collection of units that are used together. In a unit system
+there are base units and derived units. Base units are the fundamental units of
+the system and derived units are constructed from the base units. For example,
+in the SI system, the base units are the meter, kilogram, second, ampere,
+kelvin, mole, and candela. Derived units are constructed from these base units,
+for example, the newton is a derived unit of force.
+
+`unxt` provides powerful tools for defining and working with unit systems. Unit
+systems can be statically defined (useful for many tools and development
+environments) or dynamically defined (useful for interactive environments and
+Python's general dynamism). Unit systems can be extended, compared, and used for
+decomposing units on quantities. There are many more features and tools for
+working with unit systems in `unxt`.
+
+"""
+# pylint: disable=duplicate-code
+
+__all__ = (
+    "unitsystem",
+    "unitsystem_of",
+    # classes
+    "AbstractUnitSystem",
+    "UNITSYSTEMS_REGISTRY",
+    "DimensionlessUnitSystem",
+    "LTMAUnitSystem",
+    # unit system instance
+    "galactic",
+    "dimensionless",
+    "solarsystem",
+    "si",
+    "cgs",
+    # unit system alias
+    "NAMED_UNIT_SYSTEMS",
+    # unit system flags
+    "AbstractUSysFlag",
+    "StandardUSysFlag",
+    "DynamicalSimUSysFlag",
+    # functions
+    "equivalent",
+)
+
+from .setup_package import install_import_hook
+
+with install_import_hook("unxt.unitsystems"):
+    from ._src.unitsystems import (
+        NAMED_UNIT_SYSTEMS,
+        UNITSYSTEMS_REGISTRY,
+        AbstractUnitSystem,
+        AbstractUSysFlag,
+        DimensionlessUnitSystem,
+        DynamicalSimUSysFlag,
+        LTMAUnitSystem,
+        StandardUSysFlag,
+        cgs,
+        dimensionless,
+        equivalent,
+        galactic,
+        si,
+        solarsystem,
+        unitsystem,  # TODO: move unitsystem to unxt_api
+    )
+    from unxt_api import unitsystem_of
+
+# Clean up the namespace
+del install_import_hook