# Licensed under a 3-clause BSD style license - see LICENSE.rst import functools import operator import pytest import numpy as np from astropy import units as u from astropy.coordinates import ( PhysicsSphericalRepresentation, CartesianRepresentation, CylindricalRepresentation, SphericalRepresentation, UnitSphericalRepresentation, SphericalDifferential, CartesianDifferential, UnitSphericalDifferential, SphericalCosLatDifferential, UnitSphericalCosLatDifferential, PhysicsSphericalDifferential, CylindricalDifferential, RadialRepresentation, RadialDifferential, Longitude, Latitude) from astropy.coordinates.representation import DIFFERENTIAL_CLASSES from astropy.coordinates.angle_utilities import angular_separation from astropy.tests.helper import assert_quantity_allclose, quantity_allclose def assert_representation_allclose(actual, desired, rtol=1.e-7, atol=None, **kwargs): actual_xyz = actual.to_cartesian().get_xyz(xyz_axis=-1) desired_xyz = desired.to_cartesian().get_xyz(xyz_axis=-1) actual_xyz, desired_xyz = np.broadcast_arrays(actual_xyz, desired_xyz, subok=True) assert_quantity_allclose(actual_xyz, desired_xyz, rtol, atol, **kwargs) def assert_differential_allclose(actual, desired, rtol=1.e-7, **kwargs): assert actual.components == desired.components for component in actual.components: actual_c = getattr(actual, component) atol = 1.e-10 * actual_c.unit assert_quantity_allclose(actual_c, getattr(desired, component), rtol, atol, **kwargs) def representation_equal(first, second): return functools.reduce(np.logical_and, (getattr(first, component) == getattr(second, component) for component in first.components)) class TestArithmetic(): def setup(self): # Choose some specific coordinates, for which ``sum`` and ``dot`` # works out nicely. self.lon = Longitude(np.arange(0, 12.1, 2), u.hourangle) self.lat = Latitude(np.arange(-90, 91, 30), u.deg) self.distance = [5., 12., 4., 2., 4., 12., 5.] * u.kpc self.spherical = SphericalRepresentation(self.lon, self.lat, self.distance) self.unit_spherical = self.spherical.represent_as( UnitSphericalRepresentation) self.cartesian = self.spherical.to_cartesian() def test_norm_spherical(self): norm_s = self.spherical.norm() assert isinstance(norm_s, u.Quantity) # Just to be sure, test against getting object arrays. assert norm_s.dtype.kind == 'f' assert np.all(norm_s == self.distance) @pytest.mark.parametrize('representation', (PhysicsSphericalRepresentation, CartesianRepresentation, CylindricalRepresentation)) def test_norm(self, representation): in_rep = self.spherical.represent_as(representation) norm_rep = in_rep.norm() assert isinstance(norm_rep, u.Quantity) assert_quantity_allclose(norm_rep, self.distance) def test_norm_unitspherical(self): norm_rep = self.unit_spherical.norm() assert norm_rep.unit == u.dimensionless_unscaled assert np.all(norm_rep == 1. * u.dimensionless_unscaled) @pytest.mark.parametrize('representation', (SphericalRepresentation, PhysicsSphericalRepresentation, CartesianRepresentation, CylindricalRepresentation, UnitSphericalRepresentation)) def test_neg_pos(self, representation): in_rep = self.cartesian.represent_as(representation) pos_rep = +in_rep assert type(pos_rep) is type(in_rep) assert pos_rep is not in_rep assert np.all(representation_equal(pos_rep, in_rep)) neg_rep = -in_rep assert type(neg_rep) is type(in_rep) assert np.all(neg_rep.norm() == in_rep.norm()) in_rep_xyz = in_rep.to_cartesian().xyz assert_quantity_allclose(neg_rep.to_cartesian().xyz, -in_rep_xyz, atol=1.e-10*in_rep_xyz.unit) def test_mul_div_spherical(self): s0 = self.spherical / (1. * u.Myr) assert isinstance(s0, SphericalRepresentation) assert s0.distance.dtype.kind == 'f' assert np.all(s0.lon == self.spherical.lon) assert np.all(s0.lat == self.spherical.lat) assert np.all(s0.distance == self.distance / (1. * u.Myr)) s1 = (1./u.Myr) * self.spherical assert isinstance(s1, SphericalRepresentation) assert np.all(representation_equal(s1, s0)) s2 = self.spherical * np.array([[1.], [2.]]) assert isinstance(s2, SphericalRepresentation) assert s2.shape == (2, self.spherical.shape[0]) assert np.all(s2.lon == self.spherical.lon) assert np.all(s2.lat == self.spherical.lat) assert np.all(s2.distance == self.spherical.distance * np.array([[1.], [2.]])) s3 = np.array([[1.], [2.]]) * self.spherical assert isinstance(s3, SphericalRepresentation) assert np.all(representation_equal(s3, s2)) s4 = -self.spherical assert isinstance(s4, SphericalRepresentation) assert quantity_allclose(s4.to_cartesian().xyz, -self.spherical.to_cartesian().xyz, atol=1e-15*self.spherical.distance.unit) assert np.all(s4.distance == self.spherical.distance) s5 = +self.spherical assert s5 is not self.spherical assert np.all(representation_equal(s5, self.spherical)) @pytest.mark.parametrize('representation', (PhysicsSphericalRepresentation, CartesianRepresentation, CylindricalRepresentation)) def test_mul_div(self, representation): in_rep = self.spherical.represent_as(representation) r1 = in_rep / (1. * u.Myr) assert isinstance(r1, representation) for component in in_rep.components: in_rep_comp = getattr(in_rep, component) r1_comp = getattr(r1, component) if in_rep_comp.unit == self.distance.unit: assert np.all(r1_comp == in_rep_comp / (1.*u.Myr)) else: assert np.all(r1_comp == in_rep_comp) r2 = np.array([[1.], [2.]]) * in_rep assert isinstance(r2, representation) assert r2.shape == (2, in_rep.shape[0]) assert_quantity_allclose(r2.norm(), self.distance * np.array([[1.], [2.]])) r3 = -in_rep assert_representation_allclose(r3.to_cartesian(), (in_rep * -1.).to_cartesian(), atol=1e-5*u.pc) with pytest.raises(TypeError): in_rep * in_rep with pytest.raises(TypeError): dict() * in_rep def test_mul_div_unit_spherical(self): s1 = self.unit_spherical * self.distance assert isinstance(s1, SphericalRepresentation) assert np.all(s1.lon == self.unit_spherical.lon) assert np.all(s1.lat == self.unit_spherical.lat) assert np.all(s1.distance == self.spherical.distance) s2 = self.unit_spherical / u.s assert isinstance(s2, SphericalRepresentation) assert np.all(s2.lon == self.unit_spherical.lon) assert np.all(s2.lat == self.unit_spherical.lat) assert np.all(s2.distance == 1./u.s) u3 = -self.unit_spherical assert isinstance(u3, UnitSphericalRepresentation) assert_quantity_allclose(u3.lon, self.unit_spherical.lon + 180.*u.deg) assert np.all(u3.lat == -self.unit_spherical.lat) assert_quantity_allclose(u3.to_cartesian().xyz, -self.unit_spherical.to_cartesian().xyz, atol=1.e-10*u.dimensionless_unscaled) u4 = +self.unit_spherical assert isinstance(u4, UnitSphericalRepresentation) assert u4 is not self.unit_spherical assert np.all(representation_equal(u4, self.unit_spherical)) def test_add_sub_cartesian(self): c1 = self.cartesian + self.cartesian assert isinstance(c1, CartesianRepresentation) assert c1.x.dtype.kind == 'f' assert np.all(representation_equal(c1, 2. * self.cartesian)) with pytest.raises(TypeError): self.cartesian + 10.*u.m with pytest.raises(u.UnitsError): self.cartesian + (self.cartesian / u.s) c2 = self.cartesian - self.cartesian assert isinstance(c2, CartesianRepresentation) assert np.all(representation_equal( c2, CartesianRepresentation(0.*u.m, 0.*u.m, 0.*u.m))) c3 = self.cartesian - self.cartesian / 2. assert isinstance(c3, CartesianRepresentation) assert np.all(representation_equal(c3, self.cartesian / 2.)) @pytest.mark.parametrize('representation', (PhysicsSphericalRepresentation, SphericalRepresentation, CylindricalRepresentation)) def test_add_sub(self, representation): in_rep = self.cartesian.represent_as(representation) r1 = in_rep + in_rep assert isinstance(r1, representation) expected = 2. * in_rep for component in in_rep.components: assert_quantity_allclose(getattr(r1, component), getattr(expected, component)) with pytest.raises(TypeError): 10.*u.m + in_rep with pytest.raises(u.UnitsError): in_rep + (in_rep / u.s) r2 = in_rep - in_rep assert isinstance(r2, representation) assert_representation_allclose( r2.to_cartesian(), CartesianRepresentation(0.*u.m, 0.*u.m, 0.*u.m), atol=1e-15*u.kpc) r3 = in_rep - in_rep / 2. assert isinstance(r3, representation) expected = in_rep / 2. assert_representation_allclose(r3, expected) def test_add_sub_unit_spherical(self): s1 = self.unit_spherical + self.unit_spherical assert isinstance(s1, SphericalRepresentation) expected = 2. * self.unit_spherical for component in s1.components: assert_quantity_allclose(getattr(s1, component), getattr(expected, component)) with pytest.raises(TypeError): 10.*u.m - self.unit_spherical with pytest.raises(u.UnitsError): self.unit_spherical + (self.unit_spherical / u.s) s2 = self.unit_spherical - self.unit_spherical / 2. assert isinstance(s2, SphericalRepresentation) expected = self.unit_spherical / 2. for component in s2.components: assert_quantity_allclose(getattr(s2, component), getattr(expected, component)) @pytest.mark.parametrize('representation', (CartesianRepresentation, PhysicsSphericalRepresentation, SphericalRepresentation, CylindricalRepresentation)) def test_sum_mean(self, representation): in_rep = self.spherical.represent_as(representation) r_sum = in_rep.sum() assert isinstance(r_sum, representation) expected = SphericalRepresentation( 90. * u.deg, 0. * u.deg, 14. * u.kpc).represent_as(representation) for component in expected.components: exp_component = getattr(expected, component) assert_quantity_allclose(getattr(r_sum, component), exp_component, atol=1e-10*exp_component.unit) r_mean = in_rep.mean() assert isinstance(r_mean, representation) expected = expected / len(in_rep) for component in expected.components: exp_component = getattr(expected, component) assert_quantity_allclose(getattr(r_mean, component), exp_component, atol=1e-10*exp_component.unit) def test_sum_mean_unit_spherical(self): s_sum = self.unit_spherical.sum() assert isinstance(s_sum, SphericalRepresentation) expected = SphericalRepresentation( 90. * u.deg, 0. * u.deg, 3. * u.dimensionless_unscaled) for component in expected.components: exp_component = getattr(expected, component) assert_quantity_allclose(getattr(s_sum, component), exp_component, atol=1e-10*exp_component.unit) s_mean = self.unit_spherical.mean() assert isinstance(s_mean, SphericalRepresentation) expected = expected / len(self.unit_spherical) for component in expected.components: exp_component = getattr(expected, component) assert_quantity_allclose(getattr(s_mean, component), exp_component, atol=1e-10*exp_component.unit) @pytest.mark.parametrize('representation', (CartesianRepresentation, PhysicsSphericalRepresentation, SphericalRepresentation, CylindricalRepresentation)) def test_dot(self, representation): in_rep = self.cartesian.represent_as(representation) r_dot_r = in_rep.dot(in_rep) assert isinstance(r_dot_r, u.Quantity) assert r_dot_r.shape == in_rep.shape assert_quantity_allclose(np.sqrt(r_dot_r), self.distance) r_dot_r_rev = in_rep.dot(in_rep[::-1]) assert isinstance(r_dot_r_rev, u.Quantity) assert r_dot_r_rev.shape == in_rep.shape expected = [-25., -126., 2., 4., 2., -126., -25.] * u.kpc**2 assert_quantity_allclose(r_dot_r_rev, expected) for axis in 'xyz': project = CartesianRepresentation(*( (1. if axis == _axis else 0.) * u.dimensionless_unscaled for _axis in 'xyz')) assert_quantity_allclose(in_rep.dot(project), getattr(self.cartesian, axis), atol=1.*u.upc) with pytest.raises(TypeError): in_rep.dot(self.cartesian.xyz) def test_dot_unit_spherical(self): u_dot_u = self.unit_spherical.dot(self.unit_spherical) assert isinstance(u_dot_u, u.Quantity) assert u_dot_u.shape == self.unit_spherical.shape assert_quantity_allclose(u_dot_u, 1.*u.dimensionless_unscaled) cartesian = self.unit_spherical.to_cartesian() for axis in 'xyz': project = CartesianRepresentation(*( (1. if axis == _axis else 0.) * u.dimensionless_unscaled for _axis in 'xyz')) assert_quantity_allclose(self.unit_spherical.dot(project), getattr(cartesian, axis), atol=1.e-10) @pytest.mark.parametrize('representation', (CartesianRepresentation, PhysicsSphericalRepresentation, SphericalRepresentation, CylindricalRepresentation)) def test_cross(self, representation): in_rep = self.cartesian.represent_as(representation) r_cross_r = in_rep.cross(in_rep) assert isinstance(r_cross_r, representation) assert_quantity_allclose(r_cross_r.norm(), 0.*u.kpc**2, atol=1.*u.mpc**2) r_cross_r_rev = in_rep.cross(in_rep[::-1]) sep = angular_separation(self.lon, self.lat, self.lon[::-1], self.lat[::-1]) expected = self.distance * self.distance[::-1] * np.sin(sep) assert_quantity_allclose(r_cross_r_rev.norm(), expected, atol=1.*u.mpc**2) unit_vectors = CartesianRepresentation( [1., 0., 0.]*u.one, [0., 1., 0.]*u.one, [0., 0., 1.]*u.one)[:, np.newaxis] r_cross_uv = in_rep.cross(unit_vectors) assert r_cross_uv.shape == (3, 7) assert_quantity_allclose(r_cross_uv.dot(unit_vectors), 0.*u.kpc, atol=1.*u.upc) assert_quantity_allclose(r_cross_uv.dot(in_rep), 0.*u.kpc**2, atol=1.*u.mpc**2) zeros = np.zeros(len(in_rep)) * u.kpc expected = CartesianRepresentation( u.Quantity((zeros, -self.cartesian.z, self.cartesian.y)), u.Quantity((self.cartesian.z, zeros, -self.cartesian.x)), u.Quantity((-self.cartesian.y, self.cartesian.x, zeros))) # Comparison with spherical is hard since some distances are zero, # implying the angles are undefined. r_cross_uv_cartesian = r_cross_uv.to_cartesian() assert_representation_allclose(r_cross_uv_cartesian, expected, atol=1.*u.upc) # A final check, with the side benefit of ensuring __div__ and norm # work on multi-D representations. r_cross_uv_by_distance = r_cross_uv / self.distance uv_sph = unit_vectors.represent_as(UnitSphericalRepresentation) sep = angular_separation(self.lon, self.lat, uv_sph.lon, uv_sph.lat) assert_quantity_allclose(r_cross_uv_by_distance.norm(), np.sin(sep), atol=1e-9) with pytest.raises(TypeError): in_rep.cross(self.cartesian.xyz) def test_cross_unit_spherical(self): u_cross_u = self.unit_spherical.cross(self.unit_spherical) assert isinstance(u_cross_u, SphericalRepresentation) assert_quantity_allclose(u_cross_u.norm(), 0.*u.one, atol=1.e-10*u.one) u_cross_u_rev = self.unit_spherical.cross(self.unit_spherical[::-1]) assert isinstance(u_cross_u_rev, SphericalRepresentation) sep = angular_separation(self.lon, self.lat, self.lon[::-1], self.lat[::-1]) expected = np.sin(sep) assert_quantity_allclose(u_cross_u_rev.norm(), expected, atol=1.e-10*u.one) class TestUnitVectorsAndScales(): @staticmethod def check_unit_vectors(e): for v in e.values(): assert type(v) is CartesianRepresentation assert_quantity_allclose(v.norm(), 1. * u.one) return e @staticmethod def check_scale_factors(sf, rep): unit = rep.norm().unit for c, f in sf.items(): assert type(f) is u.Quantity assert (f.unit * getattr(rep, c).unit).is_equivalent(unit) def test_spherical(self): s = SphericalRepresentation(lon=[0., 6., 21.] * u.hourangle, lat=[0., -30., 85.] * u.deg, distance=[1, 2, 3] * u.kpc) e = s.unit_vectors() self.check_unit_vectors(e) sf = s.scale_factors() self.check_scale_factors(sf, s) s_lon = s + s.distance * 1e-5 * np.cos(s.lat) * e['lon'] assert_quantity_allclose(s_lon.lon, s.lon + 1e-5*u.rad, atol=1e-10*u.rad) assert_quantity_allclose(s_lon.lat, s.lat, atol=1e-10*u.rad) assert_quantity_allclose(s_lon.distance, s.distance) s_lon2 = s + 1e-5 * u.radian * sf['lon'] * e['lon'] assert_representation_allclose(s_lon2, s_lon) s_lat = s + s.distance * 1e-5 * e['lat'] assert_quantity_allclose(s_lat.lon, s.lon) assert_quantity_allclose(s_lat.lat, s.lat + 1e-5*u.rad, atol=1e-10*u.rad) assert_quantity_allclose(s_lon.distance, s.distance) s_lat2 = s + 1.e-5 * u.radian * sf['lat'] * e['lat'] assert_representation_allclose(s_lat2, s_lat) s_distance = s + 1. * u.pc * e['distance'] assert_quantity_allclose(s_distance.lon, s.lon, atol=1e-10*u.rad) assert_quantity_allclose(s_distance.lat, s.lat, atol=1e-10*u.rad) assert_quantity_allclose(s_distance.distance, s.distance + 1.*u.pc) s_distance2 = s + 1. * u.pc * sf['distance'] * e['distance'] assert_representation_allclose(s_distance2, s_distance) def test_unit_spherical(self): s = UnitSphericalRepresentation(lon=[0., 6., 21.] * u.hourangle, lat=[0., -30., 85.] * u.deg) e = s.unit_vectors() self.check_unit_vectors(e) sf = s.scale_factors() self.check_scale_factors(sf, s) s_lon = s + 1e-5 * np.cos(s.lat) * e['lon'] assert_quantity_allclose(s_lon.lon, s.lon + 1e-5*u.rad, atol=1e-10*u.rad) assert_quantity_allclose(s_lon.lat, s.lat, atol=1e-10*u.rad) s_lon2 = s + 1e-5 * u.radian * sf['lon'] * e['lon'] assert_representation_allclose(s_lon2, s_lon) s_lat = s + 1e-5 * e['lat'] assert_quantity_allclose(s_lat.lon, s.lon) assert_quantity_allclose(s_lat.lat, s.lat + 1e-5*u.rad, atol=1e-10*u.rad) s_lat2 = s + 1.e-5 * u.radian * sf['lat'] * e['lat'] assert_representation_allclose(s_lat2, s_lat) def test_radial(self): r = RadialRepresentation(10.*u.kpc) with pytest.raises(NotImplementedError): r.unit_vectors() sf = r.scale_factors() assert np.all(sf['distance'] == 1.*u.one) assert np.all(r.norm() == r.distance) with pytest.raises(TypeError): r + r def test_physical_spherical(self): s = PhysicsSphericalRepresentation(phi=[0., 6., 21.] * u.hourangle, theta=[90., 120., 5.] * u.deg, r=[1, 2, 3] * u.kpc) e = s.unit_vectors() self.check_unit_vectors(e) sf = s.scale_factors() self.check_scale_factors(sf, s) s_phi = s + s.r * 1e-5 * np.sin(s.theta) * e['phi'] assert_quantity_allclose(s_phi.phi, s.phi + 1e-5*u.rad, atol=1e-10*u.rad) assert_quantity_allclose(s_phi.theta, s.theta, atol=1e-10*u.rad) assert_quantity_allclose(s_phi.r, s.r) s_phi2 = s + 1e-5 * u.radian * sf['phi'] * e['phi'] assert_representation_allclose(s_phi2, s_phi) s_theta = s + s.r * 1e-5 * e['theta'] assert_quantity_allclose(s_theta.phi, s.phi) assert_quantity_allclose(s_theta.theta, s.theta + 1e-5*u.rad, atol=1e-10*u.rad) assert_quantity_allclose(s_theta.r, s.r) s_theta2 = s + 1.e-5 * u.radian * sf['theta'] * e['theta'] assert_representation_allclose(s_theta2, s_theta) s_r = s + 1. * u.pc * e['r'] assert_quantity_allclose(s_r.phi, s.phi, atol=1e-10*u.rad) assert_quantity_allclose(s_r.theta, s.theta, atol=1e-10*u.rad) assert_quantity_allclose(s_r.r, s.r + 1.*u.pc) s_r2 = s + 1. * u.pc * sf['r'] * e['r'] assert_representation_allclose(s_r2, s_r) def test_cartesian(self): s = CartesianRepresentation(x=[1, 2, 3] * u.pc, y=[2, 3, 4] * u.Mpc, z=[3, 4, 5] * u.kpc) e = s.unit_vectors() sf = s.scale_factors() for v, expected in zip(e.values(), ([1., 0., 0.] * u.one, [0., 1., 0.] * u.one, [0., 0., 1.] * u.one)): assert np.all(v.get_xyz(xyz_axis=-1) == expected) for f in sf.values(): assert np.all(f == 1.*u.one) def test_cylindrical(self): s = CylindricalRepresentation(rho=[1, 2, 3] * u.pc, phi=[0., 90., -45.] * u.deg, z=[3, 4, 5] * u.kpc) e = s.unit_vectors() self.check_unit_vectors(e) sf = s.scale_factors() self.check_scale_factors(sf, s) s_rho = s + 1. * u.pc * e['rho'] assert_quantity_allclose(s_rho.rho, s.rho + 1.*u.pc) assert_quantity_allclose(s_rho.phi, s.phi) assert_quantity_allclose(s_rho.z, s.z) s_rho2 = s + 1. * u.pc * sf['rho'] * e['rho'] assert_representation_allclose(s_rho2, s_rho) s_phi = s + s.rho * 1e-5 * e['phi'] assert_quantity_allclose(s_phi.rho, s.rho) assert_quantity_allclose(s_phi.phi, s.phi + 1e-5*u.rad) assert_quantity_allclose(s_phi.z, s.z) s_phi2 = s + 1e-5 * u.radian * sf['phi'] * e['phi'] assert_representation_allclose(s_phi2, s_phi) s_z = s + 1. * u.pc * e['z'] assert_quantity_allclose(s_z.rho, s.rho) assert_quantity_allclose(s_z.phi, s.phi, atol=1e-10*u.rad) assert_quantity_allclose(s_z.z, s.z + 1.*u.pc) s_z2 = s + 1. * u.pc * sf['z'] * e['z'] assert_representation_allclose(s_z2, s_z) @pytest.mark.parametrize('omit_coslat', [False, True], scope='class') class TestSphericalDifferential(): # these test cases are subclassed for SphericalCosLatDifferential, # hence some tests depend on omit_coslat. def _setup(self, omit_coslat): if omit_coslat: self.SD_cls = SphericalCosLatDifferential else: self.SD_cls = SphericalDifferential s = SphericalRepresentation(lon=[0., 6., 21.] * u.hourangle, lat=[0., -30., 85.] * u.deg, distance=[1, 2, 3] * u.kpc) self.s = s self.e = s.unit_vectors() self.sf = s.scale_factors(omit_coslat=omit_coslat) def test_name_coslat(self, omit_coslat): self._setup(omit_coslat) if omit_coslat: assert self.SD_cls is SphericalCosLatDifferential assert self.SD_cls.get_name() == 'sphericalcoslat' else: assert self.SD_cls is SphericalDifferential assert self.SD_cls.get_name() == 'spherical' assert self.SD_cls.get_name() in DIFFERENTIAL_CLASSES def test_simple_differentials(self, omit_coslat): self._setup(omit_coslat) s, e, sf = self.s, self.e, self.sf o_lon = self.SD_cls(1.*u.arcsec, 0.*u.arcsec, 0.*u.kpc) o_lonc = o_lon.to_cartesian(base=s) o_lon2 = self.SD_cls.from_cartesian(o_lonc, base=s) assert_differential_allclose(o_lon, o_lon2) # simple check by hand for first element. # lat[0] is 0, so cos(lat) term doesn't matter. assert_quantity_allclose(o_lonc[0].xyz, [0., np.pi/180./3600., 0.]*u.kpc) # check all using unit vectors and scale factors. s_lon = s + 1.*u.arcsec * sf['lon'] * e['lon'] assert_representation_allclose(o_lonc, s_lon - s, atol=1*u.npc) s_lon2 = s + o_lon assert_representation_allclose(s_lon2, s_lon, atol=1*u.npc) o_lat = self.SD_cls(0.*u.arcsec, 1.*u.arcsec, 0.*u.kpc) o_latc = o_lat.to_cartesian(base=s) assert_quantity_allclose(o_latc[0].xyz, [0., 0., np.pi/180./3600.]*u.kpc, atol=1.*u.npc) s_lat = s + 1.*u.arcsec * sf['lat'] * e['lat'] assert_representation_allclose(o_latc, s_lat - s, atol=1*u.npc) s_lat2 = s + o_lat assert_representation_allclose(s_lat2, s_lat, atol=1*u.npc) o_distance = self.SD_cls(0.*u.arcsec, 0.*u.arcsec, 1.*u.mpc) o_distancec = o_distance.to_cartesian(base=s) assert_quantity_allclose(o_distancec[0].xyz, [1e-6, 0., 0.]*u.kpc, atol=1.*u.npc) s_distance = s + 1.*u.mpc * sf['distance'] * e['distance'] assert_representation_allclose(o_distancec, s_distance - s, atol=1*u.npc) s_distance2 = s + o_distance assert_representation_allclose(s_distance2, s_distance) def test_differential_arithmetic(self, omit_coslat): self._setup(omit_coslat) s = self.s o_lon = self.SD_cls(1.*u.arcsec, 0.*u.arcsec, 0.*u.kpc) o_lon_by_2 = o_lon / 2. assert_representation_allclose(o_lon_by_2.to_cartesian(s) * 2., o_lon.to_cartesian(s), atol=1e-10*u.kpc) assert_representation_allclose(s + o_lon, s + 2 * o_lon_by_2, atol=1e-10*u.kpc) o_lon_rec = o_lon_by_2 + o_lon_by_2 assert_representation_allclose(s + o_lon, s + o_lon_rec, atol=1e-10*u.kpc) o_lon_0 = o_lon - o_lon for c in o_lon_0.components: assert np.all(getattr(o_lon_0, c) == 0.) o_lon2 = self.SD_cls(1*u.mas/u.yr, 0*u.mas/u.yr, 0*u.km/u.s) assert_quantity_allclose(o_lon2.norm(s)[0], 4.74*u.km/u.s, atol=0.01*u.km/u.s) assert_representation_allclose(o_lon2.to_cartesian(s) * 1000.*u.yr, o_lon.to_cartesian(s), atol=1e-10*u.kpc) s_off = s + o_lon s_off2 = s + o_lon2 * 1000.*u.yr assert_representation_allclose(s_off, s_off2, atol=1e-10*u.kpc) factor = 1e5 * u.radian/u.arcsec if not omit_coslat: factor = factor / np.cos(s.lat) s_off_big = s + o_lon * factor assert_representation_allclose( s_off_big, SphericalRepresentation(s.lon + 90.*u.deg, 0.*u.deg, 1e5*s.distance), atol=5.*u.kpc) o_lon3c = CartesianRepresentation(0., 4.74047, 0., unit=u.km/u.s) o_lon3 = self.SD_cls.from_cartesian(o_lon3c, base=s) expected0 = self.SD_cls(1.*u.mas/u.yr, 0.*u.mas/u.yr, 0.*u.km/u.s) assert_differential_allclose(o_lon3[0], expected0) s_off_big2 = s + o_lon3 * 1e5 * u.yr * u.radian/u.mas assert_representation_allclose( s_off_big2, SphericalRepresentation(90.*u.deg, 0.*u.deg, 1e5*u.kpc), atol=5.*u.kpc) with pytest.raises(TypeError): o_lon - s with pytest.raises(TypeError): s.to_cartesian() + o_lon def test_differential_init_errors(self, omit_coslat): self._setup(omit_coslat) s = self.s with pytest.raises(u.UnitsError): self.SD_cls(1.*u.arcsec, 0., 0.) with pytest.raises(TypeError): self.SD_cls(1.*u.arcsec, 0.*u.arcsec, 0.*u.kpc, False, False) with pytest.raises(TypeError): self.SD_cls(1.*u.arcsec, 0.*u.arcsec, 0.*u.kpc, copy=False, d_lat=0.*u.arcsec) with pytest.raises(TypeError): self.SD_cls(1.*u.arcsec, 0.*u.arcsec, 0.*u.kpc, copy=False, flying='circus') with pytest.raises(ValueError): self.SD_cls(np.ones(2)*u.arcsec, np.zeros(3)*u.arcsec, np.zeros(2)*u.kpc) with pytest.raises(u.UnitsError): self.SD_cls(1.*u.arcsec, 1.*u.s, 0.*u.kpc) with pytest.raises(u.UnitsError): self.SD_cls(1.*u.kpc, 1.*u.arcsec, 0.*u.kpc) o = self.SD_cls(1.*u.arcsec, 1.*u.arcsec, 0.*u.km/u.s) with pytest.raises(u.UnitsError): o.to_cartesian(s) with pytest.raises(AttributeError): o.d_lat = 0.*u.arcsec with pytest.raises(AttributeError): del o.d_lat o = self.SD_cls(1.*u.arcsec, 1.*u.arcsec, 0.*u.km) with pytest.raises(TypeError): o.to_cartesian() c = CartesianRepresentation(10., 0., 0., unit=u.km) with pytest.raises(TypeError): self.SD_cls.to_cartesian(c) with pytest.raises(TypeError): self.SD_cls.from_cartesian(c) with pytest.raises(TypeError): self.SD_cls.from_cartesian(c, SphericalRepresentation) @pytest.mark.parametrize('omit_coslat', [False, True], scope='class') class TestUnitSphericalDifferential(): def _setup(self, omit_coslat): if omit_coslat: self.USD_cls = UnitSphericalCosLatDifferential else: self.USD_cls = UnitSphericalDifferential s = UnitSphericalRepresentation(lon=[0., 6., 21.] * u.hourangle, lat=[0., -30., 85.] * u.deg) self.s = s self.e = s.unit_vectors() self.sf = s.scale_factors(omit_coslat=omit_coslat) def test_name_coslat(self, omit_coslat): self._setup(omit_coslat) if omit_coslat: assert self.USD_cls is UnitSphericalCosLatDifferential assert self.USD_cls.get_name() == 'unitsphericalcoslat' else: assert self.USD_cls is UnitSphericalDifferential assert self.USD_cls.get_name() == 'unitspherical' assert self.USD_cls.get_name() in DIFFERENTIAL_CLASSES def test_simple_differentials(self, omit_coslat): self._setup(omit_coslat) s, e, sf = self.s, self.e, self.sf o_lon = self.USD_cls(1.*u.arcsec, 0.*u.arcsec) o_lonc = o_lon.to_cartesian(base=s) o_lon2 = self.USD_cls.from_cartesian(o_lonc, base=s) assert_differential_allclose(o_lon, o_lon2) # simple check by hand for first element # (lat[0]=0, so works for both normal and CosLat differential) assert_quantity_allclose(o_lonc[0].xyz, [0., np.pi/180./3600., 0.]*u.one) # check all using unit vectors and scale factors. s_lon = s + 1.*u.arcsec * sf['lon'] * e['lon'] assert type(s_lon) is SphericalRepresentation assert_representation_allclose(o_lonc, s_lon - s, atol=1e-10*u.one) s_lon2 = s + o_lon assert_representation_allclose(s_lon2, s_lon, atol=1e-10*u.one) o_lat = self.USD_cls(0.*u.arcsec, 1.*u.arcsec) o_latc = o_lat.to_cartesian(base=s) assert_quantity_allclose(o_latc[0].xyz, [0., 0., np.pi/180./3600.]*u.one, atol=1e-10*u.one) s_lat = s + 1.*u.arcsec * sf['lat'] * e['lat'] assert type(s_lat) is SphericalRepresentation assert_representation_allclose(o_latc, s_lat - s, atol=1e-10*u.one) s_lat2 = s + o_lat assert_representation_allclose(s_lat2, s_lat, atol=1e-10*u.one) def test_differential_arithmetic(self, omit_coslat): self._setup(omit_coslat) s = self.s o_lon = self.USD_cls(1.*u.arcsec, 0.*u.arcsec) o_lon_by_2 = o_lon / 2. assert type(o_lon_by_2) is self.USD_cls assert_representation_allclose(o_lon_by_2.to_cartesian(s) * 2., o_lon.to_cartesian(s), atol=1e-10*u.one) s_lon = s + o_lon s_lon2 = s + 2 * o_lon_by_2 assert type(s_lon) is SphericalRepresentation assert_representation_allclose(s_lon, s_lon2, atol=1e-10*u.one) o_lon_rec = o_lon_by_2 + o_lon_by_2 assert type(o_lon_rec) is self.USD_cls assert representation_equal(o_lon, o_lon_rec) assert_representation_allclose(s + o_lon, s + o_lon_rec, atol=1e-10*u.one) o_lon_0 = o_lon - o_lon assert type(o_lon_0) is self.USD_cls for c in o_lon_0.components: assert np.all(getattr(o_lon_0, c) == 0.) o_lon2 = self.USD_cls(1.*u.mas/u.yr, 0.*u.mas/u.yr) kks = u.km/u.kpc/u.s assert_quantity_allclose(o_lon2.norm(s)[0], 4.74047*kks, atol=1e-4*kks) assert_representation_allclose(o_lon2.to_cartesian(s) * 1000.*u.yr, o_lon.to_cartesian(s), atol=1e-10*u.one) s_off = s + o_lon s_off2 = s + o_lon2 * 1000.*u.yr assert_representation_allclose(s_off, s_off2, atol=1e-10*u.one) factor = 1e5 * u.radian/u.arcsec if not omit_coslat: factor = factor / np.cos(s.lat) s_off_big = s + o_lon * factor assert_representation_allclose( s_off_big, SphericalRepresentation(s.lon + 90.*u.deg, 0.*u.deg, 1e5), atol=5.*u.one) o_lon3c = CartesianRepresentation(0., 4.74047, 0., unit=kks) # This looses information!! o_lon3 = self.USD_cls.from_cartesian(o_lon3c, base=s) expected0 = self.USD_cls(1.*u.mas/u.yr, 0.*u.mas/u.yr) assert_differential_allclose(o_lon3[0], expected0) # Part of motion kept. part_kept = s.cross(CartesianRepresentation(0, 1, 0, unit=u.one)).norm() assert_quantity_allclose(o_lon3.norm(s), 4.74047*part_kept*kks, atol=1e-10*kks) # (lat[0]=0, so works for both normal and CosLat differential) s_off_big2 = s + o_lon3 * 1e5 * u.yr * u.radian/u.mas expected0 = SphericalRepresentation(90.*u.deg, 0.*u.deg, 1e5*u.one) assert_representation_allclose(s_off_big2[0], expected0, atol=5.*u.one) def test_differential_init_errors(self, omit_coslat): self._setup(omit_coslat) with pytest.raises(u.UnitsError): self.USD_cls(0.*u.deg, 10.*u.deg/u.yr) class TestRadialDifferential(): def setup(self): s = SphericalRepresentation(lon=[0., 6., 21.] * u.hourangle, lat=[0., -30., 85.] * u.deg, distance=[1, 2, 3] * u.kpc) self.s = s self.r = s.represent_as(RadialRepresentation) self.e = s.unit_vectors() self.sf = s.scale_factors() def test_name(self): assert RadialDifferential.get_name() == 'radial' assert RadialDifferential.get_name() in DIFFERENTIAL_CLASSES def test_simple_differentials(self): r, s, e, sf = self.r, self.s, self.e, self.sf o_distance = RadialDifferential(1.*u.mpc) # Can be applied to RadialRepresentation, though not most useful. r_distance = r + o_distance assert_quantity_allclose(r_distance.distance, r.distance + o_distance.d_distance) r_distance2 = o_distance + r assert_quantity_allclose(r_distance2.distance, r.distance + o_distance.d_distance) # More sense to apply it relative to spherical representation. o_distancec = o_distance.to_cartesian(base=s) assert_quantity_allclose(o_distancec[0].xyz, [1e-6, 0., 0.]*u.kpc, atol=1.*u.npc) o_recover = RadialDifferential.from_cartesian(o_distancec, base=s) assert_quantity_allclose(o_recover.d_distance, o_distance.d_distance) s_distance = s + 1.*u.mpc * sf['distance'] * e['distance'] assert_representation_allclose(o_distancec, s_distance - s, atol=1*u.npc) s_distance2 = s + o_distance assert_representation_allclose(s_distance2, s_distance) class TestPhysicsSphericalDifferential(): """Test copied from SphericalDifferential, so less extensive.""" def setup(self): s = PhysicsSphericalRepresentation(phi=[0., 90., 315.] * u.deg, theta=[90., 120., 5.] * u.deg, r=[1, 2, 3] * u.kpc) self.s = s self.e = s.unit_vectors() self.sf = s.scale_factors() def test_name(self): assert PhysicsSphericalDifferential.get_name() == 'physicsspherical' assert PhysicsSphericalDifferential.get_name() in DIFFERENTIAL_CLASSES def test_simple_differentials(self): s, e, sf = self.s, self.e, self.sf o_phi = PhysicsSphericalDifferential(1*u.arcsec, 0*u.arcsec, 0*u.kpc) o_phic = o_phi.to_cartesian(base=s) o_phi2 = PhysicsSphericalDifferential.from_cartesian(o_phic, base=s) assert_quantity_allclose(o_phi.d_phi, o_phi2.d_phi, atol=1.*u.narcsec) assert_quantity_allclose(o_phi.d_theta, o_phi2.d_theta, atol=1.*u.narcsec) assert_quantity_allclose(o_phi.d_r, o_phi2.d_r, atol=1.*u.npc) # simple check by hand for first element. assert_quantity_allclose(o_phic[0].xyz, [0., np.pi/180./3600., 0.]*u.kpc, atol=1.*u.npc) # check all using unit vectors and scale factors. s_phi = s + 1.*u.arcsec * sf['phi'] * e['phi'] assert_representation_allclose(o_phic, s_phi - s, atol=1e-10*u.kpc) o_theta = PhysicsSphericalDifferential(0*u.arcsec, 1*u.arcsec, 0*u.kpc) o_thetac = o_theta.to_cartesian(base=s) assert_quantity_allclose(o_thetac[0].xyz, [0., 0., -np.pi/180./3600.]*u.kpc, atol=1.*u.npc) s_theta = s + 1.*u.arcsec * sf['theta'] * e['theta'] assert_representation_allclose(o_thetac, s_theta - s, atol=1e-10*u.kpc) s_theta2 = s + o_theta assert_representation_allclose(s_theta2, s_theta, atol=1e-10*u.kpc) o_r = PhysicsSphericalDifferential(0*u.arcsec, 0*u.arcsec, 1*u.mpc) o_rc = o_r.to_cartesian(base=s) assert_quantity_allclose(o_rc[0].xyz, [1e-6, 0., 0.]*u.kpc, atol=1.*u.npc) s_r = s + 1.*u.mpc * sf['r'] * e['r'] assert_representation_allclose(o_rc, s_r - s, atol=1e-10*u.kpc) s_r2 = s + o_r assert_representation_allclose(s_r2, s_r) def test_differential_init_errors(self): with pytest.raises(u.UnitsError): PhysicsSphericalDifferential(1.*u.arcsec, 0., 0.) class TestCylindricalDifferential(): """Test copied from SphericalDifferential, so less extensive.""" def setup(self): s = CylindricalRepresentation(rho=[1, 2, 3] * u.kpc, phi=[0., 90., 315.] * u.deg, z=[3, 2, 1] * u.kpc) self.s = s self.e = s.unit_vectors() self.sf = s.scale_factors() def test_name(self): assert CylindricalDifferential.get_name() == 'cylindrical' assert CylindricalDifferential.get_name() in DIFFERENTIAL_CLASSES def test_simple_differentials(self): s, e, sf = self.s, self.e, self.sf o_rho = CylindricalDifferential(1.*u.mpc, 0.*u.arcsec, 0.*u.kpc) o_rhoc = o_rho.to_cartesian(base=s) assert_quantity_allclose(o_rhoc[0].xyz, [1.e-6, 0., 0.]*u.kpc) s_rho = s + 1.*u.mpc * sf['rho'] * e['rho'] assert_representation_allclose(o_rhoc, s_rho - s, atol=1e-10*u.kpc) s_rho2 = s + o_rho assert_representation_allclose(s_rho2, s_rho) o_phi = CylindricalDifferential(0.*u.kpc, 1.*u.arcsec, 0.*u.kpc) o_phic = o_phi.to_cartesian(base=s) o_phi2 = CylindricalDifferential.from_cartesian(o_phic, base=s) assert_quantity_allclose(o_phi.d_rho, o_phi2.d_rho, atol=1.*u.npc) assert_quantity_allclose(o_phi.d_phi, o_phi2.d_phi, atol=1.*u.narcsec) assert_quantity_allclose(o_phi.d_z, o_phi2.d_z, atol=1.*u.npc) # simple check by hand for first element. assert_quantity_allclose(o_phic[0].xyz, [0., np.pi/180./3600., 0.]*u.kpc) # check all using unit vectors and scale factors. s_phi = s + 1.*u.arcsec * sf['phi'] * e['phi'] assert_representation_allclose(o_phic, s_phi - s, atol=1e-10*u.kpc) o_z = CylindricalDifferential(0.*u.kpc, 0.*u.arcsec, 1.*u.mpc) o_zc = o_z.to_cartesian(base=s) assert_quantity_allclose(o_zc[0].xyz, [0., 0., 1.e-6]*u.kpc) s_z = s + 1.*u.mpc * sf['z'] * e['z'] assert_representation_allclose(o_zc, s_z - s, atol=1e-10*u.kpc) s_z2 = s + o_z assert_representation_allclose(s_z2, s_z) def test_differential_init_errors(self): with pytest.raises(u.UnitsError): CylindricalDifferential(1.*u.pc, 1.*u.arcsec, 3.*u.km/u.s) class TestCartesianDifferential(): """Test copied from SphericalDifferential, so less extensive.""" def setup(self): s = CartesianRepresentation(x=[1, 2, 3] * u.kpc, y=[2, 3, 1] * u.kpc, z=[3, 1, 2] * u.kpc) self.s = s self.e = s.unit_vectors() self.sf = s.scale_factors() def test_name(self): assert CartesianDifferential.get_name() == 'cartesian' assert CartesianDifferential.get_name() in DIFFERENTIAL_CLASSES def test_simple_differentials(self): s, e, sf = self.s, self.e, self.sf for d, differential in ( # test different inits while we're at it. ('x', CartesianDifferential(1.*u.pc, 0.*u.pc, 0.*u.pc)), ('y', CartesianDifferential([0., 1., 0.], unit=u.pc)), ('z', CartesianDifferential(np.array([[0., 0., 1.]]) * u.pc, xyz_axis=1))): o_c = differential.to_cartesian(base=s) o_c2 = differential.to_cartesian() assert np.all(representation_equal(o_c, o_c2)) assert all(np.all(getattr(differential, 'd_'+c) == getattr(o_c, c)) for c in ('x', 'y', 'z')) differential2 = CartesianDifferential.from_cartesian(o_c) assert np.all(representation_equal(differential2, differential)) differential3 = CartesianDifferential.from_cartesian(o_c, base=o_c) assert np.all(representation_equal(differential3, differential)) s_off = s + 1.*u.pc * sf[d] * e[d] assert_representation_allclose(o_c, s_off - s, atol=1e-10*u.kpc) s_off2 = s + differential assert_representation_allclose(s_off2, s_off) def test_init_failures(self): with pytest.raises(ValueError): CartesianDifferential(1.*u.kpc/u.s, 2.*u.kpc) with pytest.raises(u.UnitsError): CartesianDifferential(1.*u.kpc/u.s, 2.*u.kpc, 3.*u.kpc) with pytest.raises(ValueError): CartesianDifferential(1.*u.kpc, 2.*u.kpc, 3.*u.kpc, xyz_axis=1) class TestDifferentialConversion(): def setup(self): self.s = SphericalRepresentation(lon=[0., 6., 21.] * u.hourangle, lat=[0., -30., 85.] * u.deg, distance=[1, 2, 3] * u.kpc) @pytest.mark.parametrize('sd_cls', [SphericalDifferential, SphericalCosLatDifferential]) def test_represent_as_own_class(self, sd_cls): so = sd_cls(1.*u.deg, 2.*u.deg, 0.1*u.kpc) so2 = so.represent_as(sd_cls) assert so2 is so def test_represent_other_coslat(self): s = self.s coslat = np.cos(s.lat) so = SphericalDifferential(1.*u.deg, 2.*u.deg, 0.1*u.kpc) so_coslat = so.represent_as(SphericalCosLatDifferential, base=s) assert_quantity_allclose(so.d_lon * coslat, so_coslat.d_lon_coslat) so2 = so_coslat.represent_as(SphericalDifferential, base=s) assert np.all(representation_equal(so2, so)) so3 = SphericalDifferential.from_representation(so_coslat, base=s) assert np.all(representation_equal(so3, so)) so_coslat2 = SphericalCosLatDifferential.from_representation(so, base=s) assert np.all(representation_equal(so_coslat2, so_coslat)) # Also test UnitSpherical us = s.represent_as(UnitSphericalRepresentation) uo = so.represent_as(UnitSphericalDifferential) uo_coslat = so.represent_as(UnitSphericalCosLatDifferential, base=s) assert_quantity_allclose(uo.d_lon * coslat, uo_coslat.d_lon_coslat) uo2 = uo_coslat.represent_as(UnitSphericalDifferential, base=us) assert np.all(representation_equal(uo2, uo)) uo3 = UnitSphericalDifferential.from_representation(uo_coslat, base=us) assert np.all(representation_equal(uo3, uo)) uo_coslat2 = UnitSphericalCosLatDifferential.from_representation( uo, base=us) assert np.all(representation_equal(uo_coslat2, uo_coslat)) uo_coslat3 = uo.represent_as(UnitSphericalCosLatDifferential, base=us) assert np.all(representation_equal(uo_coslat3, uo_coslat)) @pytest.mark.parametrize('sd_cls', [SphericalDifferential, SphericalCosLatDifferential]) @pytest.mark.parametrize('r_cls', (SphericalRepresentation, UnitSphericalRepresentation, PhysicsSphericalRepresentation, CylindricalRepresentation)) def test_represent_regular_class(self, sd_cls, r_cls): so = sd_cls(1.*u.deg, 2.*u.deg, 0.1*u.kpc) r = so.represent_as(r_cls, base=self.s) c = so.to_cartesian(self.s) r_check = c.represent_as(r_cls) assert np.all(representation_equal(r, r_check)) so2 = sd_cls.from_representation(r, base=self.s) so3 = sd_cls.from_cartesian(r.to_cartesian(), self.s) assert np.all(representation_equal(so2, so3)) @pytest.mark.parametrize('sd_cls', [SphericalDifferential, SphericalCosLatDifferential]) def test_convert_physics(self, sd_cls): # Conversion needs no base for SphericalDifferential, but does # need one (to get the latitude) for SphericalCosLatDifferential. if sd_cls is SphericalDifferential: usd_cls = UnitSphericalDifferential base_s = base_u = base_p = None else: usd_cls = UnitSphericalCosLatDifferential base_s = self.s[1] base_u = base_s.represent_as(UnitSphericalRepresentation) base_p = base_s.represent_as(PhysicsSphericalRepresentation) so = sd_cls(1.*u.deg, 2.*u.deg, 0.1*u.kpc) po = so.represent_as(PhysicsSphericalDifferential, base=base_s) so2 = sd_cls.from_representation(po, base=base_s) assert_differential_allclose(so, so2) po2 = PhysicsSphericalDifferential.from_representation(so, base=base_p) assert_differential_allclose(po, po2) so3 = po.represent_as(sd_cls, base=base_p) assert_differential_allclose(so, so3) s = self.s p = s.represent_as(PhysicsSphericalRepresentation) cso = so.to_cartesian(s[1]) cpo = po.to_cartesian(p[1]) assert_representation_allclose(cso, cpo) assert_representation_allclose(s[1] + so, p[1] + po) po2 = so.represent_as(PhysicsSphericalDifferential, base=None if base_s is None else s) assert_representation_allclose(s + so, p + po2) suo = usd_cls.from_representation(so) puo = usd_cls.from_representation(po, base=base_u) assert_differential_allclose(suo, puo) suo2 = so.represent_as(usd_cls) puo2 = po.represent_as(usd_cls, base=base_p) assert_differential_allclose(suo2, puo2) assert_differential_allclose(puo, puo2) sro = RadialDifferential.from_representation(so) pro = RadialDifferential.from_representation(po) assert representation_equal(sro, pro) sro2 = so.represent_as(RadialDifferential) pro2 = po.represent_as(RadialDifferential) assert representation_equal(sro2, pro2) assert representation_equal(pro, pro2) @pytest.mark.parametrize( ('sd_cls', 'usd_cls'), [(SphericalDifferential, UnitSphericalDifferential), (SphericalCosLatDifferential, UnitSphericalCosLatDifferential)]) def test_convert_unit_spherical_radial(self, sd_cls, usd_cls): s = self.s us = s.represent_as(UnitSphericalRepresentation) rs = s.represent_as(RadialRepresentation) assert_representation_allclose(rs * us, s) uo = usd_cls(2.*u.deg, 1.*u.deg) so = uo.represent_as(sd_cls, base=s) assert_quantity_allclose(so.d_distance, 0.*u.kpc, atol=1.*u.npc) uo2 = so.represent_as(usd_cls) assert_representation_allclose(uo.to_cartesian(us), uo2.to_cartesian(us)) so1 = sd_cls(2.*u.deg, 1.*u.deg, 5.*u.pc) uo_r = so1.represent_as(usd_cls) ro_r = so1.represent_as(RadialDifferential) assert np.all(representation_equal(uo_r, uo)) assert np.all(representation_equal(ro_r, RadialDifferential(5.*u.pc))) @pytest.mark.parametrize('sd_cls', [SphericalDifferential, SphericalCosLatDifferential]) def test_convert_cylindrial(self, sd_cls): s = self.s so = sd_cls(1.*u.deg, 2.*u.deg, 0.1*u.kpc) cyo = so.represent_as(CylindricalDifferential, base=s) cy = s.represent_as(CylindricalRepresentation) so1 = cyo.represent_as(sd_cls, base=cy) assert_representation_allclose(so.to_cartesian(s), so1.to_cartesian(s)) cyo2 = CylindricalDifferential.from_representation(so, base=cy) assert_representation_allclose(cyo2.to_cartesian(base=cy), cyo.to_cartesian(base=cy)) so2 = sd_cls.from_representation(cyo2, base=s) assert_representation_allclose(so.to_cartesian(s), so2.to_cartesian(s)) @pytest.mark.parametrize('sd_cls', [SphericalDifferential, SphericalCosLatDifferential]) def test_combinations(self, sd_cls): if sd_cls is SphericalDifferential: uo = UnitSphericalDifferential(2.*u.deg, 1.*u.deg) uo_d_lon = uo.d_lon else: uo = UnitSphericalCosLatDifferential(2.*u.deg, 1.*u.deg) uo_d_lon = uo.d_lon_coslat ro = RadialDifferential(1.*u.mpc) so1 = uo + ro so1c = sd_cls(uo_d_lon, uo.d_lat, ro.d_distance) assert np.all(representation_equal(so1, so1c)) so2 = uo - ro so2c = sd_cls(uo_d_lon, uo.d_lat, -ro.d_distance) assert np.all(representation_equal(so2, so2c)) so3 = so2 + ro so3c = sd_cls(uo_d_lon, uo.d_lat, 0.*u.kpc) assert np.all(representation_equal(so3, so3c)) so4 = so1 + ro so4c = sd_cls(uo_d_lon, uo.d_lat, 2*ro.d_distance) assert np.all(representation_equal(so4, so4c)) so5 = so1 - uo so5c = sd_cls(0*u.deg, 0.*u.deg, ro.d_distance) assert np.all(representation_equal(so5, so5c)) assert_representation_allclose(self.s + (uo+ro), self.s+so1) @pytest.mark.parametrize('op,args', [ (operator.neg, ()), (operator.pos, ()), (operator.mul, (-8.,)), (operator.truediv, ([4., 8.]*u.s,))], scope='class') class TestArithmeticWithDifferentials: def setup_class(self): self.cr = CartesianRepresentation([1, 2, 3]*u.kpc) self.cd = CartesianDifferential([.1, -.2, .3]*u.km/u.s) self.c = self.cr.with_differentials(self.cd) def test_operation_cartesian(self, op, args): ncr = op(self.c, *args) expected = (op(self.cr, *args)).with_differentials(op(self.cd, *args)) assert np.all(ncr == expected) def test_operation_radial(self, op, args): rep = self.c.represent_as(RadialRepresentation, {'s': RadialDifferential}) result = op(rep, *args) expected_distance = op(self.cr.norm(), *args) expected_rv = op((self.cr/self.cr.norm()).dot(self.cd), *args) assert u.allclose(result.distance, expected_distance) assert u.allclose(result.differentials['s'].d_distance, expected_rv) @pytest.mark.parametrize('diff_cls', [ SphericalDifferential, SphericalCosLatDifferential, PhysicsSphericalDifferential, CylindricalDifferential]) def test_operation_other(self, diff_cls, op, args): rep_cls = diff_cls.base_representation rep = self.c.represent_as(rep_cls, {'s': diff_cls}) result = op(rep, *args) expected_c = op(self.c, *args) expected = expected_c.represent_as(rep_cls, {'s': diff_cls}) # Check that we match in the representation itself. assert_representation_allclose(result, expected) assert_differential_allclose(result.differentials['s'], expected.differentials['s']) # Check that we compare correctly in cartesian as well, just to be sure. result_c = result.represent_as(CartesianRepresentation, {'s': CartesianDifferential}) assert_representation_allclose(result_c, expected_c) assert_differential_allclose(result_c.differentials['s'], expected_c.differentials['s']) @pytest.mark.parametrize('rep_cls', [ SphericalRepresentation, PhysicsSphericalRepresentation, CylindricalRepresentation]) def test_operation_cartesian_differential(self, rep_cls, op, args): rep = self.c.represent_as(rep_cls, {'s': CartesianDifferential}) result = op(rep, *args) expected_c = op(self.c, *args) expected = expected_c.represent_as(rep_cls, {'s': CartesianDifferential}) # Check that we match in the representation itself. assert_representation_allclose(result, expected) assert_differential_allclose(result.differentials['s'], expected.differentials['s']) @pytest.mark.parametrize('diff_cls', [ UnitSphericalDifferential, UnitSphericalCosLatDifferential]) def test_operation_unit_spherical(self, diff_cls, op, args): rep_cls = diff_cls.base_representation rep = self.c.represent_as(rep_cls, {'s': diff_cls}) result = op(rep, *args) if op not in (operator.neg, operator.pos): expected_cls = rep._dimensional_representation else: expected_cls = rep_cls assert type(result) is expected_cls assert type(result.differentials['s']) is diff_cls # Have lost information, so unlike above we convert our initial # unit-spherical back to Cartesian, and check that applying # the operation on that cartesian representation gives the same result. # We do not compare the output directly, since for multiplication # and division there will be sign flips in the spherical distance. expected_c = op(rep.represent_as(CartesianRepresentation, {'s': CartesianDifferential}), *args) result_c = result.represent_as(CartesianRepresentation, {'s': CartesianDifferential}) assert_representation_allclose(result_c, expected_c) assert_differential_allclose(result_c.differentials['s'], expected_c.differentials['s']) @pytest.mark.parametrize('diff_cls', [ RadialDifferential, UnitSphericalDifferential, UnitSphericalCosLatDifferential]) def test_operation_spherical_with_rv_or_pm(self, diff_cls, op, args): rep = self.c.represent_as(SphericalRepresentation, {'s': diff_cls}) result = op(rep, *args) assert type(result) is SphericalRepresentation assert type(result.differentials['s']) is diff_cls expected_c = op(rep.represent_as(CartesianRepresentation, {'s': CartesianDifferential}), *args) result_c = result.represent_as(CartesianRepresentation, {'s': CartesianDifferential}) assert_representation_allclose(result_c, expected_c) assert_differential_allclose(result_c.differentials['s'], expected_c.differentials['s']) @pytest.mark.parametrize('op,args', [ (operator.neg, ()), (operator.mul, (10.,))]) def test_operation_unitspherical_with_rv_fails(op, args): rep = UnitSphericalRepresentation( 0*u.deg, 0*u.deg, differentials={'s': RadialDifferential(10*u.km/u.s)}) with pytest.raises(ValueError, match='unit key'): op(rep, *args) @pytest.mark.parametrize('rep,dif', [ [CartesianRepresentation([1, 2, 3]*u.kpc), CartesianDifferential([.1, .2, .3]*u.km/u.s)], [SphericalRepresentation(90*u.deg, 0.*u.deg, 14.*u.kpc), SphericalDifferential(1.*u.deg, 2.*u.deg, 0.1*u.kpc)] ]) def test_arithmetic_with_differentials_fail(rep, dif): rep = rep.with_differentials(dif) with pytest.raises(TypeError): rep + rep with pytest.raises(TypeError): rep - rep with pytest.raises(TypeError): rep * rep with pytest.raises(TypeError): rep / rep