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 

266 result = task.run(diaSrc, 

267 None, 

268 self.diffim, 

269 self.exposure, 

270 self.template, 

271 # `run` sees None when the preloaded 

272 # connections have been removed. 

273 preloadedDiaObjects=None if doReloadAllApdbCatalogs else self.diaObjects, 

274 preloadedDiaSources=None if doReloadAllApdbCatalogs else self.diaSources, 

275 preloadedDiaForcedSources=( 

276 None if doReloadAllApdbCatalogs else self.diaForcedSources), 

277 band="g", 

278 idGenerator=IdGenerator(), 

279 solarSystemObjectTable=ssObjects) 

280 for subtaskName in subtasksToMock: 

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

282 assertValidOutput(task, result) 

283 # Exact type and contents of apdbMarker are undefined. 

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

285 meta = task.getFullMetadata() 

286 # Check that the expected metadata has been set. 

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

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

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

290 mainRun.assert_called_once() 

291 if doSolarSystemAssociation: 

292 ssRun.assert_called_once() 

293 else: 

294 ssRun.assert_not_called() 

295 # doReloadAllApdbCatalogs implies the DiaObject reload. 

296 if doReloadDiaObjects or doReloadAllApdbCatalogs: 

297 loadObjectsRun.assert_called_once() 

298 else: 

299 loadObjectsRun.assert_not_called() 

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

301 # tasks publish to the same Sasquatch topics. 

302 if doReloadAllApdbCatalogs: 

303 loadSourcesRun.assert_called_once() 

304 loadForcedSourcesRun.assert_called_once() 

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

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

307 else: 

308 loadSourcesRun.assert_not_called() 

309 loadForcedSourcesRun.assert_not_called() 

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

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

312 

313 def test_reloadAllApdbCatalogsRemovesPreloadedConnections(self): 

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

315 """ 

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

317 

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

319 doReloadAllApdbCatalogs=False) 

320 connections = config.connections.ConnectionsClass(config=config) 

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

322 

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

324 doReloadAllApdbCatalogs=True) 

325 connections = config.connections.ConnectionsClass(config=config) 

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

327 

328 def test_runQuantumOmitsRemovedPreloadedInputs(self): 

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

330 

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

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

333 alone when they are present. 

334 """ 

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

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

337 "preloadedDiaSources": self.diaSources, 

338 "preloadedDiaForcedSources": self.diaForcedSources} 

339 

340 def runQuantumKwargs(doReloadAllApdbCatalogs): 

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

342 doReloadAllApdbCatalogs=doReloadAllApdbCatalogs) 

343 task = DiaPipelineTask(config=config) 

344 # The quantum only carries the connections that survived. 

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

346 "diffIm": self.diffim, 

347 "exposure": self.exposure, 

348 "template": self.template, 

349 "solarSystemObjectTable": _makeMockTable()} 

350 if not doReloadAllApdbCatalogs: 

351 inputs.update(loaded) 

352 butlerQC = MagicMock() 

353 butlerQC.get.return_value = inputs 

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

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

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

357 return mockRun.call_args.kwargs 

358 

359 kwargs = runQuantumKwargs(True) 

360 for name in preloaded: 

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

362 

363 kwargs = runQuantumKwargs(False) 

364 for name in preloaded: 

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

366 

367 def test_connectionsImageType(self): 

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

369 """ 

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

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

372 "template": "DifferenceImage"}} 

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

374 with self.subTest(image_type=image_type): 

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

376 input_image_type=image_type) 

377 connections = config.connections.ConnectionsClass(config=config) 

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

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

380 

381 def test_runQuantumFutureImages(self): 

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

383 future-mode image. 

384 """ 

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

386 task = DiaPipelineTask(config=config) 

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

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

389 "solarSystemObjectTable": _makeMockTable(), 

390 "preloadedDiaObjects": self.diaObjects, 

391 "preloadedDiaSources": self.diaSources, 

392 "preloadedDiaForcedSources": self.diaForcedSources} 

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

394 inputs[name] = MagicMock() 

395 inputs[name].to_legacy.return_value = image 

396 butlerQC = MagicMock() 

397 butlerQC.get.return_value = inputs 

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

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

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

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

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

403 

404 def test_runRequiresBandAndIdGenerator(self): 

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

406 """ 

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

408 task = DiaPipelineTask(config=config) 

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

410 "legacySolarSystemTable": None, 

411 "diffIm": self.diffim, 

412 "exposure": self.exposure, 

413 "template": self.template, 

414 "band": "g", 

415 "idGenerator": IdGenerator()} 

416 

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

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

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

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

421 

422 def test_loadRefreshedDiaObjectsNoPreloaded(self): 

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

424 """ 

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

426 doReloadAllApdbCatalogs=True) 

427 task = DiaPipelineTask(config=config) 

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

429 self.diaObjects, 

430 self.diaSources, 

431 self.diaForcedSources) 

432 

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

434 

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

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

437 

438 def test_loadRefreshedDiaSources(self): 

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

440 """ 

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

442 doReloadAllApdbCatalogs=True) 

443 task = DiaPipelineTask(config=config) 

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

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

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

447 self.diaObjects, 

448 self.diaSources, 

449 self.diaForcedSources) 

450 

451 region = getRegion(self.exposure) 

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

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

454 

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

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

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

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

459 self.assertFalse(diaSources.index.has_duplicates) 

460 self.assertFalse(diaForcedSources.index.has_duplicates) 

461 

462 def test_loadRefreshedDiaSourcesNoDiaObjects(self): 

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

464 """ 

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

466 doReloadAllApdbCatalogs=True) 

467 task = DiaPipelineTask(config=config) 

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

469 region = getRegion(self.exposure) 

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

471 

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

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

474 

475 self.assertTrue(diaSources.empty) 

476 self.assertTrue(diaForcedSources.empty) 

477 

478 def test_tooManyDiaObjectsError(self): 

479 maxNewDiaObjects = 100 

480 

481 nDiaSources = maxNewDiaObjects + 1 

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

483 

484 def runAndTestWithContextManager(threshold): 

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

486 doSolarSystemAssociation=False, 

487 filterUnAssociatedSources=False, 

488 maxNewDiaObjects=threshold, 

489 ) 

490 task = DiaPipelineTask(config=config) 

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

492 else contextlib.nullcontext() 

493 with contextManager: 

494 task.associateDiaSources( 

495 diaSources, 

496 None, 

497 None, 

498 self.diaObjects, 

499 ) 

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

501 runAndTestWithContextManager(0) 

502 runAndTestWithContextManager(maxNewDiaObjects - 1) 

503 runAndTestWithContextManager(maxNewDiaObjects) 

504 runAndTestWithContextManager(maxNewDiaObjects + 1) 

505 

506 def test_createDiaObjects(self): 

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

508 """ 

509 nSources = 5 

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

511 task = DiaPipelineTask(config=config) 

512 diaSources = pd.DataFrame(data=[ 

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

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

515 "ssObjectId": 0} 

516 for idx in range(nSources)]) 

517 

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

519 skipIndex=True)) 

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

521 self.assertTrue(np.all(np.equal( 

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

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

524 self.assertTrue(np.all(np.equal( 

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

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

527 

528 def test_purgeDiaObjects(self): 

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

530 """ 

531 

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

533 task = DiaPipelineTask(config=config) 

534 exposure = makeExposure(False, False) 

535 nObj0 = 20 

536 

537 # Create diaObjects 

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

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

540 bbox = exposure.getBBox() 

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

542 bbox.grow(-size//4) 

543 exposureCut = exposure[bbox] 

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

545 buffer = 10 

546 bbox.grow(buffer) 

547 

548 def check_diaObjects(bbox, wcs, diaObjects): 

549 raVals = diaObjects.ra.to_numpy() 

550 decVals = diaObjects.dec.to_numpy() 

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

552 selector = bbox.contains(xVals, yVals) 

553 return selector 

554 

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

556 nIn0 = np.count_nonzero(selector0) 

557 nOut0 = np.count_nonzero(~selector0) 

558 self.assertEqual(nObj0, nIn0 + nOut0) 

559 

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

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

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

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

564 diaObjectIds=diaObjectIds0, buffer=buffer) 

565 diaObjectIds1 = diaObjects1["diaObjectId"] 

566 # Verify that the bounding box was not changed 

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

568 self.assertEqual(sizeCut, sizeCheck) 

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

570 nIn1 = np.count_nonzero(selector1) 

571 nOut1 = np.count_nonzero(~selector1) 

572 nObj1 = len(diaObjects1) 

573 self.assertEqual(nObj1, nIn0) 

574 # Verify that not all diaObjects were removed 

575 self.assertGreater(nObj1, 0) 

576 # Check that some diaObjects were removed 

577 self.assertLess(nObj1, nObj0) 

578 # Verify that no objects outside the bounding box remain 

579 self.assertEqual(nOut1, 0) 

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

581 self.assertEqual(nIn1, nIn0) 

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

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

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

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

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

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

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

589 # Some IDs should have been removed 

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

591 

592 def test_filterDiaObjects(self): 

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

594 Glint trail sources should also be filtered out. 

595 """ 

596 

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

598 doPackageAlerts=False, 

599 filterUnAssociatedSources=True) 

600 

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

602 doPackageAlerts=False, 

603 filterUnAssociatedSources=True, 

604 newObjectFluxField="notAFlux") 

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

606 doPackageAlerts=False, 

607 filterUnAssociatedSources=True, 

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

609 with self.assertRaises(pipeBase.InvalidQuantumError): 

610 DiaPipelineTask(config=configBadFilter) 

611 with self.assertRaises(pipeBase.InvalidQuantumError): 

612 DiaPipelineTask(config=configBadFlag) 

613 task = DiaPipelineTask(config=config) 

614 nUnassociatedDiaSources = 234 

615 

616 # Create diaSources 

617 diaSources = makeDiaSources(nUnassociatedDiaSources, 

618 np.zeros(nUnassociatedDiaSources), 

619 self.exposure, 

620 self.rng, 

621 flagList=task.config.newObjectBadFlags) 

622 reliability = self.rng.random(nUnassociatedDiaSources) 

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

624 fluxErr = np.sqrt(flux) 

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

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

627 diaSources["reliability"] = reliability 

628 diaSources["reliabilityVersion"] = '99.42' 

629 diaSources[config.newObjectFluxField] = flux 

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

631 diaSources["glint_trail"] = glint_trail 

632 badFlagName = task.config.newObjectBadFlags[0] 

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

634 nBadFlags = 20 

635 badFlags[0:nBadFlags] = True 

636 diaSources[badFlagName] = badFlags 

637 

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

639 reliabilityThreshold=None, lowSnrReliabilityThreshold=None, 

640 badFlags=None, 

641 ): 

642 

643 filterResults = task.filterSources( 

644 diaSources.copy(deep=True), 

645 snrThreshold=snrThreshold, 

646 lowReliabilitySnrThreshold=lowReliabilitySnrThreshold, 

647 reliabilityThreshold=reliabilityThreshold, 

648 lowSnrReliabilityThreshold=lowSnrReliabilityThreshold, 

649 badFlags=badFlags, 

650 ) 

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

652 nUnassociatedDiaSources) 

653 goodFlux = filterResults.goodSources[config.newObjectFluxField] 

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

655 goodSnr = np.array(goodFlux/goodFluxErr) 

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

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

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

659 goodLowSnrFlag = goodSnr < lowReliabilitySnrThreshold 

660 lowSnrReliability = goodReliability[goodLowSnrFlag] 

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

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

663 self.assertTrue(not any(glintTrailSources)) 

664 

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

666 runAndCheckFilter(diaSources, 

667 snrThreshold=0, lowReliabilitySnrThreshold=0, 

668 reliabilityThreshold=0, lowSnrReliabilityThreshold=0) 

669 runAndCheckFilter(diaSources, 

670 snrThreshold=0, lowReliabilitySnrThreshold=0, 

671 reliabilityThreshold=0, lowSnrReliabilityThreshold=0, 

672 badFlags=[badFlagName]) 

673 runAndCheckFilter(diaSources, 

674 snrThreshold=2, lowReliabilitySnrThreshold=8, 

675 reliabilityThreshold=0, lowSnrReliabilityThreshold=0) 

676 runAndCheckFilter(diaSources, 

677 snrThreshold=2, lowReliabilitySnrThreshold=8, 

678 reliabilityThreshold=0, lowSnrReliabilityThreshold=0.5) 

679 runAndCheckFilter(diaSources, 

680 snrThreshold=0, lowReliabilitySnrThreshold=0, 

681 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5) 

682 runAndCheckFilter(diaSources, 

683 snrThreshold=2, lowReliabilitySnrThreshold=8, 

684 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5) 

685 runAndCheckFilter(diaSources, 

686 snrThreshold=2, lowReliabilitySnrThreshold=8, 

687 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5, 

688 badFlags=[badFlagName]) 

689 

690 def testRunWithForcedMeasurement(self): 

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

692 

693 reliabilityThreshold = 0.5 

694 trailLengthThreshold = 1.0 

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

696 doPackageAlerts=False, 

697 forcedReliabilityThreshold=reliabilityThreshold, 

698 forcedTrailLengthThreshold=trailLengthThreshold) 

699 task = DiaPipelineTask(config=config) 

700 nDeepObjects = 20 

701 nShallowObjects = 20 

702 nGoodDiaSourcesDeep = 100 

703 nGoodDiaSourcesShallow = nShallowObjects 

704 nGoodDiaSources = nGoodDiaSourcesDeep + nGoodDiaSourcesShallow 

705 nBadDiaSources = 200 

706 

707 # Create diaObjects 

708 diaObjectsDeep = makeDiaObjects(nDeepObjects, self.exposure, self.rng, startId=1) 

709 diaObjectsShallow = makeDiaObjects(nShallowObjects, self.exposure, self.rng, startId=1 + nDeepObjects) 

710 diaObjects = task.mergeCatalogs(diaObjectsDeep, diaObjectsShallow, tableName="DiaObject") 

711 

712 diaSourcesGoodDeep = makeDiaSources(nGoodDiaSourcesDeep, diaObjectsDeep["diaObjectId"].to_numpy(), 

713 self.exposure, self.rng, startId=1, 

714 flagList=config.forcedBadFlags) 

715 diaSourcesGoodShallow = makeDiaSources(nGoodDiaSourcesShallow, 

716 diaObjectsShallow["diaObjectId"].to_numpy(), 

717 self.exposure, self.rng, startId=1 + nGoodDiaSourcesDeep, 

718 flagList=config.forcedBadFlags) 

719 

720 diaSourcesGoodDeep = convertDataFrameToSdmSchema(task.schema, diaSourcesGoodDeep, 

721 tableName="DiaSource", skipIndex=True) 

722 diaSourcesGood = task.mergeCatalogs(diaSourcesGoodDeep, diaSourcesGoodShallow, tableName="DiaSource")