autowisp.processing_steps.fit_source_extracted_psf_map — AutoWISP documentation

# Source code for autowisp.processing_steps.fit_source_extracted_psf_map

    #!/usr/bin/env python3

    """Fit a smooth dependence of source extracted PSF parameters."""

    import logging
    from collections import namedtuple

    import numpy

    from autowisp.multiprocessing_util import setup_process
    from autowisp.file_utilities import find_dr_fnames
    from autowisp import DataReductionFile
    from autowisp.fit_expression import (
        Interface as FitTermsInterface,
        iterative_fit,
    )
    from autowisp.evaluator import Evaluator
    from autowisp.processing_steps.manual_util import (
        ManualStepArgumentParser,
        ignore_progress,
    )

    _logger = logging.getLogger(__name__)

    input_type = "dr"

    def parse_command_line(*args):
        """Return the parsed command line arguments."""

        parser = ManualStepArgumentParser(
            description=__doc__,
            input_type=("" if args else input_type),
            inputs_help_extra="The DR files must already contain astrometry.",
            add_component_versions=("srcextract", "catalogue", "skytoframe"),
        )
        parser.add_argument(
            "--srcextract-only-if",
            default="True",
            help="Expression involving the header of the input images that "
            "evaluates to True/False if a particular image from the specified "
            "image collection should/should not be processed.",
        )
        parser.add_argument(
            "--srcextract-psf-params",
            nargs="+",
            default=None,
            help="List of the parameters describing PSF shapes of the extracted "
            "sources to fit a smooth dependence for. If left unspecified, smooth "
            "map will be fit for (`'S'`, `'D'`, `'K'`) or (`'fwhm'`, `'round'`, "
            "`'pa'`), whichever is available.",
        )
        parser.add_argument(
            "--srcextract-psfmap-terms",
            default="O3{x, y}*O1{phot_g_mean_mag}",
            help="An expression involving source extraction and/or catalogue "
            "variables for the terms to include in the smoothing fit.",
        )
        parser.add_argument(
            "--srcextract-psfmap-weights",
            default=None,
            type=str,
            help="An expression involving source extraction and/or catalogue "
            "variables for the weights to use for the smoothing fit.",
        )
        parser.add_argument(
            "--srcextract-psfmap-error-avg",
            default="median",
            help="How to average fitting residuals for outlier rejection. ",
        )
        parser.add_argument(
            "--srcextract-psfmap-rej-level",
            type=float,
            default=5.0,
            help="How far away from the fit should a point be before it is rejected"
            " in units of error_avg.",
        )
        parser.add_argument(
            "--srcextract-psfmap-max-rej-iter",
            type=int,
            default=20,
            help="The maximum number of rejection/re-fitting iterations to perform."
            "If the fit has not converged by then, the latest iteration is "
            "accepted. Default: %(default)s",
        )
        return parser.parse_args(*args)

    def get_predictors_and_weights(
        matched_sources, fit_terms_expression, weights_expression
    ):
        """Return the matrix of predictors to use for fitting."""
        _logger.debug("Matched columns: %s", repr(matched_sources.columns))
        if weights_expression is None:
            return (FitTermsInterface(fit_terms_expression)(matched_sources), None)
        # TODO fix matched_sources to records in Evaluator not here
        return (
            FitTermsInterface(fit_terms_expression)(matched_sources),
            Evaluator(matched_sources.to_records(index=False))(weights_expression),
        )

    def get_psf_param(matched_sources, psf_parameters):
        """Return a numpy structured array of the PSF parameters."""

        result = numpy.empty(
            len(matched_sources),
            dtype=[(param, numpy.float64) for param in psf_parameters],
        )
        for param in psf_parameters:
            _logger.debug("Setting PSF param %s", repr(param))
            _logger.debug("Matched columns: %s", repr(matched_sources.columns))
            result[param] = matched_sources[param]

        return result

    def detect_psf_parameters(matched_sources):
        """Return the default PSF parameters to fit for he given DR file."""

        for try_psf_params in [
            ("S", "D", "K"),
            ("s", "d", "k"),
            ("fwhm", "round", "pa"),
        ]:
            found_all = True
            for param in try_psf_params:
                if param not in matched_sources.columns:
                    found_all = False
            if found_all:
                return try_psf_params
        return None

    def smooth_srcextract_psf(
        dr_fname, configuration, mark_start, mark_end, **path_substitutions
    ):
        """
        Fit PSF parameters as polynomials of srcextract and catalogue info.

        Args:
            dr_fname(str):    The name of the DR file to smooth the source extracted
                parameters for.

            configuration(NamedTuple):    Configuration on how to do the smoothing.
                Should include the following attributes:

                psf_parameters([str]):    A list of the variables from the source
                    extracted datasets to smooth. If None, an attempt is made to
                    auto detect the parameters.

                fit_terms_expression(str):    A fitting terms expression defining
                    the terms to include in the fit.

                weights_expression(str):    An expression involving source
                    extraction and/or catalogue variables for the weights to use for
                    the smoothing fit.

                error_avg:    See iterative_fit().

                rej_level:    See iterative_fit().

                max_rej_iter:    See iterative_fit().

            mark_start(callable):    Invoked with the name of the DR file before the
                contents of the file are updated.

            mark_end(callable):    Invoked with the name of the DR file after
                successfully completing all updates to the DR file.

            path_substitutions:    Any substitutions required to resolve the
                path to extracted sources, catalogue sources and the
                destinationdatasets and attributes created by this method.

        Returns:
            None
        """

        with DataReductionFile(dr_fname, "r+") as dr_file:
            matched_sources = dr_file.get_matched_sources(**path_substitutions)
            predictors, weights = get_predictors_and_weights(
                matched_sources,
                configuration.fit_terms_expression,
                configuration.weights_expression,
            )
            _logger.debug(
                "Predictors (%sx%d)}: %s", *predictors.shape, repr(predictors)
            )
            if weights is not None:
                _logger.debug(
                    "Weights %s: %s", format(weights.shape), repr(weights)
                )
            else:
                _logger.debug("Not using weights")

            fit_results = {"coefficients": {}, "fit_res2": {}, "num_fit_src": {}}

            psf_parameters = configuration.psf_parameters
            if psf_parameters is None:
                psf_parameters = detect_psf_parameters(matched_sources)
            if psf_parameters is None:
                raise IOError(
                    f"Matched sources in {dr_file.filename} do not contain a full "
                    "set of either fistar or hatphot PSF parameters."
                )

            psf_param = get_psf_param(matched_sources, psf_parameters)
            for param_name in psf_parameters:
                (
                    fit_results["coefficients"][param_name],
                    fit_results["fit_res2"][param_name],
                    fit_results["num_fit_src"][param_name],
                ) = iterative_fit(
                    predictors,
                    psf_param[param_name],
                    weights=weights,
                    error_avg=configuration.error_avg,
                    rej_level=configuration.rej_level,
                    max_rej_iter=configuration.max_rej_iter,
                    fit_identifier=f"Extracted sources PSF {param_name:s} map",
                )
                if fit_results["coefficients"][param_name] is None:
                    mark_start(dr_fname)
                    mark_end(dr_fname, -2)
                    return

            mark_start(dr_fname)
            dr_file.save_source_extracted_psf_map(
                fit_results=fit_results,
                fit_configuration=configuration,
                **path_substitutions,
            )
            mark_end(dr_fname)

    def get_dr_substitutions(configuration):
        """Return the path substitutions needed to resolve input datasets."""

        return {
            what + "_version": configuration[what + "_version"]
            for what in ["srcextract", "catalogue", "skytoframe"]
        }

    def fit_source_extracted_psf_map(
        dr_collection, start_status, configuration, mark_start, mark_end
    ):
        """Fit a smooth dependence of source extraction PSF for a DR collection."""

        assert start_status is None

        # This defines a type not variable
        # pylint: disable=invalid-name
        SmoothingConfigType = namedtuple(
            "SmoothingConfigType",
            [
                "psf_parameters",
                "fit_terms_expression",
                "weights_expression",
                "error_avg",
                "rej_level",
                "max_rej_iter",
            ],
        )
        # pylint: enable=invalid-name

        smoothing_config = SmoothingConfigType(
            psf_parameters=configuration["srcextract_psf_params"],
            fit_terms_expression=configuration["srcextract_psfmap_terms"],
            weights_expression=configuration["srcextract_psfmap_weights"],
            **{
                fit_config: configuration["srcextract_psfmap_" + fit_config]
                for fit_config in ["error_avg", "rej_level", "max_rej_iter"]
            },
        )

        for dr_fname in dr_collection:
            smooth_srcextract_psf(
                dr_fname=dr_fname,
                configuration=smoothing_config,
                mark_start=mark_start,
                mark_end=mark_end,
                **get_dr_substitutions(configuration),
            )

    def cleanup_interrupted(interrupted, configuration):
        """Remove the source extracted PSF map from the given DR file."""

        path_substitutions = get_dr_substitutions(configuration)

        for dr_fname, status in interrupted:
            assert status == 0

            with DataReductionFile(dr_fname, "r+") as dr_file:
                dr_file.delete_dataset("srcextract.psf_map", **path_substitutions)

        return -1

    def has_astrometry(dr_fname, substitutions):
        """Check if the DR file contains a sky-to-frame transformation."""

        with DataReductionFile(dr_fname, mode="r") as dr_file:
            try:
                dr_file.check_for_dataset(
                    "skytoframe.coefficients", **substitutions
                )
                return True
            except IOError:
                return False

    def main():
        """Run the step from the command line."""

        configuration = parse_command_line()
        setup_process(task="main", **configuration)
        dr_substitutions = get_dr_substitutions(configuration)
        fit_source_extracted_psf_map(
            [
                dr_fname
                for dr_fname in find_dr_fnames(
                    configuration.pop("dr_files"),
                    configuration.pop("srcextract_only_if"),
                )
                if has_astrometry(dr_fname, dr_substitutions)
            ],
            None,
            configuration,
            ignore_progress,
            ignore_progress,
        )

    if __name__ == "__main__":
        main()