Coverage for tests/test_diaPipe.py: 99%

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 contextlib 

23import tempfile 

24import unittest 

25from unittest.mock import patch, MagicMock, DEFAULT 

26import warnings 

27 

28import numpy as np 

29import pandas as pd 

30import astropy.table as tb 

31import astropy.units as u 

32 

33import lsst.afw.table as afwTable 

34import lsst.dax.apdb as daxApdb 

35from lsst.meas.base import DetectorVisitIdGeneratorConfig, IdGenerator 

36import lsst.pex.config as pexConfig 

37import lsst.pipe.base as pipeBase 

38import lsst.utils.tests 

39from lsst.pipe.base.testUtils import assertValidOutput 

40 

41from lsst.ap.association import DiaPipelineTask 

42from lsst.ap.association.utils import getRegion 

43from lsst.pipe.tasks.schemaUtils import convertDataFrameToSdmSchema 

44from utils_tests import makeExposure, makeDiaObjects, makeDiaSources, makeDiaForcedSources, \ 

45 makeSolarSystemSources 

46 

47 

48def _makeMockDataFrame(): 

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

50 

51 Returns 

52 ------- 

53 mock : `unittest.mock.Mock` 

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

55 """ 

56 with warnings.catch_warnings(): 

57 # spec triggers deprecation warnings on DataFrame, but will 

58 # automatically adapt to any removals. 

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

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

61 

62 

63def _makeMockTable(): 

64 """Create a new mock of a Table. 

65 

66 Returns 

67 ------- 

68 mock : `unittest.mock.Mock` 

69 A mock guaranteed to accept all operations used by `astropy.table.Table`. 

70 """ 

71 with warnings.catch_warnings(): 

72 # spec triggers deprecation warnings on DataFrame, but will 

73 # automatically adapt to any removals. 

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

75 return MagicMock(spec=tb.Table()) 

76 

77 

78class TestDiaPipelineTask(unittest.TestCase): 

79 

80 @classmethod 

81 def _makeDefaultConfig(cls, config_file, **kwargs): 

82 config = DiaPipelineTask.ConfigClass() 

83 config.apdb_config_url = config_file 

84 config.update(**kwargs) 

85 return config 

86 

87 def setUp(self): 

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

89 rng = np.random.default_rng(1234) 

90 self.rng = rng 

91 

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

93 srcSchema = afwTable.SourceTable.makeMinimalSchema() 

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

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

96 self.srcSchema = afwTable.SourceCatalog(srcSchema) 

97 self.exposure = makeExposure(False, False) 

98 self.diffim = makeExposure(False, False) 

99 self.template = makeExposure(False, False) 

100 self.diaObjects = makeDiaObjects(20, self.exposure, rng) 

101 self.diaSources = makeDiaSources( 

102 100, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) 

103 self.diaForcedSources = makeDiaForcedSources( 

104 200, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) 

105 self.ssSources = makeSolarSystemSources( 

106 20, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) 

107 

108 sqlite_file = tempfile.NamedTemporaryFile() 

109 self.addCleanup(sqlite_file.close) 

110 self.config_file = tempfile.NamedTemporaryFile() 

111 self.addCleanup(self.config_file.close) 

112 apdb_config = daxApdb.ApdbSql.init_database(db_url=f"sqlite:///{sqlite_file.name}") 

113 apdb_config.save(self.config_file.name) 

114 

115 def testRun(self): 

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

117 """ 

118 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=True, doReloadDiaObjects=False) 

119 

120 def testRunWithSolarSystemAssociation(self): 

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

122 """ 

123 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True, doReloadDiaObjects=False) 

124 

125 def testRunWithAlerts(self): 

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

127 """ 

128 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False, doReloadDiaObjects=False) 

129 

130 def testRunWithoutAlertsOrSolarSystem(self): 

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

132 """ 

133 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False, doReloadDiaObjects=False) 

134 

135 def testRunWithReload(self): 

136 """Test running with reloading DiaObjects. 

137 """ 

138 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False, doReloadDiaObjects=True) 

139 

140 def testRunWithReloadAndSolarSystem(self): 

141 """Test running with solar system association and reloading DiaObjects. 

142 """ 

143 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True, doReloadDiaObjects=True) 

144 

145 def testRunWithReloadAndAlerts(self): 

146 """Test running with reloading DiaObjects while creating and packaging alerts. 

147 """ 

148 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False, doReloadDiaObjects=True) 

149 

150 def testRunDisableDeprecatedDoRunForcedMeasurement(self): 

151 """Test running with forced sources disabled. 

152 """ 

153 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False, doReloadDiaObjects=True, 

154 doRunForcedMeasurement=False, subtasksToMock=["diaCalculation", ] 

155 ) 

156 

157 def testRunWithReloadAllApdbCatalogs(self): 

158 """Test running with reloading the DiaSource history. 

159 """ 

160 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False, doReloadDiaObjects=True, 

161 doReloadAllApdbCatalogs=True) 

162 

163 def testRunWithReloadAllApdbCatalogsAndAlerts(self): 

164 """Test reloading the DiaSource history while packaging alerts. 

165 """ 

166 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False, doReloadDiaObjects=True, 

167 doReloadAllApdbCatalogs=True) 

168 

169 def testRunWithReloadAllApdbCatalogsOnly(self): 

170 """Test that reloading everything implies the DiaObject reload. 

171 """ 

172 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False, doReloadDiaObjects=False, 

173 doReloadAllApdbCatalogs=True) 

174 

175 def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False, 

176 doReloadDiaObjects=False, doReloadAllApdbCatalogs=False, subtasksToMock=None, **kwargs): 

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

178 """ 

179 config = self._makeDefaultConfig( 

180 config_file=self.config_file.name, 

181 doPackageAlerts=doPackageAlerts, 

182 doSolarSystemAssociation=doSolarSystemAssociation, 

183 doReloadDiaObjects=doReloadDiaObjects, 

184 doReloadAllApdbCatalogs=doReloadAllApdbCatalogs, 

185 **kwargs 

186 ) 

187 task = DiaPipelineTask(config=config) 

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

189 # true for this check otherwise. 

190 task.testDataFrameIndex = lambda x: False 

191 diaSrc = _makeMockDataFrame() 

192 ssObjects = _makeMockTable() 

193 

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

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

196 # appropriately. 

197 if subtasksToMock is None: 

198 subtasksToMock = ["diaCalculation", "diaForcedSource", ] 

199 if doPackageAlerts: 

200 subtasksToMock.append("alertPackager") 

201 else: 

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

203 

204 if not doSolarSystemAssociation: 

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

206 

207 def concatMock(_data, **_kwargs): 

208 return _makeMockDataFrame() 

209 

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

211 # return data in the correct form. 

212 def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, visitInfo, 

213 bbox, wcs): 

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

215 nAssociatedSsObjects=30, 

216 ssoAssocDiaSources=_makeMockTable(), 

217 unAssocDiaSources=_makeMockTable(), 

218 associatedSsSources=_makeMockTable(), 

219 unassociatedSsObjects=_makeMockTable()) 

220 

221 def associator_run(table, diaObjects, schema=None): 

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

223 matchedDiaSources=_makeMockDataFrame(), 

224 unAssocDiaSources=_makeMockDataFrame()) 

225 

226 def loadObjects_run(region, preloadedDiaObjects): 

227 task.metadata['loadRefreshedDiaObjectsStartUtc'] = 1.234 

228 task.metadata['loadRefreshedDiaObjectsEndUtc'] = 5.678 

229 return self.diaObjects 

230 

231 def loadSources_run(region, diaObjects, visitTime): 

232 task.metadata['loadRefreshedDiaSourcesStartUtc'] = 1.234 

233 task.metadata['loadRefreshedDiaSourcesEndUtc'] = 5.678 

234 return self.diaSources 

235 

236 def loadForcedSources_run(region, diaObjects, visitTime): 

237 task.metadata['loadRefreshedDiaForcedSourcesStartUtc'] = 1.234 

238 task.metadata['loadRefreshedDiaForcedSourcesEndUtc'] = 5.678 

239 return self.diaForcedSources 

240 

241 def updateObjectTableMock(diaObjects, diaSources): 

242 pass 

243 

244 def _selectGoodDiaObjects(diaObjectCat, mergedDiaSourceHistory): 

245 return diaObjectCat.copy(deep=True) 

246 

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

248 # execution in the test environment. 

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

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

251 patch('lsst.ap.association.diaPipe.DiaPipelineTask.updateObjectTable', 

252 side_effect=updateObjectTableMock), \ 

253 patch('lsst.ap.association.diaPipe.DiaPipelineTask._selectGoodDiaObjects', 

254 side_effect=_selectGoodDiaObjects), \ 

255 patch('lsst.ap.association.diaPipe.DiaPipelineTask.loadRefreshedDiaObjects', 

256 side_effect=loadObjects_run) as loadObjectsRun, \ 

257 patch('lsst.ap.association.diaPipe.DiaPipelineTask.loadRefreshedDiaSources', 

258 side_effect=loadSources_run) as loadSourcesRun, \ 

259 patch('lsst.ap.association.diaPipe.DiaPipelineTask.loadRefreshedDiaForcedSources', 

260 side_effect=loadForcedSources_run) as loadForcedSourcesRun, \ 

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

262 side_effect=associator_run) as mainRun, \ 

263 patch('lsst.pipe.tasks.ssoAssociation.SolarSystemAssociationTask.run', 

264 side_effect=solarSystemAssociator_run) as ssRun: 

265 if doPackageAlerts: 

266 task.alertPackager.findMissingDiaSourceColumns.return_value = [] 

267 

268 result = task.run(diaSrc, 

269 None, 

270 self.diffim, 

271 self.exposure, 

272 self.template, 

273 # `run` sees None when the preloaded 

274 # connections have been removed. 

275 preloadedDiaObjects=None if doReloadAllApdbCatalogs else self.diaObjects, 

276 preloadedDiaSources=None if doReloadAllApdbCatalogs else self.diaSources, 

277 preloadedDiaForcedSources=( 

278 None if doReloadAllApdbCatalogs else self.diaForcedSources), 

279 band="g", 

280 idGenerator=IdGenerator(), 

281 solarSystemObjectTable=ssObjects) 

282 for subtaskName in subtasksToMock: 

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

284 assertValidOutput(task, result) 

285 # Exact type and contents of apdbMarker are undefined. 

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

287 meta = task.getFullMetadata() 

288 # Check that the expected metadata has been set. 

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

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

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

292 mainRun.assert_called_once() 

293 if doSolarSystemAssociation: 

294 ssRun.assert_called_once() 

295 else: 

296 ssRun.assert_not_called() 

297 # doReloadAllApdbCatalogs implies the DiaObject reload. 

298 if doReloadDiaObjects or doReloadAllApdbCatalogs: 

299 loadObjectsRun.assert_called_once() 

300 else: 

301 loadObjectsRun.assert_not_called() 

302 # These key names are shared with LoadDiaCatalogsTask so that both 

303 # tasks publish to the same Sasquatch topics. 

304 if doReloadAllApdbCatalogs: 

305 loadSourcesRun.assert_called_once() 

306 loadForcedSourcesRun.assert_called_once() 

307 self.assertGreater(meta["diaPipe.loadDiaSourcesDuration"], 0) 

308 self.assertGreater(meta["diaPipe.loadDiaForcedSourcesDuration"], 0) 

309 else: 

310 loadSourcesRun.assert_not_called() 

311 loadForcedSourcesRun.assert_not_called() 

312 self.assertEqual(meta["diaPipe.loadDiaSourcesDuration"], -1) 

313 self.assertEqual(meta["diaPipe.loadDiaForcedSourcesDuration"], -1) 

314 

315 def test_runRejectsIncompleteDiaSourceHistory(self): 

316 """Test that DiaSource history lacking columns that alerts require 

317 stops processing before anything is written to the APDB. 

318 """ 

319 config = self._makeDefaultConfig(config_file=self.config_file.name, 

320 doPackageAlerts=True, 

321 doReloadDiaObjects=False) 

322 task = DiaPipelineTask(config=config) 

323 history = self.diaSources.drop(columns=["templateVersion", "mixedInstruments"]) 

324 with patch.object(task, "apdb") as apdb, \ 

325 self.assertRaisesRegex(ValueError, "before writing to the APDB.*: " 

326 "templateVersion, mixedInstruments\\."): 

327 task.run(_makeMockDataFrame(), 

328 None, 

329 self.diffim, 

330 self.exposure, 

331 self.template, 

332 preloadedDiaObjects=self.diaObjects, 

333 preloadedDiaSources=history, 

334 preloadedDiaForcedSources=self.diaForcedSources, 

335 band="g", 

336 idGenerator=IdGenerator(), 

337 solarSystemObjectTable=_makeMockTable()) 

338 apdb.store.assert_not_called() 

339 

340 def test_reloadAllApdbCatalogsRemovesPreloadedConnections(self): 

341 """Test that reloading drops the preloaded catalog connections. 

342 """ 

343 preloaded = {"preloadedDiaObjects", "preloadedDiaSources", "preloadedDiaForcedSources"} 

344 

345 config = self._makeDefaultConfig(config_file=self.config_file.name, 

346 doReloadAllApdbCatalogs=False) 

347 connections = config.connections.ConnectionsClass(config=config) 

348 self.assertTrue(preloaded <= set(connections.inputs)) 

349 

350 config = self._makeDefaultConfig(config_file=self.config_file.name, 

351 doReloadAllApdbCatalogs=True) 

352 connections = config.connections.ConnectionsClass(config=config) 

353 self.assertFalse(preloaded & set(connections.inputs)) 

354 

355 def test_runQuantumOmitsRemovedPreloadedInputs(self): 

356 """Test that runQuantum passes only the preloaded inputs it has. 

357 

358 `run` defaults the preloaded catalogs to `None`, so runQuantum passes 

359 just the connections that survived, and must leave the loaded catalogs 

360 alone when they are present. 

361 """ 

362 preloaded = ["preloadedDiaObjects", "preloadedDiaSources", "preloadedDiaForcedSources"] 

363 loaded = {"preloadedDiaObjects": self.diaObjects, 

364 "preloadedDiaSources": self.diaSources, 

365 "preloadedDiaForcedSources": self.diaForcedSources} 

366 

367 def runQuantumKwargs(doReloadAllApdbCatalogs): 

368 config = self._makeDefaultConfig(config_file=self.config_file.name, 

369 doReloadAllApdbCatalogs=doReloadAllApdbCatalogs) 

370 task = DiaPipelineTask(config=config) 

371 # The quantum only carries the connections that survived. 

372 inputs = {"diaSourceTable": _makeMockDataFrame(), 

373 "diffIm": self.diffim, 

374 "exposure": self.exposure, 

375 "template": self.template, 

376 "solarSystemObjectTable": _makeMockTable()} 

377 if not doReloadAllApdbCatalogs: 

378 inputs.update(loaded) 

379 butlerQC = MagicMock() 

380 butlerQC.get.return_value = inputs 

381 with patch.object(DetectorVisitIdGeneratorConfig, "apply", return_value=IdGenerator()), \ 

382 patch.object(DiaPipelineTask, "run") as mockRun: 

383 task.runQuantum(butlerQC, MagicMock(), MagicMock()) 

384 return mockRun.call_args.kwargs 

385 

386 kwargs = runQuantumKwargs(True) 

387 for name in preloaded: 

388 self.assertNotIn(name, kwargs, msg=f"{name} should be left to the default") 

389 

390 kwargs = runQuantumKwargs(False) 

391 for name in preloaded: 

392 self.assertIs(kwargs[name], loaded[name], msg=f"{name} should be the loaded catalog") 

393 

394 def test_connectionsImageType(self): 

395 """Test that input_image_type sets the storage class of each image. 

396 """ 

397 expected = {"legacy": {"diffIm": "ExposureF", "exposure": "ExposureF", "template": "ExposureF"}, 

398 "future": {"diffIm": "DifferenceImage", "exposure": "VisitImage", 

399 "template": "DifferenceImage"}} 

400 for image_type, storageClasses in expected.items(): 

401 with self.subTest(image_type=image_type): 

402 config = self._makeDefaultConfig(config_file=self.config_file.name, 

403 input_image_type=image_type) 

404 connections = config.connections.ConnectionsClass(config=config) 

405 for name, storageClass in storageClasses.items(): 

406 self.assertEqual(getattr(connections, name).storageClass, storageClass) 

407 

408 def test_runQuantumFutureImages(self): 

409 """Test that runQuantum passes run() the legacy form of each 

410 future-mode image. 

411 """ 

412 config = self._makeDefaultConfig(config_file=self.config_file.name, input_image_type="future") 

413 task = DiaPipelineTask(config=config) 

414 legacy = {"diffIm": self.diffim, "exposure": self.exposure, "template": self.template} 

415 inputs = {"diaSourceTable": _makeMockDataFrame(), 

416 "solarSystemObjectTable": _makeMockTable(), 

417 "preloadedDiaObjects": self.diaObjects, 

418 "preloadedDiaSources": self.diaSources, 

419 "preloadedDiaForcedSources": self.diaForcedSources} 

420 for name, image in legacy.items(): 

421 inputs[name] = MagicMock() 

422 inputs[name].to_legacy.return_value = image 

423 butlerQC = MagicMock() 

424 butlerQC.get.return_value = inputs 

425 with patch.object(DetectorVisitIdGeneratorConfig, "apply", return_value=IdGenerator()), \ 

426 patch.object(DiaPipelineTask, "run") as mockRun: 

427 task.runQuantum(butlerQC, MagicMock(), MagicMock()) 

428 for name, image in legacy.items(): 

429 self.assertIs(mockRun.call_args.kwargs[name], image) 

430 

431 def test_runRequiresBandAndIdGenerator(self): 

432 """Test that `run` rejects the `None` defaults on required arguments. 

433 """ 

434 config = self._makeDefaultConfig(config_file=self.config_file.name) 

435 task = DiaPipelineTask(config=config) 

436 kwargs = {"diaSourceTable": _makeMockDataFrame(), 

437 "legacySolarSystemTable": None, 

438 "diffIm": self.diffim, 

439 "exposure": self.exposure, 

440 "template": self.template, 

441 "band": "g", 

442 "idGenerator": IdGenerator()} 

443 

444 with self.assertRaisesRegex(ValueError, "band"): 

445 task.run(**(kwargs | {"band": None})) 

446 with self.assertRaisesRegex(ValueError, "idGenerator"): 

447 task.run(**(kwargs | {"idGenerator": None})) 

448 

449 def test_loadRefreshedDiaObjectsNoPreloaded(self): 

450 """Test reloading DiaObjects when no preloaded catalog is available. 

451 """ 

452 config = self._makeDefaultConfig(config_file=self.config_file.name, 

453 doReloadAllApdbCatalogs=True) 

454 task = DiaPipelineTask(config=config) 

455 task.apdb.store(self.exposure.visitInfo.date.toAstropy() - 30*u.day, 

456 self.diaObjects, 

457 self.diaSources, 

458 self.diaForcedSources) 

459 

460 diaObjects = task.loadRefreshedDiaObjects(getRegion(self.exposure)) 

461 

462 self.assertEqual(len(diaObjects), len(self.diaObjects)) 

463 self.assertEqual(diaObjects.index.name, "diaObjectId") 

464 

465 def test_loadRefreshedDiaSources(self): 

466 """Test that the DiaSource history is reloaded from the APDB. 

467 """ 

468 config = self._makeDefaultConfig(config_file=self.config_file.name, 

469 doReloadAllApdbCatalogs=True) 

470 task = DiaPipelineTask(config=config) 

471 visitTime = self.exposure.visitInfo.date.toAstropy() 

472 # Store the history as though it were observed a month earlier. 

473 task.apdb.store(visitTime - 30*u.day, 

474 self.diaObjects, 

475 self.diaSources, 

476 self.diaForcedSources) 

477 

478 region = getRegion(self.exposure) 

479 diaSources = task.loadRefreshedDiaSources(region, self.diaObjects, visitTime) 

480 diaForcedSources = task.loadRefreshedDiaForcedSources(region, self.diaObjects, visitTime) 

481 

482 self.assertEqual(len(diaSources), len(self.diaSources)) 

483 self.assertEqual(len(diaForcedSources), len(self.diaForcedSources)) 

484 self.assertEqual(list(diaSources.index.names), ["diaObjectId", "band", "diaSourceId"]) 

485 self.assertEqual(list(diaForcedSources.index.names), ["diaObjectId", "diaForcedSourceId"]) 

486 self.assertFalse(diaSources.index.has_duplicates) 

487 self.assertFalse(diaForcedSources.index.has_duplicates) 

488 

489 def test_loadRefreshedDiaSourcesNoDiaObjects(self): 

490 """Test reloading the history when no DiaObjects are in range. 

491 """ 

492 config = self._makeDefaultConfig(config_file=self.config_file.name, 

493 doReloadAllApdbCatalogs=True) 

494 task = DiaPipelineTask(config=config) 

495 emptyDiaObjects = self.diaObjects.iloc[:0] 

496 region = getRegion(self.exposure) 

497 visitTime = self.exposure.visitInfo.date.toAstropy() 

498 

499 diaSources = task.loadRefreshedDiaSources(region, emptyDiaObjects, visitTime) 

500 diaForcedSources = task.loadRefreshedDiaForcedSources(region, emptyDiaObjects, visitTime) 

501 

502 self.assertTrue(diaSources.empty) 

503 self.assertTrue(diaForcedSources.empty) 

504 

505 def test_tooManyDiaObjectsError(self): 

506 maxNewDiaObjects = 100 

507 

508 nDiaSources = maxNewDiaObjects + 1 

509 diaSources = makeDiaSources(nDiaSources, np.zeros(nDiaSources), self.exposure, self.rng) 

510 

511 def runAndTestWithContextManager(threshold): 

512 config = self._makeDefaultConfig(config_file=self.config_file.name, 

513 doSolarSystemAssociation=False, 

514 filterUnAssociatedSources=False, 

515 maxNewDiaObjects=threshold, 

516 ) 

517 task = DiaPipelineTask(config=config) 

518 contextManager = self.assertRaises(pipeBase.AlgorithmError) if nDiaSources > threshold > 0 \ 

519 else contextlib.nullcontext() 

520 with contextManager: 

521 task.associateDiaSources( 

522 diaSources, 

523 None, 

524 None, 

525 self.diaObjects, 

526 ) 

527 # Test cases at, above, and below the threshold as well as at 0. 

528 runAndTestWithContextManager(0) 

529 runAndTestWithContextManager(maxNewDiaObjects - 1) 

530 runAndTestWithContextManager(maxNewDiaObjects) 

531 runAndTestWithContextManager(maxNewDiaObjects + 1) 

532 

533 def test_createDiaObjects(self): 

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

535 """ 

536 nSources = 5 

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

538 task = DiaPipelineTask(config=config) 

539 diaSources = pd.DataFrame(data=[ 

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

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

542 "ssObjectId": 0} 

543 for idx in range(nSources)]) 

544 

545 result = task.createNewDiaObjects(convertDataFrameToSdmSchema(task.schema, diaSources, "DiaSource", 

546 skipIndex=True)) 

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

548 self.assertTrue(np.all(np.equal( 

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

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

551 self.assertTrue(np.all(np.equal( 

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

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

554 

555 def test_purgeDiaObjects(self): 

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

557 """ 

558 

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

560 task = DiaPipelineTask(config=config) 

561 exposure = makeExposure(False, False) 

562 nObj0 = 20 

563 

564 # Create diaObjects 

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

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

567 bbox = exposure.getBBox() 

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

569 bbox.grow(-size//4) 

570 exposureCut = exposure[bbox] 

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

572 buffer = 10 

573 bbox.grow(buffer) 

574 

575 def check_diaObjects(bbox, wcs, diaObjects): 

576 raVals = diaObjects.ra.to_numpy() 

577 decVals = diaObjects.dec.to_numpy() 

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

579 selector = bbox.contains(xVals, yVals) 

580 return selector 

581 

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

583 nIn0 = np.count_nonzero(selector0) 

584 nOut0 = np.count_nonzero(~selector0) 

585 self.assertEqual(nObj0, nIn0 + nOut0) 

586 

587 # Add an ID that is not in the diaObject table. It should not get removed. 

588 diaObjectIds0 = diaObjects["diaObjectId"].copy(deep=True) 

589 diaObjectIds0[max(diaObjectIds0.index) + 1] = 999 

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

591 diaObjectIds=diaObjectIds0, buffer=buffer) 

592 diaObjectIds1 = diaObjects1["diaObjectId"] 

593 # Verify that the bounding box was not changed 

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

595 self.assertEqual(sizeCut, sizeCheck) 

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

597 nIn1 = np.count_nonzero(selector1) 

598 nOut1 = np.count_nonzero(~selector1) 

599 nObj1 = len(diaObjects1) 

600 self.assertEqual(nObj1, nIn0) 

601 # Verify that not all diaObjects were removed 

602 self.assertGreater(nObj1, 0) 

603 # Check that some diaObjects were removed 

604 self.assertLess(nObj1, nObj0) 

605 # Verify that no objects outside the bounding box remain 

606 self.assertEqual(nOut1, 0) 

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

608 self.assertEqual(nIn1, nIn0) 

609 # The length of the updated object IDs should equal the number of objects 

610 # plus one, since we added an extra ID. 

611 self.assertEqual(nObj1 + 1, len(objIds)) 

612 # All of the object IDs extracted from the catalog should be in the pruned object IDs 

613 self.assertTrue(set(objIds).issuperset(diaObjectIds1)) 

614 # The pruned object IDs should contain entries that are not in the catalog 

615 self.assertFalse(set(diaObjectIds1).issuperset(objIds)) 

616 # Some IDs should have been removed 

617 self.assertLess(len(objIds), len(diaObjectIds0)) 

618 

619 def test_filterDiaObjects(self): 

620 """Unassociated diaSources that are filtered should have good reliability and SNR. 

621 Glint trail sources should also be filtered out. 

622 """ 

623 

624 config = self._makeDefaultConfig(config_file=self.config_file.name, 

625 doPackageAlerts=False, 

626 filterUnAssociatedSources=True) 

627 

628 configBadFilter = self._makeDefaultConfig(config_file=self.config_file.name, 

629 doPackageAlerts=False, 

630 filterUnAssociatedSources=True, 

631 newObjectFluxField="notAFlux") 

632 configBadFlag = self._makeDefaultConfig(config_file=self.config_file.name, 

633 doPackageAlerts=False, 

634 filterUnAssociatedSources=True, 

635 newObjectBadFlags=("junkSource", "notUsed")) 

636 with self.assertRaises(pipeBase.InvalidQuantumError): 

637 DiaPipelineTask(config=configBadFilter) 

638 with self.assertRaises(pipeBase.InvalidQuantumError): 

639 DiaPipelineTask(config=configBadFlag) 

640 task = DiaPipelineTask(config=config) 

641 nUnassociatedDiaSources = 234 

642 

643 # Create diaSources 

644 diaSources = makeDiaSources(nUnassociatedDiaSources, 

645 np.zeros(nUnassociatedDiaSources), 

646 self.exposure, 

647 self.rng, 

648 flagList=task.config.newObjectBadFlags) 

649 reliability = self.rng.random(nUnassociatedDiaSources) 

650 flux = (self.rng.random(nUnassociatedDiaSources)**2)*100 

651 fluxErr = np.sqrt(flux) 

652 glint_trail = np.zeros(nUnassociatedDiaSources, dtype=bool) 

653 glint_trail[12:16] = True # add 4 glint trail sources 

654 diaSources["reliability"] = reliability 

655 diaSources["reliabilityVersion"] = '99.42' 

656 diaSources[config.newObjectFluxField] = flux 

657 diaSources[config.newObjectFluxField + "Err"] = fluxErr 

658 diaSources["glint_trail"] = glint_trail 

659 badFlagName = task.config.newObjectBadFlags[0] 

660 badFlags = np.zeros(nUnassociatedDiaSources, dtype=bool) 

661 nBadFlags = 20 

662 badFlags[0:nBadFlags] = True 

663 diaSources[badFlagName] = badFlags 

664 

665 def runAndCheckFilter(diaSources, snrThreshold=None, lowReliabilitySnrThreshold=None, 

666 reliabilityThreshold=None, lowSnrReliabilityThreshold=None, 

667 badFlags=None, 

668 ): 

669 

670 filterResults = task.filterSources( 

671 diaSources.copy(deep=True), 

672 snrThreshold=snrThreshold, 

673 lowReliabilitySnrThreshold=lowReliabilitySnrThreshold, 

674 reliabilityThreshold=reliabilityThreshold, 

675 lowSnrReliabilityThreshold=lowSnrReliabilityThreshold, 

676 badFlags=badFlags, 

677 ) 

678 self.assertEqual(len(filterResults.goodSources) + len(filterResults.badSources), 

679 nUnassociatedDiaSources) 

680 goodFlux = filterResults.goodSources[config.newObjectFluxField] 

681 goodFluxErr = filterResults.goodSources[config.newObjectFluxField + "Err"] 

682 goodSnr = np.array(goodFlux/goodFluxErr) 

683 self.assertTrue(np.all(goodSnr > snrThreshold)) 

684 goodReliability = np.array(filterResults.goodSources["reliability"]) 

685 self.assertTrue(np.all(goodReliability > reliabilityThreshold)) 

686 goodLowSnrFlag = goodSnr < lowReliabilitySnrThreshold 

687 lowSnrReliability = goodReliability[goodLowSnrFlag] 

688 self.assertTrue(np.all(lowSnrReliability > lowSnrReliabilityThreshold)) 

689 glintTrailSources = np.array(filterResults.goodSources["glint_trail"]) 

690 self.assertTrue(not any(glintTrailSources)) 

691 

692 # No sources should be removed if the thresholds are turned off 

693 runAndCheckFilter(diaSources, 

694 snrThreshold=0, lowReliabilitySnrThreshold=0, 

695 reliabilityThreshold=0, lowSnrReliabilityThreshold=0) 

696 runAndCheckFilter(diaSources, 

697 snrThreshold=0, lowReliabilitySnrThreshold=0, 

698 reliabilityThreshold=0, lowSnrReliabilityThreshold=0, 

699 badFlags=[badFlagName]) 

700 runAndCheckFilter(diaSources, 

701 snrThreshold=2, lowReliabilitySnrThreshold=8, 

702 reliabilityThreshold=0, lowSnrReliabilityThreshold=0) 

703 runAndCheckFilter(diaSources, 

704 snrThreshold=2, lowReliabilitySnrThreshold=8, 

705 reliabilityThreshold=0, lowSnrReliabilityThreshold=0.5) 

706 runAndCheckFilter(diaSources, 

707 snrThreshold=0, lowReliabilitySnrThreshold=0, 

708 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5) 

709 runAndCheckFilter(diaSources, 

710 snrThreshold=2, lowReliabilitySnrThreshold=8, 

711 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5) 

712 runAndCheckFilter(diaSources, 

713 snrThreshold=2, lowReliabilitySnrThreshold=8, 

714 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5, 

715 badFlags=[badFlagName]) 

716 

717 def testRunWithForcedMeasurement(self): 

718 """Test running association with forced photometry.""" 

719 

720 reliabilityThreshold = 0.5 

721 trailLengthThreshold = 1.0 

722 config = self._makeDefaultConfig(config_file=self.config_file.name, 

723 doPackageAlerts=False, 

724 forcedReliabilityThreshold=reliabilityThreshold, 

725 forcedTrailLengthThreshold=trailLengthThreshold) 

726 task = DiaPipelineTask(config=config)