computeRoughPsfShapelets.py

Go to the documentation of this file.

# This file is part of meas_algorithms.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (https://www.lsst.org).
# See the COPYRIGHT file at the top-level directory of this distribution
# for details of code ownership.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
from __future__ import annotations
 
__all__ = ["ComputeRoughPsfShapeletsTask", "ComputeRoughPsfShapeletsConfig"]
 
import itertools
from typing import TYPE_CHECKING, Any, Literal
 
import numpy as np
from logging import Logger
import scipy.signal
import scipy.stats
from sklearn.covariance import MinCovDet
from sklearn.neighbors import KernelDensity
 
from lsst.afw.geom import ellipses
from lsst.afw.image import ImageD, ImageF, MaskedImageF, ExposureSummaryStats
from lsst.afw.table import Point2DKey, QuadrupoleKey, Schema, SourceCatalog
from lsst.geom import Box2D
from lsst.pex.config import Config, Field, ListField, FieldValidationError
from lsst.pipe.base import AlgorithmError, Struct, Task
from lsst.shapelet import LAGUERRE, ShapeletFunction, computeOffset
 
from ._algorithmsLib import SpanSetMoments
 
if TYPE_CHECKING:
    import matplotlib.axes
    import matplotlib.figure
    import matplotlib.image
    import matplotlib.patches
 
 
class NoStarsForShapeletsError(AlgorithmError):
    """Exception raised when any of the methods involving selection fail to
    find any usable stars (e.g. compute_raw_moments,  _threshold_with_bounds,
    _find_first_radius_mode).
    """
 
    @property
    def metadata(self) -> dict[str, Any]:
        return {}
 
 
class ComputeRoughPsfShapeletsConfig(Config):
    bad_mask_planes = ListField(
        "Mask planes to identify pixels to drop from the calculation.",
        dtype=str,
        default=["SAT", "SUSPECT", "INTRP"],
    )
    bad_pixel_max_fraction = Field(
        "Maximum fraction of a footprint's pixels that can be bad (according "
        "to bad_mask_planes) before that footprint is fully dropped.",
        dtype=float,
        default=0.25,
    )
    bad_pixel_exclusion_radius = Field(
        "If a bad pixel (according to bad_mask_planes) falls within this "
        "radius of the unweighted centroid of a footprint, drop that footprint "
        "entirely.",
        dtype=float,
        default=2.0,
    )
    max_footprint_area = Field(
        "Footprints with a pixel count larger than this threshold are dropped before computing moments.",
        dtype=int,
        default=10000,
    )
    min_snr = Field(
        "Mininum flux S/N for inclusion in the star sample.",
        dtype=float,
        default=50.0,
    )
    max_radius_factor = Field(
        "Maximium multiple of the mode of the radius distribution for inclusion in the star sample.",
        dtype=float,
        default=2.0,
    )
    max_shape_distance = Field(
        "Maximum Mahalanobis distance (distance from the center of the shape "
        "distribution in elliptical sigma units) to select a star from the candidate "
        "sample, comparing the shape of that source vs. a robust estimate of the "
        "distribution of the shapes of all sources.",
        dtype=float,
        default=3.0,
    )
    min_n_stars = Field(
        "Minimum number of stars to select.  The S/N, radius, and shape distance thresholds are "
        "relaxed as needed to meet this target.",
        dtype=int,
        default=10,
    )
    max_n_stars = Field(
        "Maximum number of stars to select.  High shape-distance sources "
        "are dropped as needed to meet this target.",
        dtype=int,
        default=20,
    )
    logarithmic_shapes = Field(
        "If True, transform the (xx, yy, xy) moments to conformal shear "
        "and log trace radius when selecting stars in order to map the ellipse "
        "parameter to a space with (-inf, inf) bounds on all quantities (but "
        "a less-linear relationship to the pixel data).",
        dtype=bool,
        default=False,
    )
    radius_mode_min = Field(
        "Minimum radius in pixels at which to start searching for the first mode. "
        "This just needs be large enough to avoid unphysically small garbage (e.g. CRs).",
        dtype=float,
        default=1.0,
    )
    radius_kde_bandwidth = Field(
        "Bandwidth of the Gaussian kernel density estimator used to find the "
        "first mode of the radius distribution.",
        dtype=float,
        default=1.0,
    )
    shapelet_order = Field(
        "Order of the shapelet expansion fit to the stars.",
        dtype=int,
        default=6,
    )
    shapelet_scale_factor = Field(
        "Scale factor to apply to the moments ellipse when computing the ellipse for the shapelet basis.",
        dtype=float,
        default=1.0,
    )
    shapelet_circular_basis = Field(
        "Whether to use a circular shapelet basis with the same moments trace instead of an elliptical one.",
        dtype=bool,
        default=True,
    )
    non_gaussian_non_atmosphere_index_list = ListField(
        doc="Index list of non-Gaussian, non-atmospheric contributors in the list of shapelet "
        "coefficients associated with shapelet_order (default accommodates an order of 6)",
        dtype=int,
        default=list(range(6, 10)) + list(range(15, 24)),
    )
 
    def validate(self) -> None:
        if self.min_n_stars > self.max_n_stars:
            raise ValueError(
                f"min_n_stars={self.min_n_stars} is greater than max_n_stars={self.max_n_stars}."
            )
        if self.shapelet_order < 0:
            raise ValueError(f"shapelet order {self.shapelet_order} must be nonnegative.")
        if self.shapelet_scale_factor <= 0.0:
            raise ValueError(f"shapelet scale factor {self.shapelet_scale_factor} must be positive.")
        # Check that the shapelet decomposition order results in at least as
        # many decomposition coefficients as the maximum index identified in
        # the non_gaussian_non_atmosphere_index_list config.
        n_coeff = 0
        for i_offset in range(self.shapelet_order + 1):
            n_coeff += (i_offset + 1)
        if max(self.non_gaussian_non_atmosphere_index_list) > n_coeff:
            raise FieldValidationError(
                ComputeRoughPsfShapeletsConfig,
                self,
                "The largest coefficient index defined in config.non_gaussian_non_atmosphere_index_list "
                f"({len(self.non_gaussian_non_atmosphere_index_list)}) is larger than the number "
                f"of coefficients expected ({n_coeff}) from the shapelet decomposition order "
                f"({self.shapelet_order})."
            )
 
 
class ComputeRoughPsfShapeletsTask(Task):
    """A task that computes a rough shapelet expansion of the PSF from a set
    of high S/N detections.
 
    Notes
    -----
    This task is expected to be run early in single-epoch processing - just
    after background subtraction and an initial high S/N detection phase, and
    before any deblending or measurement - in order to identify out-of-focus
    or otherwise bad PSFs.
 
    Given a background-subtracted `lsst.afw.image.MaskedImage`, an
    `lsst.afw.table.SourceCatalog` with footprints attached, and a random
    number generator seed, the `run` method will:
 
    - Compute the *unweighted* 0th-2nd moments of every non-child source over
      the footprint (except certain configurable masked pixels).  This is
      delegated to the `compute_raw_moments` method (which uses the C++
      `SpanSetMoments` class for the pixel-level processing).  Unweighted
      moments are used to avoid "latching onto" a small piece of PSF
      substructure, but can be much noiser than the Gaussian-weighed moments
      we usually use.
 
    - Select a "candidate" sample of sources with successfully measured
      moments that satisfy a S/N cut and a radius cut (determined from the
      first mode of the radius distribution, via kernel density estimation),
      and then use a robust covariance estimator (`scikit_learn.MinCovDet`)
      to select presumed isolated stars that are close to the center of that
      distribution, in 3-parameter shape space. This is delegated to the
      `select_stars` method.
 
    - Fit a single shapelet expansion to the selected stars.  This is
      mostly delegated to the `SpanSetMoments.fit_shapelets` method.
 
    The radial shapelet terms at 0th, 2nd, and 4th order are expected to form
    a space in which donut-shaped PSFs are well-separated from those with
    monotonic profiles.  Other terms *may* be useful in identifying other
    kinds of undesirable PSF structure.
    """
 
    ConfigClass = ComputeRoughPsfShapeletsConfig
    _DefaultName = "computeRoughPsfShapelets"
    config: ComputeRoughPsfShapeletsConfig
 
    def __init__(
        self,
        config: ComputeRoughPsfShapeletsConfig | None = None,
        *,
        schema: Schema,
        **kwargs: Any,
    ):
        super().__init__(config=config, **kwargs)
        self.schema = schema
        self.shapelets_base_name = "RoughPsfShapelets"
        self._flux_key = schema.addField(
            self.shapelets_base_name + "_flux",
            type=float,
            doc="Unweighted zeroth moment.",
        )
        self._flux_err_key = schema.addField(
            self.shapelets_base_name + "_fluxErr",
            type=float,
            doc="Uncertainty on the unweighted zeroth moment.",
        )
        self._center_key = Point2DKey.addFields(
            schema, self.shapelets_base_name, "Center from unweighted first moments.", "pixels"
        )
        self._shape_key = QuadrupoleKey.addFields(
            schema, self.shapelets_base_name, "Shape from unweighted second moments."
        )
        self._flag_key = schema.addField(
            self.shapelets_base_name + "_flag",
            type="Flag",
            doc="Flag set if the raw PSF moments were not computed.",
        )
        self._candidate_key = schema.addField(
            self.shapelets_base_name + "_candidate",
            type="Flag",
            doc="Flag set if this source passed the radius_fraction cut (see configuration).",
        )
        self._used_key = schema.addField(
            self.shapelets_base_name + "_used",
            type="Flag",
            doc=(
                "Flag set if this source passed the radius_fraction and shape_distance cuts "
                "(see configuration) and was used to fit the shapelet expansion."
            ),
        )
 
    def run(self, *, masked_image: MaskedImageF, catalog: SourceCatalog, seed: int) -> Struct:
        """Compute raw moments, select stars, and fit a shapelet expansion to
        them.
 
        Parameters
        ----------
        masked_image
            Masked image to measure on.  Must be background-subtracted.
        catalog
            Catalog of detections to extract footprints from and fill output
            columns of.  Its schema must be a superset of ``self.schema``.
        seed
            A random-number generator seed, used for the robust covariance
            estimator.
 
        Returns
        -------
        `lsst.pipe.base.Struct`
            A struct of results containing:
 
            - ``shapelet`` (`lsst.shapelet.ShapeletFunction`): A
              Gauss-Laguerre (polar shaplet) expansion of the PSF.
            - ``radial`` (`list` [`float`]): the purely radial coefficients
              of the shapelet expansion.
            - all attributes returned by the `select_stars` method.
        """
        moments = self.compute_raw_moments(masked_image=masked_image, catalog=catalog)
        result = self.select_stars(catalog, seed=seed)
        star_moments = [moments[star_id] for star_id in result.star_ids]
        result.shapelet = SpanSetMoments.fit_shapelets(
            masked_image,
            star_moments,
            self.config.shapelet_order,
            self.config.shapelet_scale_factor,
            self.config.shapelet_circular_basis,
        )
        result.shapelet.getEllipse().setCore(result.mean_shape)
        result.shapelet.changeBasisType(LAGUERRE)
        result.radial = result.shapelet.getCoefficients()[
            [computeOffset(i) for i in range(0, self.config.shapelet_order + 1, 2)]
        ]
        self.log.info("Rough PSF shapelet radial terms: %s.", result.radial)
        result = self.compute_shapelets_only_iq_score(result)
        return result
 
    def compute_raw_moments(
        self, *, masked_image: MaskedImageF, catalog: SourceCatalog
    ) -> dict[int, SpanSetMoments]:
        """Compute the unweighted moments of the footprints in a catalog.
 
        Parameters
        ----------
        masked_image
            Masked image to measure on.  Must be background-subtracted.
        catalog
            Catalog of detections to extract footprints from and fill output
            columns of.  Its schema must be a superset of ``self.schema``.
 
        Returns
        -------
        `dict` [`int`, `SpanSetMoments`]
            Objects used to construct and hold the unweighted moments and the
            pixel region used to computed them, keyed by source ID.
        """
        bitmask = masked_image.mask.getPlaneBitMask(self.config.bad_mask_planes)
        all_moments: dict[int, SpanSetMoments] = {}
        for record in catalog:
            if record.getParent() != 0:
                record.set(self._flag_key, True)
                self.log.debug("Skipping child source %s", record.getId())
                continue
            footprint_spans = record.getFootprint().getSpans()
            if footprint_spans.getArea() > self.config.max_footprint_area:
                record.set(self._flag_key, True)
                self.log.debug(
                    "Skipping source %s with footprint area %d > %d.",
                    record.getId(),
                    footprint_spans.getArea(),
                    self.config.max_footprint_area,
                )
                continue
            moments = SpanSetMoments.compute(
                record.getFootprint().getSpans(),
                masked_image=masked_image,
                bad_bitmask=bitmask,
                bad_pixel_max_fraction=self.config.bad_pixel_max_fraction,
                bad_pixel_exclusion_radius=self.config.bad_pixel_exclusion_radius,
            )
            record.set(self._flux_key, moments.flux)
            record.set(self._flux_err_key, moments.variance**0.5)
            record.set(self._center_key, moments.center)
            record.set(self._shape_key, moments.shape)
            record.set(self._flag_key, moments.any_flags_set)
            if not moments.any_flags_set:
                all_moments[record.getId()] = moments
        if all_moments:
            self.log.verbose(
                "Successfully measured raw moments for %d of %d sources.",
                len(all_moments),
                len(catalog),
            )
        else:
            raise NoStarsForShapeletsError("No raw moments could be measured.")
        return all_moments
 
    def select_stars(self, catalog: SourceCatalog, seed: int) -> Struct:
        """Select probable stars from the distribution of second moments.
 
        Parameters
        ----------
        catalog
            Catalog of detections to extract footprints from and fill output
            columns of.  Its schema must be a superset of ``self.schema``.
        seed
            A random-number generator seed, used for the robust covariance
            estimator.
 
        Returns
        -------
        `lsst.pipe.base.Struct`
            A struct of results containing:
 
            - ``star_ids`` (`numpy.ndarray`): the source IDs that are expected
              to be stars.
            - ``mean_shape`` (`lsst.afw.geom.ellipses.BaseCore`): the mean of
              the shape distribution.
            - ``shape_covariance`` (`numpy.ndarray`): the covariance of the
              distribution of shapes; a 3x3 matrix.  This uses the same
              parameterization of the shapes as ``mean_shape``.
            - ``radius_cut`` (`float`): the indended radius cut (i.e. the
              mode of the radius distribution multipled by the
              ``radius_factor`` configuration option).
            - ``radius_kde`` (`sklearn.neighbors.KernelDensity`): kernel
              density estimator on the radius distribution, used to determine
              the radius cut.
        """
        # Cut on flags and SNR first.
        indices = np.arange(len(catalog), dtype=int)[np.logical_not(catalog[self._flag_key])]
        signalToNoise = []
        for index in indices:
            if (np.isfinite(catalog[self._flux_err_key][index])
               and np.isfinite(catalog[self._flux_err_key][index])
               and catalog[self._flux_err_key][index] != 0.0):
                signalToNoise.append(catalog[self._flux_key][index] / catalog[self._flux_err_key][index])
            else:
                signalToNoise.append(np.nan)
        signalToNoise = np.array(signalToNoise)
        indices = indices[
            self._threshold_with_bounds(
                signalToNoise,
                threshold=self.config.min_snr,
                min_count=self.config.min_n_stars,
                max_count=len(catalog),
                name="S/N",
                kind=">",
            )
        ]
        # Cut on radius next.
        radii = np.zeros(indices.size, dtype=np.float64)
        for n, index in enumerate(indices):
            record = catalog[index]
            shape = record.get(self._shape_key)
            radii[n] = shape.getTraceRadius()
        radius_mode, radius_kde = self._find_first_radius_mode(radii)
        radius_cut = self.config.max_radius_factor * radius_mode
        indices = indices[
            self._threshold_with_bounds(
                radii,
                threshold=radius_cut,
                min_count=self.config.min_n_stars,
                max_count=len(catalog),
                name="radius",
                kind="<",
            )
        ]
        shape_data = np.zeros((len(indices), 3), dtype=np.float64)
        for n, index in enumerate(indices):
            record = catalog[index]
            record.set(self._candidate_key, True)
            shape = record.get(self._shape_key)
            if self.config.logarithmic_shapes:
                shape = ellipses.SeparableConformalShearLogTraceRadius(shape)
            shape_data[n, :] = shape.getParameterVector()
        shape_dist = MinCovDet(random_state=seed).fit(shape_data)
        m_distances = shape_dist.mahalanobis(shape_data)
        indices = indices[
            self._threshold_with_bounds(
                m_distances,
                threshold=self.config.max_shape_distance,
                min_count=self.config.min_n_stars,
                max_count=self.config.max_n_stars,
                name="Mahalanobis distance",
                kind="<",
            )
        ]
        for index in indices:
            catalog[index].set(self._used_key, True)
        star_ids = catalog["id"][indices]
        if self.config.logarithmic_shapes:
            mean_shape = ellipses.SeparableConformalShearLogTraceRadius(shape_dist.location_)
        else:
            mean_shape = ellipses.Quadrupole(shape_dist.location_)
        return Struct(
            star_ids=star_ids,
            mean_shape=mean_shape,
            shape_covariance=shape_dist.covariance_,
            radius_cut=radius_cut,
            radius_kde=radius_kde,
            catalog=catalog,
        )
 
    def compute_shapelets_only_iq_score(self, results: Struct) -> Struct:
        """Compute the image quality "score" that only includes the
        non-atmospheric coefficents.
 
        The "shapelet_only_iq_score" is a dimensionless image quality score
        as determined from the shapelets decomposition that only includes
        power from the non-atmospheric coefficents from the shapelet
        decomposition. The score spans the range [0.0, 1.0] with lower
        values indicating better image quality.
 
        Parameters
        ----------
        results
            Result struct returned by `run`. To be updated in-place.
 
        Returns
        -------
        `lsst.pipe.base.Struct`
            Additional results struct components include:
            - ``shapelets_summary`` (`lsst.afw.image.ExposureSummaryStats`):
              An ExposureSummaryStats object with the following parameters
              initialized:
               - ``shapeletsCoeffs`` (`list`[`float`]): The coefficients
                 from the shapelet decomposition.
               - ``nShapeletsStar`` (`int`): The number of sources used in the
                 shapelet decomposition.
               - ``shapeletsOnlyIqScore`` (`float`): The image quality score
                 which only includes power from the non-atmospheric
                 decomposition coefficients.
               - ``shapeletsStarEMedian`` (`float`): The median ellipticity
                 (sqrt(starE1**2.0 + starE2**2.0)) of the stars used in the
                 shapelet decomposition.
               - ``shapeletsStarUnNormalizedEMedian`` (`float`): The median
                 un-normalized ellipticity
                 (sqrt((starXX - starYY)**2.0 + (2.0*starXY)**2.0)) of the
                 stars used in the shapelet decomposition.
        """
        shapelets_summary = ExposureSummaryStats()
        results.non_gaussian_non_atmosphere_index_list = \
            self.config.non_gaussian_non_atmosphere_index_list
        shapelets_coeffs = results.shapelet.getCoefficients()
        shapelets_coeffs_list = [float(coeff) for coeff in shapelets_coeffs]
        shapelets_summary.shapeletsCoeffs = shapelets_coeffs_list
        shapelets_used_cat = results.catalog[
            (results.catalog[self.shapelets_base_name + "_used"])
            & (~results.catalog[self.shapelets_base_name + "_flag"])].copy(deep=True)
        shapelets_summary.nShapeletsStar = len(shapelets_used_cat)
 
        shapelets_star_xx = shapelets_used_cat[self.shapelets_base_name + "_xx"]
        shapelets_star_yy = shapelets_used_cat[self.shapelets_base_name + "_yy"]
        shapelets_star_xy = shapelets_used_cat[self.shapelets_base_name + "_xy"]
 
        # Compute a shapelets score based on the power in the
        # non-atmospheric decomposition coefficients.
        total_power = np.sum(shapelets_coeffs**2.0)
        non_atm_power = np.sum(
            shapelets_coeffs[self.config.non_gaussian_non_atmosphere_index_list]**2.0
        )
        shapelets_only_iq_score = np.where(total_power > 0, non_atm_power/total_power, 0.0)
        shapelets_summary.shapeletsOnlyIqScore = float(shapelets_only_iq_score)
 
        self.log.info("n_shapelets_star = %d; shapelets_only_iq_score = %.5f",
                      len(shapelets_used_cat), shapelets_only_iq_score)
        shapelets_star_e1 = (shapelets_star_xx - shapelets_star_yy)/(shapelets_star_xx + shapelets_star_yy)
        shapelets_star_e2 = 2*shapelets_star_xy/(shapelets_star_xx + shapelets_star_yy)
 
        shapelets_star_e = np.sqrt(shapelets_star_e1**2.0 + shapelets_star_e2**2.0)
        shapelets_star_unnormalized_e = np.sqrt(
            (shapelets_star_xx - shapelets_star_yy)**2.0 + (2.0*shapelets_star_xy)**2.0)
        shapelets_star_e_median = np.median(shapelets_star_e)
        shapelets_star_unnormalized_e_median = np.median(shapelets_star_unnormalized_e)
        shapelets_summary.shapeletsStarEMedian = float(shapelets_star_e_median)
        shapelets_summary.shapeletsStarUnNormalizedEMedian = float(shapelets_star_unnormalized_e_median)
        results.shapelets_summary = shapelets_summary
 
        return results
 
    @staticmethod
    def compute_shapelets_iq_score(shapelets_results: Struct, catalog: SourceCatalog,
                                   centroid_offset_scale: float, shapelets_base_name: str,
                                   log: Logger | None = None) -> Struct:
        """Compute the image quality "score" that includes power from the
        median centroid offset in addition to that of the non-atmospheric
        coefficents.
 
        The "shapelets_iq_score" is a dimensionless image quality score that
        includes power from the median centroid offset between those used in
        shapelet decomposition and those in the centroid slot of the provided
        source catalog, in addition to the non-atmospheric coefficents from
        the shapelet decomposition. The score spans the range [0.0, 1.0] with
        lower values indicating better image quality.  When present (non-nan)
        the shapelets_iq_score should be considered authoritative over the
        shapelets_only_iq_score.
 
        Parameters
        ----------
        shapelets_results
            Results struct as returned by `run` method of
            ComputeRoughPsfShapeletsTask. To be updated in-place.
        catalog
            Catalog of detections to extract footprints from and fill output
            columns of.  Its schema must be a superset of ``self.schema``.
            Must contain the columns added by ComputeRoughPsfShapeletsTask.
        centroid_offset_scale
            Value (in pixels) by which to scale the centroid offset to set
            and appropriate amount by which the centroid offsets can
            contribute to the computed power.
        shapelets_base_name
            Base name for the catalog columns added by
            ComputeRoughPsfShapeletsTask.
        log
            Logger to write to.
 
        Returns
        -------
        `lsst.pipe.base.Struct`
            Additional results struct components include:
            - ``shapelets_summary`` (`lsst.afw.image.ExposureSummaryStats`):
              An ExposureSummaryStats object with the following parameters
              initialized:
               - ``shapeletsIqScore`` (`float`): the image quality score which
                 includes power from the centroid shift in addition to the
                 non-atmospheric decomposition coefficients.
               - ``centroidDiffShapeletsVsSlotMedian`` (`float`): The median
                 centroid difference
                 (sqrt((slot_x - shapelet_x)**2 + (slot_y - shapelet_y)**2))
                 for sources used in the shapelet decomposition (pixels).
        """
        if shapelets_base_name + "_used" not in catalog.schema.getNames():
            raise RuntimeError("Catalog must contain columns added by ComputeRoughPsfShapeletsTask "
                               f"(e.g.{shapelets_base_name}_used).")
        shapelets_summary = shapelets_results.shapelets_summary
        shapelets_coeffs = np.array(shapelets_summary.shapeletsCoeffs)
        shapelets_used_cat = catalog[(catalog[shapelets_base_name + "_used"])
                                     & (~catalog[shapelets_base_name + "_flag"])].copy(deep=True)
        shapelets_used_finite_centroid_cat = shapelets_used_cat[(
            (np.isfinite(shapelets_used_cat["slot_Centroid_x"]))
            & (np.isfinite(shapelets_used_cat["slot_Centroid_y"])))]
        min_finite_centroid_sources = 3
        if len(shapelets_used_finite_centroid_cat) < min_finite_centroid_sources:
            if log is not None:
                log.warning("Not enough sources used in the shapelet decomposition have finite slot "
                            "centroid measurements (%d < %d). Returning without setting shapelets_iq_score.",
                            len(shapelets_used_finite_centroid_cat), min_finite_centroid_sources)
            shapelets_results.shapelets_summary = shapelets_summary
            return shapelets_results
 
        centroid_diff_shapelets_vs_slot = np.sqrt(
            (shapelets_used_finite_centroid_cat["slot_Centroid_x"]
             - shapelets_used_finite_centroid_cat[shapelets_base_name + "_x"])**2.0
            + (shapelets_used_finite_centroid_cat["slot_Centroid_y"]
               - shapelets_used_finite_centroid_cat[shapelets_base_name + "_y"])**2.0
        )
        centroid_diff_shapelets_vs_slot_median = np.nanmedian(centroid_diff_shapelets_vs_slot)
        shapelets_summary.centroidDiffShapeletsVsSlotMedian = float(centroid_diff_shapelets_vs_slot_median)
 
        # Use a fixed number of pixels (one roughly estimated to be the
        # largest shift possible considering the maximum footprint size
        # imposed in ComputeRoughPsfShapeletsTask has been determined to
        # work well), to scale the centroid offset between the slot value
        # and that determined by ComputeRoughPsfShapeletsTask.
        centroid_diff_scaled = centroid_diff_shapelets_vs_slot/centroid_offset_scale
        centroid_diff_scaled_median = np.median(centroid_diff_scaled)
 
        # Compute a shapelets score that includes power from the centroid shift
        # in addition to the non-atmospheric decomposition coefficients.
        total_power_plus_scaled = (np.sum(shapelets_coeffs**2.0)
                                   + centroid_diff_scaled_median**2.0)
        non_atm_power_plus_scaled = (
            np.sum(shapelets_coeffs[shapelets_results.non_gaussian_non_atmosphere_index_list]**2.0)
            + centroid_diff_scaled_median**2.0)
 
        shapelets_iq_score = np.where(
            total_power_plus_scaled > 0, non_atm_power_plus_scaled/total_power_plus_scaled, 0.0
        )
        shapelets_summary.shapeletsIqScore = float(shapelets_iq_score)
 
        if log is not None:
            log.info("shapelets_iq_score = %.5f; centroid_diff_shapelets_vs_slot_median = %.3f (pixels), "
                     "centroid_diff_scaled_median (by %.1f pixels) = %.3f",
                     shapelets_iq_score, centroid_diff_shapelets_vs_slot_median,
                     centroid_offset_scale, centroid_diff_scaled_median)
 
        shapelets_results.shapelets_summary = shapelets_summary
        return shapelets_results
 
    def plot_selection(
        self, figure: matplotlib.figure.Figure, *, catalog: SourceCatalog, results: Struct
    ) -> None:
        """Create plots of the shape distribution space used to select stars.
 
        Parameters
        ----------
        figure
            Matplotlib figure to plot to.
        catalog
            Catalog of sources with columns populated by the `run` method (at
            least through the `select_stars` step).
        results
            Result struct returned by `run` or `select_stars`.
        """
        from matplotlib.lines import Line2D
 
        shape_data = np.zeros((len(catalog), 3), dtype=np.float64)
        radii = np.zeros(len(catalog), dtype=np.float64)
        for n, record in enumerate(catalog):
            if record[self._flag_key]:
                continue
            shape = record[self._shape_key]
            if self.config.logarithmic_shapes:
                shape = ellipses.SeparableConformalShearLogTraceRadius(shape)
            shape_data[n, :] = shape.getParameterVector()
            radii[n] = shape.getTraceRadius()
        used_mask = catalog[self._used_key]
        candidate_mask = np.logical_and(catalog[self._candidate_key], np.logical_not(used_mask))
        measured_mask = np.logical_and(
            np.logical_not(catalog[self._flag_key]), np.logical_not(catalog[self._candidate_key])
        )
        # Set up the axes.
        axes = figure.subplot_mosaic(
            [
                ["radius", "radius", "radius"],
                ["hist0", ".", "."],
                ["scatter01", "hist1", "."],
                ["scatter02", "scatter12", "hist2"],
            ],
            gridspec_kw=dict(bottom=0.2, top=1.0),
        )
        axes["scatter01"].sharex(axes["hist0"])
        axes["scatter02"].sharex(axes["hist0"])
        axes["scatter12"].sharex(axes["hist1"])
        axes["scatter12"].sharey(axes["scatter02"])
        for tk in itertools.chain(
            axes["hist0"].get_xticklabels(),
            axes["scatter01"].get_xticklabels(),
            axes["hist1"].get_xticklabels(),
            axes["scatter12"].get_yticklabels(),
        ):
            tk.set_visible(False)
        # Move hist y axes to the outside.
        for ax in [axes["hist0"], axes["hist1"], axes["hist2"]]:
            ax.yaxis.set_label_position("right")
            ax.yaxis.tick_right()
        # Add labels to outer axes.
        names = ["Ixx", "Iyy", "Ixy"] if not self.config.logarithmic_shapes else ["𝜂1", "𝜂2", "ln(r)"]
        axes["scatter02"].set_xlabel(names[0])
        axes["scatter12"].set_xlabel(names[1])
        axes["hist2"].set_xlabel(names[2])
        axes["scatter01"].set_ylabel(names[1])
        axes["scatter02"].set_ylabel(names[2])
        # Make the plots.
        mu = results.mean_shape.getParameterVector()
        sigma = np.diagonal(results.shape_covariance) ** 0.5
        lower_bounds = [max(mu[i] - 3 * sigma[i], min(shape_data[:, i])) for i in range(3)]
        upper_bounds = [min(mu[i] + 3 * sigma[i], max(shape_data[:, i])) for i in range(3)]
        grids = [np.linspace(lower_bounds[i], upper_bounds[i], 50) for i in range(3)]
        axes["radius"].axvline(results.radius_cut, color="k")
        for color, mask, alpha in [
            ("grey", measured_mask, 0.5),
            ("blue", candidate_mask, 0.75),
            ("green", used_mask, 1.0),
        ]:
            axes["radius"].hist(
                radii[mask],
                color=color,
                alpha=alpha,
                bins=16,
                histtype="step",
                range=(radii.min(), 15.0),
                density=True,
            )
            for i in range(3):
                axes[f"hist{i}"].hist(
                    shape_data[mask, i],
                    bins=16,
                    range=(lower_bounds[i], upper_bounds[i]),
                    density=True,
                    color=color,
                    histtype="step",
                    alpha=alpha,
                )
                for j in range(3):
                    if (ax := axes.get(f"scatter{i}{j}")) is not None:
                        ax.scatter(
                            shape_data[mask, i],
                            shape_data[mask, j],
                            c=color,
                            s=4,
                            alpha=alpha,
                            edgecolors=None,
                        )
        for i in range(3):
            axes[f"hist{i}"].plot(
                grids[i], scipy.stats.norm.pdf(grids[i], loc=mu[i], scale=sigma[i]), "k", alpha=0.5
            )
            axes[f"hist{i}"].set_xlim(lower_bounds[i], upper_bounds[i])
            for j in range(3):
                if (ax := axes.get(f"scatter{i}{j}")) is not None:
                    sigma_ellipse = ellipses.Quadrupole(
                        results.shape_covariance[i, i],
                        results.shape_covariance[j, j],
                        results.shape_covariance[i, j],
                    )
                    for factor in [1, 2, 3]:
                        self._draw_ellipse(
                            ax,
                            sigma_ellipse,
                            x=mu[i],
                            y=mu[j],
                            scale=factor,
                            fill=False,
                            edgecolor="k",
                            alpha=0.5,
                        )
                        ax.set_xlim(lower_bounds[i], upper_bounds[i])
                        ax.set_ylim(lower_bounds[j], upper_bounds[j])
        figure.legend(
            [
                Line2D([], [], color="green", alpha=1.0),
                Line2D([], [], color="blue", alpha=0.75),
                Line2D([], [], color="gray", alpha=0.5),
            ],
            [
                f"{self.shapelets_base_name}_used ({used_mask.sum()})",
                f"{self.shapelets_base_name}_candidate & ~{self.shapelets_base_name}_used "
                "({candidate_mask.sum()})",
                f"~{self.shapelets_base_name}_flag & ~{self.shapelets_base_name}_candidate "
                f"({measured_mask.sum()})",
            ],
            loc="lower center",
        )
        return figure
 
    def plot_shapelets(
        self,
        figure: matplotlib.figure.Figure,
        *,
        image: ImageF,
        catalog: SourceCatalog,
        results: Struct,
        n_stars: int = 3,
        stamp_size: float = 2.0,
    ) -> None:
        """Create data/model/residual plots of stars and the shapelet model.
 
        Parameters
        ----------
        figure
            Matplotlib figure to plot to.
        image
            The image the stars were measured on.
        catalog
            Catalog of sources with columns populated by the `run` method .
        results
            Result struct returned by `run`.
        n_stars
            Number of stars to include.
        stamp_size
            Stamp size in inches.
        """
        from matplotlib.colors import Normalize
 
        width = stamp_size * 3 + 1.5
        figure.set_size_inches(w=width, h=stamp_size * n_stars)
        axes = figure.subplot_mosaic(
            [
                ["image_cbar", f"d{star_id}", f"m{star_id}", f"r{star_id}", "res_cbar"]
                for star_id in results.star_ids[:n_stars]
            ],
            gridspec_kw=dict(
                wspace=0.01, hspace=0.01, left=0.5 / width, right=1.0 - 0.5 / width, bottom=0.01, top=0.99
            ),
            width_ratios=[0.25, stamp_size, stamp_size, stamp_size, 0.25],
        )
        for name, ax in axes.items():
            if not name.endswith("cbar"):
                ax.axis("off")
        norm: Normalize | None = None
        res_norm: Normalize | None = None
        for star_id in results.star_ids[:n_stars]:
            record = catalog.find(star_id)
            star_bbox = record.getFootprint().getBBox()
            star_model = ImageD(star_bbox)
            star_center = record[self._center_key]
            star_ellipse = ellipses.Ellipse(ellipses.Axes(record[self._shape_key]), star_center)
            star_shapelet = ShapeletFunction(results.shapelet)
            star_shapelet.setEllipse(star_ellipse)
            star_shapelet.evaluate().addToImage(star_model)
            if norm is None:
                norm = Normalize(vmin=star_model.array.min(), vmax=star_model.array.max())
            star_image = image[star_bbox].clone()
            star_image /= record[self._flux_key]
            self._draw_image(axes[f"d{star_id}"], star_image, norm=norm, cmap="YlGnBu")
            self._draw_ellipse(axes[f"d{star_id}"], star_ellipse, fill=False, edgecolor="blue", alpha=0.5)
            image_plot = self._draw_image(axes[f"m{star_id}"], star_model, norm=norm, cmap="YlGnBu")
            self._draw_ellipse(axes[f"m{star_id}"], star_ellipse, fill=False, edgecolor="blue", alpha=0.5)
            star_image -= star_model.convertF()
            amax = np.abs(star_image.array).max()
            if res_norm is None:
                res_norm = Normalize(vmin=-amax, vmax=amax)
            res_plot = self._draw_image(axes[f"r{star_id}"], star_image, norm=res_norm, cmap="RdBu")
            self._draw_ellipse(axes[f"r{star_id}"], star_ellipse, fill=False, edgecolor="blue", alpha=0.5)
        figure.colorbar(image_plot, cax=axes["image_cbar"], location="left")
        figure.colorbar(res_plot, cax=axes["res_cbar"], location="right")
        return figure
 
    def _threshold_with_bounds(
        self,
        values: np.ndarray,
        threshold: float,
        min_count: int,
        max_count: int,
        name: str,
        kind: Literal["<", ">"],
    ) -> np.ndarray:
        """Return the indices of an array that satisfy an inequality
        and/or lower and upper bounds on the number of indices returned.
 
        Parameters
        ----------
        values
            Array of values to threshold on.
        threshold
            Threshold value that selected elements must be above or below.
        min_count
            The minimum number of indices returned.  When thresholding would
            yield fewer than this number, the threshold is ignored.  Note that
            the number of indices may still be less than this if the size of
            ``values`` is less than this.
        max_count
            The maximum number of indices returned.
        name
            Name of the quantity being thresholded, for log messages.
        kind
            Whether the threshold is a upper bound (``<``) or lower bound
            (``>``).  This also sets how values are ranked when they are added
            or dropped to satisfy the count constraints.
 
        Returns
        -------
        indices
            Indices into ``values``.
        """
        if min_count > len(values):
            raise NoStarsForShapeletsError(
                f"Not enough sources ({len(values)}) for {name} cut that must yield at least {min_count}."
            )
        sorter = values.argsort()
        n = np.searchsorted(values[sorter], threshold)
        if kind == ">":
            sorter = sorter[::-1]
            n = len(sorter) - n
        if n < min_count:
            self.log.verbose(
                "Applying a %s %s %f cut yields only %d sources; keeping the top %d (%s %s %f) instead.",
                name,
                kind,
                threshold,
                n,
                min_count,
                name,
                kind,
                values[sorter[min_count - 1]],
            )
            n = min_count
        elif n > max_count:
            self.log.verbose(
                "%d sources have %s %s %f; keeping only the top %d (%s %s %f) instead.",
                n,
                name,
                kind,
                threshold,
                max_count,
                name,
                kind,
                values[sorter[max_count - 1]],
            )
            n = max_count
        else:
            self.log.verbose("Keeping %d sources with %s %s %f.", n, name, kind, threshold)
        return sorter[:n]
 
    def _find_first_radius_mode(self, radii: np.ndarray) -> tuple[float, KernelDensity]:
        """Find the first peak in a 1-d distribution of radii."""
        kde = KernelDensity(bandwidth=self.config.radius_kde_bandwidth).fit(radii.reshape(-1, 1))
        sorted_radii = radii.copy()
        sorted_radii.sort()
        sorted_radii = sorted_radii[sorted_radii.searchsorted(self.config.radius_mode_min):]
        scores = kde.score_samples(sorted_radii.reshape(-1, 1))
        peaks, _ = scipy.signal.find_peaks(scores)
        if not peaks.size:
            raise NoStarsForShapeletsError("Radius distribute has no mode.")
        return sorted_radii[peaks.min()], kde
 
    @staticmethod
    def _draw_ellipse(
        axes: matplotlib.axes.Axes,
        ellipse: ellipses.BaseCore | ellipses.Ellipse,
        *,
        x: float | None = None,
        y: float | None = None,
        scale: float = 1.0,
        **kwargs: Any,
    ) -> matplotlib.patches.Ellipse:
        from matplotlib.patches import Ellipse as EllipsePatch
 
        if isinstance(ellipse, ellipses.Ellipse):
            if x is None:
                x = ellipse.getCenter().getX()
            if y is None:
                y = ellipse.getCenter().getY()
            ellipse = ellipse.getCore()
        else:
            if x is None:
                x = 0.0
            if y is None:
                y = 0.0
        ellipse = ellipses.Axes(ellipse)
        patch = EllipsePatch(
            (x, y),
            ellipse.getA() * 2 * scale,  # factor of 2 for radius->diameter
            ellipse.getB() * 2 * scale,
            angle=ellipse.getTheta() * 180.0 / np.pi,
            **kwargs,
        )
        axes.add_patch(patch)
        return patch
 
    @staticmethod
    def _draw_image(axes: matplotlib.axes.Axes, image: ImageF, **kwargs: Any) -> matplotlib.image.AxesImage:
        fp_bbox = Box2D(image.getBBox())
        return axes.imshow(
            image.array,
            interpolation="nearest",
            origin="lower",
            aspect="equal",
            extent=(fp_bbox.x.min, fp_bbox.x.max, fp_bbox.y.min, fp_bbox.y.max),
            **kwargs,
        )