brightStarStamps.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/>.
 
"""Collection of small images (stamps), each centered on a bright star."""
 
__all__ = ["BrightStarStamp", "BrightStarStamps"]
 
import logging
from collections.abc import Collection
from dataclasses import dataclass
from functools import reduce
from operator import ior
 
import numpy as np
from lsst.afw.geom import SpanSet, Stencil
from lsst.afw.image import MaskedImageF
from lsst.afw.math import Property, StatisticsControl, makeStatistics, stringToStatisticsProperty
from lsst.afw.table.io import Persistable
from lsst.geom import Point2I
 
from .stamps import AbstractStamp, Stamps, readFitsWithOptions
 
logger = logging.getLogger(__name__)
 
 
@dataclass
class BrightStarStamp(AbstractStamp):
    """Single stamp centered on a bright star, normalized by its annularFlux.
 
    Parameters
    ----------
    stamp_im : `~lsst.afw.image.MaskedImage`
        Pixel data for this postage stamp
    gaiaGMag : `float`
        Gaia G magnitude for the object in this stamp
    gaiaId : `int`
        Gaia object identifier
    position : `~lsst.geom.Point2I`
        Origin of the stamps in its origin exposure (pixels)
    archive_element : `~lsst.afw.table.io.Persistable` or None, optional
        Archive element (e.g. Transform or WCS) associated with this stamp.
    annularFlux : `float` or None, optional
        Flux in an annulus around the object
    """
 
    stamp_im: MaskedImageF
    gaiaGMag: float
    gaiaId: int
    position: Point2I
    archive_element: Persistable | None = None
    annularFlux: float | None = None
    minValidAnnulusFraction: float = 0.0
 
    @classmethod
    def factory(cls, stamp_im, metadata, idx, archive_element=None, minValidAnnulusFraction=0.0):
        """This method is needed to service the FITS reader. We need a standard
        interface to construct objects like this. Parameters needed to
        construct this object are passed in via a metadata dictionary and then
        passed to the constructor of this class.  This particular factory
        method requires keys: G_MAGS, GAIA_IDS, and ANNULAR_FLUXES. They should
        each point to lists of values.
 
        Parameters
        ----------
        stamp_im : `~lsst.afw.image.MaskedImage`
            Pixel data to pass to the constructor
        metadata : `dict`
            Dictionary containing the information
            needed by the constructor.
        idx : `int`
            Index into the lists in ``metadata``
        archive_element : `~lsst.afw.table.io.Persistable` or None, optional
            Archive element (e.g. Transform or WCS) associated with this stamp.
        minValidAnnulusFraction : `float`, optional
            The fraction of valid pixels within the normalization annulus of a
            star.
 
        Returns
        -------
        brightstarstamp : `BrightStarStamp`
            An instance of this class
        """
        if "X0S" in metadata and "Y0S" in metadata:
            x0 = metadata.getArray("X0S")[idx]
            y0 = metadata.getArray("Y0S")[idx]
            position = Point2I(x0, y0)
        else:
            position = None
        return cls(
            stamp_im=stamp_im,
            gaiaGMag=metadata.getArray("G_MAGS")[idx],
            gaiaId=metadata.getArray("GAIA_IDS")[idx],
            position=position,
            archive_element=archive_element,
            annularFlux=metadata.getArray("ANNULAR_FLUXES")[idx],
            minValidAnnulusFraction=minValidAnnulusFraction,
        )
 
    def measureAndNormalize(
        self,
        annulus: SpanSet,
        statsControl: StatisticsControl = StatisticsControl(),
        statsFlag: Property = stringToStatisticsProperty("MEAN"),
        badMaskPlanes: Collection[str] = ("BAD", "SAT", "NO_DATA"),
    ):
        """Compute "annularFlux", the integrated flux within an annulus
        around an object's center, and normalize it.
 
        Since the center of bright stars are saturated and/or heavily affected
        by ghosts, we measure their flux in an annulus with a large enough
        inner radius to avoid the most severe ghosts and contain enough
        non-saturated pixels.
 
        Parameters
        ----------
        annulus : `~lsst.afw.geom.spanSet.SpanSet`
            SpanSet containing the annulus to use for normalization.
        statsControl : `~lsst.afw.math.statistics.StatisticsControl`, optional
            StatisticsControl to be used when computing flux over all pixels
            within the annulus.
        statsFlag : `~lsst.afw.math.statistics.Property`, optional
            statsFlag to be passed on to ``afwMath.makeStatistics`` to compute
            annularFlux. Defaults to a simple MEAN.
        badMaskPlanes : `collections.abc.Collection` [`str`]
            Collection of mask planes to ignore when computing annularFlux.
        """
        stampSize = self.stamp_imstamp_im.getDimensions()
        # create image: same pixel values within annulus, NO_DATA elsewhere
        maskPlaneDict = self.stamp_imstamp_im.mask.getMaskPlaneDict()
        annulusImage = MaskedImageF(stampSize, planeDict=maskPlaneDict)
        annulusMask = annulusImage.mask
        annulusMask.array[:] = 2 ** maskPlaneDict["NO_DATA"]
        annulus.copyMaskedImage(self.stamp_imstamp_im, annulusImage)
        # set mask planes to be ignored
        andMask = reduce(ior, (annulusMask.getPlaneBitMask(bm) for bm in badMaskPlanes))
        statsControl.setAndMask(andMask)
 
        annulusStat = makeStatistics(annulusImage, statsFlag, statsControl)
        # determining the number of valid (unmasked) pixels within the annulus
        unMasked = annulusMask.array.size - np.count_nonzero(annulusMask.array)
        # should we have some messages printed here?
        logger.info(
            "The Star's annulus contains %s valid pixels and the annulus itself contains %s pixels.",
            unMasked,
            annulus.getArea(),
        )
        if unMasked > (annulus.getArea() * self.minValidAnnulusFractionminValidAnnulusFraction):
            # compute annularFlux
            self.annularFluxannularFlux = annulusStat.getValue()
            logger.info("Annular flux is: %s", self.annularFluxannularFlux)
        else:
            raise RuntimeError(
                f"Less than {self.minValidAnnulusFraction * 100}% of pixels within the annulus are valid."
            )
        if np.isnan(self.annularFluxannularFlux):
            raise RuntimeError("Annular flux computation failed, likely because there are no valid pixels.")
        if self.annularFluxannularFlux < 0:
            raise RuntimeError("The annular flux is negative. The stamp can not be normalized!")
        # normalize stamps
        self.stamp_imstamp_im.image.array /= self.annularFluxannularFlux
        return None
 
 
class BrightStarStamps(Stamps):
    """Collection of bright star stamps and associated metadata.
 
    Parameters
    ----------
    starStamps : `collections.abc.Sequence` [`BrightStarStamp`]
        Sequence of star stamps. Cannot contain both normalized and
        unnormalized stamps.
    innerRadius : `int`, optional
        Inner radius value, in pixels. This and ``outerRadius`` define the
        annulus used to compute the ``"annularFlux"`` values within each
        ``starStamp``. Must be provided if ``normalize`` is True.
    outerRadius : `int`, optional
        Outer radius value, in pixels. This and ``innerRadius`` define the
        annulus used to compute the ``"annularFlux"`` values within each
        ``starStamp``. Must be provided if ``normalize`` is True.
    nb90Rots : `int`, optional
        Number of 90 degree rotations required to compensate for detector
        orientation.
    metadata : `~lsst.daf.base.PropertyList`, optional
        Metadata associated with the bright stars.
    use_mask : `bool`
        If `True` read and write mask data. Default `True`.
    use_variance : `bool`
        If ``True`` read and write variance data. Default ``False``.
    use_archive : `bool`
        If ``True`` read and write an Archive that contains a Persistable
        associated with each stamp. In the case of bright stars, this is
        usually a ``TransformPoint2ToPoint2``, used to warp each stamp
        to the same pixel grid before stacking.
 
    Raises
    ------
    ValueError
        Raised if one of the star stamps provided does not contain the
        required keys.
    AttributeError
        Raised if there is a mix-and-match of normalized and unnormalized
        stamps, stamps normalized with different annulus definitions, or if
        stamps are to be normalized but annular radii were not provided.
 
    Notes
    -----
    A butler can be used to read only a part of the stamps, specified by a
    bbox:
 
    >>> starSubregions = butler.get(
            "brightStarStamps",
            dataId,
            parameters={"bbox": bbox}
        )
    """
 
    def __init__(
        self,
        starStamps,
        innerRadius=None,
        outerRadius=None,
        nb90Rots=None,
        metadata=None,
        use_mask=True,
        use_variance=False,
        use_archive=False,
    ):
        super().__init__(starStamps, metadata, use_mask, use_variance, use_archive)
        # Ensure stamps contain a flux measurement if and only if they are
        # already expected to be normalized
        self._checkNormalization(False, innerRadius, outerRadius)
        self._innerRadius, self._outerRadius = innerRadius, outerRadius
        if innerRadius is not None and outerRadius is not None:
            self.normalized = True
        else:
            self.normalized = False
        self.nb90Rots = nb90Rots
 
    @classmethod
    def initAndNormalize(
        cls,
        starStamps,
        innerRadius,
        outerRadius,
        nb90Rots=None,
        metadata=None,
        use_mask=True,
        use_variance=False,
        use_archive=False,
        imCenter=None,
        discardNanFluxObjects=True,
        statsControl=StatisticsControl(),
        statsFlag=stringToStatisticsProperty("MEAN"),
        badMaskPlanes=("BAD", "SAT", "NO_DATA"),
    ):
        """Normalize a set of bright star stamps and initialize a
        BrightStarStamps instance.
 
        Since the center of bright stars are saturated and/or heavily affected
        by ghosts, we measure their flux in an annulus with a large enough
        inner radius to avoid the most severe ghosts and contain enough
        non-saturated pixels.
 
        Parameters
        ----------
        starStamps : `collections.abc.Sequence` [`BrightStarStamp`]
            Sequence of star stamps. Cannot contain both normalized and
            unnormalized stamps.
        innerRadius : `int`
            Inner radius value, in pixels. This and ``outerRadius`` define the
            annulus used to compute the ``"annularFlux"`` values within each
            ``starStamp``.
        outerRadius : `int`
            Outer radius value, in pixels. This and ``innerRadius`` define the
            annulus used to compute the ``"annularFlux"`` values within each
            ``starStamp``.
        nb90Rots : `int`, optional
            Number of 90 degree rotations required to compensate for detector
            orientation.
        metadata : `~lsst.daf.base.PropertyList`, optional
            Metadata associated with the bright stars.
        use_mask : `bool`
            If `True` read and write mask data. Default `True`.
        use_variance : `bool`
            If ``True`` read and write variance data. Default ``False``.
        use_archive : `bool`
            If ``True`` read and write an Archive that contains a Persistable
            associated with each stamp. In the case of bright stars, this is
            usually a ``TransformPoint2ToPoint2``, used to warp each stamp
            to the same pixel grid before stacking.
        imCenter : `collections.abc.Sequence`, optional
            Center of the object, in pixels. If not provided, the center of the
            first stamp's pixel grid will be used.
        discardNanFluxObjects : `bool`
            Whether objects with NaN annular flux should be discarded.
            If False, these objects will not be normalized.
        statsControl : `~lsst.afw.math.statistics.StatisticsControl`, optional
            StatisticsControl to be used when computing flux over all pixels
            within the annulus.
        statsFlag : `~lsst.afw.math.statistics.Property`, optional
            statsFlag to be passed on to ``~lsst.afw.math.makeStatistics`` to
            compute annularFlux. Defaults to a simple MEAN.
        badMaskPlanes : `collections.abc.Collection` [`str`]
            Collection of mask planes to ignore when computing annularFlux.
 
        Raises
        ------
        ValueError
            Raised if one of the star stamps provided does not contain the
            required keys.
        AttributeError
            Raised if there is a mix-and-match of normalized and unnormalized
            stamps, stamps normalized with different annulus definitions, or if
            stamps are to be normalized but annular radii were not provided.
        """
        if imCenter is None:
            stampSize = starStamps[0].stamp_im.getDimensions()
            imCenter = stampSize[0] // 2, stampSize[1] // 2
        # Create SpanSet of annulus
        outerCircle = SpanSet.fromShape(outerRadius, Stencil.CIRCLE, offset=imCenter)
        innerCircle = SpanSet.fromShape(innerRadius, Stencil.CIRCLE, offset=imCenter)
        annulus = outerCircle.intersectNot(innerCircle)
        # Initialize (unnormalized) brightStarStamps instance
        bss = cls(
            starStamps,
            innerRadius=None,
            outerRadius=None,
            nb90Rots=nb90Rots,
            metadata=metadata,
            use_mask=use_mask,
            use_variance=use_variance,
            use_archive=use_archive,
        )
        # Ensure no stamps had already been normalized.
        bss._checkNormalization(True, innerRadius, outerRadius)
        bss._innerRadius, bss._outerRadius = innerRadius, outerRadius
        # Create a list to contain rejected stamps.
        rejects = []
        # Apply normalization
        for stamp in bss._stamps:
            try:
                stamp.measureAndNormalize(
                    annulus, statsControl=statsControl, statsFlag=statsFlag, badMaskPlanes=badMaskPlanes
                )
            except RuntimeError as err:
                logger.error(err)
                # Optionally keep NaN flux objects, for bookkeeping purposes,
                # and to avoid having to re-find and redo the preprocessing
                # steps needed before bright stars can be subtracted.
                if discardNanFluxObjects:
                    rejects.append(stamp)
                else:
                    stamp.annularFlux = np.nan
        # Remove rejected stamps.
        if discardNanFluxObjects:
            for reject in rejects:
                bss._stamps.remove(reject)
        bss.normalized = True
        return bss
 
    def _refresh_metadata(self):
        """Refresh metadata. Should be called before writing the object out."""
        # add full list of positions, Gaia magnitudes, IDs and annularFlxes to
        # shared metadata
        self._metadata["G_MAGS"] = self.getMagnitudes()
        self._metadata["GAIA_IDS"] = self.getGaiaIds()
        positions = self.getPositions()
        self._metadata["X0S"] = [xy0[0] for xy0 in positions]
        self._metadata["Y0S"] = [xy0[1] for xy0 in positions]
        self._metadata["ANNULAR_FLUXES"] = self.getAnnularFluxes()
        self._metadata["NORMALIZED"] = self.normalized
        self._metadata["INNER_RADIUS"] = self._innerRadius
        self._metadata["OUTER_RADIUS"] = self._outerRadius
        if self.nb90Rots is not None:
            self._metadata["NB_90_ROTS"] = self.nb90Rots
        return None
 
    @classmethod
    def readFits(cls, filename):
        """Build an instance of this class from a file.
 
        Parameters
        ----------
        filename : `str`
            Name of the file to read
        """
        return cls.readFitsWithOptionsreadFitsWithOptionsreadFitsWithOptions(filename, None)
 
    @classmethod
    def readFitsWithOptions(cls, filename, options):
        """Build an instance of this class with options.
 
        Parameters
        ----------
        filename : `str`
            Name of the file to read
        options : `PropertyList`
            Collection of metadata parameters
        """
        stamps, metadata = readFitsWithOptions(filename, BrightStarStamp.factory, options)
        nb90Rots = metadata["NB_90_ROTS"] if "NB_90_ROTS" in metadata else None
        if metadata["NORMALIZED"]:
            return cls(
                stamps,
                innerRadius=metadata["INNER_RADIUS"],
                outerRadius=metadata["OUTER_RADIUS"],
                nb90Rots=nb90Rots,
                metadata=metadata,
                use_mask=metadata["HAS_MASK"],
                use_variance=metadata["HAS_VARIANCE"],
                use_archive=metadata["HAS_ARCHIVE"],
            )
        else:
            return cls(
                stamps,
                nb90Rots=nb90Rots,
                metadata=metadata,
                use_mask=metadata["HAS_MASK"],
                use_variance=metadata["HAS_VARIANCE"],
                use_archive=metadata["HAS_ARCHIVE"],
            )
 
    def append(self, item, innerRadius=None, outerRadius=None):
        """Add an additional bright star stamp.
 
        Parameters
        ----------
        item : `BrightStarStamp`
            Bright star stamp to append.
        innerRadius : `int`, optional
            Inner radius value, in pixels. This and ``outerRadius`` define the
            annulus used to compute the ``"annularFlux"`` values within each
            ``BrightStarStamp``.
        outerRadius : `int`, optional
            Outer radius value, in pixels. This and ``innerRadius`` define the
            annulus used to compute the ``"annularFlux"`` values within each
            ``BrightStarStamp``.
        """
        if not isinstance(item, BrightStarStamp):
            raise ValueError(f"Can only add instances of BrightStarStamp, got {type(item)}.")
        if (item.annularFlux is None) == self.normalized:
            raise AttributeError(
                "Trying to append an unnormalized stamp to a normalized BrightStarStamps "
                "instance, or vice-versa."
            )
        else:
            self._checkRadius(innerRadius, outerRadius)
        self._stamps.append(item)
        return None
 
    def extend(self, bss):
        """Extend BrightStarStamps instance by appending elements from another
        instance.
 
        Parameters
        ----------
        bss : `BrightStarStamps`
            Other instance to concatenate.
        """
        if not isinstance(bss, BrightStarStamps):
            raise ValueError(f"Can only extend with a BrightStarStamps object. Got {type(bss)}.")
        self._checkRadius(bss._innerRadius, bss._outerRadius)
        self._stamps += bss._stamps
 
    def getMagnitudes(self):
        """Retrieve Gaia G magnitudes for each star.
 
        Returns
        -------
        gaiaGMags : `list` [`float`]
        """
        return [stamp.gaiaGMag for stamp in self._stamps]
 
    def getGaiaIds(self):
        """Retrieve Gaia IDs for each star.
 
        Returns
        -------
        gaiaIds : `list` [`int`]
        """
        return [stamp.gaiaId for stamp in self._stamps]
 
    def getAnnularFluxes(self):
        """Retrieve normalization factors for each star.
 
        These are computed by integrating the flux in annulus centered on the
        bright star, far enough from center to be beyond most severe ghosts and
        saturation. The inner and outer radii that define the annulus can be
        recovered from the metadata.
 
        Returns
        -------
        annularFluxes : `list` [`float`]
        """
        return [stamp.annularFlux for stamp in self._stamps]
 
    def selectByMag(self, magMin=None, magMax=None):
        """Return the subset of bright star stamps for objects with specified
        magnitude cuts (in Gaia G).
 
        Parameters
        ----------
        magMin : `float`, optional
            Keep only stars fainter than this value.
        magMax : `float`, optional
            Keep only stars brighter than this value.
        """
        subset = [
            stamp
            for stamp in self._stamps
            if (magMin is None or stamp.gaiaGMag > magMin) and (magMax is None or stamp.gaiaGMag < magMax)
        ]
        # This is an optimization to save looping over the init argument when
        # it is already guaranteed to be the correct type
        instance = BrightStarStamps(
            (), innerRadius=self._innerRadius, outerRadius=self._outerRadius, metadata=self._metadata
        )
        instance._stamps = subset
        return instance
 
    def _checkRadius(self, innerRadius, outerRadius):
        """Ensure provided annulus radius is consistent with that already
        present in the instance, or with arguments passed on at initialization.
        """
        if innerRadius != self._innerRadius or outerRadius != self._outerRadius:
            raise AttributeError(
                "Trying to mix stamps normalized with annulus radii "
                f"{innerRadius, outerRadius} with those of BrightStarStamp instance\n"
                f"(computed with annular radii {self._innerRadius, self._outerRadius})."
            )
 
    def _checkNormalization(self, normalize, innerRadius, outerRadius):
        """Ensure there is no mixing of normalized and unnormalized stars, and
        that, if requested, normalization can be performed.
        """
        noneFluxCount = self.getAnnularFluxes().count(None)
        nStamps = len(self)
        nFluxVals = nStamps - noneFluxCount
        if noneFluxCount and noneFluxCount < nStamps:
            # at least one stamp contains an annularFlux value (i.e. has been
            # normalized), but not all of them do
            raise AttributeError(
                f"Only {nFluxVals} stamps contain an annularFlux value.\nAll stamps in a "
                "BrightStarStamps instance must either be normalized with the same annulus "
                "definition, or none of them can contain an annularFlux value."
            )
        elif normalize:
            # stamps are to be normalized; ensure annular radii are specified
            # and they have no annularFlux
            if innerRadius is None or outerRadius is None:
                raise AttributeError(
                    "For stamps to be normalized (normalize=True), please provide a valid "
                    "value (in pixels) for both innerRadius and outerRadius."
                )
            elif noneFluxCount < nStamps:
                raise AttributeError(
                    f"{nFluxVals} stamps already contain an annularFlux value. For stamps to "
                    "be normalized, all their annularFlux must be None."
                )
        elif innerRadius is not None and outerRadius is not None:
            # Radii provided, but normalize=False; check that stamps
            # already contain annularFluxes
            if noneFluxCount:
                raise AttributeError(
                    f"{noneFluxCount} stamps contain no annularFlux, but annular radius "
                    "values were provided and normalize=False.\nTo normalize stamps, set "
                    "normalize to True."
                )
        else:
            # At least one radius value is missing; ensure no stamps have
            # already been normalized
            if nFluxVals:
                raise AttributeError(
                    f"{nFluxVals} stamps contain an annularFlux value. If stamps have "
                    "been normalized, the innerRadius and outerRadius values used must "
                    "be provided."
                )
        return None