Coverage for tests/test_diaPipe.py: 13%

1# This file is part of ap_association.

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.

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 contextlib

23import tempfile

24import unittest

25from unittest.mock import patch, MagicMock, DEFAULT

26import warnings

28import numpy as np

29import pandas as pd

30import astropy.table as tb

32import lsst.afw.table as afwTable

33import lsst.dax.apdb as daxApdb

34from lsst.meas.base import IdGenerator

35import lsst.pex.config as pexConfig

36import lsst.pipe.base as pipeBase

37import lsst.utils.tests

38from lsst.pipe.base.testUtils import assertValidOutput

40from lsst.ap.association import DiaPipelineTask

41from lsst.pipe.tasks.schemaUtils import convertDataFrameToSdmSchema

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

43 makeSolarSystemSources

46def _makeMockDataFrame():

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

49 Returns

50 -------

51 mock : `unittest.mock.Mock`

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

53 """

54 with warnings.catch_warnings():

55 # spec triggers deprecation warnings on DataFrame, but will

56 # automatically adapt to any removals.

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

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

61def _makeMockTable():

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

64 Returns

65 -------

66 mock : `unittest.mock.Mock`

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

68 """

69 with warnings.catch_warnings():

70 # spec triggers deprecation warnings on DataFrame, but will

71 # automatically adapt to any removals.

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

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

76class TestDiaPipelineTask(unittest.TestCase):

78 @classmethod

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

80 config = DiaPipelineTask.ConfigClass()

81 config.apdb_config_url = config_file

82 config.update(**kwargs)

83 return config

85 def setUp(self):

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

87 rng = np.random.default_rng(1234)

88 self.rng = rng

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

91 srcSchema = afwTable.SourceTable.makeMinimalSchema()

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

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

94 self.srcSchema = afwTable.SourceCatalog(srcSchema)

95 self.exposure = makeExposure(False, False)

96 self.diffim = makeExposure(False, False)

97 self.template = makeExposure(False, False)

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

99 self.diaSources = makeDiaSources(

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

101 self.diaForcedSources = makeDiaForcedSources(

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

103 self.ssSources = makeSolarSystemSources(

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

105

106 sqlite_file = tempfile.NamedTemporaryFile()

107 self.addCleanup(sqlite_file.close)

108 self.config_file = tempfile.NamedTemporaryFile()

109 self.addCleanup(self.config_file.close)

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

111 apdb_config.save(self.config_file.name)

112

113 def testRun(self):

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

115 """

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

117

118 def testRunWithSolarSystemAssociation(self):

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

120 """

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

122

123 def testRunWithAlerts(self):

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

125 """

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

127

128 def testRunWithoutAlertsOrSolarSystem(self):

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

130 """

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

132

133 def testRunWithReload(self):

134 """Test running with reloading DiaObjects.

135 """

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

137

138 def testRunWithReloadAndSolarSystem(self):

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

140 """

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

142

143 def testRunWithReloadAndAlerts(self):

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

145 """

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

147

148 def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False, doRunForcedMeasurement=False,

149 doReloadDiaObjects=False):

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

151 """

152 config = self._makeDefaultConfig(

153 config_file=self.config_file.name,

154 doPackageAlerts=doPackageAlerts,

155 doSolarSystemAssociation=doSolarSystemAssociation,

156 doRunForcedMeasurement=doRunForcedMeasurement,

157 doReloadDiaObjects=doReloadDiaObjects,

158 )

159 task = DiaPipelineTask(config=config)

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

161 # true for this check otherwise.

162 task.testDataFrameIndex = lambda x: False

163 diaSrc = _makeMockDataFrame()

164 ssObjects = _makeMockTable()

165

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

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

168 # appropriately.

169 subtasksToMock = [

170 "diaCalculation",

171 ]

172 if doPackageAlerts:

173 subtasksToMock.append("alertPackager")

174 else:

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

176

177 if not doSolarSystemAssociation:

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

179

180 def concatMock(_data, **_kwargs):

181 return _makeMockDataFrame()

182

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

184 # return data in the correct form.

185 def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, visitInfo,

186 bbox, wcs):

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

188 nAssociatedSsObjects=30,

189 ssoAssocDiaSources=_makeMockTable(),

190 unAssocDiaSources=_makeMockTable(),

191 associatedSsSources=_makeMockTable(),

192 unassociatedSsObjects=_makeMockTable())

193

194 def associator_run(table, diaObjects):

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

196 matchedDiaSources=_makeMockDataFrame(),

197 unAssocDiaSources=_makeMockDataFrame())

198

199 def loadObjects_run(region, preloadedDiaObjects):

200 task.metadata['loadRefreshedDiaObjectsStartUtc'] = 1.234

201 task.metadata['loadRefreshedDiaObjectsEndUtc'] = 5.678

202 return self.diaObjects

203

204 def updateObjectTableMock(diaObjects, diaSources):

205 pass

206

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

208 # execution in the test environment.

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

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

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

212 side_effect=updateObjectTableMock), \

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

214 side_effect=loadObjects_run), \

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

216 side_effect=associator_run) as mainRun, \

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

218 side_effect=solarSystemAssociator_run) as ssRun:

219

220 result = task.run(diaSrc,

221 None,

222 self.diffim,

223 self.exposure,

224 self.template,

225 preloadedDiaObjects=self.diaObjects,

226 preloadedDiaSources=self.diaSources,

227 preloadedDiaForcedSources=self.diaForcedSources,

228 band="g",

229 idGenerator=IdGenerator(),

230 solarSystemObjectTable=ssObjects)

231 for subtaskName in subtasksToMock:

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

233 assertValidOutput(task, result)

234 # Exact type and contents of apdbMarker are undefined.

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

236 meta = task.getFullMetadata()

237 # Check that the expected metadata has been set.

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

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

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

241 mainRun.assert_called_once()

242 if doSolarSystemAssociation:

243 ssRun.assert_called_once()

244 else:

245 ssRun.assert_not_called()

246

247 def test_tooManyDiaObjectsError(self):

248 maxNewDiaObjects = 100

249

250 nDiaSources = maxNewDiaObjects + 1

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

252

253 def runAndTestWithContextManager(threshold):

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

255 doSolarSystemAssociation=False,

256 filterUnAssociatedSources=False,

257 maxNewDiaObjects=threshold,

258 )

259 task = DiaPipelineTask(config=config)

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

261 else contextlib.nullcontext()

262 with contextManager:

263 task.associateDiaSources(

264 diaSources,

265 None,

266 None,

267 self.diaObjects,

268 )

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

270 runAndTestWithContextManager(0)

271 runAndTestWithContextManager(maxNewDiaObjects - 1)

272 runAndTestWithContextManager(maxNewDiaObjects)

273 runAndTestWithContextManager(maxNewDiaObjects + 1)

274

275 def test_createDiaObjects(self):

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

277 """

278 nSources = 5

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

280 task = DiaPipelineTask(config=config)

281 diaSources = pd.DataFrame(data=[

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

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

284 "ssObjectId": 0}

285 for idx in range(nSources)])

286

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

288 skipIndex=True))

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

290 self.assertTrue(np.all(np.equal(

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

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

293 self.assertTrue(np.all(np.equal(

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

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

296

297 def test_purgeDiaObjects(self):

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

299 """

300

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

302 task = DiaPipelineTask(config=config)

303 exposure = makeExposure(False, False)

304 nObj0 = 20

305

306 # Create diaObjects

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

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

309 bbox = exposure.getBBox()

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

311 bbox.grow(-size//4)

312 exposureCut = exposure[bbox]

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

314 buffer = 10

315 bbox.grow(buffer)

316

317 def check_diaObjects(bbox, wcs, diaObjects):

318 raVals = diaObjects.ra.to_numpy()

319 decVals = diaObjects.dec.to_numpy()

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

321 selector = bbox.contains(xVals, yVals)

322 return selector

323

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

325 nIn0 = np.count_nonzero(selector0)

326 nOut0 = np.count_nonzero(~selector0)

327 self.assertEqual(nObj0, nIn0 + nOut0)

328

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

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

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

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

333 diaObjectIds=diaObjectIds0, buffer=buffer)

334 diaObjectIds1 = diaObjects1["diaObjectId"]

335 # Verify that the bounding box was not changed

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

337 self.assertEqual(sizeCut, sizeCheck)

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

339 nIn1 = np.count_nonzero(selector1)

340 nOut1 = np.count_nonzero(~selector1)

341 nObj1 = len(diaObjects1)

342 self.assertEqual(nObj1, nIn0)

343 # Verify that not all diaObjects were removed

344 self.assertGreater(nObj1, 0)

345 # Check that some diaObjects were removed

346 self.assertLess(nObj1, nObj0)

347 # Verify that no objects outside the bounding box remain

348 self.assertEqual(nOut1, 0)

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

350 self.assertEqual(nIn1, nIn0)

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

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

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

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

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

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

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

358 # Some IDs should have been removed

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

360

361 def test_filterDiaObjects(self):

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

363 Glint trail sources should also be filtered out.

364 """

365

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

367 doPackageAlerts=False,

368 filterUnAssociatedSources=True)

369

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

371 doPackageAlerts=False,

372 filterUnAssociatedSources=True,

373 newObjectFluxField="notAFlux")

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

375 doPackageAlerts=False,

376 filterUnAssociatedSources=True,

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

378 with self.assertRaises(pipeBase.InvalidQuantumError):

379 DiaPipelineTask(config=configBadFilter)

380 with self.assertRaises(pipeBase.InvalidQuantumError):

381 DiaPipelineTask(config=configBadFlag)

382 task = DiaPipelineTask(config=config)

383 nUnassociatedDiaSources = 234

384

385 # Create diaSources

386 diaSources = makeDiaSources(nUnassociatedDiaSources,

387 np.zeros(nUnassociatedDiaSources),

388 self.exposure,

389 self.rng,

390 flagList=task.config.newObjectBadFlags)

391 reliability = self.rng.random(nUnassociatedDiaSources)

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

393 fluxErr = np.sqrt(flux)

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

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

396 diaSources["reliability"] = reliability

397 diaSources[config.newObjectFluxField] = flux

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

399 diaSources["glint_trail"] = glint_trail

400 badFlagName = task.config.newObjectBadFlags[0]

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

402 nBadFlags = 20

403 badFlags[0:nBadFlags] = True

404 diaSources[badFlagName] = badFlags

405

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

407 reliabilityThreshold=None, lowSnrReliabilityThreshold=None,

408 badFlags=None,

409 ):

410

411 filterResults = task.filterSources(

412 diaSources.copy(deep=True),

413 snrThreshold=snrThreshold,

414 lowReliabilitySnrThreshold=lowReliabilitySnrThreshold,

415 reliabilityThreshold=reliabilityThreshold,

416 lowSnrReliabilityThreshold=lowSnrReliabilityThreshold,

417 badFlags=badFlags,

418 )

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

420 nUnassociatedDiaSources)

421 goodFlux = filterResults.goodSources[config.newObjectFluxField]

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

423 goodSnr = np.array(goodFlux/goodFluxErr)

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

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

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

427 goodLowSnrFlag = goodSnr < lowReliabilitySnrThreshold

428 lowSnrReliability = goodReliability[goodLowSnrFlag]

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

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

431 self.assertTrue(not any(glintTrailSources))

432

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

434 runAndCheckFilter(diaSources,

435 snrThreshold=0, lowReliabilitySnrThreshold=0,

436 reliabilityThreshold=0, lowSnrReliabilityThreshold=0)

437 runAndCheckFilter(diaSources,

438 snrThreshold=0, lowReliabilitySnrThreshold=0,

439 reliabilityThreshold=0, lowSnrReliabilityThreshold=0,

440 badFlags=[badFlagName])

441 runAndCheckFilter(diaSources,

442 snrThreshold=2, lowReliabilitySnrThreshold=8,

443 reliabilityThreshold=0, lowSnrReliabilityThreshold=0)

444 runAndCheckFilter(diaSources,

445 snrThreshold=2, lowReliabilitySnrThreshold=8,

446 reliabilityThreshold=0, lowSnrReliabilityThreshold=0.5)

447 runAndCheckFilter(diaSources,

448 snrThreshold=0, lowReliabilitySnrThreshold=0,

449 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5)

450 runAndCheckFilter(diaSources,

451 snrThreshold=2, lowReliabilitySnrThreshold=8,

452 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5)

453 runAndCheckFilter(diaSources,

454 snrThreshold=2, lowReliabilitySnrThreshold=8,

455 reliabilityThreshold=0.1, lowSnrReliabilityThreshold=0.5,

456 badFlags=[badFlagName])

457

458 def testRunWithForcedMeasurement(self):

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

460

461 reliabilityThreshold = 0.5

462 trailLengthThreshold = 1.0

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

464 doPackageAlerts=False,

465 doRunForcedMeasurement=True,

466 forcedReliabilityThreshold=reliabilityThreshold,

467 forcedTrailLengthThreshold=trailLengthThreshold)

468 task = DiaPipelineTask(config=config)

469 nDeepObjects = 20

470 nShallowObjects = 20

471 nGoodDiaSourcesDeep = 100

472 nGoodDiaSourcesShallow = nShallowObjects

473 nGoodDiaSources = nGoodDiaSourcesDeep + nGoodDiaSourcesShallow

474 nBadDiaSources = 200

475

476 # Create diaObjects

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

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

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

480

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

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

483 flagList=config.forcedBadFlags)

484 diaSourcesGoodShallow = makeDiaSources(nGoodDiaSourcesShallow,

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

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

487 flagList=config.forcedBadFlags)

488

489 diaSourcesGoodDeep = convertDataFrameToSdmSchema(task.schema, diaSourcesGoodDeep,

490 tableName="DiaSource", skipIndex=True)

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

492

493 diaSourcesBad = makeDiaSources(nBadDiaSources, diaObjects["diaObjectId"].to_numpy(), self.exposure,

494 self.rng, randomizeObjects=True, startId=1 + nGoodDiaSources,

495 flagList=config.forcedBadFlags)

496 diaSourcesGood['reliability'] = (self.rng.random(nGoodDiaSources)*(1 - reliabilityThreshold)

497 + reliabilityThreshold)

498 diaSourcesGood['trailLength'] = self.rng.random(nGoodDiaSources)*trailLengthThreshold

499 diaSourcesBad = convertDataFrameToSdmSchema(task.schema, diaSourcesBad,

500 tableName="DiaSource", skipIndex=True)

501

502 # Set some "bad" diaSources to have bad reliability, some good

503 diaSourcesBad['reliability'] = self.rng.random(nBadDiaSources)

504 # Set some "bad" diaSources to have too long trail lengths, some acceptible

505 diaSourcesBad['trailLength'] = self.rng.random(nBadDiaSources)*trailLengthThreshold + 0.5

506 missingBadFlags = ((diaSourcesBad['reliability'] > reliabilityThreshold)

507 & (diaSourcesBad['trailLength'] < trailLengthThreshold)

508 )

509 for badFlag in config.forcedBadFlags:

510 # Set a fraction of the bad diaSources to have each flag, assigned randomly

511 diaSourcesBad[badFlag] = self.rng.random(nBadDiaSources) > 0.8

512 missingBadFlags &= ~diaSourcesBad[badFlag]

513 # Catch any "bad" diaSources that have not yet been flagged

514 if np.any(missingBadFlags):

515 diaSourcesBad[badFlag] |= missingBadFlags

516 diaObjectsForcedGoodOnly = task._selectGoodDiaObjects(diaObjects, diaSourcesGood)

517 diaObjectsForcedBadOnly = task._selectGoodDiaObjects(diaObjects, diaSourcesBad)

518 # None of the diaObjects should be selected if only given the bad diaSources

519 self.assertTrue(diaObjectsForcedBadOnly.empty)

520

521 diaSources = task.mergeCatalogs(diaSourcesGood, diaSourcesBad, tableName="DiaSource")

522 diaObjectsForced = task._selectGoodDiaObjects(diaObjects, diaSources)

523 # The number of diaObjects selected should be the same regardless of

524 # whether the bad diaSources are included

525 self.assertEqual(len(diaObjectsForced), len(diaObjectsForcedGoodOnly))

526 # All of the deep diaObjects should be selected, and none of the shallow

527 self.assertEqual(len(diaObjectsForced), nDeepObjects)

528 fSrc = task.runForcedMeasurement(diaObjectsForced, diaObjectsForced, self.exposure, self.exposure,

529 IdGenerator())

530 self.assertEqual(set(fSrc['diaObjectId']), set(diaObjectsDeep['diaObjectId']))

531

532 def test_selectGoodObjects(self):

533 """Test the diaObject selection funtion used for forced photometry.

534 """

535 reliabilityThreshold = 0.5

536 trailLengthThreshold = 1.0

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

538 doPackageAlerts=False,

539 doRunForcedMeasurement=True,

540 forcedReliabilityThreshold=reliabilityThreshold,

541 forcedTrailLengthThreshold=trailLengthThreshold)

542 task = DiaPipelineTask(config=config)

543 nObjects = 20

544 nDiaSources = 100

545

546 # Create diaObjects

547 diaObjects = makeDiaObjects(nObjects, self.exposure, self.rng, startId=1)

548

549 diaSources = makeDiaSources(nDiaSources, diaObjects["diaObjectId"].to_numpy(),

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

551 flagList=config.forcedBadFlags)

552 # Since flagList is not specified, the columns will be missing, and added and filled with NaNs

553 # when put through convertDataFrameToSdmSchema.

554 diaSourcesMissingColumns = makeDiaSources(nDiaSources, diaObjects["diaObjectId"].to_numpy(),

555 self.exposure, self.rng, startId=1 + nDiaSources)

556 # Should raise an error if run before adding the required columns

557 with self.assertRaises(RuntimeError):

558 task._selectGoodDiaObjects(diaObjects, diaSources)

559

560 diaSources['reliability'] = (self.rng.random(nDiaSources)*(1 - reliabilityThreshold)

561 + reliabilityThreshold)

562 diaSources['trailLength'] = self.rng.random(nDiaSources)*trailLengthThreshold

563

564 diaSourcesMissingColumns['reliability'] = (self.rng.random(nDiaSources)*(1 - reliabilityThreshold)

565 + reliabilityThreshold)

566 diaSourcesMissingColumns['trailLength'] = self.rng.random(nDiaSources)*trailLengthThreshold

567

568 diaSourcesNotMatched = makeDiaSources(nDiaSources, diaObjects["diaObjectId"].to_numpy() + 999,

569 self.exposure, self.rng, startId=1)

570

571 diaSourcesNotMatched['reliability'] = (self.rng.random(nDiaSources)*(1 - reliabilityThreshold)

572 + reliabilityThreshold)

573 diaSourcesNotMatched['trailLength'] = self.rng.random(nDiaSources)*trailLengthThreshold

574

575 diaSources = convertDataFrameToSdmSchema(task.schema, diaSources,

576 tableName="DiaSource", skipIndex=True)

577 diaSourcesMissingColumns = convertDataFrameToSdmSchema(task.schema, diaSourcesMissingColumns,

578 tableName="DiaSource", skipIndex=True)

579 diaSourcesNotMatched = convertDataFrameToSdmSchema(task.schema, diaSourcesNotMatched,

580 tableName="DiaSource", skipIndex=True)

581

582 # Run the method with missing flags

583 diaObjectsSelected = task._selectGoodDiaObjects(diaObjects, diaSources)

584 diaObjectsBadSelection = task._selectGoodDiaObjects(diaObjects, diaSourcesNotMatched)

585 diaObjectsNaNSelection = task._selectGoodDiaObjects(diaObjects, diaSourcesMissingColumns)

586 # The matched catalog of good diaSources should not drop any diaObjects

587 self.assertTrue(diaObjects.equals(diaObjectsSelected))

588 # The mis-matched catalog of good diaSources should drop every diaObject

589 self.assertTrue(diaObjectsBadSelection.empty)

590 # All diaObjects should be dropped for the diaSource catalog that has all NaN values for the flags

591 self.assertTrue(diaObjectsNaNSelection.empty)

592

593 def test_selectGoodObjectsWithWrongFlags(self):

594 """Test the diaObject selection funtion used for forced photometry.

595 """

596 reliabilityThreshold = 0.5

597 trailLengthThreshold = 1.0

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

599 doPackageAlerts=False,

600 doRunForcedMeasurement=True,

601 forcedReliabilityThreshold=reliabilityThreshold,

602 forcedTrailLengthThreshold=trailLengthThreshold,

603 forcedBadFlags=['foo', 'bar'])

604 task = DiaPipelineTask(config=config)

605 nObjects = 20

606 nDiaSources = 100

607

608 # Create diaObjects

609 diaObjects = makeDiaObjects(nObjects, self.exposure, self.rng, startId=1)

610

611 diaSources = makeDiaSources(nDiaSources, diaObjects["diaObjectId"].to_numpy(),

612 self.exposure, self.rng, startId=1)

613

614 diaSources['reliability'] = (self.rng.random(nDiaSources)*(1 - reliabilityThreshold)

615 + reliabilityThreshold)

616 diaSources['trailLength'] = self.rng.random(nDiaSources)*trailLengthThreshold

617

618 diaSourcesNotMatched = makeDiaSources(nDiaSources, diaObjects["diaObjectId"].to_numpy() + 999,

619 self.exposure, self.rng, startId=1)

620

621 diaSourcesNotMatched['reliability'] = (self.rng.random(nDiaSources)*(1 - reliabilityThreshold)

622 + reliabilityThreshold)

623 diaSourcesNotMatched['trailLength'] = self.rng.random(nDiaSources)*trailLengthThreshold

624

625 diaSources = convertDataFrameToSdmSchema(task.schema, diaSources,

626 tableName="DiaSource", skipIndex=True)

627 diaSourcesNotMatched = convertDataFrameToSdmSchema(task.schema, diaSourcesNotMatched,

628 tableName="DiaSource", skipIndex=True)

629

630 # Verify that the flags are in fact missing

631 for flag in config.forcedBadFlags:

632 self.assertNotIn(flag, diaSources.columns)

633 # Run the method with missing flags

634 diaObjectsSelected = task._selectGoodDiaObjects(diaObjects, diaSources)

635 diaObjectsBadSelection = task._selectGoodDiaObjects(diaObjects, diaSourcesNotMatched)

636 # The matched catalog of good diaSources should not drop any diaObjects

637 self.assertTrue(diaObjects.equals(diaObjectsSelected))

638 # The mis-matched catalog of good diaSources should drop every diaObject

639 self.assertTrue(diaObjectsBadSelection.empty)

640

641 def test_mergeEmptyCatalog(self):

642 """Test that a catalog is unchanged if it is merged with an empty

643 catalog.

644 """

645 diaSourcesBase = self.diaSources

646

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

648 task = DiaPipelineTask(config=config)

649 # Include some but not all columns that should be in diaSourcesBase, and some that are mis-matched

650 diaSourcesEmpty = pd.DataFrame(columns=["ra", "dec", "foo"])

651 diaSourcesTest = task.mergeCatalogs(diaSourcesBase, diaSourcesEmpty, tableName="DiaSource")

652 self.assertTrue(diaSourcesBase.equals(diaSourcesTest))

653

654 def test_mergeCatalogs(self):

655 """Test that a merged catalog is concatenated correctly.

656 """

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

658 task = DiaPipelineTask(config=config)

659

660 diaSourcesBase = convertDataFrameToSdmSchema(task.schema, self.diaSources, "DiaSource",

661 skipIndex=True)

662 nBase = len(diaSourcesBase)

663 nNew = int(nBase/2)

664

665 diaSourcesNew = makeDiaSources(nNew, self.diaObjects["diaObjectId"].to_numpy(), self.exposure,

666 self.rng)

667 diaSourcesNew = convertDataFrameToSdmSchema(task.schema, diaSourcesNew, "DiaSource", skipIndex=True)

668 diaSourcesTest = task.mergeCatalogs(diaSourcesBase, diaSourcesNew, tableName="DiaSource")

669 self.assertEqual(len(diaSourcesTest), nBase + nNew)

670 diaSourcesExtract1 = diaSourcesTest.iloc[:nBase]

671 diaSourcesExtract2 = diaSourcesTest.iloc[nBase:]

672

673 pd.testing.assert_frame_equal(diaSourcesBase, diaSourcesExtract1)

674 pd.testing.assert_frame_equal(diaSourcesNew, diaSourcesExtract2)

675

676 def test_updateObjectTable(self):

677 """Test that the diaObject record is updated with the number of

678 diaSources.

679 """

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

681 task = DiaPipelineTask(config=config)

682 nObjects = 20

683 nSrcPerObject = 10

684 nExtraSources = 5

685 nSources = nSrcPerObject*nObjects + nExtraSources

686 expectedSourcesPerObject = nSrcPerObject*np.ones(nObjects)

687 expectedSourcesPerObject[:nExtraSources] += 1

688 diaObjects = makeDiaObjects(nObjects, self.exposure, self.rng)

689 diaSources = makeDiaSources(nSources, diaObjects["diaObjectId"].to_numpy(), self.exposure, self.rng)

690 updatedDiaObjects = task.updateObjectTable(diaObjects, diaSources)

691 self.assertTrue(np.all(updatedDiaObjects.nDiaSources.values == expectedSourcesPerObject))

692

693

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

695 pass

696

697

698def setup_module(module):

699 lsst.utils.tests.init()

700

701

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

703 lsst.utils.tests.init()

704 unittest.main()

Coverage for tests / test_diaPipe.py: 13%

345 statements