feat: q as vec for p represent_as

pyproject.toml

@@ -87,6 +87,7 @@
   filterwarnings = [
     "error",
     "ignore:ast\\.Str is deprecated and will be removed in Python 3\\.14:DeprecationWarning",
+    "ignore:Explicitly requested dtype <class 'jax.numpy\\.float64'> requested in astype is not available"
   ]
   log_cli_level = "INFO"
   minversion = "6.0"

src/coordinax/_transform.py

@@ -30,46 +30,111 @@
     (
         Abstract1DVectorDifferential,
         type[Abstract1DVectorDifferential],  # type: ignore[misc]
-        AbstractVector,
+        AbstractVector | Quantity["length"],
     ),
     (
         Abstract2DVectorDifferential,
         type[Abstract2DVectorDifferential],  # type: ignore[misc]
-        AbstractVector,
+        AbstractVector | Quantity["length"],
     ),
     (
         Abstract3DVectorDifferential,
         type[Abstract3DVectorDifferential],  # type: ignore[misc]
-        AbstractVector,
+        AbstractVector | Quantity["length"],
     ),
 )
 def represent_as(
     current: AbstractVectorDifferential,
     target: type[AbstractVectorDifferential],
-    position: AbstractVector,
+    position: AbstractVector | Quantity["length"],
     /,
     **kwargs: Any,
 ) -> AbstractVectorDifferential:
-    """Abstract3DVectorDifferential -> Cartesian -> Abstract3DVectorDifferential.
+    """AbstractVectorDifferential -> Cartesian -> AbstractVectorDifferential.
+
+    This is the base case for the transformation of vector differentials.
+
+    Parameters
+    ----------
+    current : AbstractVectorDifferential
+        The vector differential to transform.
+    target : type[AbstractVectorDifferential]
+        The target type of the vector differential.
+    position : AbstractVector
+        The position vector used to transform the differential.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Examples
+    --------
+    >>> import coordinax as cx
+    >>> from unxt import Quantity
+
+    Let's start in 1D:
+
+    >>> q = cx.Cartesian1DVector(x=Quantity(1.0, "km"))
+    >>> p = cx.CartesianDifferential1D(d_x=Quantity(1.0, "km/s"))
+    >>> cx.represent_as(p, cx.RadialDifferential, q)
+    RadialDifferential( d_r=Quantity[...]( value=f32[], unit=Unit("km / s") ) )
+
+    Now in 2D:
+
+    >>> q = cx.Cartesian2DVector.constructor(Quantity([1.0, 2.0], "km"))
+    >>> p = cx.CartesianDifferential2D.constructor(Quantity([1.0, 2.0], "km/s"))
+    >>> cx.represent_as(p, cx.PolarDifferential, q)
+    PolarDifferential(
+      d_r=Quantity[...]( value=f32[], unit=Unit("km / s") ),
+      d_phi=Quantity[...]( value=f32[], unit=Unit("rad / s") )
+    )
+
+    And in 3D:
+
+    >>> q = cx.Cartesian3DVector.constructor(Quantity([1.0, 2.0, 3.0], "km"))
+    >>> p = cx.CartesianDifferential3D.constructor(Quantity([1.0, 2.0, 3.0], "km/s"))
+    >>> cx.represent_as(p, cx.SphericalDifferential, q)
+    SphericalDifferential(
+      d_r=Quantity[...]( value=f32[], unit=Unit("km / s") ),
+      d_theta=Quantity[...]( value=f32[], unit=Unit("rad / s") ),
+      d_phi=Quantity[...]( value=f32[], unit=Unit("rad / s") )
+    )
+
+    If given a position as a Quantity, it will be converted to the appropriate
+    Cartesian vector:
+
+    >>> p = cx.CartesianDifferential3D.constructor(Quantity([1.0, 2.0, 3.0], "km/s"))
+    >>> cx.represent_as(p, cx.SphericalDifferential, Quantity([1.0, 2.0, 3.0], "km"))
+    SphericalDifferential(
+      d_r=Quantity[...]( value=f32[], unit=Unit("km / s") ),
+      d_theta=Quantity[...]( value=f32[], unit=Unit("rad / s") ),
+      d_phi=Quantity[...]( value=f32[], unit=Unit("rad / s") )
+    )
 
-    This is the base case for the transformation of 1D vector differentials.
     """
     # TODO: not require the shape munging / support more shapes
     shape = current.shape
     flat_shape = prod(shape)
-    position = position.reshape(flat_shape)  # flattened
+
+    # Parse the position to an AbstractVector
+    if isinstance(position, AbstractVector):
+        posvec = position
+    else:  # Q -> Cart<X>D
+        posvec = current.integral_cls._cartesian_cls.constructor(  # noqa: SLF001
+            position
+        )
+
+    posvec = posvec.reshape(flat_shape)  # flattened
 
     # Start by transforming the position to the type required by the
     # differential to construct the Jacobian.
-    current_position = represent_as(position, current.integral_cls, **kwargs)
+    current_pos = represent_as(posvec, current.integral_cls, **kwargs)
 
     # Takes the Jacobian through the representation transformation function.  This
     # returns a representation of the target type, where the value of each field the
     # corresponding row of the Jacobian. The value of the field is a Quantity with
     # the correct numerator unit (of the Jacobian row). The value is a Vector of the
     # original type, with fields that are the columns of that row, but with only the
     # denomicator's units.
-    jac_nested_vecs = jac_rep_as(current_position, target.integral_cls)
+    jac_nested_vecs = jac_rep_as(current_pos, target.integral_cls)
 
     # This changes the Jacobian to be a dictionary of each row, with the value
     # being that row's column as a dictionary, now with the correct units for

tests/test_d2.py

@@ -1,7 +1,6 @@
 """Test :mod:`coordinax._d2`."""
 
 import astropy.units as u
-import jax.numpy as jnp
 import pytest
 
 import quaxed.array_api as xp
@@ -83,7 +82,7 @@ def test_cartesian2d_to_polar(self, vector):
 
         assert isinstance(polar, PolarVector)
         assert qnp.array_equal(polar.r, qnp.hypot(vector.x, vector.y))
-        assert jnp.allclose(
+        assert qnp.allclose(
             polar.phi.value,
             xp.asarray([1.3734008, 1.2490457, 1.1659045, 1.1071488]),
         )
@@ -115,7 +114,7 @@ def test_cartesian2d_to_spherical(self, vector):
 
         assert isinstance(spherical, SphericalVector)
         assert qnp.array_equal(spherical.r, qnp.hypot(vector.x, vector.y))
-        assert jnp.allclose(
+        assert qnp.allclose(
             spherical.phi.to_value(u.rad),
             xp.asarray([1.3734008, 1.2490457, 1.1659045, 1.1071488]),
         )
@@ -159,7 +158,7 @@ def test_polar_to_cartesian1d(self, vector):
         cart1d = vector.represent_as(Cartesian1DVector)
 
         assert isinstance(cart1d, Cartesian1DVector)
-        assert jnp.allclose(
+        assert qnp.allclose(
             cart1d.x.to_value(u.kpc), xp.asarray([1.0, 1.0806047, -1.2484405, -3.95997])
         )
         assert qnp.array_equal(cart1d.x, vector.r * xp.cos(vector.phi))
@@ -180,11 +179,11 @@ def test_polar_to_cartesian2d(self, vector):
         assert qnp.array_equal(
             cart2d.x, Quantity([1.0, 1.0806046, -1.2484405, -3.95997], u.kpc)
         )
-        assert jnp.allclose(cart2d.x.value, (vector.r * xp.cos(vector.phi)).value)
+        assert qnp.allclose(cart2d.x.value, (vector.r * xp.cos(vector.phi)).value)
         assert qnp.array_equal(
             cart2d.y, Quantity([0.0, 1.6829419, 2.7278922, 0.56448], u.kpc)
         )
-        assert jnp.allclose(cart2d.y.value, (vector.r * xp.sin(vector.phi)).value)
+        assert qnp.allclose(cart2d.y.value, (vector.r * xp.sin(vector.phi)).value)
 
     def test_polar_to_polar(self, vector):
         """Test ``coordinax.represent_as(PolarVector)``."""