Coverage for tests/test_diaPipe.py: 17%
257 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-15 01:05 -0800
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, doReloadDiaObjects=False)
117 def testRunWithSolarSystemAssociation(self):
118 """Test running while creating and packaging alerts.
119 """
120 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True, doReloadDiaObjects=False)
122 def testRunWithAlerts(self):
123 """Test running while creating and packaging alerts.
124 """
125 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False, doReloadDiaObjects=False)
127 def testRunWithoutAlertsOrSolarSystem(self):
128 """Test running without creating and packaging alerts.
129 """
130 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False, doReloadDiaObjects=False)
132 def testRunWithReload(self):
133 """Test running with reloading DiaObjects.
134 """
135 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False, doReloadDiaObjects=True)
137 def testRunWithReloadAndSolarSystem(self):
138 """Test running with solar system association and reloading DiaObjects.
139 """
140 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True, doReloadDiaObjects=True)
142 def testRunWithReloadAndAlerts(self):
143 """Test running with reloading DiaObjects while creating and packaging alerts.
144 """
145 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False, doReloadDiaObjects=True)
147 def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False, doRunForcedMeasurement=False,
148 doReloadDiaObjects=False):
149 """Test the normal workflow of each ap_pipe step.
150 """
151 config = self._makeDefaultConfig(
152 config_file=self.config_file.name,
153 doPackageAlerts=doPackageAlerts,
154 doSolarSystemAssociation=doSolarSystemAssociation,
155 doRunForcedMeasurement=doRunForcedMeasurement,
156 doReloadDiaObjects=doReloadDiaObjects,
157 )
158 task = DiaPipelineTask(config=config)
159 # Set DataFrame index testing to always return False. Mocks return
160 # true for this check otherwise.
161 task.testDataFrameIndex = lambda x: False
162 diffIm = Mock(spec=afwImage.ExposureF)
163 template = Mock(spec=afwImage.ExposureF)
164 diaSrc = _makeMockDataFrame()
165 ssObjects = _makeMockTable()
167 # Each of these subtasks should be called once during diaPipe
168 # execution. We use mocks here to check they are being executed
169 # appropriately.
170 subtasksToMock = [
171 "diaCalculation",
172 ]
173 if doPackageAlerts:
174 subtasksToMock.append("alertPackager")
175 else:
176 self.assertFalse(hasattr(task, "alertPackager"))
178 if not doSolarSystemAssociation:
179 self.assertFalse(hasattr(task, "solarSystemAssociator"))
181 def concatMock(_data, **_kwargs):
182 return _makeMockDataFrame()
184 # Mock out the run() methods of these Tasks to ensure they
185 # return data in the correct form.
186 def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, visitInfo,
187 bbox, wcs):
188 return lsst.pipe.base.Struct(nTotalSsObjects=42,
189 nAssociatedSsObjects=30,
190 ssoAssocDiaSources=_makeMockTable(),
191 unAssocDiaSources=_makeMockTable(),
192 associatedSsSources=_makeMockTable(),
193 unassociatedSsObjects=_makeMockTable())
195 def associator_run(table, diaObjects):
196 return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3,
197 matchedDiaSources=_makeMockDataFrame(),
198 unAssocDiaSources=_makeMockDataFrame())
200 def loadObjects_run(region, preloadedDiaObjects):
201 task.metadata['loadRefreshedDiaObjectsStartUtc'] = 1.234
202 task.metadata['loadRefreshedDiaObjectsEndUtc'] = 5.678
203 return self.diaObjects
205 def updateObjectTableMock(diaObjects, diaSources):
206 pass
208 # apdb isn't a subtask, but still needs to be mocked out for correct
209 # execution in the test environment.
210 with patch.multiple(task, **{task: DEFAULT for task in subtasksToMock + ["apdb"]}), \
211 patch('lsst.ap.association.diaPipe.pd.concat', side_effect=concatMock), \
212 patch('lsst.ap.association.diaPipe.DiaPipelineTask.updateObjectTable',
213 side_effect=updateObjectTableMock), \
214 patch('lsst.ap.association.diaPipe.DiaPipelineTask.loadRefreshedDiaObjects',
215 side_effect=loadObjects_run), \
216 patch('lsst.ap.association.association.AssociationTask.run',
217 side_effect=associator_run) as mainRun, \
218 patch('lsst.pipe.tasks.ssoAssociation.SolarSystemAssociationTask.run',
219 side_effect=solarSystemAssociator_run) as ssRun:
221 result = task.run(diaSrc,
222 None,
223 diffIm,
224 self.exposure,
225 template,
226 preloadedDiaObjects=self.diaObjects,
227 preloadedDiaSources=self.diaSources,
228 preloadedDiaForcedSources=self.diaForcedSources,
229 band="g",
230 idGenerator=IdGenerator(),
231 solarSystemObjectTable=ssObjects)
232 for subtaskName in subtasksToMock:
233 getattr(task, subtaskName).run.assert_called_once()
234 assertValidOutput(task, result)
235 # Exact type and contents of apdbMarker are undefined.
236 self.assertIsInstance(result.apdbMarker, pexConfig.Config)
237 meta = task.getFullMetadata()
238 # Check that the expected metadata has been set.
239 self.assertEqual(meta["diaPipe.numUpdatedDiaObjects"], 2)
240 self.assertEqual(meta["diaPipe.numUnassociatedDiaObjects"], 3)
241 # and that associators ran once or not at all.
242 mainRun.assert_called_once()
243 if doSolarSystemAssociation:
244 ssRun.assert_called_once()
245 else:
246 ssRun.assert_not_called()
248 def test_tooManyDiaObjectsError(self):
249 maxNewDiaObjects = 100
251 nDiaSources = maxNewDiaObjects + 1
252 diaSources = makeDiaSources(nDiaSources, np.zeros(nDiaSources), self.exposure, self.rng)
254 def runAndTestWithContextManager(threshold):
255 config = self._makeDefaultConfig(config_file=self.config_file.name,
256 doSolarSystemAssociation=False,
257 filterUnAssociatedSources=False,
258 maxNewDiaObjects=threshold,
259 )
260 task = DiaPipelineTask(config=config)
261 contextManager = self.assertRaises(pipeBase.AlgorithmError) if nDiaSources > threshold > 0 \
262 else contextlib.nullcontext()
263 with contextManager:
264 task.associateDiaSources(
265 diaSources,
266 None,
267 None,
268 self.diaObjects,
269 )
270 # Test cases at, above, and below the threshold as well as at 0.
271 runAndTestWithContextManager(0)
272 runAndTestWithContextManager(maxNewDiaObjects - 1)
273 runAndTestWithContextManager(maxNewDiaObjects)
274 runAndTestWithContextManager(maxNewDiaObjects + 1)
276 def test_createDiaObjects(self):
277 """Test that creating new DiaObjects works as expected.
278 """
279 nSources = 5
280 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
281 task = DiaPipelineTask(config=config)
282 diaSources = pd.DataFrame(data=[
283 {"ra": 0.04*idx, "dec": 0.04*idx,
284 "diaSourceId": idx + 1 + nSources, "diaObjectId": 0,
285 "ssObjectId": 0}
286 for idx in range(nSources)])
288 result = task.createNewDiaObjects(convertDataFrameToSdmSchema(task.schema, diaSources, "DiaSource"))
289 self.assertEqual(nSources, len(result.newDiaObjects))
290 self.assertTrue(np.all(np.equal(
291 result.diaSources["diaObjectId"].to_numpy(),
292 result.diaSources["diaSourceId"].to_numpy())))
293 self.assertTrue(np.all(np.equal(
294 result.newDiaObjects["diaObjectId"].to_numpy(),
295 result.diaSources["diaSourceId"].to_numpy())))
297 def test_purgeDiaObjects(self):
298 """Remove diaOjects that are outside an image's bounding box.
299 """
301 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
302 task = DiaPipelineTask(config=config)
303 exposure = makeExposure(False, False)
304 nObj0 = 20
306 # Create diaObjects
307 diaObjects = makeDiaObjects(nObj0, exposure, self.rng)
308 # Shrink the bounding box so that some of the diaObjects will be outside
309 bbox = exposure.getBBox()
310 size = np.minimum(bbox.getHeight(), bbox.getWidth())
311 bbox.grow(-size//4)
312 exposureCut = exposure[bbox]
313 sizeCut = np.minimum(bbox.getHeight(), bbox.getWidth())
314 buffer = 10
315 bbox.grow(buffer)
317 def check_diaObjects(bbox, wcs, diaObjects):
318 raVals = diaObjects.ra.to_numpy()
319 decVals = diaObjects.dec.to_numpy()
320 xVals, yVals = wcs.skyToPixelArray(raVals, decVals, degrees=True)
321 selector = bbox.contains(xVals, yVals)
322 return selector
324 selector0 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects)
325 nIn0 = np.count_nonzero(selector0)
326 nOut0 = np.count_nonzero(~selector0)
327 self.assertEqual(nObj0, nIn0 + nOut0)
329 # Add an ID that is not in the diaObject table. It should not get removed.
330 diaObjectIds0 = diaObjects["diaObjectId"].copy(deep=True)
331 diaObjectIds0[max(diaObjectIds0.index) + 1] = 999
332 diaObjects1, objIds = task.purgeDiaObjects(exposureCut.getBBox(), exposureCut.getWcs(), diaObjects,
333 diaObjectIds=diaObjectIds0, buffer=buffer)
334 diaObjectIds1 = diaObjects1["diaObjectId"]
335 # Verify that the bounding box was not changed
336 sizeCheck = np.minimum(exposureCut.getBBox().getHeight(), exposureCut.getBBox().getWidth())
337 self.assertEqual(sizeCut, sizeCheck)
338 selector1 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects1)
339 nIn1 = np.count_nonzero(selector1)
340 nOut1 = np.count_nonzero(~selector1)
341 nObj1 = len(diaObjects1)
342 self.assertEqual(nObj1, nIn0)
343 # Verify that not all diaObjects were removed
344 self.assertGreater(nObj1, 0)
345 # Check that some diaObjects were removed
346 self.assertLess(nObj1, nObj0)
347 # Verify that no objects outside the bounding box remain
348 self.assertEqual(nOut1, 0)
349 # Verify that no objects inside the bounding box were removed
350 self.assertEqual(nIn1, nIn0)
351 # The length of the updated object IDs should equal the number of objects
352 # plus one, since we added an extra ID.
353 self.assertEqual(nObj1 + 1, len(objIds))
354 # All of the object IDs extracted from the catalog should be in the pruned object IDs
355 self.assertTrue(set(objIds).issuperset(diaObjectIds1))
356 # The pruned object IDs should contain entries that are not in the catalog
357 self.assertFalse(set(diaObjectIds1).issuperset(objIds))
358 # Some IDs should have been removed
359 self.assertLess(len(objIds), len(diaObjectIds0))
361 def test_filterDiaObjects(self):
362 """Unassociated diaSources that are filtered should have good reliability and SNR.
363 Glint trail sources should also be filtered out.
364 """
366 config = self._makeDefaultConfig(config_file=self.config_file.name,
367 doPackageAlerts=False,
368 filterUnAssociatedSources=True)
370 configBadFilter = self._makeDefaultConfig(config_file=self.config_file.name,
371 doPackageAlerts=False,
372 filterUnAssociatedSources=True,
373 newObjectFluxField="notAFlux")
374 configBadFlag = self._makeDefaultConfig(config_file=self.config_file.name,
375 doPackageAlerts=False,
376 filterUnAssociatedSources=True,
377 newObjectBadFlags=("junkSource", "notUsed"))
378 with self.assertRaises(pipeBase.InvalidQuantumError):
379 DiaPipelineTask(config=configBadFilter)
380 with self.assertRaises(pipeBase.InvalidQuantumError):
381 DiaPipelineTask(config=configBadFlag)
382 task = DiaPipelineTask(config=config)
383 nUnassociatedDiaSources = 234
385 # Create diaSources
386 diaSources = makeDiaSources(nUnassociatedDiaSources,
387 np.zeros(nUnassociatedDiaSources),
388 self.exposure,
389 self.rng,
390 flagList=task.config.newObjectBadFlags)
391 reliability = self.rng.random(nUnassociatedDiaSources)
392 flux = (self.rng.random(nUnassociatedDiaSources)**2)*100
393 fluxErr = np.sqrt(flux)
394 glint_trail = np.zeros(nUnassociatedDiaSources, dtype=bool)
395 glint_trail[12:16] = True # add 4 glint trail sources
396 diaSources["reliability"] = reliability
397 diaSources[config.newObjectFluxField] = flux
398 diaSources[config.newObjectFluxField + "Err"] = fluxErr
399 diaSources["glint_trail"] = glint_trail
400 badFlagName = task.config.newObjectBadFlags[0]
401 badFlags = np.zeros(nUnassociatedDiaSources, dtype=bool)
402 nBadFlags = 20
403 badFlags[0:nBadFlags] = True
404 diaSources[badFlagName] = badFlags
406 def runAndCheckFilter(diaSources, snrThreshold=None, lowReliabilitySnrThreshold=None,
407 reliabilityThreshold=None, lowSnrReliabilityThreshold=None,
408 badFlags=None,
409 ):
411 filterResults = task.filterSources(
412 diaSources.copy(deep=True),
413 snrThreshold=snrThreshold,
414 lowReliabilitySnrThreshold=lowReliabilitySnrThreshold,
415 reliabilityThreshold=reliabilityThreshold,
416 lowSnrReliabilityThreshold=lowSnrReliabilityThreshold,
417 badFlags=badFlags,
418 )
419 self.assertEqual(len(filterResults.goodSources) + len(filterResults.badSources),
420 nUnassociatedDiaSources)
421 goodFlux = filterResults.goodSources[config.newObjectFluxField]
422 goodFluxErr = filterResults.goodSources[config.newObjectFluxField + "Err"]
423 goodSnr = np.array(goodFlux/goodFluxErr)
424 self.assertTrue(np.all(goodSnr > snrThreshold))
425 goodReliability = np.array(filterResults.goodSources["reliability"])
426 self.assertTrue(np.all(goodReliability > reliabilityThreshold))
427 goodLowSnrFlag = goodSnr < lowReliabilitySnrThreshold
428 lowSnrReliability = goodReliability[goodLowSnrFlag]
429 self.assertTrue(np.all(lowSnrReliability > lowSnrReliabilityThreshold))
430 glintTrailSources = np.array(filterResults.goodSources["glint_trail"])
431 self.assertTrue(not any(glintTrailSources))
433 # No sources should be removed if the thresholds are turned off
434 runAndCheckFilter(diaSources,
435 snrThreshold=0, lowReliabilitySnrThreshold=0,
436 reliabilityThreshold=0, lowSnrReliabilityThreshold=0)
437 runAndCheckFilter(diaSources,
438 snrThreshold=0, lowReliabilitySnrThreshold=0,
439 reliabilityThreshold=0, lowSnrReliabilityThreshold=0,
440 badFlags=[badFlagName])
441 runAndCheckFilter(diaSources,
442 snrThreshold=2, lowReliabilitySnrThreshold=8,
443 reliabilityThreshold=0, lowSnrReliabilityThreshold=0)
444 runAndCheckFilter(diaSources,
445 snrThreshold=2, lowReliabilitySnrThreshold=8,
446 reliabilityThreshold=0, lowSnrReliabilityThreshold=0.5)
447 runAndCheckFilter(diaSources,
448 snrThreshold=0, lowReliabilitySnrThreshold=0,
449 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5)
450 runAndCheckFilter(diaSources,
451 snrThreshold=2, lowReliabilitySnrThreshold=8,
452 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5)
453 runAndCheckFilter(diaSources,
454 snrThreshold=2, lowReliabilitySnrThreshold=8,
455 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5,
456 badFlags=[badFlagName])
458 def test_mergeEmptyCatalog(self):
459 """Test that a catalog is unchanged if it is merged with an empty
460 catalog.
461 """
462 diaSourcesBase = self.diaSources
464 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
465 task = DiaPipelineTask(config=config)
466 # Include some but not all columns that should be in diaSourcesBase, and some that are mis-matched
467 diaSourcesEmpty = pd.DataFrame(columns=["ra", "dec", "foo"])
468 diaSourcesTest = task.mergeCatalogs(diaSourcesBase, diaSourcesEmpty, tableName="DiaSource")
469 self.assertTrue(diaSourcesBase.equals(diaSourcesTest))
471 def test_mergeCatalogs(self):
472 """Test that a merged catalog is concatenated correctly.
473 """
474 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
475 task = DiaPipelineTask(config=config)
477 diaSourcesBase = convertDataFrameToSdmSchema(task.schema, self.diaSources, "DiaSource")
478 nBase = len(diaSourcesBase)
479 nNew = int(nBase/2)
481 diaSourcesNew = makeDiaSources(nNew, self.diaObjects["diaObjectId"].to_numpy(), self.exposure,
482 self.rng)
483 diaSourcesNew = convertDataFrameToSdmSchema(task.schema, diaSourcesNew, "DiaSource")
484 diaSourcesTest = task.mergeCatalogs(diaSourcesBase, diaSourcesNew, tableName="DiaSource")
485 self.assertEqual(len(diaSourcesTest), nBase + nNew)
486 diaSourcesExtract1 = diaSourcesTest.iloc[:nBase]
487 diaSourcesExtract2 = diaSourcesTest.iloc[nBase:]
489 pd.testing.assert_frame_equal(diaSourcesBase, diaSourcesExtract1)
490 pd.testing.assert_frame_equal(diaSourcesNew, diaSourcesExtract2)
492 def test_updateObjectTable(self):
493 """Test that the diaObject record is updated with the number of
494 diaSources.
495 """
496 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False)
497 task = DiaPipelineTask(config=config)
498 nObjects = 20
499 nSrcPerObject = 10
500 nExtraSources = 5
501 nSources = nSrcPerObject*nObjects + nExtraSources
502 expectedSourcesPerObject = nSrcPerObject*np.ones(nObjects)
503 expectedSourcesPerObject[:nExtraSources] += 1
504 diaObjects = makeDiaObjects(nObjects, self.exposure, self.rng)
505 diaSources = makeDiaSources(nSources, diaObjects["diaObjectId"].to_numpy(), self.exposure, self.rng)
506 updatedDiaObjects = task.updateObjectTable(diaObjects, diaSources)
507 self.assertTrue(np.all(updatedDiaObjects.nDiaSources.values == expectedSourcesPerObject))
510class MemoryTester(lsst.utils.tests.MemoryTestCase):
511 pass
514def setup_module(module):
515 lsst.utils.tests.init()
518if __name__ == "__main__": 518 ↛ 519line 518 didn't jump to line 519 because the condition on line 518 was never true
519 lsst.utils.tests.init()
520 unittest.main()