Merge pull request #6289 from adrn/coordinate/repr-fix

__repr__ fix for coordinate frames with velocities

astropy/coordinates/baseframe.py

@@ -803,6 +803,9 @@ def represent_as(self, base, s='base', in_frame_units=False):
                 data = data.__class__(copy=False, **datakwargs)
 
             if differential_cls:
+                # the original differential
+                data_diff = self.data.differentials['s']
+
                 # If the new differential is known to this frame and has a
                 # defined set of names and units, then use that.
                 new_attrs = self.representation_info.get(differential_cls)
@@ -811,8 +814,23 @@ def represent_as(self, base, s='base', in_frame_units=False):
                                       for comp in diff.components)
                     for comp, new_attr_unit in zip(diff.components,
                                                    new_attrs['units']):
-                        if (new_attr_unit and
-                                hasattr(self._data.differentials['s'], comp)):
+                        # Some special-casing to treat a situation where the
+                        # input data has a UnitSphericalDifferential or a
+                        # RadialDifferential. It is re-represented to the
+                        # frame's differential class (which might be, e.g., a
+                        # dimensional Differential), so we don't want to try to
+                        # convert the empty component units
+                        if (isinstance(data_diff,
+                                       (r.UnitSphericalDifferential,
+                                        r.UnitSphericalCosLatDifferential)) and
+                                comp not in data_diff.__class__.attr_classes):
+                            continue
+
+                        elif (isinstance(data_diff, r.RadialDifferential) and
+                              comp not in data_diff.__class__.attr_classes):
+                            continue
+
+                        if new_attr_unit and hasattr(diff, comp):
                             diffkwargs[comp] = diffkwargs[comp].to(new_attr_unit)
 
                     diff = diff.__class__(copy=False, **diffkwargs)
@@ -1007,11 +1025,22 @@ def _data_repr(self):
         if self.representation:
             if (issubclass(self.representation, r.SphericalRepresentation) and
                     isinstance(self.data, r.UnitSphericalRepresentation)):
-                data = self.represent_as(self.data.__class__,
-                                         in_frame_units=True)
+                rep_cls = self.data.__class__
+            else:
+                rep_cls = self.representation
+
+            if 's' in self.data.differentials:
+                dif_cls = self.get_representation_cls('s')
+                dif_data = self.data.differentials['s']
+                if isinstance(dif_data, (r.UnitSphericalDifferential,
+                                         r.UnitSphericalCosLatDifferential,
+                                         r.RadialDifferential)):
+                    dif_cls = dif_data.__class__
+
             else:
-                data = self.represent_as(self.representation,
-                                         in_frame_units=True)
+                dif_cls = None
+
+            data = self.represent_as(rep_cls, dif_cls, in_frame_units=True)
 
             data_repr = repr(data)
             for nmpref, nmrepr in self.representation_component_names.items():
@@ -1028,11 +1057,10 @@ def _data_repr(self):
         else:
             data_repr = 'Data:\n' + data_repr
 
-        if (getattr(self.data, 'differentials', None) and
-                's' in self.data.differentials):
+        if 's' in self.data.differentials:
             data_repr_spl = data_repr.split('\n')
             if 'has differentials' in data_repr_spl[-1]:
-                diffrepr = repr(self.data.differentials['s']).split('\n')
+                diffrepr = repr(data.differentials['s']).split('\n')
                 if diffrepr[0].startswith('<'):
                     diffrepr[0] = ' ' + ' '.join(diffrepr[0].split(' ')[1:])
                 for frm_nm, rep_nm in self.get_representation_component_names('s').items():

astropy/coordinates/representation.py

@@ -2350,6 +2350,10 @@ class UnitSphericalDifferential(BaseSphericalDifferential):
     """
     base_representation = UnitSphericalRepresentation
 
+    @classproperty
+    def _dimensional_differential(cls):
+        return SphericalDifferential
+
     def __init__(self, d_lon, d_lat, copy=True):
         super(UnitSphericalDifferential,
               self).__init__(d_lon, d_lat, copy=copy)
@@ -2551,6 +2555,10 @@ class UnitSphericalCosLatDifferential(BaseSphericalCosLatDifferential):
     attr_classes = OrderedDict([('d_lon_coslat', u.Quantity),
                                 ('d_lat', u.Quantity)])
 
+    @classproperty
+    def _dimensional_differential(cls):
+        return SphericalCosLatDifferential
+
     def __init__(self, d_lon_coslat, d_lat, copy=True):
         super(UnitSphericalCosLatDifferential,
               self).__init__(d_lon_coslat, d_lat, copy=copy)

astropy/coordinates/tests/test_frames.py

@@ -205,9 +205,12 @@ def test_frame_repr_vels():
 
     i = ICRS(ra=1*u.deg, dec=2*u.deg,
              pm_ra_cosdec=1*u.marcsec/u.yr, pm_dec=2*u.marcsec/u.yr)
+
+    # unit comes out as mas/yr because of the preferred units defined in the
+    # frame RepresentationMapping
     assert repr(i) == ('<ICRS Coordinate: (ra, dec) in deg\n'
                        '    ( 1.,  2.)\n'
-                       ' (pm_ra_cosdec, pm_dec) in marcsec / yr\n'
+                       ' (pm_ra_cosdec, pm_dec) in mas / yr\n'
                        '    ( 1.,  2.)>')
 
 

astropy/coordinates/tests/test_frames_with_velocity.py

@@ -56,10 +56,26 @@ def test_all_arg_options(kwargs):
     # access to the relevant attributes from the resulting object
     icrs = ICRS(**kwargs)
     gal = icrs.transform_to(Galactic)
+    repr_gal = repr(gal)
 
     for k in kwargs:
         getattr(icrs, k)
 
+    if 'pm_ra_cosdec' in kwargs: # should have both
+        assert 'pm_l_cosb' in repr_gal
+        assert 'pm_b' in repr_gal
+        assert 'mas / yr' in repr_gal
+
+        if 'radial_velocity' not in kwargs:
+            assert 'radial_velocity' not in repr_gal
+
+    if 'radial_velocity' in kwargs:
+        assert 'radial_velocity' in repr_gal
+        assert 'km / s' in repr_gal
+
+        if 'pm_ra_cosdec' not in kwargs:
+            assert 'pm_l_cosb' not in repr_gal
+            assert 'pm_b' not in repr_gal
 
 @pytest.mark.parametrize('cls,lon,lat', [
     [bf.ICRS, 'ra', 'dec'], [bf.FK4, 'ra', 'dec'], [bf.FK4NoETerms, 'ra', 'dec'],

astropy/coordinates/transformations.py

@@ -1020,7 +1020,7 @@ def _apply_transform(self, fromcoord, matrix, offset):
             if offset is not None and 's' in offset.differentials:
                 veldiff = veldiff + offset.differentials['s']
 
-            newrep = newrep.with_differentials(veldiff)
+            newrep = newrep.with_differentials({'s': veldiff})
 
         if isinstance(fromcoord.data, UnitSphericalRepresentation):
             # Special-case this because otherwise the return object will think
@@ -1052,6 +1052,21 @@ def _apply_transform(self, fromcoord, matrix, offset):
             newrep = newrep.represent_as(fromcoord.data.__class__)  # drops diffs
             newrep = newrep.with_differentials(diffs)
 
+        elif has_velocity and unit_vel_diff:
+            # Here, we're in the case where the representation is not
+            # UnitSpherical, but the differential *is* one of the UnitSpherical
+            # types. We have to convert back to that differential class or the
+            # resulting frame will think it has a valid radial_velocity. This
+            # can probably be cleaned up: we currently have to go through the
+            # dimensional version of the differential before representing as the
+            # unit differential so that the units work out (the distance length
+            # unit shouldn't appear in the resulting proper motions)
+
+            diff_cls = fromcoord.data.differentials['s'].__class__
+            newrep = newrep.represent_as(fromcoord.data.__class__,
+                                         diff_cls._dimensional_differential)
+            newrep = newrep.represent_as(fromcoord.data.__class__, diff_cls)
+
         # We pulled the radial differential off of the representation
         # earlier, so now we need to put it back. But, in order to do that, we
         # have to turn the representation into a repr that is compatible with