fix self-transform and add examples

GalacticDynamics · Nov 25, 2024 · a6e2c09 · a6e2c09
1 parent 12f4b6f
commit a6e2c09
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 31 deletions.
diff --git a/src/coordinax/_src/vectors/d3/transform.py b/src/coordinax/_src/vectors/d3/transform.py
@@ -21,7 +21,6 @@
 from .mathspherical import MathSphericalPos, MathSphericalVel
 from .spherical import SphericalPos, SphericalVel
 from .spheroidal import ProlateSpheroidalPos, ProlateSpheroidalVel
-from coordinax._src.distance import AbstractDistance
 from coordinax._src.vectors.base import AbstractPos
 
 ###############################################################################
@@ -603,7 +602,8 @@ def represent_as(
     ...     Delta=Quantity(0.5, "kpc")
     ... )
     >>> print(cx.represent_as(vec, cx.CylindricalPos))
-    TODO: add
+    <CylindricalPos (rho[kpc], phi[deg], z[kpc])
+        [ 0.387 90.     0.894]>
 
     """
     Delta2 = current.Delta**2
@@ -622,9 +622,7 @@ def represent_as(
 def represent_as(
     current: CylindricalPos,
     target: type[ProlateSpheroidalPos],
-    *,
-    Delta: AbstractDistance | Quantity["length"],  # noqa: N803
-    **kwargs: Any,
+    Delta: Quantity["length"],  # noqa: N803
 ) -> ProlateSpheroidalPos:
     """CylindricalPos -> ProlateSpheroidalPos.
 
@@ -639,8 +637,8 @@ def represent_as(
     ...     z=Quantity(1, "kpc")
     ... )
     >>> print(vec.represent_as(cx.ProlateSpheroidalPos, Delta=Quantity(0.5, "kpc")))
-    <ProlateSpheroidalPos (mu[kpc2], nu[kpc2], phi[deg], Delta[kpc])
-        [ 2.133  0.117 90.     0.5  ]>
+    <ProlateSpheroidalPos (mu[kpc2], nu[kpc2], phi[deg])
+        [ 2.133  0.117 90.   ]>
 
     """
     R2 = current.rho**2
@@ -693,7 +691,24 @@ def represent_as(
 def represent_as(
     current: ProlateSpheroidalPos, target: type[CartesianPos3D], /, **kwargs: Any
 ) -> CartesianPos3D:
-    """ProlateSpheroidalPos -> CartesianPos3D."""
+    """ProlateSpheroidalPos -> CartesianPos3D.
+
+    Examples
+    --------
+    >>> from unxt import Quantity
+    >>> import coordinax as cx
+
+    >>> vec = cx.ProlateSpheroidalPos(
+    ...     mu=Quantity(1., "kpc2"),
+    ...     nu=Quantity(0.2, "kpc2"),
+    ...     phi=Quantity(90, "deg"),
+    ...     Delta=Quantity(0.5, "kpc")
+    ... )
+    >>> print(cx.represent_as(vec, cx.CartesianPos3D))
+    <CartesianPos3D (x[kpc], y[kpc], z[kpc])
+        [-1.693e-08  3.873e-01  8.944e-01]>
+
+    """
     cyl = represent_as(current, CylindricalPos)
     return represent_as(cyl, target)
 
@@ -702,34 +717,64 @@ def represent_as(
 def represent_as(
     current: CartesianPos3D,
     target: type[ProlateSpheroidalPos],
-    *,
-    Delta: AbstractDistance | Quantity["length"],  # noqa: N803
-    **kwargs: Any,
+    Delta: Quantity["length"],  # noqa: N803
 ) -> ProlateSpheroidalPos:
     """CartesianPos3D -> ProlateSpheroidalPos."""
     cyl = represent_as(current, CylindricalPos)
-    return represent_as(cyl, target, Delta=Delta)
+    return represent_as(cyl, target, Delta)
 
 
 @dispatch
 def represent_as(
-    current: ProlateSpheroidalPos, target: type[ProlateSpheroidalPos], /
+    current: ProlateSpheroidalPos, target: type[ProlateSpheroidalPos]
 ) -> ProlateSpheroidalPos:
-    """ProlateSpheroidalPos -> ProlateSpheroidalPos."""
+    """ProlateSpheroidalPos -> ProlateSpheroidalPos.
+
+    Examples
+    --------
+    >>> from unxt import Quantity
+    >>> import coordinax as cx
+
+    >>> vec = cx.ProlateSpheroidalPos(
+    ...     mu=Quantity(1., "kpc2"),
+    ...     nu=Quantity(0.2, "kpc2"),
+    ...     phi=Quantity(90, "deg"),
+    ...     Delta=Quantity(0.5, "kpc")
+    ... )
+    >>> print(cx.represent_as(vec, cx.ProlateSpheroidalPos))
+    <ProlateSpheroidalPos...>
+
+    """
     return current
 
 
 @dispatch
 def represent_as(
     current: ProlateSpheroidalPos,
     target: type[ProlateSpheroidalPos],
-    *,
-    Delta: AbstractDistance | Quantity["length"],  # noqa: N803
+    Delta: Quantity["length"],  # noqa: N803
+    /,
     **kwargs: Any,
 ) -> ProlateSpheroidalPos:
-    """ProlateSpheroidalPos -> ProlateSpheroidalPos."""
+    """ProlateSpheroidalPos -> ProlateSpheroidalPos.
+
+    Examples
+    --------
+    >>> from unxt import Quantity
+    >>> import coordinax as cx
+
+    >>> vec = cx.ProlateSpheroidalPos(
+    ...     mu=Quantity(1., "kpc2"),
+    ...     nu=Quantity(0.2, "kpc2"),
+    ...     phi=Quantity(90, "deg"),
+    ...     Delta=Quantity(0.5, "kpc")
+    ... )
+    >>> print(cx.represent_as(vec, cx.ProlateSpheroidalPos, Delta=Quantity(0.5, "kpc")))
+    <ProlateSpheroidalPos...>
+
+    """
     cyl = represent_as(current, CylindricalPos)
-    return represent_as(cyl, target, Delta=Delta)
+    return represent_as(cyl, target, Delta)
 
 
 # =============================================================================

diff --git a/tests/test_d3.py b/tests/test_d3.py
@@ -598,25 +598,20 @@ def test_prolatespheroidal_to_spherical(self, vector):
     def test_prolatespheroidal_to_prolatespheroidal(self, vector):
         """Test ``coordinax.represent_as(ProlateSpheroidalPos)``."""
         # Jit can copy
-        newvec = vector.represent_as(cx.ProlateSpheroidalPos, Delta=vector.Delta)
-        assert jnp.array_equal(newvec.mu, vector.mu)
+        newvec = vector.represent_as(cx.ProlateSpheroidalPos, vector.Delta)
+        assert jnp.allclose(newvec.mu.value, vector.mu.value)
         assert jnp.allclose(newvec.nu.value, vector.nu.value)
         assert jnp.array_equal(newvec.phi, vector.phi)
 
         # With a different focal length, should not be the same:
-        newvec = vector.represent_as(
-            cx.ProlateSpheroidalPos, Delta=u.Quantity(0.7, "kpc")
-        )
-        # assert not jnp.array_equal(newvec.mu, vector.mu)
-        assert not jnp.array_equal(newvec.nu, vector.nu)
-        # assert not jnp.array_equal(newvec.phi, vector.phi)
+        newvec = vector.represent_as(cx.ProlateSpheroidalPos, u.Quantity(0.5, "kpc"))
+        assert not jnp.allclose(newvec.mu.value, vector.mu.value)
+        assert not jnp.allclose(newvec.nu.value, vector.nu.value)
+        assert jnp.array_equal(newvec.phi, vector.phi)
 
         # The normal `represent_as` method should return the same object
-        # TODO: this fails because (I think) the output gets wrapped in a call of the
-        # initializer cx.ProlateSpheroidalPos(...), which then does not have the
-        # Delta=... argument
-        # newvec = cx.represent_as(vector, cx.ProlateSpheroidalPos)
-        # assert newvec is vector
+        newvec = cx.represent_as(vector, cx.ProlateSpheroidalPos)
+        assert newvec is vector
 
 
 class AbstractVel3DTest(AbstractVelTest):