diff --git a/src/vector/_d3/compat.py b/src/vector/_d3/compat.py index 75b49b86..900773ab 100644 --- a/src/vector/_d3/compat.py +++ b/src/vector/_d3/compat.py @@ -8,7 +8,7 @@ import astropy.units as apyu from jax_quantity import Quantity from jaxtyping import Shaped -from plum import conversion_method +from plum import conversion_method, convert from vector._utils import dataclass_values, full_shaped @@ -53,9 +53,9 @@ def vec_diff_to_q( def cart3_to_apycart3(obj: Cartesian3DVector, /) -> apyc.CartesianRepresentation: """`vector.Cartesian3DVector` -> `astropy.CartesianRepresentation`.""" return apyc.CartesianRepresentation( - x=obj.x.as_type(apyu.Quantity), - y=obj.y.as_type(apyu.Quantity), - z=obj.z.as_type(apyu.Quantity), + x=convert(obj.x, apyu.Quantity), + y=convert(obj.y, apyu.Quantity), + z=convert(obj.z, apyu.Quantity), ) @@ -76,9 +76,9 @@ def apycart3_to_cart3(obj: apyc.CartesianRepresentation, /) -> Cartesian3DVector def sph_to_apysph(obj: SphericalVector, /) -> apyc.PhysicsSphericalRepresentation: """`vector.SphericalVector` -> `astropy.PhysicsSphericalRepresentation`.""" return apyc.PhysicsSphericalRepresentation( - r=obj.r.as_type(apyu.Quantity), - phi=obj.phi.as_type(apyu.Quantity), - theta=obj.theta.as_type(apyu.Quantity), + r=convert(obj.r, apyu.Quantity), + phi=convert(obj.phi, apyu.Quantity), + theta=convert(obj.theta, apyu.Quantity), ) @@ -99,9 +99,9 @@ def apysph_to_sph(obj: apyc.PhysicsSphericalRepresentation, /) -> SphericalVecto def cyl_to_apycyl(obj: CylindricalVector, /) -> apyc.CylindricalRepresentation: """`vector.CylindricalVector` -> `astropy.CylindricalRepresentation`.""" return apyc.CylindricalRepresentation( - rho=obj.rho.as_type(apyu.Quantity), - phi=obj.phi.as_type(apyu.Quantity), - z=obj.z.as_type(apyu.Quantity), + rho=convert(obj.rho, apyu.Quantity), + phi=convert(obj.phi, apyu.Quantity), + z=convert(obj.z, apyu.Quantity), ) @@ -123,9 +123,9 @@ def diffcart3_to_apycart3( ) -> apyc.CartesianDifferential: """`vector.CartesianDifferential3D` -> `astropy.CartesianDifferential`.""" return apyc.CartesianDifferential( - d_x=obj.d_x.as_type(apyu.Quantity), - d_y=obj.d_y.as_type(apyu.Quantity), - d_z=obj.d_z.as_type(apyu.Quantity), + d_x=convert(obj.d_x, apyu.Quantity), + d_y=convert(obj.d_y, apyu.Quantity), + d_z=convert(obj.d_z, apyu.Quantity), ) @@ -152,9 +152,9 @@ def diffsph_to_apysph( ) -> apyc.PhysicsSphericalDifferential: """`vector.SphericalDifferential` -> `astropy.PhysicsSphericalDifferential`.""" return apyc.PhysicsSphericalDifferential( - d_r=obj.d_r.as_type(apyu.Quantity), - d_phi=obj.d_phi.as_type(apyu.Quantity), - d_theta=obj.d_theta.as_type(apyu.Quantity), + d_r=convert(obj.d_r, apyu.Quantity), + d_phi=convert(obj.d_phi, apyu.Quantity), + d_theta=convert(obj.d_theta, apyu.Quantity), ) @@ -179,9 +179,9 @@ def apysph_to_diffsph( def diffcyl_to_apycyl(obj: CylindricalDifferential, /) -> apyc.CylindricalDifferential: """`vector.CylindricalDifferential` -> `astropy.CylindricalDifferential`.""" return apyc.CylindricalDifferential( - d_rho=obj.d_rho.as_type(apyu.Quantity), - d_phi=obj.d_phi.as_type(apyu.Quantity), - d_z=obj.d_z.as_type(apyu.Quantity), + d_rho=convert(obj.d_rho, apyu.Quantity), + d_phi=convert(obj.d_phi, apyu.Quantity), + d_z=convert(obj.d_z, apyu.Quantity), ) diff --git a/tests/test_d3.py b/tests/test_d3.py index 38cbecc6..11943ce7 100644 --- a/tests/test_d3.py +++ b/tests/test_d3.py @@ -120,9 +120,9 @@ def test_cartesian3d_to_cartesian3d_astropy(self, vector, apyvector): """Test Astropy equivalence.""" newvec = vector.represent_as(Cartesian3DVector) - assert np.allclose(newvec.x.as_type(u.Quantity), apyvector.x) - assert np.allclose(newvec.y.as_type(u.Quantity), apyvector.y) - assert np.allclose(newvec.z.as_type(u.Quantity), apyvector.z) + assert np.allclose(convert(newvec.x, u.Quantity), apyvector.x) + assert np.allclose(convert(newvec.y, u.Quantity), apyvector.y) + assert np.allclose(convert(newvec.z, u.Quantity), apyvector.z) def test_cartesian3d_to_spherical(self, vector): """Test ``vector.represent_as(SphericalVector)``.""" @@ -145,9 +145,9 @@ def test_cartesian3d_to_spherical_astropy(self, vector, apyvector): sph = vector.represent_as(SphericalVector) apysph = apyvector.represent_as(apyc.PhysicsSphericalRepresentation) - assert np.allclose(sph.r.as_type(u.Quantity), apysph.r) - assert np.allclose(sph.theta.as_type(u.Quantity), apysph.theta) - assert np.allclose(sph.phi.as_type(u.Quantity), apysph.phi) + assert np.allclose(convert(sph.r, u.Quantity), apysph.r) + assert np.allclose(convert(sph.theta, u.Quantity), apysph.theta) + assert np.allclose(convert(sph.phi, u.Quantity), apysph.phi) def test_cartesian3d_to_cylindrical(self, vector): """Test ``vector.represent_as(CylindricalVector)``.""" @@ -166,9 +166,9 @@ def test_cartesian3d_to_cylindrical_astropy(self, vector, apyvector): cyl = vector.represent_as(CylindricalVector) apycyl = apyvector.represent_as(apyc.CylindricalRepresentation) - assert np.allclose(cyl.rho.as_type(u.Quantity), apycyl.rho) - assert np.allclose(cyl.z.as_type(u.Quantity), apycyl.z) - assert np.allclose(cyl.phi.as_type(u.Quantity), apycyl.phi) + assert np.allclose(convert(cyl.rho, u.Quantity), apycyl.rho) + assert np.allclose(convert(cyl.z, u.Quantity), apycyl.z) + assert np.allclose(convert(cyl.phi, u.Quantity), apycyl.phi) class TestSphericalVector: @@ -264,9 +264,9 @@ def test_spherical_to_cartesian3d_astropy(self, vector, apyvector): cart3d = vector.represent_as(Cartesian3DVector) apycart3 = apyvector.represent_as(apyc.CartesianRepresentation) - assert np.allclose(cart3d.x.as_type(u.Quantity), apycart3.x) - assert np.allclose(cart3d.y.as_type(u.Quantity), apycart3.y) - assert np.allclose(cart3d.z.as_type(u.Quantity), apycart3.z) + assert np.allclose(convert(cart3d.x, u.Quantity), apycart3.x) + assert np.allclose(convert(cart3d.y, u.Quantity), apycart3.y) + assert np.allclose(convert(cart3d.z, u.Quantity), apycart3.z) def test_spherical_to_spherical(self, vector): """Test ``vector.represent_as(SphericalVector)``.""" @@ -283,9 +283,9 @@ def test_spherical_to_spherical_astropy(self, vector, apyvector): sph = vector.represent_as(SphericalVector) apysph = apyvector.represent_as(apyc.PhysicsSphericalRepresentation) - assert np.allclose(sph.r.as_type(u.Quantity), apysph.r) - assert np.allclose(sph.theta.as_type(u.Quantity), apysph.theta) - assert np.allclose(sph.phi.as_type(u.Quantity), apysph.phi) + assert np.allclose(convert(sph.r, u.Quantity), apysph.r) + assert np.allclose(convert(sph.theta, u.Quantity), apysph.theta) + assert np.allclose(convert(sph.phi, u.Quantity), apysph.phi) def test_spherical_to_cylindrical(self, vector): """Test ``vector.represent_as(CylindricalVector)``.""" @@ -308,11 +308,11 @@ def test_spherical_to_cylindrical_astropy(self, vector, apyvector): cyl = vector.represent_as(CylindricalVector, z=Quantity([9, 10, 11, 12], u.m)) apycyl = apyvector.represent_as(apyc.CylindricalRepresentation) - assert np.allclose(cyl.rho.as_type(u.Quantity), apycyl.rho) - assert np.allclose(cyl.z.as_type(u.Quantity), apycyl.z) + assert np.allclose(convert(cyl.rho, u.Quantity), apycyl.rho) + assert np.allclose(convert(cyl.z, u.Quantity), apycyl.z) with pytest.raises(AssertionError): # TODO: Fix this - assert np.allclose(cyl.phi.as_type(u.Quantity), apycyl.phi) + assert np.allclose(convert(cyl.phi, u.Quantity), apycyl.phi) class TestCylindricalVector: @@ -405,9 +405,9 @@ def test_cylindrical_to_cartesian3d_astropy(self, vector, apyvector): cart3d = vector.represent_as(Cartesian3DVector) apycart3 = apyvector.represent_as(apyc.CartesianRepresentation) - assert np.allclose(cart3d.x.as_type(u.Quantity), apycart3.x) - assert np.allclose(cart3d.y.as_type(u.Quantity), apycart3.y) - assert np.allclose(cart3d.z.as_type(u.Quantity), apycart3.z) + assert np.allclose(convert(cart3d.x, u.Quantity), apycart3.x) + assert np.allclose(convert(cart3d.y, u.Quantity), apycart3.y) + assert np.allclose(convert(cart3d.z, u.Quantity), apycart3.z) def test_cylindrical_to_spherical(self, vector): """Test ``vector.represent_as(SphericalVector)``.""" @@ -422,9 +422,9 @@ def test_cylindrical_to_spherical_astropy(self, vector, apyvector): """Test Astropy equivalence.""" sph = vector.represent_as(SphericalVector) apysph = apyvector.represent_as(apyc.PhysicsSphericalRepresentation) - assert np.allclose(sph.r.as_type(u.Quantity), apysph.r) - assert np.allclose(sph.theta.as_type(u.Quantity), apysph.theta) - assert np.allclose(sph.phi.as_type(u.Quantity), apysph.phi) + assert np.allclose(convert(sph.r, u.Quantity), apysph.r) + assert np.allclose(convert(sph.theta, u.Quantity), apysph.theta) + assert np.allclose(convert(sph.phi, u.Quantity), apysph.phi) def test_cylindrical_to_cylindrical(self, vector): """Test ``vector.represent_as(CylindricalVector)``.""" @@ -441,9 +441,9 @@ def test_cylindrical_to_cylindrical_astropy(self, vector, apyvector): cyl = vector.represent_as(CylindricalVector) apycyl = apyvector.represent_as(apyc.CylindricalRepresentation) - assert np.allclose(cyl.rho.as_type(u.Quantity), apycyl.rho) - assert np.allclose(cyl.z.as_type(u.Quantity), apycyl.z) - assert np.allclose(cyl.phi.as_type(u.Quantity), apycyl.phi) + assert np.allclose(convert(cyl.rho, u.Quantity), apycyl.rho) + assert np.allclose(convert(cyl.z, u.Quantity), apycyl.z) + assert np.allclose(convert(cyl.phi, u.Quantity), apycyl.phi) class Abstract3DVectorDifferentialTest(AbstractVectorDifferentialTest): @@ -538,9 +538,9 @@ def test_cartesian3d_to_cartesian3d_astropy( cart3 = difntl.represent_as(CartesianDifferential3D, vector) apycart3 = apydifntl.represent_as(apyc.CartesianDifferential, apyvector) - assert np.allclose(cart3.d_x.as_type(u.Quantity), apycart3.d_x) - assert np.allclose(cart3.d_y.as_type(u.Quantity), apycart3.d_y) - assert np.allclose(cart3.d_z.as_type(u.Quantity), apycart3.d_z) + assert np.allclose(convert(cart3.d_x, u.Quantity), apycart3.d_x) + assert np.allclose(convert(cart3.d_y, u.Quantity), apycart3.d_y) + assert np.allclose(convert(cart3.d_z, u.Quantity), apycart3.d_z) def test_cartesian3d_to_spherical(self, difntl, vector): """Test ``vector.represent_as(SphericalDifferential)``.""" @@ -566,16 +566,16 @@ def test_cartesian3d_to_spherical_astropy( sph = difntl.represent_as(SphericalDifferential, vector) apysph = apydifntl.represent_as(apyc.PhysicsSphericalDifferential, apyvector) - assert np.allclose(sph.d_r.as_type(u.Quantity), apysph.d_r) + assert np.allclose(convert(sph.d_r, u.Quantity), apysph.d_r) with pytest.raises(AssertionError): # TODO: fixme assert np.allclose( - sph.d_theta.as_type(u.Quantity).to(u.mas / u.Myr), + convert(sph.d_theta, u.Quantity).to(u.mas / u.Myr), apysph.d_theta.to(u.mas / u.Myr), atol=1e-9, ) with pytest.raises(AssertionError): # TODO: fixme assert np.allclose( - sph.d_phi.as_type(u.Quantity).to(u.mas / u.Myr), + convert(sph.d_phi, u.Quantity).to(u.mas / u.Myr), apysph.d_phi.to(u.mas / u.Myr), atol=1e-7, ) @@ -600,10 +600,10 @@ def test_cartesian3d_to_spherical_astropy( """Test Astropy equivalence.""" cyl = difntl.represent_as(CylindricalDifferential, vector) apycyl = apydifntl.represent_as(apyc.CylindricalDifferential, apyvector) - assert np.allclose(cyl.d_rho.as_type(u.Quantity), apycyl.d_rho) - assert np.allclose(cyl.d_z.as_type(u.Quantity), apycyl.d_z) + assert np.allclose(convert(cyl.d_rho, u.Quantity), apycyl.d_rho) + assert np.allclose(convert(cyl.d_z, u.Quantity), apycyl.d_z) with pytest.raises(AssertionError): # TODO: fixme - assert np.allclose(cyl.d_phi.as_type(u.Quantity), apycyl.d_phi) + assert np.allclose(convert(cyl.d_phi, u.Quantity), apycyl.d_phi) class TestSphericalDifferential(Abstract3DVectorDifferentialTest): @@ -704,9 +704,9 @@ def test_spherical_to_cartesian3d_astropy( cart3d = difntl.represent_as(CartesianDifferential3D, vector) apycart3 = apydifntl.represent_as(apyc.CartesianDifferential, apyvector) - assert np.allclose(cart3d.d_x.as_type(u.Quantity), apycart3.d_x) - assert np.allclose(cart3d.d_y.as_type(u.Quantity), apycart3.d_y) - assert np.allclose(cart3d.d_z.as_type(u.Quantity), apycart3.d_z) + assert np.allclose(convert(cart3d.d_x, u.Quantity), apycart3.d_x) + assert np.allclose(convert(cart3d.d_y, u.Quantity), apycart3.d_y) + assert np.allclose(convert(cart3d.d_z, u.Quantity), apycart3.d_z) def test_spherical_to_spherical(self, difntl, vector): """Test ``vector.represent_as(SphericalDifferential)``.""" @@ -722,9 +722,9 @@ def test_spherical_to_spherical_astropy(self, difntl, vector, apydifntl, apyvect """Test Astropy equivalence.""" sph = difntl.represent_as(SphericalDifferential, vector) apysph = apydifntl.represent_as(apyc.PhysicsSphericalDifferential, apyvector) - assert np.allclose(sph.d_r.as_type(u.Quantity), apysph.d_r) - assert np.allclose(sph.d_theta.as_type(u.Quantity), apysph.d_theta) - assert np.allclose(sph.d_phi.as_type(u.Quantity), apysph.d_phi) + assert np.allclose(convert(sph.d_r, u.Quantity), apysph.d_r) + assert np.allclose(convert(sph.d_theta, u.Quantity), apysph.d_theta) + assert np.allclose(convert(sph.d_phi, u.Quantity), apysph.d_phi) def test_spherical_to_cylindrical(self, difntl, vector): """Test ``vector.represent_as(CylindricalDifferential)``.""" @@ -747,9 +747,9 @@ def test_spherical_to_cylindrical_astropy( """Test Astropy equivalence.""" cyl = difntl.represent_as(CylindricalDifferential, vector) apycyl = apydifntl.represent_as(apyc.CylindricalDifferential, apyvector) - assert np.allclose(cyl.d_rho.as_type(u.Quantity), apycyl.d_rho) - assert np.allclose(cyl.d_z.as_type(u.Quantity), apycyl.d_z) - assert np.allclose(cyl.d_phi.as_type(u.Quantity), apycyl.d_phi) + assert np.allclose(convert(cyl.d_rho, u.Quantity), apycyl.d_rho) + assert np.allclose(convert(cyl.d_z, u.Quantity), apycyl.d_z) + assert np.allclose(convert(cyl.d_phi, u.Quantity), apycyl.d_phi) class TestCylindricalDifferential(Abstract3DVectorDifferentialTest): @@ -838,9 +838,9 @@ def test_cylindrical_to_cartesian3d(self, difntl, vector, apydifntl, apyvector): assert array_equal(cart3d.d_z, Quantity([9, 10, 11, 12], u.km / u.s)) apycart3 = apydifntl.represent_as(apyc.CartesianDifferential, apyvector) - assert np.allclose(cart3d.d_x.as_type(u.Quantity), apycart3.d_x) - assert np.allclose(cart3d.d_y.as_type(u.Quantity), apycart3.d_y) - assert np.allclose(cart3d.d_z.as_type(u.Quantity), apycart3.d_z) + assert np.allclose(convert(cart3d.d_x, u.Quantity), apycart3.d_x) + assert np.allclose(convert(cart3d.d_y, u.Quantity), apycart3.d_y) + assert np.allclose(convert(cart3d.d_z, u.Quantity), apycart3.d_z) def test_cylindrical_to_spherical(self, difntl, vector): """Test ``vector.represent_as(SphericalDifferential)``.""" @@ -863,10 +863,10 @@ def test_cylindrical_to_spherical_astropy( """Test Astropy equivalence.""" sph = difntl.represent_as(SphericalDifferential, vector) apysph = apydifntl.represent_as(apyc.PhysicsSphericalDifferential, apyvector) - assert np.allclose(sph.d_r.as_type(u.Quantity), apysph.d_r) + assert np.allclose(convert(sph.d_r, u.Quantity), apysph.d_r) with pytest.raises(AssertionError): - assert np.allclose(sph.d_theta.as_type(u.Quantity), apysph.d_theta) - assert np.allclose(sph.d_phi.as_type(u.Quantity), apysph.d_phi) + assert np.allclose(convert(sph.d_theta, u.Quantity), apysph.d_theta) + assert np.allclose(convert(sph.d_phi, u.Quantity), apysph.d_phi) def test_cylindrical_to_cylindrical(self, difntl, vector): """Test ``vector.represent_as(CylindricalDifferential)``.""" @@ -882,6 +882,6 @@ def test_cylindrical_to_cylindrical(self, difntl, vector, apydifntl, apyvector): """Test Astropy equivalence.""" cyl = difntl.represent_as(CylindricalDifferential, vector) apycyl = apydifntl.represent_as(apyc.CylindricalDifferential, apyvector) - assert np.allclose(cyl.d_rho.as_type(u.Quantity), apycyl.d_rho) - assert np.allclose(cyl.d_z.as_type(u.Quantity), apycyl.d_z) - assert np.allclose(cyl.d_phi.as_type(u.Quantity), apycyl.d_phi) + assert np.allclose(convert(cyl.d_rho, u.Quantity), apycyl.d_rho) + assert np.allclose(convert(cyl.d_z, u.Quantity), apycyl.d_z) + assert np.allclose(convert(cyl.d_phi, u.Quantity), apycyl.d_phi)