DM-31619: Fully incorporate SSP object association in DiaPipe

python/lsst/ap/association/association.py

@@ -77,9 +77,13 @@ def run(self,
         result : `lsst.pipe.base.Struct`
             Results struct with components.
 
-            - ``"diaSources"`` : Full set of diaSources after matching. Matched
+            - ``"matchedDiaSources"`` : DiaSources that were matched. Matched
               Sources have their diaObjectId updated and set to the id of the
               diaObject they were matched to. (`pandas.DataFrame`)
+            - ``"unAssocDiaSources"`` : DiaSources that were not matched.
+              Unassociated sources have their diaObject set to 0 as they
+              were not associated with any existing DiaObjects.
+              (`pandas.DataFrame`)
             - ``"nUpdatedDiaObjects"`` : Number of DiaObjects that were
               matched to new DiaSources. (`int`)
             - ``"nUnassociatedDiaObjects"`` : Number of DiaObjects that were
@@ -88,14 +92,18 @@ def run(self,
         diaSources = self.check_dia_source_radec(diaSources)
         if len(diaObjects) == 0:
             return pipeBase.Struct(
-                diaSources=diaSources,
+                matchedDiaSources=pd.DataFrame(columns=diaSources.columns),
+                unAssocDiaSources=diaSources,
                 nUpdatedDiaObjects=0,
                 nUnassociatedDiaObjects=0)
 
         matchResult = self.associate_sources(diaObjects, diaSources)
 
+        mask = matchResult.diaSources["diaObjectId"] != 0
+
         return pipeBase.Struct(
-            diaSources=matchResult.diaSources,
+            matchedDiaSources=matchResult.diaSources[mask].reset_index(drop=True),
+            unAssocDiaSources=matchResult.diaSources[~mask].reset_index(drop=True),
             nUpdatedDiaObjects=matchResult.nUpdatedDiaObjects,
             nUnassociatedDiaObjects=matchResult.nUnassociatedDiaObjects)
 

python/lsst/ap/association/packageAlerts.py

@@ -120,6 +120,10 @@ def run(self,
         templatePhotoCalib = template.getPhotoCalib()
         for srcIndex, diaSource in diaSourceCat.iterrows():
             # Get all diaSources for the associated diaObject.
+            # TODO: DM-31992 skip DiaSources associated with Solar System
+            #       Objects for now.
+            if srcIndex[0] == 0:
+                continue
             diaObject = diaObjectCat.loc[srcIndex[0]]
             if diaObject["nDiaSources"] > 1:
                 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]

python/lsst/ap/association/skyBotEphemerisQuery.py

@@ -128,7 +128,7 @@ def run(self, visitInfos, visit):
             datasetType=self.config.connections.visitInfos)
 
         # Midpoint time of the exposure in MJD
-        expMidPointEPOCH = visitInfo.date.get(system=DateTime.EPOCH)
+        expMidPointEPOCH = visitInfo.date.get(system=DateTime.JD)
 
         # Boresight of the exposure on sky.
         expCenter = visitInfo.boresightRaDec

python/lsst/ap/association/ssoAssociation.py

@@ -96,8 +96,8 @@ def run(self, diaSourceCatalog, solarSystemObjects):
 
         assocMask = diaSourceCatalog["ssObjectId"] != 0
         return pipeBase.Struct(
-            ssoAssocDiaSources=diaSourceCatalog[assocMask],
-            unAssocDiaSources=diaSourceCatalog[~assocMask])
+            ssoAssocDiaSources=diaSourceCatalog[assocMask].reset_index(drop=True),
+            unAssocDiaSources=diaSourceCatalog[~assocMask].reset_index(drop=True))
 
     def _radec_to_xyz(self, ras, decs):
         """Convert input ra/dec coordinates to spherical unit-vectors.

tests/test_association_task.py

@@ -63,9 +63,16 @@ def test_run(self):
 
         self.assertEqual(results.nUpdatedDiaObjects, len(self.diaObjects) - 1)
         self.assertEqual(results.nUnassociatedDiaObjects, 1)
+        self.assertEqual(len(results.matchedDiaSources),
+                         len(self.diaObjects) - 1)
+        self.assertEqual(len(results.unAssocDiaSources), 1)
         for test_obj_id, expected_obj_id in zip(
-                results.diaSources["diaObjectId"].to_numpy(),
-                [0, 1, 2, 3, 4]):
+                results.matchedDiaSources["diaObjectId"].to_numpy(),
+                [1, 2, 3, 4]):
+            self.assertEqual(test_obj_id, expected_obj_id)
+        for test_obj_id, expected_obj_id in zip(
+                results.unAssocDiaSources["diaObjectId"].to_numpy(),
+                [0]):
             self.assertEqual(test_obj_id, expected_obj_id)
 
     def test_run_no_existing_objects(self):
@@ -77,7 +84,8 @@ def test_run_no_existing_objects(self):
             pd.DataFrame(columns=["ra", "decl", "diaObjectId"]))
         self.assertEqual(results.nUpdatedDiaObjects, 0)
         self.assertEqual(results.nUnassociatedDiaObjects, 0)
-        self.assertTrue(np.all(results.diaSources["diaObjectId"] == 0))
+        self.assertEqual(len(results.matchedDiaSources), 0)
+        self.assertTrue(np.all(results.unAssocDiaSources["diaObjectId"] == 0))
 
     def test_associate_sources(self):
         """Test the performance of the associate_sources method in

tests/test_diaPipe.py

@@ -27,18 +27,21 @@
 import lsst.afw.table as afwTable
 import lsst.pipe.base as pipeBase
 import lsst.utils.tests
-from unittest.mock import patch, Mock, DEFAULT
+from unittest.mock import patch, Mock, MagicMock, DEFAULT
 
 from lsst.ap.association import DiaPipelineTask
 
 
 class TestDiaPipelineTask(unittest.TestCase):
 
     @classmethod
-    def _makeDefaultConfig(cls, doPackageAlerts=False):
+    def _makeDefaultConfig(cls,
+                           doPackageAlerts=False,
+                           doSolarSystemAssociation=False):
         config = DiaPipelineTask.ConfigClass()
         config.apdb.db_url = "sqlite://"
         config.doPackageAlerts = doPackageAlerts
+        config.doSolarSystemAssociation = doSolarSystemAssociation
         return config
 
     def setUp(self):
@@ -51,28 +54,41 @@ def setUp(self):
     def tearDown(self):
         pass
 
+    def testRun(self):
+        """Test running while creating and packaging alerts.
+        """
+        self._testRun(doPackageAlerts=True, doSolarSystemAssociation=True)
+
+    def testRunWithSolarSystemAssociation(self):
+        """Test running while creating and packaging alerts.
+        """
+        self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True)
+
     def testRunWithAlerts(self):
         """Test running while creating and packaging alerts.
         """
-        self._testRun(True)
+        self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False)
 
-    def testRunWithoutAlerts(self):
+    def testRunWithoutAlertsOrSolarSystem(self):
         """Test running without creating and packaging alerts.
         """
-        self._testRun(False)
+        self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False)
 
-    def _testRun(self, doPackageAlerts=False):
+    def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False):
         """Test the normal workflow of each ap_pipe step.
         """
-        config = self._makeDefaultConfig(doPackageAlerts=doPackageAlerts)
+        config = self._makeDefaultConfig(
+            doPackageAlerts=doPackageAlerts,
+            doSolarSystemAssociation=doSolarSystemAssociation)
         task = DiaPipelineTask(config=config)
         # Set DataFrame index testing to always return False. Mocks return
         # true for this check otherwise.
         task.testDataFrameIndex = lambda x: False
         diffIm = Mock(spec=afwImage.ExposureF)
         exposure = Mock(spec=afwImage.ExposureF)
         template = Mock(spec=afwImage.ExposureF)
-        diaSrc = Mock(sepc=pd.DataFrame)
+        diaSrc = MagicMock(spec=pd.DataFrame())
+        ssObjects = MagicMock(spec=pd.DataFrame())
         ccdExposureIdBits = 32
 
         # Each of these subtasks should be called once during diaPipe
@@ -89,21 +105,31 @@ def _testRun(self, doPackageAlerts=False):
         else:
             self.assertFalse(hasattr(task, "alertPackager"))
 
+        if doSolarSystemAssociation:
+            subtasksToMock.append("solarSystemAssociator")
+        else:
+            self.assertFalse(hasattr(task, "solarSystemAssociator"))
+
         # apdb isn't a subtask, but still needs to be mocked out for correct
         # execution in the test environment.
+        def concatMock(data):
+            return MagicMock(spec=pd.DataFrame)
         with patch.multiple(
             task, **{task: DEFAULT for task in subtasksToMock + ["apdb"]}
         ):
-            result = task.run(diaSrc,
-                              diffIm,
-                              exposure,
-                              template,
-                              ccdExposureIdBits,
-                              "g")
-            for subtaskName in subtasksToMock:
-                getattr(task, subtaskName).run.assert_called_once()
-            pipeBase.testUtils.assertValidOutput(task, result)
-            self.assertEqual(result.apdbMarker.db_url, "sqlite://")
+            with patch('lsst.ap.association.diaPipe.pd.concat',
+                       new=concatMock):
+                result = task.run(diaSrc,
+                                  ssObjects,
+                                  diffIm,
+                                  exposure,
+                                  template,
+                                  ccdExposureIdBits,
+                                  "g")
+                for subtaskName in subtasksToMock:
+                    getattr(task, subtaskName).run.assert_called_once()
+                pipeBase.testUtils.assertValidOutput(task, result)
+                self.assertEqual(result.apdbMarker.db_url, "sqlite://")
 
     def test_createDiaObjects(self):
         """Test that creating new DiaObjects works as expected.

tests/test_packageAlerts.py

@@ -91,7 +91,7 @@ def makeDiaObjects(nObjects, exposure):
         newObject = {"ra": coord.getRa().asDegrees(),
                      "decl": coord.getDec().asDegrees(),
                      "radecTai": midPointTaiMJD,
-                     "diaObjectId": idx,
+                     "diaObjectId": idx + 1,
                      "pmParallaxNdata": 0,
                      "nearbyObj1": 0,
                      "nearbyObj2": 0,
@@ -146,7 +146,7 @@ def makeDiaSources(nSources, diaObjectIds, exposure):
                      "ssObjectId": 0,
                      "parentDiaSourceId": 0,
                      "prv_procOrder": 0,
-                     "diaSourceId": idx,
+                     "diaSourceId": idx + 1,
                      "midPointTai": midPointTaiMJD + 1.0 * idx,
                      "filterName": exposure.getFilterLabel().bandLabel,
                      "psNdata": 0,
@@ -183,7 +183,7 @@ def makeDiaForcedSources(nSources, diaObjectIds, exposure):
     for idx in range(nSources):
         objId = diaObjectIds[idx % len(diaObjectIds)]
         # Put together the minimum values for the alert.
-        data.append({"diaForcedSourceId": idx,
+        data.append({"diaForcedSourceId": idx + 1,
                      "ccdVisitId": ccdVisitId + idx,
                      "diaObjectId": objId,
                      "midPointTai": midPointTaiMJD + 1.0 * idx,
@@ -295,10 +295,10 @@ def setUp(self):
         diaSourceHistory["programId"] = 0
 
         self.diaSources = diaSourceHistory.loc[
-            [(0, "g", 8), (1, "g", 9)], :]
+            [(1, "g", 9), (2, "g", 10)], :]
         self.diaSources["bboxSize"] = self.cutoutSize
-        self.diaSourceHistory = diaSourceHistory.drop(labels=[(0, "g", 8),
-                                                              (1, "g", 9)])
+        self.diaSourceHistory = diaSourceHistory.drop(labels=[(1, "g", 9),
+                                                              (2, "g", 10)])
 
         self.cutoutWcs = wcs.WCS(naxis=2)
         self.cutoutWcs.wcs.crpix = [self.center[0], self.center[1]]