Drop diaObjects that are off the image after association

python/lsst/ap/association/diaPipe.py

@@ -455,8 +455,9 @@ def run(self,
             solarSystemObjectTable = legacySolarSystemTable
 
         if not preloadedDiaObjects.empty:
-            diaObjects = self.purgeDiaObjects(diffIm.getBBox(), diffIm.getWcs(), preloadedDiaObjects,
-                                              buffer=self.config.imagePixelMargin)
+            # Include a small buffer outside the image so that we can associate sources near the edge
+            diaObjects, _ = self.purgeDiaObjects(diffIm.getBBox(), diffIm.getWcs(), preloadedDiaObjects,
+                                                 buffer=self.config.imagePixelMargin)
         else:
             diaObjects = preloadedDiaObjects
 
@@ -467,7 +468,7 @@ def run(self,
 
         # Merge associated diaSources
         mergedDiaSourceHistory, mergedDiaObjects, updatedDiaObjectIds = self.mergeAssociatedCatalogs(
-            preloadedDiaSources, associatedDiaSources, diaObjects, newDiaObjects
+            preloadedDiaSources, associatedDiaSources, diaObjects, newDiaObjects, diffIm
         )
 
         # Compute DiaObject Summary statistics from their full DiaSource
@@ -648,7 +649,8 @@ def associateDiaSources(self, diaSourceTable, solarSystemObjectTable, diffIm, di
         return (associatedDiaSources, createResults.newDiaObjects)
 
     @timeMethod
-    def mergeAssociatedCatalogs(self, preloadedDiaSources, associatedDiaSources, diaObjects, newDiaObjects):
+    def mergeAssociatedCatalogs(self, preloadedDiaSources, associatedDiaSources, diaObjects, newDiaObjects,
+                                diffIm):
         """Merge the associated diaSource and diaObjects to their previous history.
 
         Parameters
@@ -697,7 +699,13 @@ def mergeAssociatedCatalogs(self, preloadedDiaSources, associatedDiaSources, dia
                 sort=True)
         else:
             mergedDiaObjects = diaObjects
-        if self.testDataFrameIndex(diaObjects):
+
+        # Exclude any objects that are off the image after association.
+        mergedDiaObjects, updatedDiaObjectIds = self.purgeDiaObjects(diffIm.getBBox(), diffIm.getWcs(),
+                                                                     mergedDiaObjects,
+                                                                     diaObjectIds=updatedDiaObjectIds,
+                                                                     buffer=-1)
+        if self.testDataFrameIndex(mergedDiaObjects):
             raise RuntimeError(
                 "Duplicate DiaObjects created after association. This is "
                 "likely due to re-running data with an already populated "
@@ -840,7 +848,7 @@ def _add_association_meta_data(self,
         self.metadata.add('numTotalSolarSystemObjects', nTotalSsObjects)
         self.metadata.add('numAssociatedSsObjects', nAssociatedSsObjects)
 
-    def purgeDiaObjects(self, bbox, wcs, diaObjCat, buffer=0):
+    def purgeDiaObjects(self, bbox, wcs, diaObjCat, diaObjectIds=None, buffer=0):
         """Drop diaObjects that are outside the exposure bounding box.
 
         Parameters
@@ -868,12 +876,21 @@ def purgeDiaObjects(self, bbox, wcs, diaObjCat, buffer=0):
             selector = bbox.contains(xVals, yVals)
             nPurged = np.sum(~selector)
             if nPurged > 0:
+                if diaObjectIds is not None:
+                    # We also need to drop any of the associated IDs if this runs after association
+                    purgedIds = diaObjCat[~selector].diaObjectId
+                    diaObjectIds = diaObjectIds[~np.isin(diaObjectIds, purgedIds)]
+                    self.log.info("Dropped %i diaObjects that were outside the bbox "
+                                  "after association, leaving %i in the catalog",
+                                  nPurged, len(diaObjCat) - nPurged)
+                else:
+                    self.log.info("Dropped %i diaObjects that were outside the padded bbox "
+                                  "before association, leaving %i in the catalog",
+                                  nPurged, len(diaObjCat) - nPurged)
                 diaObjCat = diaObjCat[selector].copy()
-                self.log.info(f"Dropped {nPurged} diaObjects that were outside the bbox "
-                              f"leaving {len(diaObjCat)} in the catalog")
         except Exception as e:
             self.log.warning("Error attempting to check diaObject history: %s", e)
-        return diaObjCat
+        return diaObjCat, diaObjectIds
 
     def mergeCatalogs(self, originalCatalog, newCatalog, catalogName):
         """Combine two catalogs, ensuring that the columns of the new catalog

tests/test_diaPipe.py

@@ -282,8 +282,12 @@ def check_diaObjects(bbox, wcs, diaObjects):
         nOut0 = np.count_nonzero(~selector0)
         self.assertEqual(nObj0, nIn0 + nOut0)
 
-        diaObjects1 = task.purgeDiaObjects(exposureCut.getBBox(), exposureCut.getWcs(), diaObjects,
-                                           buffer=buffer)
+        # Add an ID that is not in the diaObject table. It should not get removed.
+        diaObjectIds0 = diaObjects["diaObjectId"].copy(deep=True)
+        diaObjectIds0[max(diaObjectIds0.index) + 1] = 999
+        diaObjects1, objIds = task.purgeDiaObjects(exposureCut.getBBox(), exposureCut.getWcs(), diaObjects,
+                                                   diaObjectIds=diaObjectIds0, buffer=buffer)
+        diaObjectIds1 = diaObjects1["diaObjectId"]
         # Verify that the bounding box was not changed
         sizeCheck = np.minimum(exposureCut.getBBox().getHeight(), exposureCut.getBBox().getWidth())
         self.assertEqual(sizeCut, sizeCheck)
@@ -300,6 +304,15 @@ def check_diaObjects(bbox, wcs, diaObjects):
         self.assertEqual(nOut1, 0)
         # Verify that no objects inside the bounding box were removed
         self.assertEqual(nIn1, nIn0)
+        # The length of the updated object IDs should equal the number of objects
+        # plus one, since we added an extra ID.
+        self.assertEqual(nObj1 + 1, len(objIds))
+        # All of the object IDs extracted from the catalog should be in the pruned object IDs
+        self.assertTrue(set(objIds).issuperset(diaObjectIds1))
+        # The pruned object IDs should contain entries that are not in the catalog
+        self.assertFalse(set(diaObjectIds1).issuperset(objIds))
+        # Some IDs should have been removed
+        self.assertLess(len(objIds), len(diaObjectIds0))
 
     def test_mergeEmptyCatalog(self):
         """Test that a catalog is unchanged if it is merged with an empty