Coverage for tests/test_diaPipe.py: 20%

157 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-07 03:43 -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/>. 

21 

22import tempfile 

23import unittest 

24from unittest.mock import patch, Mock, MagicMock, DEFAULT 

25import warnings 

26 

27import numpy as np 

28import pandas as pd 

29 

30import lsst.afw.image as afwImage 

31import lsst.afw.table as afwTable 

32import lsst.dax.apdb as daxApdb 

33import lsst.pex.config as pexConfig 

34import lsst.utils.tests 

35from lsst.pipe.base.testUtils import assertValidOutput 

36 

37from lsst.ap.association import DiaPipelineTask 

38from utils_tests import makeExposure, makeDiaObjects 

39 

40 

41def _makeMockDataFrame(): 

42 """Create a new mock of a DataFrame. 

43 

44 Returns 

45 ------- 

46 mock : `unittest.mock.Mock` 

47 A mock guaranteed to accept all operations used by `pandas.DataFrame`. 

48 """ 

49 with warnings.catch_warnings(): 

50 # spec triggers deprecation warnings on DataFrame, but will 

51 # automatically adapt to any removals. 

52 warnings.simplefilter("ignore", category=DeprecationWarning) 

53 return MagicMock(spec=pd.DataFrame()) 

54 

55 

56class TestDiaPipelineTask(unittest.TestCase): 

57 

58 @classmethod 

59 def _makeDefaultConfig(cls, 

60 config_file, 

61 doPackageAlerts=False, 

62 doSolarSystemAssociation=False): 

63 config = DiaPipelineTask.ConfigClass() 

64 config.doConfigureApdb = False 

65 config.apdb_config_url = config_file 

66 config.doPackageAlerts = doPackageAlerts 

67 config.doSolarSystemAssociation = doSolarSystemAssociation 

68 return config 

69 

70 def setUp(self): 

71 # Create an instance of random generator with fixed seed. 

72 rng = np.random.default_rng(1234) 

73 self.rng = rng 

74 

75 # schemas are persisted in both Gen 2 and Gen 3 butler as prototypical catalogs 

76 srcSchema = afwTable.SourceTable.makeMinimalSchema() 

77 srcSchema.addField("base_PixelFlags_flag", type="Flag") 

78 srcSchema.addField("base_PixelFlags_flag_offimage", type="Flag") 

79 self.srcSchema = afwTable.SourceCatalog(srcSchema) 

80 

81 apdb_config = daxApdb.ApdbSql.init_database(db_url="sqlite://") 

82 self.config_file = tempfile.NamedTemporaryFile() 

83 self.addCleanup(self.config_file.close) 

84 apdb_config.save(self.config_file.name) 

85 

86 # TODO: remove on DM-43419 

87 def testConfigApdbNestedOk(self): 

88 config = DiaPipelineTask.ConfigClass() 

89 config.doConfigureApdb = True 

90 with self.assertWarns(FutureWarning): 

91 config.apdb.db_url = "sqlite://" 

92 config.freeze() 

93 config.validate() 

94 

95 # TODO: remove on DM-43419 

96 def testConfigApdbNestedInvalid(self): 

97 config = DiaPipelineTask.ConfigClass() 

98 config.doConfigureApdb = True 

99 # Don't set db_url 

100 config.freeze() 

101 with self.assertRaises(pexConfig.FieldValidationError): 

102 config.validate() 

103 

104 # TODO: remove on DM-43419 

105 def testConfigApdbFileOk(self): 

106 config = DiaPipelineTask.ConfigClass() 

107 config.doConfigureApdb = False 

108 config.apdb_config_url = "some/file/path.yaml" 

109 config.freeze() 

110 config.validate() 

111 

112 # TODO: remove on DM-43419 

113 def testConfigApdbFileInvalid(self): 

114 config = DiaPipelineTask.ConfigClass() 

115 config.doConfigureApdb = False 

116 # Don't set apdb_config_url 

117 config.freeze() 

118 with self.assertRaises(pexConfig.FieldValidationError): 

119 config.validate() 

120 

121 def testRun(self): 

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

123 """ 

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

125 

126 def testRunWithSolarSystemAssociation(self): 

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

128 """ 

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

130 

131 def testRunWithAlerts(self): 

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

133 """ 

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

135 

136 def testRunWithoutAlertsOrSolarSystem(self): 

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

138 """ 

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

140 

141 def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False): 

142 """Test the normal workflow of each ap_pipe step. 

143 """ 

144 config = self._makeDefaultConfig( 

145 config_file=self.config_file.name, 

146 doPackageAlerts=doPackageAlerts, 

147 doSolarSystemAssociation=doSolarSystemAssociation) 

148 task = DiaPipelineTask(config=config) 

149 # Set DataFrame index testing to always return False. Mocks return 

150 # true for this check otherwise. 

151 task.testDataFrameIndex = lambda x: False 

152 diffIm = Mock(spec=afwImage.ExposureF) 

153 exposure = Mock(spec=afwImage.ExposureF) 

154 template = Mock(spec=afwImage.ExposureF) 

155 diaSrc = _makeMockDataFrame() 

156 ssObjects = _makeMockDataFrame() 

157 ccdExposureIdBits = 32 

158 

159 # Each of these subtasks should be called once during diaPipe 

160 # execution. We use mocks here to check they are being executed 

161 # appropriately. 

162 subtasksToMock = [ 

163 "diaCatalogLoader", 

164 "diaCalculation", 

165 "diaForcedSource", 

166 ] 

167 if doPackageAlerts: 

168 subtasksToMock.append("alertPackager") 

169 else: 

170 self.assertFalse(hasattr(task, "alertPackager")) 

171 

172 if not doSolarSystemAssociation: 

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

174 

175 def concatMock(_data, **_kwargs): 

176 return _makeMockDataFrame() 

177 

178 # Mock out the run() methods of these two Tasks to ensure they 

179 # return data in the correct form. 

180 def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, diffIm): 

181 return lsst.pipe.base.Struct(nTotalSsObjects=42, 

182 nAssociatedSsObjects=30, 

183 ssoAssocDiaSources=_makeMockDataFrame(), 

184 unAssocDiaSources=_makeMockDataFrame()) 

185 

186 def associator_run(table, diaObjects, exposure_time=None): 

187 return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3, 

188 matchedDiaSources=_makeMockDataFrame(), 

189 unAssocDiaSources=_makeMockDataFrame(), 

190 longTrailedSources=None) 

191 

192 # apdb isn't a subtask, but still needs to be mocked out for correct 

193 # execution in the test environment. 

194 with patch.multiple(task, **{task: DEFAULT for task in subtasksToMock + ["apdb"]}), \ 

195 patch('lsst.ap.association.diaPipe.pd.concat', side_effect=concatMock), \ 

196 patch('lsst.ap.association.association.AssociationTask.run', 

197 side_effect=associator_run) as mainRun, \ 

198 patch('lsst.ap.association.ssoAssociation.SolarSystemAssociationTask.run', 

199 side_effect=solarSystemAssociator_run) as ssRun: 

200 

201 result = task.run(diaSrc, 

202 ssObjects, 

203 diffIm, 

204 exposure, 

205 template, 

206 ccdExposureIdBits, 

207 "g") 

208 for subtaskName in subtasksToMock: 

209 getattr(task, subtaskName).run.assert_called_once() 

210 assertValidOutput(task, result) 

211 # Exact type and contents of apdbMarker are undefined. 

212 self.assertIsInstance(result.apdbMarker, pexConfig.Config) 

213 meta = task.getFullMetadata() 

214 # Check that the expected metadata has been set. 

215 self.assertEqual(meta["diaPipe.numUpdatedDiaObjects"], 2) 

216 self.assertEqual(meta["diaPipe.numUnassociatedDiaObjects"], 3) 

217 # and that associators ran once or not at all. 

218 mainRun.assert_called_once() 

219 if doSolarSystemAssociation: 

220 ssRun.assert_called_once() 

221 else: 

222 ssRun.assert_not_called() 

223 

224 def test_createDiaObjects(self): 

225 """Test that creating new DiaObjects works as expected. 

226 """ 

227 nSources = 5 

228 diaSources = pd.DataFrame(data=[ 

229 {"ra": 0.04*idx, "dec": 0.04*idx, 

230 "diaSourceId": idx + 1 + nSources, "diaObjectId": 0, 

231 "ssObjectId": 0} 

232 for idx in range(nSources)]) 

233 

234 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False) 

235 task = DiaPipelineTask(config=config) 

236 result = task.createNewDiaObjects(diaSources) 

237 self.assertEqual(nSources, len(result.newDiaObjects)) 

238 self.assertTrue(np.all(np.equal( 

239 result.diaSources["diaObjectId"].to_numpy(), 

240 result.diaSources["diaSourceId"].to_numpy()))) 

241 self.assertTrue(np.all(np.equal( 

242 result.newDiaObjects["diaObjectId"].to_numpy(), 

243 result.diaSources["diaSourceId"].to_numpy()))) 

244 

245 def test_purgeDiaObjects(self): 

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

247 """ 

248 

249 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False) 

250 task = DiaPipelineTask(config=config) 

251 exposure = makeExposure(False, False) 

252 nObj0 = 20 

253 

254 # Create diaObjects 

255 diaObjects = makeDiaObjects(nObj0, exposure, self.rng) 

256 # Shrink the bounding box so that some of the diaObjects will be outside 

257 bbox = exposure.getBBox() 

258 size = np.minimum(bbox.getHeight(), bbox.getWidth()) 

259 bbox.grow(-size//4) 

260 exposureCut = exposure[bbox] 

261 sizeCut = np.minimum(bbox.getHeight(), bbox.getWidth()) 

262 buffer = 10 

263 bbox.grow(buffer) 

264 

265 def check_diaObjects(bbox, wcs, diaObjects): 

266 raVals = diaObjects.ra.to_numpy() 

267 decVals = diaObjects.dec.to_numpy() 

268 xVals, yVals = wcs.skyToPixelArray(raVals, decVals, degrees=True) 

269 selector = bbox.contains(xVals, yVals) 

270 return selector 

271 

272 selector0 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects) 

273 nIn0 = np.count_nonzero(selector0) 

274 nOut0 = np.count_nonzero(~selector0) 

275 self.assertEqual(nObj0, nIn0 + nOut0) 

276 

277 diaObjects1 = task.purgeDiaObjects(exposureCut.getBBox(), exposureCut.getWcs(), diaObjects, 

278 buffer=buffer) 

279 # Verify that the bounding box was not changed 

280 sizeCheck = np.minimum(exposureCut.getBBox().getHeight(), exposureCut.getBBox().getWidth()) 

281 self.assertEqual(sizeCut, sizeCheck) 

282 selector1 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects1) 

283 nIn1 = np.count_nonzero(selector1) 

284 nOut1 = np.count_nonzero(~selector1) 

285 nObj1 = len(diaObjects1) 

286 self.assertEqual(nObj1, nIn0) 

287 # Verify that not all diaObjects were removed 

288 self.assertGreater(nObj1, 0) 

289 # Check that some diaObjects were removed 

290 self.assertLess(nObj1, nObj0) 

291 # Verify that no objects outside the bounding box remain 

292 self.assertEqual(nOut1, 0) 

293 # Verify that no objects inside the bounding box were removed 

294 self.assertEqual(nIn1, nIn0) 

295 

296 

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

298 pass 

299 

300 

301def setup_module(module): 

302 lsst.utils.tests.init() 

303 

304 

305if __name__ == "__main__": 305 ↛ 306line 305 didn't jump to line 306, because the condition on line 305 was never true

306 lsst.utils.tests.init() 

307 unittest.main()