Coverage for tests/test_diaPipe.py: 20%

123 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-03 12:06 +0000

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/>. 

21 

22import unittest 

23import numpy as np 

24import pandas as pd 

25 

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 

33 

34from lsst.ap.association import DiaPipelineTask 

35 

36 

37class TestDiaPipelineTask(unittest.TestCase): 

38 

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 

48 

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) 

55 

56 def tearDown(self): 

57 pass 

58 

59 def testRun(self): 

60 """Test running while creating and packaging alerts. 

61 """ 

62 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=True) 

63 

64 def testRunWithSolarSystemAssociation(self): 

65 """Test running while creating and packaging alerts. 

66 """ 

67 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True) 

68 

69 def testRunWithAlerts(self): 

70 """Test running while creating and packaging alerts. 

71 """ 

72 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False) 

73 

74 def testRunWithoutAlertsOrSolarSystem(self): 

75 """Test running without creating and packaging alerts. 

76 """ 

77 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False) 

78 

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 

95 

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")) 

108 

109 if not doSolarSystemAssociation: 

110 self.assertFalse(hasattr(task, "solarSystemAssociator")) 

111 

112 def concatMock(_data, **_kwargs): 

113 return MagicMock(spec=pd.DataFrame) 

114 

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())) 

123 

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) 

130 

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): 

140 

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) 

163 

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)]) 

173 

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()))) 

184 

185 def test_purgeDiaObjects(self): 

186 """Remove diaOjects that are outside an image's bounding box. 

187 """ 

188 

189 config = self._makeDefaultConfig(doPackageAlerts=False) 

190 task = DiaPipelineTask(config=config) 

191 exposure = makeExposure(False, False) 

192 nObj0 = 20 

193 

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) 

204 

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 

211 

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) 

216 

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) 

235 

236 

237class MemoryTester(lsst.utils.tests.MemoryTestCase): 

238 pass 

239 

240 

241def setup_module(module): 

242 lsst.utils.tests.init() 

243 

244 

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()