Merge pull request #621 from lsst/tickets/DM-30229

DM-30229: Add convenience methods to get PSF shape quantities

python/lsst/afw/cameraGeom/_camera.py

@@ -26,7 +26,7 @@
 from ._cameraGeom import Camera, FOCAL_PLANE
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Camera:  # noqa: F811
 
     def getPupilFactory(self, visitInfo, pupilSize, npix, **kwargs):

python/lsst/afw/cameraGeom/_detector.py

@@ -35,7 +35,7 @@
                            DetectorTypeValNameDict.items()}
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class DetectorBase:  # noqa: F811
     def __iter__(self):
         return (self[i] for i in range(len(self)))

python/lsst/afw/detection/_footprintMerge.py

@@ -25,7 +25,7 @@
 from ._detection import FootprintMergeList
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class FootprintMergeList:  # noqa: F811
     def getMergedSourceCatalog(self, catalogs, filters,
                                peakDist, schema, idFactory, samePeakDist):

python/lsst/afw/fits/_fitsContinued.py

@@ -26,7 +26,7 @@
                     compressionAlgorithmToString, scalingAlgorithmToString)
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Fits:  # noqa: F811
     def __enter__(self):
         return self
@@ -35,20 +35,20 @@ def __exit__(self, cls, exc, traceback):
         self.closeFile()
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class ImageWriteOptions:  # noqa: F811
     def __repr__(self):
         return f"{self.__class__.__name__}(compression={self.compression!r}, scaling={self.scaling!r})"
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class ImageCompressionOptions:  # noqa: F811
     def __repr__(self):
         return (f"{self.__class__.__name__}(algorithm={compressionAlgorithmToString(self.algorithm)!r}, "
                 f"tiles={self.tiles.tolist()!r}, quantizeLevel={self.quantizeLevel:f})")
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class ImageScalingOptions:  # noqa: F811
     def __repr__(self):
         return (f"{self.__class__.__name__}(algorithm={scalingAlgorithmToString(self.algorithm)!r}, "

python/lsst/afw/geom/ellipses/_axes.py

@@ -4,7 +4,7 @@
 from ._ellipses import Axes
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Axes:  # noqa: F811
     def __repr__(self):
         return f"Axes(a={self.getA()!r}, b={self.getB()!r}, theta={self.getTheta()!r})"

python/lsst/afw/geom/ellipses/_ellipse.py

@@ -4,7 +4,7 @@
 from ._ellipses import Ellipse
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Ellipse:  # noqa: F811
     def __repr__(self):
         return f"Ellipse({self.getCore()!r}, {self.getCenter()!r})"

python/lsst/afw/geom/ellipses/_quadrupole.py

@@ -4,7 +4,7 @@
 from ._ellipses import Quadrupole
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Quadrupole:  # noqa: F811
     def __repr__(self):
         return f"Quadrupole(ixx={self.getIxx()!r}, iyy={self.getIyy()!r}, ixy={self.getIxy()!r})"

python/lsst/afw/geom/polygon.py

@@ -25,7 +25,7 @@
 from ._geom import Polygon
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Polygon:  # noqa: F811
     def __repr__(self):
         return f"{self.__class__.__name__}({[p for p in self.getVertices()]})"

python/lsst/afw/image/apCorrMap/_apCorrMapContinued.py

@@ -24,7 +24,7 @@
 __all__ = ['ApCorrMap']
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class ApCorrMap:  # noqa: F811
 
     def keys(self):

python/lsst/afw/math/_background.py

@@ -25,7 +25,7 @@
 __all__ = []  # import this module only for its side effects
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Background:  # noqa: F811
     def __reduce__(self):
         """Pickling"""

python/lsst/afw/math/_chebyshevBoundedField.py

@@ -27,7 +27,7 @@
 __all__ = []  # import this module only for its side effects
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class ChebyshevBoundedField:  # noqa: F811
     @classmethod
     def approximate(cls, boundedField,

python/lsst/afw/table/_aliasMap.py

@@ -25,7 +25,7 @@
 from ._table import AliasMap
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class AliasMap:  # noqa: F811
 
     def keys(self):

python/lsst/afw/table/_base.py

@@ -28,7 +28,7 @@
 __all__ = ["Catalog"]
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class BaseRecord:  # noqa: F811
 
     def extract(self, *patterns, **kwargs):

python/lsst/afw/table/_baseColumnView.py

@@ -31,7 +31,7 @@
 # base class, so we use the same naming convention we use for those.
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class _BaseColumnViewBase:  # noqa: F811
 
     def getBits(self, keys=None):

python/lsst/afw/table/_schema.py

@@ -87,7 +87,7 @@ def _registerInstantiations(abc, types):
 SchemaItem.alias(float, _SchemaItem["D"])
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class Schema:  # noqa: F811
 
     def getOrderedNames(self):

python/lsst/afw/table/_schemaMapper.py

@@ -28,7 +28,7 @@
 from ._table import SchemaMapper
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class SchemaMapper:  # noqa: F811
 
     def addOutputField(self, field, type=None, doc=None, units="", size=None,

python/lsst/afw/table/_source.py

@@ -22,8 +22,9 @@
 __all__ = []
 
 from lsst.utils import continueClass
+from lsst.pex.exceptions import LogicError
 from ._base import Catalog
-from ._table import SourceCatalog, SourceTable
+from ._table import QuadrupoleKey, SourceCatalog, SourceColumnView, SourceRecord, SourceTable
 
 Catalog.register("Source", SourceCatalog)
 
@@ -106,3 +107,66 @@ def _getChildrenWithoutChecking(parent):
             return (_getChildrenWithoutChecking(p) for p in parent)
         except TypeError:
             return _getChildrenWithoutChecking(parent)
+
+
+@continueClass
+class SourceRecord:  # noqa: F811
+
+    def getPsfShape(self):
+        # Catch the KeyError and raise LogicError from `pex.exception` for
+        # consistent behavior with similar C++ methods (getIxx, getIyy, etc.)
+        try:
+            return QuadrupoleKey(self.schema["slot_PsfShape"]).get(self)
+        except KeyError:
+            raise LogicError("Key is not valid (if this is a SourceRecord, make sure slot aliases have been "
+                             "set up)") from None
+
+    def _getPsfShapeComponent(self, suffix):
+        # Catch the KeyError and raise LogicError from `pex.exception` for
+        # consistent behavior with similar C++ methods (getIxx, getIyy, etc.)
+        try:
+            return self["slot_PsfShape_" + suffix]
+        except KeyError:
+            raise LogicError("Key is not valid (if this is a SourceRecord, make sure slot aliases have been "
+                             "set up)") from None
+
+    def getPsfIxx(self):
+        return self._getPsfShapeComponent("xx")
+
+    def getPsfIyy(self):
+        return self._getPsfShapeComponent("yy")
+
+    def getPsfIxy(self):
+        return self._getPsfShapeComponent("xy")
+
+    def getPsfShapeFlag(self):
+        return self._getPsfShapeComponent("flag")
+
+
+@continueClass
+class SourceColumnView:  # noqa: F811
+
+    def _getPsfShapeComponent(self, suffix):
+        # Catch the KeyError and raise LogicError from `pex.exception` for
+        # consistent behavior with similar C++ methods (getIxx, getIyy, etc.)
+        try:
+            return self["slot_PsfShape_" + suffix]
+        except KeyError:
+            raise LogicError("Key is not valid (if this is a SourceCatalog, make sure slot aliases have been "
+                             "set up)") from None
+
+    def getPsfIxx(self):
+        return self._getPsfShapeComponent("xx")
+
+    def getPsfIyy(self):
+        return self._getPsfShapeComponent("yy")
+
+    def getPsfIxy(self):
+        return self._getPsfShapeComponent("xy")
+
+
+@continueClass
+class SourceTable:  # noqa: F811
+
+    def definePsfShape(self, name):
+        self.schema.getAliasMap().set("slot_PsfShape", name)

python/lsst/afw/typehandling/_SimpleGenericMap.py

@@ -67,7 +67,7 @@ def fromkeys(cls, iterable, value=None):
 _oldInit = SimpleGenericMapS.__init__
 
 
-@continueClass  # noqa: F811 (FIXME: remove for py 3.8+)
+@continueClass
 class SimpleGenericMapS:  # noqa: F811
     def __init__(self, source=None, **kwargs):
         _oldInit(self)

tests/test_sourceTable.py

@@ -79,6 +79,9 @@ def fillRecord(self, record):
         record.set(self.xxErrKey, np.random.randn())
         record.set(self.yyErrKey, np.random.randn())
         record.set(self.xyErrKey, np.random.randn())
+        record.set(self.psfShapeKey.getIxx(), np.random.randn())
+        record.set(self.psfShapeKey.getIyy(), np.random.randn())
+        record.set(self.psfShapeKey.getIxy(), np.random.randn())
         record.set(self.fluxFlagKey, np.random.randn() > 0)
         record.set(self.centroidFlagKey, np.random.randn() > 0)
         record.set(self.shapeFlagKey, np.random.randn() > 0)
@@ -105,6 +108,10 @@ def setUp(self):
         self.yyErrKey = self.schema.addField("c_yyErr", type="F")
         self.shapeFlagKey = self.schema.addField("c_flag", type="Flag")
 
+        self.psfShapeKey = lsst.afw.table.QuadrupoleKey.addFields(
+            self.schema, "d", "", lsst.afw.table.CoordinateType.PIXEL)
+        self.psfShapeFlagKey = self.schema.addField("d_flag", type="Flag")
+
         self.table = lsst.afw.table.SourceTable.make(self.schema)
         self.catalog = lsst.afw.table.SourceCatalog(self.table)
         self.record = self.catalog.addNew()
@@ -145,11 +152,15 @@ def checkCanonical(self):
         self.assertFloatsAlmostEqual(
             math.fabs(self.record.get(self.xyErrKey)),
             math.sqrt(self.record.getShapeErr()[2, 2]), rtol=1e-6)
+        self.assertEqual(self.table.getSchema().getAliasMap().get("slot_PsfShape"), "d")
+        self.assertEqual(self.psfShapeKey.get(self.record),
+                         self.record.getPsfShape())
 
     def testPersisted(self):
         self.table.definePsfFlux("a")
         self.table.defineCentroid("b")
         self.table.defineShape("c")
+        self.table.definePsfShape("d")
         with lsst.utils.tests.getTempFilePath(".fits") as filename:
             self.catalog.writeFits(filename)
             catalog = lsst.afw.table.SourceCatalog.readFits(filename)
@@ -180,11 +191,15 @@ def testPersisted(self):
             self.assertFloatsAlmostEqual(
                 math.fabs(self.record.get(self.xyErrKey)),
                 math.sqrt(self.record.getShapeErr()[2, 2]), rtol=1e-6)
+            psfShape = self.psfShapeKey.get(self.record)
+            self.assertEqual(table.getSchema().getAliasMap().get("slot_PsfShape"), "d")
+            self.assertEqual(psfShape, record.getPsfShape())
 
     def testCanonical2(self):
         self.table.definePsfFlux("a")
         self.table.defineCentroid("b")
         self.table.defineShape("c")
+        self.table.definePsfShape("d")
         self.checkCanonical()
 
     def testPickle(self):
@@ -242,18 +257,23 @@ def testColumnView(self):
         self.table.definePsfFlux("a")
         self.table.defineCentroid("b")
         self.table.defineShape("c")
+        self.table.definePsfShape("d")
         self.assertFloatsEqual(cols2["a_instFlux"], cols2.getPsfInstFlux())
         self.assertFloatsEqual(cols2["a_instFluxErr"], cols2.getPsfInstFluxErr())
         self.assertFloatsEqual(cols2["b_x"], cols2.getX())
         self.assertFloatsEqual(cols2["b_y"], cols2.getY())
         self.assertFloatsEqual(cols2["c_xx"], cols2.getIxx())
         self.assertFloatsEqual(cols2["c_yy"], cols2.getIyy())
         self.assertFloatsEqual(cols2["c_xy"], cols2.getIxy())
+        self.assertFloatsEqual(cols2["d_xx"], cols2.getPsfIxx())
+        self.assertFloatsEqual(cols2["d_yy"], cols2.getPsfIyy())
+        self.assertFloatsEqual(cols2["d_xy"], cols2.getPsfIxy())
 
         # Trying to access slots which have been removed should raise.
         self.catalog.table.schema.getAliasMap().erase("slot_Centroid")
         self.catalog.table.schema.getAliasMap().erase("slot_Shape")
-        for quantity in ["X", "Y", "Ixx", "Iyy", "Ixy"]:
+        self.catalog.table.schema.getAliasMap().erase("slot_PsfShape")
+        for quantity in ["X", "Y", "Ixx", "Iyy", "Ixy", "PsfIxx", "PsfIyy", "PsfIxy"]:
             with self.assertRaises(lsst.pex.exceptions.LogicError):
                 getattr(self.catalog, f"get{quantity}")()