Coverage for tests/test_diaPipe.py: 16%
248 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-10 02:09 -0700
1# This file is part of ap_association.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22import contextlib
23import tempfile
24import unittest
25from unittest.mock import patch, Mock, MagicMock, DEFAULT
26import warnings
28import numpy as np
29import pandas as pd
30import astropy.table as tb
32import lsst.afw.image as afwImage
33import lsst.afw.table as afwTable
34import lsst.dax.apdb as daxApdb
35from lsst.meas.base import IdGenerator
36import lsst.pex.config as pexConfig
37import lsst.pipe.base as pipeBase
38import lsst.utils.tests
39from lsst.pipe.base.testUtils import assertValidOutput
41from lsst.ap.association import DiaPipelineTask
42from lsst.ap.association.utils import convertDataFrameToSdmSchema
43from utils_tests import makeExposure, makeDiaObjects, makeDiaSources, makeDiaForcedSources, \
44 makeSolarSystemSources
47def _makeMockDataFrame():
48 """Create a new mock of a DataFrame.
50 Returns
51 -------
52 mock : `unittest.mock.Mock`
53 A mock guaranteed to accept all operations used by `pandas.DataFrame`.
54 """
55 with warnings.catch_warnings():
56 # spec triggers deprecation warnings on DataFrame, but will
57 # automatically adapt to any removals.
58 warnings.simplefilter("ignore", category=DeprecationWarning)
59 return MagicMock(spec=pd.DataFrame())
62def _makeMockTable():
63 """Create a new mock of a Table.
65 Returns
66 -------
67 mock : `unittest.mock.Mock`
68 A mock guaranteed to accept all operations used by `astropy.table.Table`.
69 """
70 with warnings.catch_warnings():
71 # spec triggers deprecation warnings on DataFrame, but will
72 # automatically adapt to any removals.
73 warnings.simplefilter("ignore", category=DeprecationWarning)
74 return MagicMock(spec=tb.Table())
77class TestDiaPipelineTask(unittest.TestCase):
79 @classmethod
80 def _makeDefaultConfig(cls, config_file, **kwargs):
81 config = DiaPipelineTask.ConfigClass()
82 config.apdb_config_url = config_file
83 config.update(**kwargs)
84 return config
86 def setUp(self):
87 # Create an instance of random generator with fixed seed.
88 rng = np.random.default_rng(1234)
89 self.rng = rng
91 # schemas are persisted in both Gen 2 and Gen 3 butler as prototypical catalogs
92 srcSchema = afwTable.SourceTable.makeMinimalSchema()
93 srcSchema.addField("base_PixelFlags_flag", type="Flag")
94 srcSchema.addField("base_PixelFlags_flag_offimage", type="Flag")
95 self.srcSchema = afwTable.SourceCatalog(srcSchema)
96 self.exposure = makeExposure(False, False)
97 self.diaObjects = makeDiaObjects(20, self.exposure, rng)
98 self.diaSources = makeDiaSources(
99 100, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng)
100 self.diaForcedSources = makeDiaForcedSources(
101 200, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng)
102 self.ssSources = makeSolarSystemSources(
103 20, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng)
105 sqlite_file = tempfile.NamedTemporaryFile()
106 self.addCleanup(sqlite_file.close)
107 self.config_file = tempfile.NamedTemporaryFile()
108 self.addCleanup(self.config_file.close)
109 apdb_config = daxApdb.ApdbSql.init_database(db_url=f"sqlite:///{sqlite_file.name}")
110 apdb_config.save(self.config_file.name)
112 def testRun(self):
113 """Test running while creating and packaging alerts.
114 """
115 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=True)
117 def testRunWithSolarSystemAssociation(self):
118 """Test running while creating and packaging alerts.
119 """
120 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True)
122 def testRunWithAlerts(self):
123 """Test running while creating and packaging alerts.
124 """
125 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False)
127 def testRunWithoutAlertsOrSolarSystem(self):
128 """Test running without creating and packaging alerts.
129 """
130 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False)
132 def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False, doRunForcedMeasurement=False):
133 """Test the normal workflow of each ap_pipe step.
134 """
135 config = self._makeDefaultConfig(
136 config_file=self.config_file.name,
137 doPackageAlerts=doPackageAlerts,
138 doSolarSystemAssociation=doSolarSystemAssociation,
139 doRunForcedMeasurement=doRunForcedMeasurement,
140 )
141 task = DiaPipelineTask(config=config)
142 # Set DataFrame index testing to always return False. Mocks return
143 # true for this check otherwise.
144 task.testDataFrameIndex = lambda x: False
145 diffIm = Mock(spec=afwImage.ExposureF)
146 exposure = Mock(spec=afwImage.ExposureF)
147 template = Mock(spec=afwImage.ExposureF)
148 diaSrc = _makeMockDataFrame()
149 ssObjects = _makeMockTable()
151 # Each of these subtasks should be called once during diaPipe
152 # execution. We use mocks here to check they are being executed
153 # appropriately.
154 subtasksToMock = [
155 "diaCalculation",
156 ]
157 if doPackageAlerts:
158 subtasksToMock.append("alertPackager")
159 else:
160 self.assertFalse(hasattr(task, "alertPackager"))
162 if not doSolarSystemAssociation:
163 self.assertFalse(hasattr(task, "solarSystemAssociator"))
165 def concatMock(_data, **_kwargs):
166 return _makeMockDataFrame()
168 # Mock out the run() methods of these two Tasks to ensure they
169 # return data in the correct form.
170 def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, visitInfo,
171 bbox, wcs):
172 return lsst.pipe.base.Struct(nTotalSsObjects=42,
173 nAssociatedSsObjects=30,
174 ssoAssocDiaSources=_makeMockTable(),
175 unAssocDiaSources=_makeMockTable(),
176 associatedSsSources=_makeMockTable(),
177 unassociatedSsObjects=_makeMockTable())
179 def associator_run(table, diaObjects):
180 return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3,
181 matchedDiaSources=_makeMockDataFrame(),
182 unAssocDiaSources=_makeMockDataFrame())
184 def updateObjectTableMock(diaObjects, diaSources):
185 pass
187 # apdb isn't a subtask, but still needs to be mocked out for correct
188 # execution in the test environment.
189 with patch.multiple(task, **{task: DEFAULT for task in subtasksToMock + ["apdb"]}), \
190 patch('lsst.ap.association.diaPipe.pd.concat', side_effect=concatMock), \
191 patch('lsst.ap.association.diaPipe.DiaPipelineTask.updateObjectTable',
192 side_effect=updateObjectTableMock), \
193 patch('lsst.ap.association.association.AssociationTask.run',
194 side_effect=associator_run) as mainRun, \
195 patch('lsst.pipe.tasks.ssoAssociation.SolarSystemAssociationTask.run',
196 side_effect=solarSystemAssociator_run) as ssRun:
198 result = task.run(diaSrc,
199 None,
200 diffIm,
201 exposure,
202 template,
203 preloadedDiaObjects=self.diaObjects,
204 preloadedDiaSources=self.diaSources,
205 preloadedDiaForcedSources=self.diaForcedSources,
206 band="g",
207 idGenerator=IdGenerator(),
208 solarSystemObjectTable=ssObjects)
209 for subtaskName in subtasksToMock:
210 getattr(task, subtaskName).run.assert_called_once()
211 assertValidOutput(task, result)
212 # Exact type and contents of apdbMarker are undefined.
213 self.assertIsInstance(result.apdbMarker, pexConfig.Config)
214 meta = task.getFullMetadata()
215 # Check that the expected metadata has been set.
216 self.assertEqual(meta["diaPipe.numUpdatedDiaObjects"], 2)
217 self.assertEqual(meta["diaPipe.numUnassociatedDiaObjects"], 3)
218 # and that associators ran once or not at all.
219 mainRun.assert_called_once()
220 if doSolarSystemAssociation:
221 ssRun.assert_called_once()
222 else:
223 ssRun.assert_not_called()
225 def test_tooManyDiaObjectsError(self):
226 maxNewDiaObjects = 100
228 nDiaSources = maxNewDiaObjects + 1
229 diaSources = makeDiaSources(nDiaSources, np.zeros(nDiaSources), self.exposure, self.rng)
231 def runAndTestWithContextManager(threshold):
232 config = self._makeDefaultConfig(config_file=self.config_file.name,
233 doSolarSystemAssociation=False,
234 filterUnAssociatedSources=False,
235 maxNewDiaObjects=threshold,
236 )
237 task = DiaPipelineTask(config=config)
238 contextManager = self.assertRaises(pipeBase.AlgorithmError) if nDiaSources > threshold > 0 \
239 else contextlib.nullcontext()
240 with contextManager:
241 task.associateDiaSources(
242 diaSources,
243 None,
244 None,
245 self.diaObjects,
246 )
247 # Test cases at, above, and below the threshold as well as at 0.
248 runAndTestWithContextManager(0)
249 runAndTestWithContextManager(maxNewDiaObjects - 1)
250 runAndTestWithContextManager(maxNewDiaObjects)
251 runAndTestWithContextManager(maxNewDiaObjects + 1)
253 def test_createDiaObjects(self):
254 """Test that creating new DiaObjects works as expected.
255 """
256 nSources = 5
257 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
258 task = DiaPipelineTask(config=config)
259 diaSources = pd.DataFrame(data=[
260 {"ra": 0.04*idx, "dec": 0.04*idx,
261 "diaSourceId": idx + 1 + nSources, "diaObjectId": 0,
262 "ssObjectId": 0}
263 for idx in range(nSources)])
265 result = task.createNewDiaObjects(convertDataFrameToSdmSchema(task.schema, diaSources, "DiaSource"))
266 self.assertEqual(nSources, len(result.newDiaObjects))
267 self.assertTrue(np.all(np.equal(
268 result.diaSources["diaObjectId"].to_numpy(),
269 result.diaSources["diaSourceId"].to_numpy())))
270 self.assertTrue(np.all(np.equal(
271 result.newDiaObjects["diaObjectId"].to_numpy(),
272 result.diaSources["diaSourceId"].to_numpy())))
274 def test_purgeDiaObjects(self):
275 """Remove diaOjects that are outside an image's bounding box.
276 """
278 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
279 task = DiaPipelineTask(config=config)
280 exposure = makeExposure(False, False)
281 nObj0 = 20
283 # Create diaObjects
284 diaObjects = makeDiaObjects(nObj0, exposure, self.rng)
285 # Shrink the bounding box so that some of the diaObjects will be outside
286 bbox = exposure.getBBox()
287 size = np.minimum(bbox.getHeight(), bbox.getWidth())
288 bbox.grow(-size//4)
289 exposureCut = exposure[bbox]
290 sizeCut = np.minimum(bbox.getHeight(), bbox.getWidth())
291 buffer = 10
292 bbox.grow(buffer)
294 def check_diaObjects(bbox, wcs, diaObjects):
295 raVals = diaObjects.ra.to_numpy()
296 decVals = diaObjects.dec.to_numpy()
297 xVals, yVals = wcs.skyToPixelArray(raVals, decVals, degrees=True)
298 selector = bbox.contains(xVals, yVals)
299 return selector
301 selector0 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects)
302 nIn0 = np.count_nonzero(selector0)
303 nOut0 = np.count_nonzero(~selector0)
304 self.assertEqual(nObj0, nIn0 + nOut0)
306 # Add an ID that is not in the diaObject table. It should not get removed.
307 diaObjectIds0 = diaObjects["diaObjectId"].copy(deep=True)
308 diaObjectIds0[max(diaObjectIds0.index) + 1] = 999
309 diaObjects1, objIds = task.purgeDiaObjects(exposureCut.getBBox(), exposureCut.getWcs(), diaObjects,
310 diaObjectIds=diaObjectIds0, buffer=buffer)
311 diaObjectIds1 = diaObjects1["diaObjectId"]
312 # Verify that the bounding box was not changed
313 sizeCheck = np.minimum(exposureCut.getBBox().getHeight(), exposureCut.getBBox().getWidth())
314 self.assertEqual(sizeCut, sizeCheck)
315 selector1 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects1)
316 nIn1 = np.count_nonzero(selector1)
317 nOut1 = np.count_nonzero(~selector1)
318 nObj1 = len(diaObjects1)
319 self.assertEqual(nObj1, nIn0)
320 # Verify that not all diaObjects were removed
321 self.assertGreater(nObj1, 0)
322 # Check that some diaObjects were removed
323 self.assertLess(nObj1, nObj0)
324 # Verify that no objects outside the bounding box remain
325 self.assertEqual(nOut1, 0)
326 # Verify that no objects inside the bounding box were removed
327 self.assertEqual(nIn1, nIn0)
328 # The length of the updated object IDs should equal the number of objects
329 # plus one, since we added an extra ID.
330 self.assertEqual(nObj1 + 1, len(objIds))
331 # All of the object IDs extracted from the catalog should be in the pruned object IDs
332 self.assertTrue(set(objIds).issuperset(diaObjectIds1))
333 # The pruned object IDs should contain entries that are not in the catalog
334 self.assertFalse(set(diaObjectIds1).issuperset(objIds))
335 # Some IDs should have been removed
336 self.assertLess(len(objIds), len(diaObjectIds0))
338 def test_filterDiaObjects(self):
339 """Unassociated diaSources that are filtered should have good reliability and SNR.
340 Glint trail sources should also be filtered out.
341 """
343 config = self._makeDefaultConfig(config_file=self.config_file.name,
344 doPackageAlerts=False,
345 filterUnAssociatedSources=True)
347 configBadFilter = self._makeDefaultConfig(config_file=self.config_file.name,
348 doPackageAlerts=False,
349 filterUnAssociatedSources=True,
350 newObjectFluxField="notAFlux")
351 configBadFlag = self._makeDefaultConfig(config_file=self.config_file.name,
352 doPackageAlerts=False,
353 filterUnAssociatedSources=True,
354 newObjectBadFlags=("junkSource", "notUsed"))
355 with self.assertRaises(pipeBase.InvalidQuantumError):
356 DiaPipelineTask(config=configBadFilter)
357 with self.assertRaises(pipeBase.InvalidQuantumError):
358 DiaPipelineTask(config=configBadFlag)
359 task = DiaPipelineTask(config=config)
360 nUnassociatedDiaSources = 234
362 # Create diaSources
363 diaSources = makeDiaSources(nUnassociatedDiaSources,
364 np.zeros(nUnassociatedDiaSources),
365 self.exposure,
366 self.rng,
367 flagList=task.config.newObjectBadFlags)
368 reliability = self.rng.random(nUnassociatedDiaSources)
369 flux = (self.rng.random(nUnassociatedDiaSources)**2)*100
370 fluxErr = np.sqrt(flux)
371 glint_trail = np.zeros(nUnassociatedDiaSources, dtype=bool)
372 glint_trail[12:16] = True # add 4 glint trail sources
373 diaSources["reliability"] = reliability
374 diaSources[config.newObjectFluxField] = flux
375 diaSources[config.newObjectFluxField + "Err"] = fluxErr
376 diaSources["glint_trail"] = glint_trail
377 badFlagName = task.config.newObjectBadFlags[0]
378 badFlags = np.zeros(nUnassociatedDiaSources, dtype=bool)
379 nBadFlags = 20
380 badFlags[0:nBadFlags] = True
381 diaSources[badFlagName] = badFlags
383 def runAndCheckFilter(diaSources, snrThreshold=None, lowReliabilitySnrThreshold=None,
384 reliabilityThreshold=None, lowSnrReliabilityThreshold=None,
385 badFlags=None,
386 ):
388 filterResults = task.filterSources(
389 diaSources.copy(deep=True),
390 snrThreshold=snrThreshold,
391 lowReliabilitySnrThreshold=lowReliabilitySnrThreshold,
392 reliabilityThreshold=reliabilityThreshold,
393 lowSnrReliabilityThreshold=lowSnrReliabilityThreshold,
394 badFlags=badFlags,
395 )
396 self.assertEqual(len(filterResults.goodSources) + len(filterResults.badSources),
397 nUnassociatedDiaSources)
398 goodFlux = filterResults.goodSources[config.newObjectFluxField]
399 goodFluxErr = filterResults.goodSources[config.newObjectFluxField + "Err"]
400 goodSnr = np.array(goodFlux/goodFluxErr)
401 self.assertTrue(np.all(goodSnr > snrThreshold))
402 goodReliability = np.array(filterResults.goodSources["reliability"])
403 self.assertTrue(np.all(goodReliability > reliabilityThreshold))
404 goodLowSnrFlag = goodSnr < lowReliabilitySnrThreshold
405 lowSnrReliability = goodReliability[goodLowSnrFlag]
406 self.assertTrue(np.all(lowSnrReliability > lowSnrReliabilityThreshold))
407 glintTrailSources = np.array(filterResults.goodSources["glint_trail"])
408 self.assertTrue(not any(glintTrailSources))
410 # No sources should be removed if the thresholds are turned off
411 runAndCheckFilter(diaSources,
412 snrThreshold=0, lowReliabilitySnrThreshold=0,
413 reliabilityThreshold=0, lowSnrReliabilityThreshold=0)
414 runAndCheckFilter(diaSources,
415 snrThreshold=0, lowReliabilitySnrThreshold=0,
416 reliabilityThreshold=0, lowSnrReliabilityThreshold=0,
417 badFlags=[badFlagName])
418 runAndCheckFilter(diaSources,
419 snrThreshold=2, lowReliabilitySnrThreshold=8,
420 reliabilityThreshold=0, lowSnrReliabilityThreshold=0)
421 runAndCheckFilter(diaSources,
422 snrThreshold=2, lowReliabilitySnrThreshold=8,
423 reliabilityThreshold=0, lowSnrReliabilityThreshold=0.5)
424 runAndCheckFilter(diaSources,
425 snrThreshold=0, lowReliabilitySnrThreshold=0,
426 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5)
427 runAndCheckFilter(diaSources,
428 snrThreshold=2, lowReliabilitySnrThreshold=8,
429 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5)
430 runAndCheckFilter(diaSources,
431 snrThreshold=2, lowReliabilitySnrThreshold=8,
432 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5,
433 badFlags=[badFlagName])
435 def test_mergeEmptyCatalog(self):
436 """Test that a catalog is unchanged if it is merged with an empty
437 catalog.
438 """
439 diaSourcesBase = self.diaSources
441 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
442 task = DiaPipelineTask(config=config)
443 # Include some but not all columns that should be in diaSourcesBase, and some that are mis-matched
444 diaSourcesEmpty = pd.DataFrame(columns=["ra", "dec", "foo"])
445 diaSourcesTest = task.mergeCatalogs(diaSourcesBase, diaSourcesEmpty, tableName="DiaSource")
446 self.assertTrue(diaSourcesBase.equals(diaSourcesTest))
448 def test_mergeCatalogs(self):
449 """Test that a merged catalog is concatenated correctly.
450 """
451 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
452 task = DiaPipelineTask(config=config)
454 diaSourcesBase = convertDataFrameToSdmSchema(task.schema, self.diaSources, "DiaSource")
455 nBase = len(diaSourcesBase)
456 nNew = int(nBase/2)
458 diaSourcesNew = makeDiaSources(nNew, self.diaObjects["diaObjectId"].to_numpy(), self.exposure,
459 self.rng)
460 diaSourcesNew = convertDataFrameToSdmSchema(task.schema, diaSourcesNew, "DiaSource")
461 diaSourcesTest = task.mergeCatalogs(diaSourcesBase, diaSourcesNew, tableName="DiaSource")
462 self.assertEqual(len(diaSourcesTest), nBase + nNew)
463 diaSourcesExtract1 = diaSourcesTest.iloc[:nBase]
464 diaSourcesExtract2 = diaSourcesTest.iloc[nBase:]
466 pd.testing.assert_frame_equal(diaSourcesBase, diaSourcesExtract1)
467 pd.testing.assert_frame_equal(diaSourcesNew, diaSourcesExtract2)
469 def test_updateObjectTable(self):
470 """Test that the diaObject record is updated with the number of
471 diaSources.
472 """
473 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
474 task = DiaPipelineTask(config=config)
475 nObjects = 20
476 nSrcPerObject = 10
477 nExtraSources = 5
478 nSources = nSrcPerObject*nObjects + nExtraSources
479 expectedSourcesPerObject = nSrcPerObject*np.ones(nObjects)
480 expectedSourcesPerObject[:nExtraSources] += 1
481 diaObjects = makeDiaObjects(nObjects, self.exposure, self.rng)
482 diaSources = makeDiaSources(nSources, diaObjects["diaObjectId"].to_numpy(), self.exposure, self.rng)
483 updatedDiaObjects = task.updateObjectTable(diaObjects, diaSources)
484 self.assertTrue(np.all(updatedDiaObjects.nDiaSources.values == expectedSourcesPerObject))
487class MemoryTester(lsst.utils.tests.MemoryTestCase):
488 pass
491def setup_module(module):
492 lsst.utils.tests.init()
495if __name__ == "__main__": 495 ↛ 496line 495 didn't jump to line 496 because the condition on line 495 was never true
496 lsst.utils.tests.init()
497 unittest.main()