Merge pull request #298 from lsst/tickets/DM-48928

DM-48928: Fix serialization issue with LombScarglePeriodogramMulti plugin

Author: jeffcarlin

python/lsst/meas/base/compensatedGaussian/_compensatedTophat.py

@@ -44,7 +44,7 @@ class SingleFrameCompensatedTophatFluxConfig(SingleFramePluginConfig):
         doc="The aperture radii (in pixels) to measure the top-hats.",
         dtype=int,
         minLength=1,
-        default=[12,],
+        default=[12, ],
     )
     inner_scale = RangeField(
         doc="Inner background annulus scale (relative to aperture).",

python/lsst/meas/base/diaCalculationPlugins.py

@@ -316,26 +316,36 @@ def calculate(self,
         **kwargs : `dict`
             Unused kwargs that are always passed to a plugin.
         """
-        n_bands = len(diaSources["band"].unique())
+
+        bands_arr = diaSources['band'].unique().values
+        unique_bands = np.unique(np.concatenate(bands_arr))
         # Check and initialize output columns in diaObjects.
         if (periodCol := "multiPeriod") not in diaObjects.columns:
             diaObjects[periodCol] = np.nan
         if (powerCol := "multiPower") not in diaObjects.columns:
             diaObjects[powerCol] = np.nan
         if (fapCol := "multiFap") not in diaObjects.columns:
             diaObjects[fapCol] = np.nan
-        if (ampCol := "multiAmp") not in diaObjects.columns:
-            diaObjects[ampCol] = pd.Series([np.nan]*n_bands, dtype="object")
-        if (phaseCol := "multiPhase") not in diaObjects.columns:
-            diaObjects[phaseCol] = pd.Series([np.nan]*n_bands, dtype="object")
-
-        def _calculate_period_multi(df, min_detections=9, oversampling_factor=5, nyquist_factor=100):
+        ampCol = "multiAmp"
+        phaseCol = "multiPhase"
+        for i in range(len(unique_bands)):
+            ampCol_band = f"{unique_bands[i]}_{ampCol}"
+            if ampCol_band not in diaObjects.columns:
+                diaObjects[ampCol_band] = np.nan
+            phaseCol_band = f"{unique_bands[i]}_{phaseCol}"
+            if phaseCol_band not in diaObjects.columns:
+                diaObjects[phaseCol_band] = np.nan
+
+        def _calculate_period_multi(df, all_unique_bands,
+                                    min_detections=9, oversampling_factor=5, nyquist_factor=100):
             """Calculate the multi-band Lomb-Scargle periodogram.
 
             Parameters
             ----------
             df : `pandas.DataFrame`
                 The input DataFrame.
+            all_unique_bands : `list` of `str`
+                List of all bands present in the diaSource table that is being worked on.
             min_detections : `int`, optional
                 The minimum number of detections, including all bands.
             oversampling_factor : `int`, optional
@@ -352,11 +362,14 @@ def _calculate_period_multi(df, min_detections=9, oversampling_factor=5, nyquist
                                       np.isnan(df["midpointMjdTai"]))]
 
             if (len(tmpDf)) < min_detections:
-                return pd.Series({periodCol: np.nan,
-                                  powerCol: np.nan,
-                                  fapCol: np.nan,
-                                  ampCol: pd.Series([np.nan]*n_bands, dtype="object"),
-                                  phaseCol: pd.Series([np.nan]*n_bands, dtype="object")})
+                pd_tab_nodet = pd.Series({periodCol: np.nan,
+                                          powerCol: np.nan,
+                                          fapCol: np.nan})
+                for band in all_unique_bands:
+                    pd_tab_nodet[f"{band}_{ampCol}"] = np.nan
+                    pd_tab_nodet[f"{band}_{phaseCol}"] = np.nan
+
+                return pd_tab_nodet
 
             time = tmpDf["midpointMjdTai"].to_numpy()
             flux = tmpDf["psfFlux"].to_numpy()
@@ -378,15 +391,29 @@ def _calculate_period_multi(df, min_detections=9, oversampling_factor=5, nyquist
 
             pd_tab = pd.Series({periodCol: period[np.argmax(power)],
                                 powerCol: np.max(power),
-                                fapCol: fap_estimate,
-                                ampCol: params_table_new[0],
-                                phaseCol: params_table_new[1]
+                                fapCol: fap_estimate
                                 })
 
+            # Initialize the per-band amplitude/phase columns as NaNs
+            for band in all_unique_bands:
+                pd_tab[f"{band}_{ampCol}"] = np.nan
+                pd_tab[f"{band}_{phaseCol}"] = np.nan
+
+            # Populate the values of only the bands that have data for this diaSource
+            unique_bands = np.unique(bands)
+            for i in range(len(unique_bands)):
+                pd_tab[f"{unique_bands[i]}_{ampCol}"] = params_table_new[0][i]
+                pd_tab[f"{unique_bands[i]}_{phaseCol}"] = params_table_new[1][i]
+
             return pd_tab
 
-        diaObjects.loc[:, [periodCol, powerCol, fapCol, ampCol, phaseCol]
-                       ] = diaSources.apply(_calculate_period_multi)
+        columns_list = [periodCol, powerCol, fapCol]
+        for i in range(len(unique_bands)):
+            columns_list.append(f"{unique_bands[i]}_{ampCol}")
+            columns_list.append(f"{unique_bands[i]}_{phaseCol}")
+
+        diaObjects.loc[:, columns_list
+                       ] = diaSources.apply(_calculate_period_multi, unique_bands)
 
 
 class MeanDiaPositionConfig(DiaObjectCalculationPluginConfig):

tests/test_diaCalculationPlugins.py

@@ -97,7 +97,7 @@ def run_multi_plugin(diaObjectCat, diaSourceCat, band, plugin):
         Input object catalog to store data into and read from.
     diaSourcesCat : `pandas.DataFrame`
         DiaSource catalog to read data from and groupby on.
-    fitlerName : `str`
+    filterName : `str`
         String name of the filter to process.
     plugin : `lsst.ap.association.DiaCalculationPlugin`
         Plugin to run.
@@ -121,6 +121,33 @@ def run_multi_plugin(diaObjectCat, diaSourceCat, band, plugin):
                      band=band)
 
 
+def run_multiband_plugin(diaObjectCat, diaSourceCat, plugin):
+    """Wrapper for running multi plugins.
+
+    Reproduces some of the behavior of `lsst.ap.association.DiaCalcuation.run`
+
+    Parameters
+    ----------
+    diaObjectCat : `pandas.DataFrame`
+        Input object catalog to store data into and read from.
+    diaSourcesCat : `pandas.DataFrame`
+        DiaSource catalog to read data from and groupby on.
+    plugin : `lsst.ap.association.DiaCalculationPlugin`
+        Plugin to run.
+    """
+    diaObjectCat.set_index("diaObjectId", inplace=True, drop=False)
+    diaSourceCat.set_index(
+        ["diaObjectId", "band", "diaSourceId"],
+        inplace=True,
+        drop=False)
+
+    diaSourcesGB = diaSourceCat.groupby(level=0)
+
+    plugin.calculate(diaObjects=diaObjectCat,
+                     diaSources=diaSourcesGB,
+                     )
+
+
 def make_diaObject_table(objId, plugin, default_value=None, band=None):
     """Create a minimal diaObject table with columns required for the plugin
 
@@ -958,18 +985,18 @@ def testCalculate(self):
                                              "ap_lombScarglePeriodogramMulti",
                                              None)
 
-        run_multi_plugin(diaObjects, diaSources, "u", plugin)
+        run_multiband_plugin(diaObjects, diaSources, plugin)
         self.assertAlmostEqual(diaObjects.at[objId, "multiPeriod"], 10, delta=0.04)
         self.assertAlmostEqual(diaObjects.at[objId, "multiPower"], 1, delta=1e-2)
         # This implementation of LS returns a normalized power < 1.
         self.assertLess(diaObjects.at[objId, "multiPower"], 1)
         self.assertAlmostEqual(diaObjects.at[objId, "multiFap"], 0, delta=0.04)
         # Note: The below values are empirical, but seem reasonable, and
-        # test that we get an array with one value per band.
-        self.assertFloatsAlmostEqual(np.array(diaObjects.at[objId, "multiAmp"]),
-                                     np.array([0.029, 0.029]), atol=1e-3)
-        self.assertFloatsAlmostEqual(np.array(diaObjects.at[objId, "multiPhase"]),
-                                     np.array([1., -2.]), rtol=6e-2)
+        # test that we get values for each band.
+        self.assertAlmostEqual(diaObjects.at[objId, "u_multiAmp"], 0.029, delta=0.01)
+        self.assertAlmostEqual(diaObjects.at[objId, "g_multiAmp"], 0.029, delta=0.01)
+        self.assertAlmostEqual(diaObjects.at[objId, "u_multiPhase"], -2.0, delta=0.2)
+        self.assertAlmostEqual(diaObjects.at[objId, "g_multiPhase"], 1.0, delta=0.1)
 
     def testCalculateTwoSources(self):
         """Test Mulitband Lomb Scargle Periodogram with 2 sources (minimum