Coverage for tests/test_diaPipe.py: 20%
123 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-19 02:24 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-19 02:24 -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 unittest
23import numpy as np
24import pandas as pd
26import lsst.afw.image as afwImage
27import lsst.afw.table as afwTable
28from lsst.pipe.base.testUtils import assertValidOutput
29from utils_tests import makeExposure, makeDiaObjects
30import lsst.utils.tests
31import lsst.utils.timer
32from unittest.mock import patch, Mock, MagicMock, DEFAULT
34from lsst.ap.association import DiaPipelineTask
37class TestDiaPipelineTask(unittest.TestCase):
39 @classmethod
40 def _makeDefaultConfig(cls,
41 doPackageAlerts=False,
42 doSolarSystemAssociation=False):
43 config = DiaPipelineTask.ConfigClass()
44 config.apdb.db_url = "sqlite://"
45 config.doPackageAlerts = doPackageAlerts
46 config.doSolarSystemAssociation = doSolarSystemAssociation
47 return config
49 def setUp(self):
50 # schemas are persisted in both Gen 2 and Gen 3 butler as prototypical catalogs
51 srcSchema = afwTable.SourceTable.makeMinimalSchema()
52 srcSchema.addField("base_PixelFlags_flag", type="Flag")
53 srcSchema.addField("base_PixelFlags_flag_offimage", type="Flag")
54 self.srcSchema = afwTable.SourceCatalog(srcSchema)
56 def tearDown(self):
57 pass
59 def testRun(self):
60 """Test running while creating and packaging alerts.
61 """
62 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=True)
64 def testRunWithSolarSystemAssociation(self):
65 """Test running while creating and packaging alerts.
66 """
67 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True)
69 def testRunWithAlerts(self):
70 """Test running while creating and packaging alerts.
71 """
72 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False)
74 def testRunWithoutAlertsOrSolarSystem(self):
75 """Test running without creating and packaging alerts.
76 """
77 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False)
79 def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False):
80 """Test the normal workflow of each ap_pipe step.
81 """
82 config = self._makeDefaultConfig(
83 doPackageAlerts=doPackageAlerts,
84 doSolarSystemAssociation=doSolarSystemAssociation)
85 task = DiaPipelineTask(config=config)
86 # Set DataFrame index testing to always return False. Mocks return
87 # true for this check otherwise.
88 task.testDataFrameIndex = lambda x: False
89 diffIm = Mock(spec=afwImage.ExposureF)
90 exposure = Mock(spec=afwImage.ExposureF)
91 template = Mock(spec=afwImage.ExposureF)
92 diaSrc = MagicMock(spec=pd.DataFrame())
93 ssObjects = MagicMock(spec=pd.DataFrame())
94 ccdExposureIdBits = 32
96 # Each of these subtasks should be called once during diaPipe
97 # execution. We use mocks here to check they are being executed
98 # appropriately.
99 subtasksToMock = [
100 "diaCatalogLoader",
101 "diaCalculation",
102 "diaForcedSource",
103 ]
104 if doPackageAlerts:
105 subtasksToMock.append("alertPackager")
106 else:
107 self.assertFalse(hasattr(task, "alertPackager"))
109 if not doSolarSystemAssociation:
110 self.assertFalse(hasattr(task, "solarSystemAssociator"))
112 def concatMock(_data, **_kwargs):
113 return MagicMock(spec=pd.DataFrame)
115 # Mock out the run() methods of these two Tasks to ensure they
116 # return data in the correct form.
117 @lsst.utils.timer.timeMethod
118 def solarSystemAssociator_run(self, unAssocDiaSources, solarSystemObjectTable, diffIm):
119 return lsst.pipe.base.Struct(nTotalSsObjects=42,
120 nAssociatedSsObjects=30,
121 ssoAssocDiaSources=MagicMock(spec=pd.DataFrame()),
122 unAssocDiaSources=MagicMock(spec=pd.DataFrame()))
124 @lsst.utils.timer.timeMethod
125 def associator_run(self, table, diaObjects, exposure_time=None):
126 return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3,
127 matchedDiaSources=MagicMock(spec=pd.DataFrame()),
128 unAssocDiaSources=MagicMock(spec=pd.DataFrame()),
129 longTrailedSources=None)
131 # apdb isn't a subtask, but still needs to be mocked out for correct
132 # execution in the test environment.
133 with patch.multiple(
134 task, **{task: DEFAULT for task in subtasksToMock + ["apdb"]}
135 ):
136 with patch('lsst.ap.association.diaPipe.pd.concat', new=concatMock), \
137 patch('lsst.ap.association.association.AssociationTask.run', new=associator_run), \
138 patch('lsst.ap.association.ssoAssociation.SolarSystemAssociationTask.run',
139 new=solarSystemAssociator_run):
141 result = task.run(diaSrc,
142 ssObjects,
143 diffIm,
144 exposure,
145 template,
146 ccdExposureIdBits,
147 "g")
148 for subtaskName in subtasksToMock:
149 getattr(task, subtaskName).run.assert_called_once()
150 assertValidOutput(task, result)
151 self.assertEqual(result.apdbMarker.db_url, "sqlite://")
152 meta = task.getFullMetadata()
153 # Check that the expected metadata has been set.
154 self.assertEqual(meta["diaPipe.numUpdatedDiaObjects"], 2)
155 self.assertEqual(meta["diaPipe.numUnassociatedDiaObjects"], 3)
156 # and that associators ran once or not at all.
157 self.assertEqual(len(meta.getArray("diaPipe:associator.associator_runEndUtc")), 1)
158 if doSolarSystemAssociation:
159 self.assertEqual(len(meta.getArray("diaPipe:solarSystemAssociator."
160 "solarSystemAssociator_runEndUtc")), 1)
161 else:
162 self.assertNotIn("diaPipe:solarSystemAssociator", meta)
164 def test_createDiaObjects(self):
165 """Test that creating new DiaObjects works as expected.
166 """
167 nSources = 5
168 diaSources = pd.DataFrame(data=[
169 {"ra": 0.04*idx, "dec": 0.04*idx,
170 "diaSourceId": idx + 1 + nSources, "diaObjectId": 0,
171 "ssObjectId": 0}
172 for idx in range(nSources)])
174 config = self._makeDefaultConfig(doPackageAlerts=False)
175 task = DiaPipelineTask(config=config)
176 result = task.createNewDiaObjects(diaSources)
177 self.assertEqual(nSources, len(result.newDiaObjects))
178 self.assertTrue(np.all(np.equal(
179 result.diaSources["diaObjectId"].to_numpy(),
180 result.diaSources["diaSourceId"].to_numpy())))
181 self.assertTrue(np.all(np.equal(
182 result.newDiaObjects["diaObjectId"].to_numpy(),
183 result.diaSources["diaSourceId"].to_numpy())))
185 def test_purgeDiaObjects(self):
186 """Remove diaOjects that are outside an image's bounding box.
187 """
189 config = self._makeDefaultConfig(doPackageAlerts=False)
190 task = DiaPipelineTask(config=config)
191 exposure = makeExposure(False, False)
192 nObj0 = 20
194 # Create diaObjects
195 diaObjects = makeDiaObjects(nObj0, exposure)
196 # Shrink the bounding box so that some of the diaObjects will be outside
197 bbox = exposure.getBBox()
198 size = np.minimum(bbox.getHeight(), bbox.getWidth())
199 bbox.grow(-size//4)
200 exposureCut = exposure[bbox]
201 sizeCut = np.minimum(bbox.getHeight(), bbox.getWidth())
202 buffer = 10
203 bbox.grow(buffer)
205 def check_diaObjects(bbox, wcs, diaObjects):
206 raVals = diaObjects.ra.to_numpy()
207 decVals = diaObjects.dec.to_numpy()
208 xVals, yVals = wcs.skyToPixelArray(raVals, decVals, degrees=True)
209 selector = bbox.contains(xVals, yVals)
210 return selector
212 selector0 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects)
213 nIn0 = np.count_nonzero(selector0)
214 nOut0 = np.count_nonzero(~selector0)
215 self.assertEqual(nObj0, nIn0 + nOut0)
217 diaObjects1 = task.purgeDiaObjects(exposureCut.getBBox(), exposureCut.getWcs(), diaObjects,
218 buffer=buffer)
219 # Verify that the bounding box was not changed
220 sizeCheck = np.minimum(exposureCut.getBBox().getHeight(), exposureCut.getBBox().getWidth())
221 self.assertEqual(sizeCut, sizeCheck)
222 selector1 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects1)
223 nIn1 = np.count_nonzero(selector1)
224 nOut1 = np.count_nonzero(~selector1)
225 nObj1 = len(diaObjects1)
226 self.assertEqual(nObj1, nIn0)
227 # Verify that not all diaObjects were removed
228 self.assertGreater(nObj1, 0)
229 # Check that some diaObjects were removed
230 self.assertLess(nObj1, nObj0)
231 # Verify that no objects outside the bounding box remain
232 self.assertEqual(nOut1, 0)
233 # Verify that no objects inside the bounding box were removed
234 self.assertEqual(nIn1, nIn0)
237class MemoryTester(lsst.utils.tests.MemoryTestCase):
238 pass
241def setup_module(module):
242 lsst.utils.tests.init()
245if __name__ == "__main__": 245 ↛ 246line 245 didn't jump to line 246, because the condition on line 245 was never true
246 lsst.utils.tests.init()
247 unittest.main()