Skip to main content
`Ctrl`+`K`

# Source code for astropy.coordinates.transformations.function

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

    """Function-based coordinate transformations.

    These are transformations that cannot be represented as an affine transformation.
    """

    from contextlib import suppress
    from inspect import signature
    from warnings import warn

    from astropy import units as u
    from astropy.coordinates.transformations.base import CoordinateTransform
    from astropy.utils.exceptions import AstropyWarning

    __all__ = ["FunctionTransform", "FunctionTransformWithFiniteDifference"]

    class FunctionTransform(CoordinateTransform):
        """
        A coordinate transformation defined by a function that accepts a
        coordinate object and returns the transformed coordinate object.

        Parameters
        ----------
        func : callable
            The transformation function. Should have a call signature
            ``func(formcoord, toframe)``. Note that, unlike
            `CoordinateTransform.__call__`, ``toframe`` is assumed to be of type
            ``tosys`` for this function.
        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 ``func`` is not callable.
        ValueError
            If ``func`` cannot accept two arguments.

        """

        def __init__(self, func, fromsys, tosys, priority=1, register_graph=None):
            if not callable(func):
                raise TypeError("func must be callable")

            with suppress(TypeError):
                sig = signature(func)
                kinds = [x.kind for x in sig.parameters.values()]
                if (
                    len(x for x in kinds if x == sig.POSITIONAL_ONLY) != 2
                    and sig.VAR_POSITIONAL not in kinds
                ):
                    raise ValueError("provided function does not accept two arguments")

            self.func = func

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

        def __call__(self, fromcoord, toframe):
            res = self.func(fromcoord, toframe)
            if not isinstance(res, self.tosys):
                raise TypeError(
                    f"the transformation function yielded {res} but "
                    f"should have been of type {self.tosys}"
                )
            if fromcoord.data.differentials and not res.data.differentials:
                warn(
                    "Applied a FunctionTransform to a coordinate frame with "
                    "differentials, but the FunctionTransform does not handle "
                    "differentials, so they have been dropped.",
                    AstropyWarning,
                )
            return res

    class FunctionTransformWithFiniteDifference(FunctionTransform):
        r"""Transformation based on functions using finite difference for velocities.

        A coordinate transformation that works like a
        `~astropy.coordinates.FunctionTransform`, but computes velocity shifts
        based on the finite-difference relative to one of the frame attributes.
        Note that the transform function should *not* change the differential at
        all in this case, as any differentials will be overridden.

        When a differential is in the from coordinate, the finite difference
        calculation has two components. The first part is simple the existing
        differential, but re-orientation (using finite-difference techniques) to
        point in the direction the velocity vector has in the *new* frame. The
        second component is the "induced" velocity.  That is, the velocity
        intrinsic to the frame itself, estimated by shifting the frame using the
        ``finite_difference_frameattr_name`` frame attribute a small amount
        (``finite_difference_dt``) in time and re-calculating the position.

        Parameters
        ----------
        finite_difference_frameattr_name : str or None
            The name of the frame attribute on the frames to use for the finite
            difference.  Both the to and the from frame will be checked for this
            attribute, but only one needs to have it. If None, no velocity
            component induced from the frame itself will be included - only the
            re-orientation of any existing differential.
        finite_difference_dt : `~astropy.units.Quantity` ['time'] or callable
            If a quantity, this is the size of the differential used to do the
            finite difference.  If a callable, should accept
            ``(fromcoord, toframe)`` and return the ``dt`` value.
        symmetric_finite_difference : bool
            If True, the finite difference is computed as
            :math:`\frac{x(t + \Delta t / 2) - x(t + \Delta t / 2)}{\Delta t}`, or
            if False, :math:`\frac{x(t + \Delta t) - x(t)}{\Delta t}`.  The latter
            case has slightly better performance (and more stable finite difference
            behavior).

        All other parameters are identical to the initializer for
        `~astropy.coordinates.FunctionTransform`.

        """

        def __init__(
            self,
            func,
            fromsys,
            tosys,
            priority=1,
            register_graph=None,
            finite_difference_frameattr_name="obstime",
            finite_difference_dt=1 * u.second,
            symmetric_finite_difference=True,
        ):
            super().__init__(func, fromsys, tosys, priority, register_graph)
            self.finite_difference_frameattr_name = finite_difference_frameattr_name
            self.finite_difference_dt = finite_difference_dt
            self.symmetric_finite_difference = symmetric_finite_difference

        @property
        def finite_difference_frameattr_name(self):
            return self._finite_difference_frameattr_name

        @finite_difference_frameattr_name.setter
        def finite_difference_frameattr_name(self, value):
            if value is None:
                self._diff_attr_in_fromsys = self._diff_attr_in_tosys = False
            else:
                diff_attr_in_fromsys = value in self.fromsys.frame_attributes
                diff_attr_in_tosys = value in self.tosys.frame_attributes
                if diff_attr_in_fromsys or diff_attr_in_tosys:
                    self._diff_attr_in_fromsys = diff_attr_in_fromsys
                    self._diff_attr_in_tosys = diff_attr_in_tosys
                else:
                    raise ValueError(
                        f"Frame attribute name {value} is not a frame attribute of"
                        f" {self.fromsys} or {self.tosys}"
                    )
            self._finite_difference_frameattr_name = value

        def __call__(self, fromcoord, toframe):
            from astropy.coordinates.representation import (
                CartesianDifferential,
                CartesianRepresentation,
            )

            supcall = self.func
            if not fromcoord.data.differentials:
                return supcall(fromcoord, toframe)
            # this is the finite difference case

            if callable(self.finite_difference_dt):
                dt = self.finite_difference_dt(fromcoord, toframe)
            else:
                dt = self.finite_difference_dt
            halfdt = dt / 2

            from_diffless = fromcoord.realize_frame(fromcoord.data.without_differentials())
            reprwithoutdiff = supcall(from_diffless, toframe)

            # first we use the existing differential to compute an offset due to
            # the already-existing velocity, but in the new frame
            fromcoord_cart = fromcoord.cartesian
            if self.symmetric_finite_difference:
                fwdxyz = (
                    fromcoord_cart.xyz + fromcoord_cart.differentials["s"].d_xyz * halfdt
                )
                fwd = supcall(
                    fromcoord.realize_frame(CartesianRepresentation(fwdxyz)), toframe
                )
                backxyz = (
                    fromcoord_cart.xyz - fromcoord_cart.differentials["s"].d_xyz * halfdt
                )
                back = supcall(
                    fromcoord.realize_frame(CartesianRepresentation(backxyz)), toframe
                )
            else:
                fwdxyz = fromcoord_cart.xyz + fromcoord_cart.differentials["s"].d_xyz * dt
                fwd = supcall(
                    fromcoord.realize_frame(CartesianRepresentation(fwdxyz)), toframe
                )
                back = reprwithoutdiff
            diffxyz = (fwd.cartesian - back.cartesian).xyz / dt

            # now we compute the "induced" velocities due to any movement in
            # the frame itself over time
            attrname = self.finite_difference_frameattr_name
            if attrname is not None:
                if self.symmetric_finite_difference:
                    if self._diff_attr_in_fromsys:
                        kws = {attrname: getattr(from_diffless, attrname) + halfdt}
                        from_diffless_fwd = from_diffless.replicate(**kws)
                    else:
                        from_diffless_fwd = from_diffless
                    if self._diff_attr_in_tosys:
                        kws = {attrname: getattr(toframe, attrname) + halfdt}
                        fwd_frame = toframe.replicate_without_data(**kws)
                    else:
                        fwd_frame = toframe
                    fwd = supcall(from_diffless_fwd, fwd_frame)

                    if self._diff_attr_in_fromsys:
                        kws = {attrname: getattr(from_diffless, attrname) - halfdt}
                        from_diffless_back = from_diffless.replicate(**kws)
                    else:
                        from_diffless_back = from_diffless
                    if self._diff_attr_in_tosys:
                        kws = {attrname: getattr(toframe, attrname) - halfdt}
                        back_frame = toframe.replicate_without_data(**kws)
                    else:
                        back_frame = toframe
                    back = supcall(from_diffless_back, back_frame)
                else:
                    if self._diff_attr_in_fromsys:
                        kws = {attrname: getattr(from_diffless, attrname) + dt}
                        from_diffless_fwd = from_diffless.replicate(**kws)
                    else:
                        from_diffless_fwd = from_diffless
                    if self._diff_attr_in_tosys:
                        kws = {attrname: getattr(toframe, attrname) + dt}
                        fwd_frame = toframe.replicate_without_data(**kws)
                    else:
                        fwd_frame = toframe
                    fwd = supcall(from_diffless_fwd, fwd_frame)
                    back = reprwithoutdiff

                diffxyz += (fwd.cartesian - back.cartesian).xyz / dt

            newdiff = CartesianDifferential(diffxyz)
            reprwithdiff = reprwithoutdiff.data.to_cartesian().with_differentials(newdiff)
            return reprwithoutdiff.realize_frame(reprwithdiff)