Skip to main content
`Ctrl`+`K`

# Source code for astropy.coordinates.transformations.affine

    # Licensed under a 3-clause BSD style license - see LICENSE.rst

    """Affine coordinate transformations."""

    from abc import abstractmethod

    import numpy as np

    from astropy.coordinates.representation import (
        CartesianDifferential,
        RadialDifferential,
        SphericalCosLatDifferential,
        SphericalDifferential,
        UnitSphericalRepresentation,
    )
    from astropy.coordinates.transformations.base import CoordinateTransform

    __all__ = [
        "AffineTransform",
        "BaseAffineTransform",
        "DynamicMatrixTransform",
        "StaticMatrixTransform",
    ]

    class BaseAffineTransform(CoordinateTransform):
        """Base class for common functionality between the ``AffineTransform``-type
        subclasses.

        This base class is needed because `~astropy.coordinates.AffineTransform`
        and the matrix transform classes share the ``__call__()`` method, but
        differ in how they generate the affine parameters.
        `~astropy.coordinates.StaticMatrixTransform` passes in a matrix stored as a
        class attribute, and both of the matrix transforms pass in ``None`` for the
        offset. Hence, user subclasses would likely want to subclass this (rather
        than `~astropy.coordinates.AffineTransform`) if they want to provide
        alternative transformations using this machinery.

        """

        def _apply_transform(self, fromcoord, matrix, offset):
            data = fromcoord.data
            has_velocity = "s" in data.differentials

            # Bail out if no transform is actually requested
            if matrix is None and offset is None:
                return data

            # list of unit differentials
            _unit_diffs = (
                SphericalDifferential._unit_differential,
                SphericalCosLatDifferential._unit_differential,
            )
            unit_vel_diff = has_velocity and isinstance(
                data.differentials["s"], _unit_diffs
            )
            rad_vel_diff = has_velocity and isinstance(
                data.differentials["s"], RadialDifferential
            )

            # Some initial checking to short-circuit doing any re-representation if
            # we're going to fail anyways:
            if isinstance(data, UnitSphericalRepresentation) and offset is not None:
                raise TypeError(
                    "Position information stored on coordinate frame "
                    "is insufficient to do a full-space position "
                    f"transformation (representation class: {data.__class__})"
                )

            if (
                has_velocity
                and (unit_vel_diff or rad_vel_diff)
                and offset is not None
                and "s" in offset.differentials
            ):
                # Coordinate has a velocity, but it is not a full-space velocity
                # that we need to do a velocity offset
                raise TypeError(
                    "Velocity information stored on coordinate frame is insufficient to do"
                    " a full-space velocity transformation (differential class:"
                    f" {data.differentials['s'].__class__})"
                )

            if len(data.differentials) > 1:
                # We should never get here because the frame initializer shouldn't
                # allow more differentials, but this just adds protection for
                # subclasses that somehow skip the checks
                raise ValueError(
                    "Representation passed to AffineTransform contains multiple associated"
                    " differentials. Only a single differential with velocity units is"
                    f" presently supported (differentials: {data.differentials})."
                )

            # If the representation is a UnitSphericalRepresentation, and this is
            # just a MatrixTransform, we have to try to turn the differential into a
            # Unit version of the differential (if no radial velocity) or a
            # sphericaldifferential with zero proper motion (if only a radial
            # velocity) so that the matrix operation works
            if (
                has_velocity
                and isinstance(data, UnitSphericalRepresentation)
                and not (unit_vel_diff or rad_vel_diff)
            ):
                # retrieve just velocity differential
                unit_diff = data.differentials["s"].represent_as(
                    data.differentials["s"]._unit_differential, data
                )
                data = data.with_differentials({"s": unit_diff})  # updates key

            # If it's a RadialDifferential, we flat-out ignore the differentials
            # This is because, by this point (past the validation above), we can
            # only possibly be doing a rotation-only transformation, and that
            # won't change the radial differential. We later add it back in
            elif rad_vel_diff:
                data = data.without_differentials()

            # Convert the representation and differentials to cartesian without
            # having them attached to a frame
            rep = data.to_cartesian()
            diffs = {
                k: diff.represent_as(CartesianDifferential, data)
                for k, diff in data.differentials.items()
            }
            rep = rep.with_differentials(diffs)

            # Only do transform if matrix is specified. This is for speed in
            # transformations that only specify an offset (e.g., LSR)
            if matrix is not None:
                # Note: this applies to both representation and differentials
                rep = rep.transform(matrix)

            # TODO: if we decide to allow arithmetic between representations that
            # contain differentials, this can be tidied up
            newrep = rep.without_differentials()
            if offset is not None:
                newrep += offset.without_differentials()

            # We need a velocity (time derivative) and, for now, are strict: the
            # representation can only contain a velocity differential and no others.
            if has_velocity and not rad_vel_diff:
                veldiff = rep.differentials["s"]  # already in Cartesian form

                if offset is not None and "s" in offset.differentials:
                    veldiff += offset.differentials["s"]

                newrep = newrep.with_differentials({"s": veldiff})

            if isinstance(fromcoord.data, UnitSphericalRepresentation):
                # Special-case this because otherwise the return object will think
                # it has a valid distance with the default return (a
                # CartesianRepresentation instance)

                if has_velocity and not unit_vel_diff and not rad_vel_diff:
                    # We have to first represent as the Unit types we converted to,
                    # then put the d_distance information back in to the
                    # differentials and re-represent as their original forms
                    newdiff = newrep.differentials["s"]
                    _unit_cls = fromcoord.data.differentials["s"]._unit_differential
                    newdiff = newdiff.represent_as(_unit_cls, newrep)

                    kwargs = {comp: getattr(newdiff, comp) for comp in newdiff.components}
                    kwargs["d_distance"] = fromcoord.data.differentials["s"].d_distance
                    diffs = {
                        "s": type(fromcoord.data.differentials["s"])(copy=False, **kwargs)
                    }

                elif has_velocity and unit_vel_diff:
                    newdiff = newrep.differentials["s"].represent_as(
                        fromcoord.data.differentials["s"].__class__, newrep
                    )
                    diffs = {"s": newdiff}

                else:
                    diffs = newrep.differentials

                newrep = newrep.represent_as(type(fromcoord.data)).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(
                    type(fromcoord.data), diff_cls._dimensional_differential
                ).represent_as(type(fromcoord.data), 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
            # having a RadialDifferential
            if has_velocity and rad_vel_diff:
                newrep = newrep.represent_as(fromcoord.data.__class__)
                newrep = newrep.with_differentials({"s": fromcoord.data.differentials["s"]})

            return newrep

        def __call__(self, fromcoord, toframe):
            params = self._affine_params(fromcoord, toframe)
            newrep = self._apply_transform(fromcoord, *params)
            return toframe.realize_frame(newrep)

        @abstractmethod
        def _affine_params(self, fromcoord, toframe):
            pass

    class AffineTransform(BaseAffineTransform):
        """
        A coordinate transformation specified as a function that yields a 3 x 3
        cartesian transformation matrix and a tuple of displacement vectors.

        See `~astropy.coordinates.Galactocentric` for
        an example.

        Parameters
        ----------
        transform_func : callable
            A callable that has the signature ``transform_func(fromcoord, toframe)``
            and returns: a (3, 3) matrix that operates on ``fromcoord`` in a
            Cartesian representation, and a ``CartesianRepresentation`` with
            (optionally) an attached velocity ``CartesianDifferential`` to represent
            a translation and offset in velocity to apply after the matrix
            operation.
        fromsys : class
            The coordinate frame class to start from.
        tosys : class
            The coordinate frame class to transform into.
        priority : float or int
            The priority if this transform when finding the shortest
            coordinate transform path - large numbers are lower priorities.
        register_graph : `~astropy.coordinates.TransformGraph` or None
            A graph to register this transformation with on creation, or
            `None` to leave it unregistered.

        Raises
        ------
        TypeError
            If ``transform_func`` is not callable

        """

        def __init__(self, transform_func, fromsys, tosys, priority=1, register_graph=None):
            if not callable(transform_func):
                raise TypeError("transform_func is not callable")
            self.transform_func = transform_func

            super().__init__(
                fromsys, tosys, priority=priority, register_graph=register_graph
            )

        def _affine_params(self, fromcoord, toframe):
            return self.transform_func(fromcoord, toframe)

    class StaticMatrixTransform(BaseAffineTransform):
        """
        A coordinate transformation defined as a 3 x 3 cartesian
        transformation matrix.

        This is distinct from DynamicMatrixTransform in that this kind of matrix is
        independent of frame attributes.  That is, it depends *only* on the class of
        the frame.

        Parameters
        ----------
        matrix : array-like or callable
            A 3 x 3 matrix for transforming 3-vectors. In most cases will
            be unitary (although this is not strictly required). If a callable,
            will be called *with no arguments* to get the matrix.
        fromsys : class
            The coordinate frame class to start from.
        tosys : class
            The coordinate frame class to transform into.
        priority : float or int
            The priority if this transform when finding the shortest
            coordinate transform path - large numbers are lower priorities.
        register_graph : `~astropy.coordinates.TransformGraph` or None
            A graph to register this transformation with on creation, or
            `None` to leave it unregistered.

        Raises
        ------
        ValueError
            If the matrix is not 3 x 3

        """

        def __init__(self, matrix, fromsys, tosys, priority=1, register_graph=None):
            if callable(matrix):
                matrix = matrix()
            self.matrix = np.array(matrix)

            if self.matrix.shape != (3, 3):
                raise ValueError("Provided matrix is not 3 x 3")

            super().__init__(
                fromsys, tosys, priority=priority, register_graph=register_graph
            )

        def _affine_params(self, fromcoord, toframe):
            return self.matrix, None

    class DynamicMatrixTransform(BaseAffineTransform):
        """
        A coordinate transformation specified as a function that yields a
        3 x 3 cartesian transformation matrix.

        This is similar to, but distinct from StaticMatrixTransform, in that the
        matrix for this class might depend on frame attributes.

        Parameters
        ----------
        matrix_func : callable
            A callable that has the signature ``matrix_func(fromcoord, toframe)`` and
            returns a 3 x 3 matrix that converts ``fromcoord`` in a cartesian
            representation to the new coordinate system.
        fromsys : class
            The coordinate frame class to start from.
        tosys : class
            The coordinate frame class to transform into.
        priority : float or int
            The priority if this transform when finding the shortest
            coordinate transform path - large numbers are lower priorities.
        register_graph : `~astropy.coordinates.TransformGraph` or None
            A graph to register this transformation with on creation, or
            `None` to leave it unregistered.

        Raises
        ------
        TypeError
            If ``matrix_func`` is not callable

        """

        def __init__(self, matrix_func, fromsys, tosys, priority=1, register_graph=None):
            if not callable(matrix_func):
                raise TypeError("matrix_func is not callable")
            self.matrix_func = matrix_func

            super().__init__(
                fromsys, tosys, priority=priority, register_graph=register_graph
            )

        def _affine_params(self, fromcoord, toframe):
            return self.matrix_func(fromcoord, toframe), None