diff --git a/src/unxt/quantity.pyi b/src/unxt/quantity.pyi new file mode 100644 index 00000000..8a2e6469 --- /dev/null +++ b/src/unxt/quantity.pyi @@ -0,0 +1,178 @@ +from typing import ( + Any, + Generic, + LiteralString, + Self, + TypeGuard, + TypeVar, + final, + overload, +) + +import jax +import numpy as np +import optype as op +from jaxtyping import Array, ArrayLike, Bool +from numpy.typing import NDArray +from quax import ArrayValue + +from unxt.units import AbstractUnit +from unxt.unitsystems import AbstractUnitSystem + +_T = TypeVar("_T", bound=LiteralString) + +@final +class StaticValue: + def __init__(self, array: op.numpy.CanArrayND[Any, Any], /) -> None: ... + @property + def array(self) -> NDArray[Any]: ... + @classmethod + def from_(cls, *args: Any, **kwargs: Any) -> Self: ... + def __eq__(self, other: object, /) -> bool | NDArray[np.bool_]: ... # type: ignore[override] + def __ne__(self, other: object, /) -> bool | NDArray[np.bool_]: ... # type: ignore[override] + def __lt__(self, other: op.CanLessThan[Any], /) -> NDArray[np.bool_]: ... + def __le__(self, other: op.CanLessThanOrEqual[Any], /) -> NDArray[np.bool_]: ... + def __gt__(self, other: op.CanGreaterThan[Any], /) -> NDArray[np.bool_]: ... + def __ge__(self, other: op.CanGreaterThanOrEqual[Any], /) -> NDArray[np.bool_]: ... + def __hash__(self) -> int: ... + def __add__(self, other: op.CanAdd[Any]) -> Any: ... + def __radd__(self, other: op.CanAdd[Any]) -> Any: ... + def __sub__(self, other: op.CanSub[Any]) -> Any: ... + def __rsub__(self, other: op.CanSub[Any]) -> Any: ... + def __mul__(self, other: op.CanMul[Any]) -> Any: ... + def __rmul__(self, other: op.CanMul[Any]) -> Any: ... + def __truediv__(self, other: op.CanTrueDiv[Any]) -> Any: ... + def __rtruediv__(self, other: op.CanTrueDiv[Any]) -> Any: ... + def __floordiv__(self, other: op.CanFloorDiv[Any]) -> Any: ... + def __rfloordiv__(self, other: op.CanFloorDiv[Any]) -> Any: ... + def __pow__(self, other: op.CanPow[Any]) -> Any: ... + def __rpow__(self, other: op.CanPow[Any]) -> Any: ... + def __mod__(self, other: op.CanMod[Any]) -> Any: ... + def __rmod__(self, other: op.CanMod[Any]) -> Any: ... + def __matmul__(self, other: op.CanMatmul[Any]) -> Any: ... + def __rmatmul__(self, other: op.CanMatmul[Any]) -> Any: ... + def __neg__(self) -> Self: ... + def __pos__(self) -> Self: ... + def __abs__(self) -> Self: ... + def __len__(self) -> int: ... + def __iter__(self) -> Any: ... + def __getitem__(self, key: Any) -> Any: ... + +class AllowValue: ... + +class AbstractQuantity(ArrayValue): + # Core public attributes + value: Array | StaticValue + unit: AbstractUnit + + # Implements quax.ArrayValue abstract methods + def aval(self) -> jax.core.ShapedArray: ... + def materialise(self) -> Array: ... + + # Core public methods + def uconvert(self, unit: Any) -> Self: ... + def ustrip(self) -> Any: ... + + # NumpyBinaryOpsMixin + def __add__(self, other: Any) -> Self: ... + def __radd__(self, other: Any) -> Self: ... + def __sub__(self, other: Any) -> Self: ... + def __rsub__(self, other: Any) -> Self: ... + def __mul__(self, other: Any) -> AbstractQuantity: ... + def __rmul__(self, other: Any) -> AbstractQuantity: ... + def __matmul__(self, other: Any) -> AbstractQuantity: ... + def __rmatmul__(self, other: Any) -> AbstractQuantity: ... + def __truediv__(self, other: Any) -> AbstractQuantity: ... + def __rtruediv__(self, other: Any) -> AbstractQuantity: ... + def __floordiv__(self, other: Any) -> AbstractQuantity: ... + def __rfloordiv__(self, other: Any) -> AbstractQuantity: ... + def __mod__(self, other: Any) -> AbstractQuantity: ... + def __rmod__(self, other: Any) -> AbstractQuantity: ... + def __pow__(self, other: Any) -> AbstractQuantity: ... + def __rpow__(self, other: Any) -> AbstractQuantity: ... + + # NumpyComparisonMixin + def __eq__(self, other: object) -> Bool[Array, ...]: ... + def __ne__(self, other: object) -> Bool[Array, ...]: ... # type: ignore[override] + def __lt__(self, other: Any) -> Bool[Array, ...]: ... + def __le__(self, other: Any) -> Bool[Array, ...]: ... + def __gt__(self, other: Any) -> Bool[Array, ...]: ... + def __ge__(self, other: Any) -> Bool[Array, ...]: ... + + # NumpyUnaryMixin + def __neg__(self) -> Self: ... + def __pos__(self) -> Self: ... + def __abs__(self) -> Self: ... + + # NumpyRoundMixin / TruncMixin / FloorMixin / CeilMixin + def __round__(self, ndigits: int = ...) -> Self: ... + def __trunc__(self) -> Self: ... + def __floor__(self) -> Self: ... + def __ceil__(self) -> Self: ... + @classmethod + def from_(cls, *args: Any, **kwargs: Any) -> Self: ... + def __len__(self) -> int: ... + def __iter__(self) -> Any: ... + def __getitem__(self, key: Any) -> Self: ... + +class AbstractAngle(AbstractQuantity): + def wrap_to( + self, /, min: AbstractQuantity, max: AbstractQuantity + ) -> AbstractAngle: ... + +class AbstractParametricQuantity(AbstractQuantity, Generic[_T]): ... + +class Quantity(AbstractParametricQuantity[_T]): + def __init__(self, value: ArrayLike, unit: str | AbstractUnit) -> None: ... + +Q = Quantity + +class BareQuantity(AbstractQuantity): + def __init__(self, value: ArrayLike, unit: str | AbstractUnit) -> None: ... + +class StaticQuantity(AbstractParametricQuantity[_T]): + value: StaticValue + def __init__( + self, value: np.ndarray | StaticValue, unit: str | AbstractUnit + ) -> None: ... + +class Angle(AbstractAngle): + def __init__(self, value: ArrayLike, unit: str | AbstractUnit) -> None: ... + +def is_any_quantity(obj: Any, /) -> TypeGuard[AbstractQuantity]: ... +def convert_to_quantity_value(obj: Any, /) -> Any: ... + +# ustrip: no-unit form, unit-first form, and AllowValue flag forms +@overload +def ustrip(x: StaticQuantity, /) -> np.ndarray: ... +@overload +def ustrip(x: AbstractQuantity, /) -> Array | np.ndarray: ... +@overload +def ustrip( + u: str | AbstractUnit | AbstractUnitSystem, x: AbstractQuantity, / +) -> Array | np.ndarray: ... +@overload +def ustrip(flag: type[AllowValue], x: Any, /) -> Any: ... +@overload +def ustrip( + flag: type[AllowValue], u: str | AbstractUnit | AbstractUnitSystem, x: Any, / +) -> Any: ... +def uconvert( + u: str | AbstractUnit | AbstractUnitSystem, x: AbstractQuantity, / +) -> AbstractQuantity: ... +@overload +def uconvert_value( + uto: str | AbstractUnit, ufrom: str | AbstractUnit, x: ArrayLike, / +) -> ArrayLike: ... +@overload +def uconvert_value( + uto: AbstractUnitSystem, ufrom: str | AbstractUnit, x: ArrayLike, / +) -> ArrayLike: ... +@overload +def uconvert_value( + uto: Any, ufrom: Any, x: AbstractQuantity, / +) -> AbstractQuantity: ... +def is_unit_convertible(to_unit: Any, from_: Any, /) -> bool: ... +def wrap_to( + angle: AbstractQuantity, min: AbstractQuantity, max: AbstractQuantity, / +) -> AbstractQuantity: ...