From 54985ab2393f8322bc3c86dd2bf458344304b63d Mon Sep 17 00:00:00 2001
From: John Parejko <parejkoj@uw.edu>
Date: Wed, 27 Nov 2024 15:02:32 -0800
Subject: [PATCH] Add check on psf size to MeasurePsfTask

This will catch bad psfs returned from both psfex and piff.

Unfortunately, there are no tests for MeasurePsf, so there's not a good
place to put a test of this new exception (e.g. with a mocked-bad psf).
I checked it with one of the calibrateImage tests, removing the `not`
and seeing that the appropriate exception was raised.
---
 python/lsst/pipe/tasks/measurePsf.py | 32 ++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/python/lsst/pipe/tasks/measurePsf.py b/python/lsst/pipe/tasks/measurePsf.py
index 51f41e856..721772e72 100644
--- a/python/lsst/pipe/tasks/measurePsf.py
+++ b/python/lsst/pipe/tasks/measurePsf.py
@@ -21,6 +21,8 @@
 
 __all__ = ["MeasurePsfConfig", "MeasurePsfTask"]
 
+import numpy as np
+
 import lsst.afw.display as afwDisplay
 import lsst.afw.math as afwMath
 import lsst.meas.algorithms as measAlg
@@ -31,6 +33,27 @@
 from lsst.utils.timer import timeMethod
 
 
+class NonfinitePsfShapeError(pipeBase.AlgorithmError):
+    """Raised if the radius of the fitted psf model is non-finite.
+
+    Parameters
+    ----------
+    psf_size : `float`
+        Fitted size of the failing PSF model.
+    """
+    def __init__(self, psf_size) -> None:
+        self._psf_size = psf_size
+        super().__init__(
+            f"Failed to determine PSF: fitter returned model with non-finite PSF size ({psf_size} pixels)."
+        )
+
+    @property
+    def metadata(self) -> dict:
+        return {
+            "psf_size": self._psf_size,
+        }
+
+
 class MeasurePsfConfig(pexConfig.Config):
     starSelector = measAlg.sourceSelectorRegistry.makeField(
         "Star selection algorithm",
@@ -199,6 +222,11 @@ def run(self, exposure, sources, expId=0, matches=None):
             ``cellSet``
                 An `lsst.afw.math.SpatialCellSet` containing the PSF candidates
                 as returned by the psf determiner.
+
+        Raises
+        ------
+        NonfinitePsfShapeError
+            If the new PSF has NaN or Inf width.
         """
         self.log.info("Measuring PSF")
 
@@ -242,6 +270,10 @@ def run(self, exposure, sources, expId=0, matches=None):
                                                        flagKey=self.usedKey)
         self.log.info("PSF determination using %d/%d stars.",
                       self.metadata.getScalar("numGoodStars"), self.metadata.getScalar("numAvailStars"))
+        if not np.isfinite((psfSize := psf.computeShape(psf.getAveragePosition()).getDeterminantRadius())):
+            raise NonfinitePsfShapeError(psf_size=psfSize)
+        else:
+            self.log.info("Fitted PSF size: %f pixels", psfSize)
 
         exposure.setPsf(psf)