vector norm convenience method

Signed-off-by: nstarman <[email protected]>

pyproject.toml

@@ -98,7 +98,10 @@
     ]
 
 [tool.mypy]
-  disable_error_code = ["no-redef"]
+  disable_error_code = [
+    "no-redef",  # for plum-dispatch
+    "name-defined",  # for jaxtyping
+  ]
   disallow_incomplete_defs = false
   disallow_untyped_defs = false
   enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]

src/vector/_base.py

@@ -136,6 +136,12 @@ def represent_as(self, target: type[VT], /, *args: Any, **kwargs: Any) -> VT:
 
         return represent_as(self, target, **kwargs)
 
+    @abstractmethod
+    def norm(self) -> Quantity:
+        """Return the norm of the vector."""
+        # TODO: make a generic method that works on all dimensions
+        raise NotImplementedError
+
 
 class AbstractVectorDifferential(AbstractVectorBase):
     """Abstract representation of vector differentials in different systems."""

src/vector/_d1/base.py

@@ -3,14 +3,25 @@
 __all__ = ["Abstract1DVector", "Abstract1DVectorDifferential"]
 
 
+from functools import partial
+
 import equinox as eqx
+import jax
+from jax_quantity import Quantity
 
 from vector._base import AbstractVector, AbstractVectorDifferential
 
 
 class Abstract1DVector(AbstractVector):
     """Abstract representation of 1D coordinates in different systems."""
 
+    @partial(jax.jit)
+    def norm(self) -> Quantity["length"]:
+        """Return the norm of the vector."""
+        from .builtin import Cartesian1DVector  # pylint: disable=C0415
+
+        return self.represent_as(Cartesian1DVector).norm()
+
 
 class Abstract1DVectorDifferential(AbstractVectorDifferential):
     """Abstract representation of 1D differentials in different systems."""

src/vector/_d1/builtin.py

@@ -9,9 +9,12 @@
     "RadialDifferential",
 ]
 
+from functools import partial
 from typing import ClassVar, final
 
+import array_api_jax_compat as xp
 import equinox as eqx
+import jax
 
 from vector._checks import check_r_non_negative
 from vector._typing import BatchableLength, BatchableSpeed
@@ -30,6 +33,11 @@ class Cartesian1DVector(Abstract1DVector):
     x: BatchableLength = eqx.field(converter=converter_quantity_array)
     r"""X coordinate :math:`x \in (-\infty,+\infty)`."""
 
+    @partial(jax.jit)
+    def norm(self) -> BatchableLength:
+        """Return the norm of the vector."""
+        return xp.abs(self.x)
+
 
 @final
 class RadialVector(Abstract1DVector):

src/vector/_d2/base.py

@@ -3,14 +3,25 @@
 __all__ = ["Abstract2DVector", "Abstract2DVectorDifferential"]
 
 
+from functools import partial
+
 import equinox as eqx
+import jax
+from jax_quantity import Quantity
 
 from vector._base import AbstractVector, AbstractVectorDifferential
 
 
 class Abstract2DVector(AbstractVector):
     """Abstract representation of 2D coordinates in different systems."""
 
+    @partial(jax.jit)
+    def norm(self) -> Quantity["length"]:
+        """Return the norm of the vector."""
+        from .builtin import Cartesian2DVector  # pylint: disable=C0415
+
+        return self.represent_as(Cartesian2DVector).norm()
+
 
 class Abstract2DVectorDifferential(AbstractVectorDifferential):
     """Abstract representation of 2D vector differentials."""

src/vector/_d2/builtin.py

@@ -11,9 +11,12 @@
     "PolarDifferential",
 ]
 
+from functools import partial
 from typing import ClassVar, final
 
+import array_api_jax_compat as xp
 import equinox as eqx
+import jax
 
 from vector._checks import check_phi_range, check_r_non_negative
 from vector._typing import (
@@ -40,6 +43,11 @@ class Cartesian2DVector(Abstract2DVector):
     y: BatchableLength = eqx.field(converter=converter_quantity_array)
     r"""Y coordinate :math:`y \in (-\infty,+\infty)`."""
 
+    @partial(jax.jit)
+    def norm(self) -> BatchableLength:
+        """Return the norm of the vector."""
+        return xp.sqrt(self.x**2 + self.y**2)
+
 
 @final
 class PolarVector(Abstract2DVector):
@@ -59,6 +67,11 @@ def __check_init__(self) -> None:
         check_r_non_negative(self.r)
         check_phi_range(self.phi)
 
+    @partial(jax.jit)
+    def norm(self) -> BatchableLength:
+        """Return the norm of the vector."""
+        return self.r
+
 
 # class LnPolarVector(Abstract2DVector):
 #     """Log-polar vector representation."""

src/vector/_d3/base.py

@@ -3,14 +3,25 @@
 __all__ = ["Abstract3DVector", "Abstract3DVectorDifferential"]
 
 
+from functools import partial
+
 import equinox as eqx
+import jax
+from jax_quantity import Quantity
 
 from vector._base import AbstractVector, AbstractVectorDifferential
 
 
 class Abstract3DVector(AbstractVector):
     """Abstract representation of 3D coordinates in different systems."""
 
+    @partial(jax.jit)
+    def norm(self) -> Quantity["length"]:
+        """Return the norm of the vector."""
+        from .builtin import Cartesian3DVector  # pylint: disable=C0415
+
+        return self.represent_as(Cartesian3DVector).norm()
+
 
 class Abstract3DVectorDifferential(AbstractVectorDifferential):
     """Abstract representation of 3D vector differentials."""

src/vector/_d3/builtin.py

@@ -11,9 +11,12 @@
     "CylindricalDifferential",
 ]
 
+from functools import partial
 from typing import ClassVar, final
 
+import array_api_jax_compat as xp
 import equinox as eqx
+import jax
 
 from vector._checks import check_phi_range, check_r_non_negative, check_theta_range
 from vector._typing import (
@@ -43,6 +46,11 @@ class Cartesian3DVector(Abstract3DVector):
     z: BatchableLength = eqx.field(converter=converter_quantity_array)
     r"""Z coordinate :math:`z \in (-\infty,+\infty)`."""
 
+    @partial(jax.jit)
+    def norm(self) -> BatchableLength:
+        """Return the norm of the vector."""
+        return xp.sqrt(self.x**2 + self.y**2 + self.z**2)
+
 
 @final
 class SphericalVector(Abstract3DVector):
@@ -63,6 +71,11 @@ def __check_init__(self) -> None:
         check_theta_range(self.theta)
         check_phi_range(self.phi)
 
+    @partial(jax.jit)
+    def norm(self) -> BatchableLength:
+        """Return the norm of the vector."""
+        return self.r
+
 
 @final
 class CylindricalVector(Abstract3DVector):
@@ -82,6 +95,11 @@ def __check_init__(self) -> None:
         check_r_non_negative(self.rho)
         check_phi_range(self.phi)
 
+    @partial(jax.jit)
+    def norm(self) -> BatchableLength:
+        """Return the norm of the vector."""
+        return xp.sqrt(self.rho**2 + self.z**2)
+
 
 ##############################################################################
 # Differential