From a6e2c09b36fa67a9287e8d3589e9865b62267c08 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 25 Nov 2024 14:14:11 -0500 Subject: [PATCH] fix self-transform and add examples --- src/coordinax/_src/vectors/d3/transform.py | 81 +++++++++++++++++----- tests/test_d3.py | 21 +++--- 2 files changed, 71 insertions(+), 31 deletions(-) diff --git a/src/coordinax/_src/vectors/d3/transform.py b/src/coordinax/_src/vectors/d3/transform.py index b1ade82..9a75b41 100644 --- 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 + """ 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"))) - + """ 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)) + + + """ cyl = represent_as(current, CylindricalPos) return represent_as(cyl, target) @@ -702,20 +717,34 @@ 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)) + + + """ return current @@ -723,13 +752,29 @@ def represent_as( 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"))) + + + """ 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 index eec2f9b..1c21648 100644 --- 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):