From a2fcae83cece90ac310f9bcd08763117b44b009a Mon Sep 17 00:00:00 2001 From: Nathaniel Starkman Date: Fri, 13 Dec 2024 17:06:36 -0500 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(vecs):=20scaled=20cartesian=20?= =?UTF-8?q?3D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Nathaniel Starkman --- src/coordinax/_src/vectors/d3/__init__.py | 10 + .../_src/vectors/d3/scaled_cartesian.py | 223 ++++++++++++++++++ 2 files changed, 233 insertions(+) create mode 100644 src/coordinax/_src/vectors/d3/scaled_cartesian.py diff --git a/src/coordinax/_src/vectors/d3/__init__.py b/src/coordinax/_src/vectors/d3/__init__.py index fcffbbe4..4c33390f 100644 --- a/src/coordinax/_src/vectors/d3/__init__.py +++ b/src/coordinax/_src/vectors/d3/__init__.py @@ -35,10 +35,20 @@ "ProlateSpheroidalPos", "ProlateSpheroidalVel", "ProlateSpheroidalAcc", + # Scaled Cartesian + "ScaledCartesianPos3D", + "ScaledCartesianVel3D", + "ScaledCartesianAcc3D", # Generic "CartesianGeneric3D", ] +from scaled_cartesian import ( + ScaledCartesianAcc3D, + ScaledCartesianPos3D, + ScaledCartesianVel3D, +) + from .base import AbstractAcc3D, AbstractPos3D, AbstractVel3D from .base_spherical import ( AbstractSphericalAcc, diff --git a/src/coordinax/_src/vectors/d3/scaled_cartesian.py b/src/coordinax/_src/vectors/d3/scaled_cartesian.py new file mode 100644 index 00000000..90120f8a --- /dev/null +++ b/src/coordinax/_src/vectors/d3/scaled_cartesian.py @@ -0,0 +1,223 @@ +"""Built-in vector classes.""" + +__all__ = [ + "ScaledCartesianPos3D", + "ScaledCartesianVel3D", + "ScaledCartesianAcc3D", +] + +from functools import partial +from typing import final +from typing_extensions import override + +import equinox as eqx +import jax +from plum import dispatch + +from unxt import Quantity +from unxt.quantity import AbstractQuantity + +import coordinax._src.typing as ct +from .base import AbstractAcc3D, AbstractPos3D, AbstractVel3D +from .cartesian import CartesianPos3D +from .generic import CartesianGeneric3D +from coordinax._src.distances import ( + AbstractDistance, + BatchableDistance, + Distance, +) +from coordinax._src.utils import classproperty + +##################################################################### +# Position + + +@final +class ScaledCartesianPos3D(AbstractPos3D): + """Cartesian vector representation.""" + + r: BatchableDistance = eqx.field( + converter=Unless(AbstractDistance, partial(Distance.from_, dtype=float)) + ) + + s: Quantity["dimensionless"] = eqx.field(converter=Quantity[""].from_) + r"""Scaled x coordinate :math:`x \in (-\infty,+\infty)`.""" + + t: Quantity["dimensionless"] = eqx.field(converter=Quantity[""].from_) + r"""Scaled y coordinate :math:`y \in (-\infty,+\infty)`.""" + + u: Quantity["dimensionless"] = eqx.field(converter=Quantity[""].from_) + r"""Scaled z coordinate :math:`z \in (-\infty,+\infty)`.""" + + @override + @classproperty + @classmethod + def differential_cls(cls) -> type["ScaledCartesianVel3D"]: + """Return the differential of the class.""" + return ScaledCartesianVel3D + + @partial(jax.jit) + def norm(self) -> BatchableDistance: + return self.r + + +# ===================================================== +# Constructors + + +@ScaledCartesianPos3D.from_.dispatch # type: ignore[attr-defined, misc] +def from_( + cls: type[ScaledCartesianPos3D], + obj: AbstractQuantity, # TODO: Shaped[AbstractQuantity, "*batch 3"] + /, +) -> ScaledCartesianPos3D: + """Construct a scaled 3D Cartesian position. + + Examples + -------- + >>> import unxt as u + >>> import coordinax as cx + + >>> vec = cx.ScaledCartesianPos3D.from_(Quantity([1, 2, 3], "m")) + >>> vec + CartesianPos3D( + r=Distance(value=f32[], unit=Unit("m")), + s=Quantity[...](value=f32[], unit=Unit("")), + t=Quantity[...](value=f32[], unit=Unit("")), + u=Quantity[...](value=f32[], unit=Unit("")) + ) + + """ + cart = CartesianPos3D.from_(obj) + return cart.vconvert(cls) + + +# ===================================================== +# Functions + + +# from coordinax.vectors.funcs +@dispatch # type: ignore[misc] +@partial(eqx.filter_jit, inline=True) +def normalize_vector(obj: ScaledCartesianPos3D, /) -> ScaledCartesianPos3D: + """Return the norm of the vector. + + This has length 1. + + .. note:: + + The unit vector is dimensionless, even if the input vector has units. + This is because the unit vector is a ratio of two quantities: each + component and the norm of the vector. + + Returns + ------- + CartesianGeneric3D + The norm of the vector. + + Examples + -------- + >>> import coordinax as cx + >>> q = cx.ScaledCartesianPos3D.from_([1, 2, 3], "km") + >>> cx.vecs.normalize_vector(q) + CartesianGeneric3D( + x=Quantity[...]( value=f32[], unit=Unit(dimensionless) ), + y=Quantity[...]( value=f32[], unit=Unit(dimensionless) ), + z=Quantity[...]( value=f32[], unit=Unit(dimensionless) ) + ) + + """ + return CartesianGeneric3D(x=obj.x / obj.r, y=obj.y / obj.r, z=obj.z / obj.r) + + +##################################################################### +# Velocity + + +@final +class ScaledCartesianVel3D(AbstractVel3D): + """Scaled 3D Cartesian velocity.""" + + d_r: ct.BatchableSpeed = eqx.field(converter=Quantity["speed"].from_) + r"""speed :math:`dx/dt \in [-\infty, \infty].""" + + d_s: ct.BatchableSpeed = eqx.field(converter=Quantity["frequency"].from_) + r"""s speed :math:`dx/dt \in [-\infty, \infty] [frequency].""" + + d_t: ct.BatchableSpeed = eqx.field(converter=Quantity["frequency"].from_) + r"""s speed :math:`dx/dt \in [-\infty, \infty] [frequency].""" + + d_u: ct.BatchableSpeed = eqx.field(converter=Quantity["frequency"].from_) + r"""Z speed :math:`dz/dt \in [-\infty, \infty] [frequency].""" + + @override + @classproperty + @classmethod + def integral_cls(cls) -> type[ScaledCartesianPos3D]: + return ScaledCartesianPos3D + + @override + @classproperty + @classmethod + def differential_cls(cls) -> type["ScaledCartesianAcc3D"]: + return ScaledCartesianAcc3D + + @partial(eqx.filter_jit, inline=True) + def norm(self, _: AbstractPos3D | None = None, /) -> ct.BatchableSpeed: + """Return the norm of the vector. + + Examples + -------- + >>> import coordinax as cx + >>> c = cx.CartesianVel3D.from_([1, 2, 3], "km/s") + >>> c.norm() + Quantity['speed'](Array(3.7416575, dtype=float32), unit='km / s') + + """ + return self.d_r + + +##################################################################### +# Acceleration + + +@final +class ScaledCartesianAcc3D(AbstractAcc3D): + """Scaled 3D Cartesian acceleration.""" + + d2_r: ct.BatchableAcc = eqx.field(converter=Quantity["acceleration"].from_) + r"""scaled radial acceleration :math:`d^2x/dt^2 \in [-\infty, \infty].""" + + d2_s: ct.BatchableAcc = eqx.field(converter=Quantity["frequency drift"].from_) + r"""s acceleration :math:`d^2s/dt^2 \in [-\infty, \infty] [time^-2].""" + + d2_t: ct.BatchableAcc = eqx.field(converter=Quantity["frequency drift"].from_) + r"""t acceleration :math:`d^2t/dt^2 \in [-\infty, \infty] [time^-2].""" + + d2_u: ct.BatchableAcc = eqx.field(converter=Quantity["frequency drift"].from_) + r"""u acceleration :math:`d^2u/dt^2 \in [-\infty, \infty] [time^-2].""" + + @classproperty + @classmethod + def integral_cls(cls) -> type[ScaledCartesianVel3D]: + return ScaledCartesianVel3D + + # ----------------------------------------------------- + # Methods + + @override + @partial(jax.jit, inline=True) + def norm( + self, _: AbstractVel3D | None = None, __: AbstractPos3D | None = None, / + ) -> ct.BatchableAcc: + """Return the norm of the vector. + + Examples + -------- + >>> import coordinax as cx + >>> c = cx.vecs.CartesianAcc3D.from_([1, 2, 3], "km/s2") + >>> c.norm() + Quantity['acceleration'](Array(3.7416575, dtype=float32), unit='km / s2') + + """ + return self.d2_r