DM-21166: Create DiaCalculation plugins that replicate current ap_association behavior

python/lsst/ap/association/diaCalculationPlugins.py

@@ -22,14 +22,16 @@
 import numpy as np
 
 import lsst.geom as geom
+import lsst.pex.config as pexConfig
 
 from .diaCalculation import (
     DiaObjectCalculationPluginConfig,
     DiaObjectCalculationPlugin)
 from lsst.meas.base.pluginRegistry import register
 
 __all__ = ("MeanDiaPositionConfig", "MeanDiaPosition",
-           "WeightedMeanDiaPsFluxConfig", "WeightedMeanDiaPsFlux")
+           "WeightedMeanDiaPsFluxConfig", "WeightedMeanDiaPsFlux",
+           "PercentileDiaPsFlux", "PercentileDiaPsFluxConfig")
 
 
 class MeanDiaPositionConfig(DiaObjectCalculationPluginConfig):
@@ -92,10 +94,12 @@ class WeightedMeanDiaPsFluxConfig(DiaObjectCalculationPluginConfig):
 class WeightedMeanDiaPsFlux(DiaObjectCalculationPlugin):
     """Compute the weighted mean and mean error on the point source fluxes
     of the DiaSource measured on the difference image.
+
+    Additionally store number of usable data points.
     """
 
     ConfigClass = WeightedMeanDiaPsFluxConfig
-    outputCols = ["PSFluxMean", "PSFluxMeanErr"]
+    outputCols = ["PSFluxMean", "PSFluxMeanErr", "PSFFluxNdata"]
 
     @classmethod
     def getExecutionOrder(cls):
@@ -128,13 +132,90 @@ def calculate(self,
                                  filterDiaFluxes["psFluxErr"] ** 2)
             fluxMean /= tot_weight
             fluxMeanErr = np.sqrt(1 / tot_weight)
+            nFluxData = np.sum(np.isfinite(filterDiaFluxes["psFlux"]))
         else:
             fluxMean = np.nan
             fluxMeanErr = np.nan
-
+            nFluxData = 0
         if np.isfinite(fluxMean) and np.isfinite(fluxMeanErr):
-            diaObject["%sPSFluxMean" % filterName] = fluxMean
-            diaObject["%sPSFluxMeanErr" % filterName] = fluxMeanErr
+            diaObject["{}PSFluxMean".format(filterName)] = fluxMean
+            diaObject["{}PSFluxMeanErr".format(filterName)] = fluxMeanErr
+            diaObject["{}PSFluxNdata".format(filterName)] = nFluxData
+        else:
+            self.fail(diaObject, filterName)
+
+    def fail(self, diaObject, filterName, error=None):
+        """Set diaObject position values to nan.
+
+        Parameters
+        ----------
+        diaObject : `dict`
+            Summary object to store values in.
+        filterName : `str`
+            Simple name of the filter for the flux being calculated.
+        error : `BaseException`
+            Error to pass.
+        """
+        diaObject["{}PSFluxMean".format(filterName)] = np.nan
+        diaObject["{}PSFluxMeanErr".format(filterName)] = np.nan
+        diaObject["{}PSFluxNdata".format(filterName)] = 0
+
+
+class PercentileDiaPsFluxConfig(DiaObjectCalculationPluginConfig):
+
+    percentiles = pexConfig.ListField(
+        dtype=int,
+        default=[5, 25, 50, 75, 95],
+        doc="Percentiles to calculate to compute values for. Should be "
+            "integer values."
+    )
+
+
+@register("ap_percentileFlux")
+class PercentileDiaPsFlux(DiaObjectCalculationPlugin):
+    """
+    """
+
+    ConfigClass = PercentileDiaPsFluxConfig
+    outputCols = []
+
+    def __init__(self, config, name, metadata):
+        DiaObjectCalculationPlugin.__init__(self, config, name)
+        self.outputCols = ["PSFluxPercentile{:02d}".format(percent)
+                           for percent in self.config.percentiles]
+
+    @classmethod
+    def getExecutionOrder(cls):
+        return cls.DEFAULT_CATALOGCALCULATION
+
+    def calculate(self,
+                  diaObject,
+                  diaSources,
+                  filterDiaFluxes,
+                  filterName,
+                  **kwargs):
+        """Compute the weighted mean and mean error of the point source flux.
+
+        Parameters
+        ----------
+        diaObject : `dict`
+            Summary object to store values in.
+        diaSources : `pandas.DataFrame`
+            DataFrame representing all diaSources associated with this
+            diaObject.
+        filterDiaFluxes : `pandas.DataFrame`
+            DataFrame representing diaSources associated with this
+            diaObject that are observed in the band pass ``filterName``.
+        filterName : `str`
+            Simple, string name of the filter for the fluxb eing calculated.
+        """
+        if len(filterDiaFluxes) > 0:
+            pTiles = np.nanpercentile(filterDiaFluxes["psFlux"],
+                                      self.config.percentiles)
+            for pTile, tilePercent in zip(pTiles, self.config.precentiles):
+                diaObject[
+                    "{}PSFluxPercentile{:02d}".format(filterName,
+                                                      tilePercent)] = pTile
         else:
             self.fail(diaObject, filterName)
 
@@ -150,5 +231,6 @@ def fail(self, diaObject, filterName, error=None):
         error : `BaseException`
             Error to pass.
         """
-        diaObject["%sPSFluxMean" % filterName] = np.nan
-        diaObject["%sPSFluxMeanErr" % filterName] = np.nan
+        for pTile in self.config.precentiles:
+            diaObject["{}PSFluxPercentile{:02d}".format(filterName,
+                                                        pTile)] = np.nan

tests/test_dia_calculation_plugins.py

@@ -24,18 +24,16 @@
 import unittest
 
 from lsst.ap.association import (
-    MeanDiaPosition,
-    MeanDiaPositionConfig,
-    WeightedMeanDiaPsFlux,
-    WeightedMeanDiaPsFluxConfig)
+    MeanDiaPosition, MeanDiaPositionConfig,
+    WeightedMeanDiaPsFlux, WeightedMeanDiaPsFluxConfig,
+    PercentileDiaPsFlux, PercentileDiaPsFluxConfig)
 import lsst.utils.tests
 
 
 class TestMeanPosition(unittest.TestCase):
 
     def testCalculate(self):
-        """Test that forced source catalogs are successfully created and have
-        sensible values.
+        """Test mean position calculation.
         """
         n_sources = 10
 
@@ -80,8 +78,7 @@ def testCalculate(self):
 class TestWeightedMeanDiaPsFlux(unittest.TestCase):
 
     def testCalculate(self):
-        """Test that forced source catalogs are successfully created and have
-        sensible values.
+        """Test mean value calculation.
         """
         n_sources = 10
         diaObject = dict()
@@ -95,19 +92,61 @@ def testCalculate(self):
 
         self.assertAlmostEqual(diaObject["uPSFluxMean"], 0.0)
         self.assertAlmostEqual(diaObject["uPSFluxMeanErr"], np.sqrt(1 / n_sources))
+        self.assertEqual(diaObject["uPSFluxNdata"], n_sources)
 
         diaObject = dict()
         mean_flux.calculate(diaObject, [], [], "g")
 
         self.assertTrue(np.isnan(diaObject["gPSFluxMean"]))
         self.assertTrue(np.isnan(diaObject["gPSFluxMeanErr"]))
+        self.assertEqual(diaObject["gPSFluxNdata"], 0)
 
         diaObject = dict()
         diaSources.loc[4, "psFlux"] = np.nan
         mean_flux.calculate(diaObject, diaSources, diaSources, "r")
 
         self.assertTrue(~np.isnan(diaObject["rPSFluxMean"]))
         self.assertTrue(~np.isnan(diaObject["rPSFluxMeanErr"]))
+        self.assertEqual(diaObject["rPSFluxNdata"], n_sources - 1)
+
+
+class TestPercentileDiaPsFlux(unittest.TestCase):
+
+    def testCalculate(self):
+        """Test flux percentile calculation.
+        """
+        n_sources = 10
+        diaObject = dict()
+        fluxes = np.linspace(-1, 1, n_sources)
+        diaSources = pd.DataFrame(data={"psFlux": fluxes,
+                                        "psFluxErr": np.ones(n_sources)})
+
+        pTileFlux = PercentileDiaPsFlux(PercentileDiaPsFluxConfig(),
+                                        "ap_percentileFlux",
+                                        None)
+        pTileFlux.calculate(diaObject, diaSources, diaSources, "u")
+        for pTile, testVal in zip(pTileFlux.config.percentiles,
+                                  np.nanpercentile(fluxes)):
+            self.assertAlmostEqual(
+                diaObject["uPSFluxPercentile{:02d}".format(pTile)],
+                testVal)
+
+        diaObject = dict()
+        pTileFlux.calculate(diaObject, [], [], "g")
+        for pTile, testVal in zip(pTileFlux.config.percentiles,
+                                  np.nanpercentile(fluxes)):
+            self.assertTrue(np.isnan(
+                diaObject["uPSFluxPercentile{:02d}".format(pTile)]))
+
+        diaObject = dict()
+        diaSources.loc[4, "psFlux"] = np.nan
+        fluxes[4] = np.nan
+        pTileFlux.calculate(diaObject, diaSources, diaSources, "r")
+        for pTile, testVal in zip(pTileFlux.config.percentiles,
+                                  np.nanpercentile(fluxes)):
+            self.assertAlmostEqual(
+                diaObject["uPSFluxPercentile{:02d}".format(pTile)],
+                testVal)
 
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):