astropy / astropy

Astronomy and astrophysics core library
https://www.astropy.org
BSD 3-Clause "New" or "Revised" License
4.41k stars 1.77k forks source link

Enhance in-place modification of coordinate Frame or SkyCoord #10151

Open taldcroft opened 4 years ago

taldcroft commented 4 years ago

Description

Currently the __setitem__ support for BaseFrame and SkyCoord require that frame attributes are exactly the same in the object being updated and the right hand side value. Some WIP code that would help in making things somewhat more flexible is shown below and was removed from #9857 in 8b161db70. See also https://github.com/astropy/astropy/pull/9857#pullrequestreview-385950530.

For instance there is no reason that one should not be able to do:

sc1 = SkyCoord([1]*u.deg, [2]*u.deg, obstime='J2001')
sc2 = SkyCoord([3]*u.deg, [4]*u.deg, obstime='J2010')
t1 = Table([sc1])
t2 = Table([sc2])
table.vstack([t1, t2])

Additional context

@contextlib.contextmanager
def _set_time_writeable_true(obj):
    """Context manager to allow __setitem__ to operate on frame attributes.

    Not used in current implementation which just insists that they are
    equivalent but see commit 4e3f99943.
    """
    is_time = isinstance(obj, Time)
    try:
        if is_time:
            writeable = obj.writeable
            obj.writeable = True
        yield
    finally:
        if is_time:
            obj.writeable = writeable

def _frame_attribute_setitem(frame, attr, item, value):
    """Set a frame attribute ``frame.attr[item] = value```

    This is not used in the current implementation of __setitem__ but here in
    the WIP for reference of what might be possible.

    The trick is that frame attributes can be scalars and they might need to be
    broadcast to the SkyCoord shape.  For example if a SkyCoord ``sc1`` has
    shape=(3,) with scalar equinox='B1999', and ``sc2`` has equinox='B2015',
    then setting ``sc1[1] = sc2[2]`` needs to make sc1.equinox into shape=(3,)
    and result in ['B1999', 'B2015', 'B1999'].

    It isn't obvious if this can cause problems, so don't do that without
    some discussion.
    """
    _attr = '_' + attr
    frame_attr = getattr(frame, _attr)
    value_attr = getattr(value, _attr)
    frame_attr_shape = getattr(frame_attr, 'shape', ())

    # If all the same then do nothing
    if np.all(frame_attr == value_attr):
        return

    # At this point we need to ensure frame_attr has the same shape as frame
    # because ``item`` is referenced to frame shape and we will set frame_attr[item].
    if frame_attr_shape == ():
        frame_attr = frame_attr.take(np.zeros(frame.shape))
        setattr(frame, _attr, frame_attr)

    elif frame_attr_shape != frame.shape:
        # How to broadcast a ShapedLikeNDArray subclass like Time. E.g.
        # an obstime attribute with shape (3,) and the frame has shape (10, 3).
        raise NotImplementedError('help me @mhvk!')

    # Finally do the setting
    frame_attr[item] = value_attr
mhvk commented 4 years ago

Probably good to restate my general worry, which is that whether or not set values propagate back to parents will become rather random. E.g.,

sc1 = SkyCoord([1, 0, 1]*u.deg, [2, 3, 4]*u.deg, obstime='J2001')
sc2 = sc1[:2]
sc2[0] = SkyCoord(0*u.deg, 0.*.deg, obstime='J2010')

will, with the current implementation of __getitem__, change the coordinates in sc1[0] as well, but it will not change obstime (since that is a scalar, and by setting one element in sc2, sc2.obstime will have to be turned into an array). I fear this will lead to very hard to debug problems.

That said, I think in principle it is OK if obstime is already an array in the first place.

Also, I think there is no problem when copies of everything are being made anyway. In fact, there is no reason not to have the general np.stack and np.concatenate work with SkyCoord (using __array_ufunc__ - #8610) - all those can work, I think, via SkyCoordInfo.new_like).