Source code for autowisp.astrometry.anmatch_transformation

"""Define a class to apply sky-to-frame transformations found by anmatch."""

import os.path
import numpy

from autowisp.astrometry.map_projections import gnomonic_projection




class AnmatchTransformation:
    """A clas that applies transformations generated by anmatch."""



    def __init__(self, trans_file=None):
        """
        Parse the transformation to apply from the given file.

        Args:
            trans_file:    The transformation file contaning the transformation
                to apply. If this is the name of an existing file, it is opened
                and parsed, otherwise this should be a sequence of the lines
                contained in the transfomation file.

        Returns:
            None
        """

        self.coefficients = {"x": [], "y": []}
        if trans_file is not None:
            if os.path.exists(trans_file):
                with open(trans_file, encoding="ascii") as trans_lines:
                    self.read_anmatch_transformation(trans_lines)
            else:
                self.read_anmatch_transformation(trans_file)




    def read_anmatch_transformation(self, trans_lines):
        """
        Parse the transformation to apply from the given file.

        Args:
            trans_lines:    A sequence of the individual lines of the trans file
                defining the transformation to apply.

        Returns:
            None
        """

        for line in trans_lines:
            if line[0] == "#":
                if line[1:].strip().startswith("2MASS:"):
                    # ra is perfectly reasonable for sky position
                    # pylint: disable=invalid-name
                    ra, dec = map(float, line[1:].strip().split()[1:3])
                    # pylint: enable=invalid-name
                    self.center = {"RA": ra, "Dec": dec}
            else:
                key, value = map(str.strip, line.split("="))
                if key == "type":
                    assert value == "polynomial"
                elif key == "order":
                    self.order = int(value)
                elif key == "offset":
                    assert value == "0, 0"
                elif key == "scale":
                    assert value == "1"
                elif key == "basisshift":
                    assert value == "0, 0"
                elif (
                    key[0] == "d" and key[1] in ["x", "y"] and key[2:] == "fit"
                ):
                    self.coefficients[key[1]] = [
                        float(word.strip()) for word in value.split(",")
                    ]

        assert (
            len(self.coefficients["x"])
            == (self.order + 1) * (self.order + 2) // 2
        )
        assert (
            len(self.coefficients["y"])
            == (self.order + 1) * (self.order + 2) // 2
        )




    def evaluate_polynomial(self, projected, transformed):
        """
        Apply the currently defined polynomial transfomation to (xi, eta).

        Args:
            projected:    The projected (c.f. gnomonic_projection()) coordinates
                to transform. Should have `'xi'` and `'eta'` keys containing the
                projected coordinates to transform in degrees.

            transformed:    A numpy array with `'x'` and `'y'` fields to fill
                with the transformed coordinates (in pixels).

        Returns:
            None
        """

        transformed["x"] = transformed["y"] = 0.0

        xi_to_order = 1.0
        order_factorial = 1.0
        coef_index = 0
        for order in range(self.order + 1):

            xi_term = numpy.copy(xi_to_order)
            eta_term = 1.0

            xi_factor = order_factorial
            eta_factor = 1.0

            for eta_power in range(order + 1):
                for out_var in ["x", "y"]:
                    transformed[out_var] += (
                        self.coefficients[out_var][coef_index]
                        * xi_term
                        * eta_term
                        / xi_factor
                        / eta_factor
                    )

                coef_index += 1
                if eta_power < order:
                    xi_term /= projected["xi"]
                    eta_term *= projected["eta"]

                    xi_factor /= order - eta_power
                    eta_factor *= eta_power + 1

            order_factorial *= order + 1
            xi_to_order *= projected["xi"]




    def __call__(self, sources, save_intermediate=False):
        """
        Return the x and y coordinates to which the given (ra, dec) project to.

        Args:
            sources:    The sky position to project to frame coordinates. Should
                have keys named `'RA'` and `'Dec'` in degrees.

            save_intermediate:    Should the intermediate coordinates be saved.
                Those are the coordinates after apllying the sky projection and
                before applying the polynomial transformation.

        Returns:
            numpy.array(dtype=[('x', float), ('y', float)]):
                The position in pixels relative to the lower left frame of the
                corner to which the given sky position projects under the
                currently defined transformation, as a numpy array with fields
                named `'x'` and `'y'`. If `save_intermediate` is True, the
                result also contains fields called `'xi'` and `'eta'` containing
                the intermediate coordinates."""

        try:
            projected_shape = sources["RA"].shape
        except TypeError:
            projected_shape = (1,)

        # pylint false positive
        # pylint: disable=no-member
        projected_dtype = [("x", numpy.float64), ("y", numpy.float64)]
        intermediate_dtype = [("xi", numpy.float64), ("eta", numpy.float64)]
        # pylint: enable=no-member
        if save_intermediate:
            projected_dtype.extend(intermediate_dtype)

        projected = numpy.empty(projected_shape, dtype=projected_dtype)
        intermediate = (
            projected
            if save_intermediate
            else numpy.empty(projected_shape, intermediate_dtype)
        )

        gnomonic_projection(sources, intermediate, **self.center)
        self.evaluate_polynomial(intermediate, projected)
        return projected