Coverage for python/lsst/pipe/tasks/matchDiffimSourceInjected.py: 71%

154 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-16 09:04 +0000

1# This file is part of ap_pipe. 

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 

22__all__ = ["MatchInjectedToDiaSourceTask", 

23 "MatchInjectedToDiaSourceConfig", 

24 "MatchInjectedToAssocDiaSourceTask", 

25 "MatchInjectedToAssocDiaSourceConfig"] 

26 

27import astropy.units as u 

28from astropy.table import join, vstack 

29import numpy as np 

30from scipy.spatial import cKDTree 

31 

32from lsst.afw import table as afwTable 

33from lsst import geom as lsstGeom 

34import lsst.pex.config as pexConfig 

35from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct 

36import lsst.pipe.base.connectionTypes as connTypes 

37from lsst.meas.base import ForcedMeasurementTask, ForcedMeasurementConfig 

38 

39 

40class MatchInjectedToDiaSourceConnections( 

41 PipelineTaskConnections, 

42 defaultTemplates={"coaddName": "deep"}, 

43 dimensions=("instrument", 

44 "visit", 

45 "detector")): 

46 injectionCat = connTypes.Input( 

47 doc="Catalog of sources injected in the images.", 

48 name="VisitDetectorFakeSourceCat", 

49 storageClass="ArrowAstropy", 

50 dimensions=("instrument", "visit", "detector"), 

51 ) 

52 diffIm = connTypes.Input( 

53 doc="Difference image on which the DiaSources were detected.", 

54 name="{coaddName}Diff_differenceExp", 

55 storageClass="ExposureF", 

56 dimensions=("instrument", "visit", "detector"), 

57 ) 

58 diaSources = connTypes.Input( 

59 doc="A DiaSource catalog to match against fakeCat.", 

60 name="{coaddName}Diff_diaSrc", 

61 storageClass="SourceCatalog", 

62 dimensions=("instrument", "visit", "detector"), 

63 ) 

64 matchDiaSources = connTypes.Output( 

65 doc="A catalog of those fakeCat sources that have a match in " 

66 "diaSrc. The schema is the union of the schemas for " 

67 "``fakeCat`` and ``diaSrc``.", 

68 name="{coaddName}Diff_matchDiaSrc", 

69 storageClass="ArrowAstropy", 

70 dimensions=("instrument", "visit", "detector"), 

71 ) 

72 

73 

74class MatchInjectedToDiaSourceConfig( 

75 PipelineTaskConfig, 

76 pipelineConnections=MatchInjectedToDiaSourceConnections): 

77 """Config for MatchFakesTask. 

78 """ 

79 matchDistanceArcseconds = pexConfig.RangeField( 

80 doc="Distance in arcseconds to match within.", 

81 dtype=float, 

82 default=0.5, 

83 min=0, 

84 max=10, 

85 ) 

86 doMatchVisit = pexConfig.Field( 

87 dtype=bool, 

88 default=True, 

89 doc="Match visit to trim the fakeCat" 

90 ) 

91 trimBuffer = pexConfig.Field( 

92 doc="Size of the pixel buffer surrounding the image." 

93 "Only those fake sources with a centroid" 

94 "falling within the image+buffer region will be considered matches.", 

95 dtype=int, 

96 default=50, 

97 ) 

98 doForcedMeasurement = pexConfig.Field( 

99 dtype=bool, 

100 default=True, 

101 doc="Force measurement of the fakes at the injection locations." 

102 ) 

103 forcedMeasurement = pexConfig.ConfigurableField( 

104 target=ForcedMeasurementTask, 

105 doc="Task to force photometer difference image at injection locations.", 

106 ) 

107 

108 

109class MatchInjectedToDiaSourceTask(PipelineTask): 

110 

111 _DefaultName = "matchInjectedToDiaSource" 

112 ConfigClass = MatchInjectedToDiaSourceConfig 

113 

114 def run(self, injectionCat, diffIm, diaSources): 

115 """Match injected sources to detected diaSources within a difference image bound. 

116 

117 Parameters 

118 ---------- 

119 injectionCat : `astropy.table.Table` 

120 Table of catalog of synthetic sources to match to detected diaSources. 

121 diffIm : `lsst.afw.image.Exposure` 

122 Difference image where ``diaSources`` were detected. 

123 diaSources : `afw.table.SourceCatalog` 

124 Catalog of difference image sources detected in ``diffIm``. 

125 Returns 

126 ------- 

127 result : `lsst.pipe.base.Struct` 

128 Results struct with components. 

129 

130 - ``matchedDiaSources`` : Fakes matched to input diaSources. Has 

131 length of ``injectionCat``. (`astropy.table.Table`) 

132 """ 

133 

134 if self.config.doMatchVisit: 

135 fakeCat = self._trimFakeCat(injectionCat, diffIm) 

136 else: 

137 fakeCat = injectionCat 

138 if self.config.doForcedMeasurement: 138 ↛ 139line 138 didn't jump to line 139 because the condition on line 138 was never true

139 self._estimateFakesSNR(fakeCat, diffIm) 

140 

141 # Split the fake catalog into the initial injections and the variable sources themselves, 

142 # which are generated as duplicates of the initial injections with a twin_id column. 

143 # We then match only the initial injections to the diaSources, 

144 # and then add back in the variable sources by matching them to their twins 

145 initialFakeCat, variableDoublesFakeCat = self._splitVariables(fakeCat) 

146 matchedFakes = self._processFakes(initialFakeCat, diaSources) 

147 fullMatchedFakes = self._add_variables_to_matched(matchedFakes, variableDoublesFakeCat) 

148 

149 return Struct(matchDiaSources=fullMatchedFakes) 

150 

151 def _estimateFakesSNR(self, injectionCat, diffIm): 

152 """Estimate the signal-to-noise ratio of the fakes in the given catalog. 

153 

154 Parameters 

155 ---------- 

156 injectionCat : `astropy.table.Table` 

157 Catalog of synthetic sources to estimate the S/N of. **This table 

158 will be modified in place**. 

159 diffIm : `lsst.afw.image.Exposure` 

160 Difference image where the sources were detected. 

161 """ 

162 # Create a schema for the forced measurement task 

163 schema = afwTable.SourceTable.makeMinimalSchema() 

164 schema.addField("x", "D", "x position in image.", units="pixel") 

165 schema.addField("y", "D", "y position in image.", units="pixel") 

166 schema.addField("deblend_nChild", "I", "Need for minimal forced phot schema") 

167 

168 pluginList = [ 

169 "base_PixelFlags", 

170 "base_SdssCentroid", 

171 "base_CircularApertureFlux", 

172 "base_PsfFlux", 

173 "base_LocalBackground" 

174 ] 

175 forcedMeasConfig = ForcedMeasurementConfig(plugins=pluginList) 

176 forcedMeasConfig.slots.centroid = 'base_SdssCentroid' 

177 forcedMeasConfig.slots.shape = None 

178 

179 # Create the forced measurement task 

180 forcedMeas = ForcedMeasurementTask(schema, config=forcedMeasConfig) 

181 

182 # Specify the columns to copy from the input catalog to the output catalog 

183 forcedMeas.copyColumns = {"coord_ra": "ra", "coord_dec": "dec"} 

184 

185 # Create an afw table from the input catalog 

186 outputCatalog = afwTable.SourceCatalog(schema) 

187 outputCatalog.reserve(len(injectionCat)) 

188 for row in injectionCat: 

189 outputRecord = outputCatalog.addNew() 

190 outputRecord.setId(row['injection_id']) 

191 outputRecord.setCoord(lsstGeom.SpherePoint(row["ra"], row["dec"], lsstGeom.degrees)) 

192 outputRecord.set("x", row["x"]) 

193 outputRecord.set("y", row["y"]) 

194 

195 # Generate the forced measurement catalog 

196 forcedSources = forcedMeas.generateMeasCat(diffIm, outputCatalog, diffIm.getWcs()) 

197 # Attach the PSF shape footprints to the forced measurement catalog 

198 forcedMeas.attachPsfShapeFootprints(forcedSources, diffIm) 

199 

200 # Copy the x and y positions from the forced measurement catalog back 

201 # to the input catalog 

202 for src, tgt in zip(forcedSources, outputCatalog): 

203 src.set('base_SdssCentroid_x', tgt['x']) 

204 src.set('base_SdssCentroid_y', tgt['y']) 

205 

206 # Define the centroid for the forced measurement catalog 

207 forcedSources.defineCentroid('base_SdssCentroid') 

208 # Run the forced measurement task 

209 forcedMeas.run(forcedSources, diffIm, outputCatalog, diffIm.getWcs()) 

210 # Convert the forced measurement catalog to an astropy table 

211 forcedSources_table = forcedSources.asAstropy() 

212 

213 # Add the forced measurement columns to the input catalog 

214 for column in forcedSources_table.columns: 

215 if "Flux" in column or "flag" in column: 

216 injectionCat["forced_"+column] = forcedSources_table[column] 

217 

218 # Add the SNR columns to the input catalog 

219 for column in injectionCat.colnames: 

220 if column.endswith("instFlux"): 

221 flux = np.abs(injectionCat[column]) 

222 fluxErr = injectionCat[column+"Err"].copy() 

223 fluxErr = np.where( 

224 (fluxErr <= 0) | (np.isnan(fluxErr)), np.nanmax(fluxErr), fluxErr) 

225 

226 injectionCat[column+"_SNR"] = flux / fluxErr 

227 

228 def _processFakes(self, injectedCat, diaSources): 

229 """Match fakes to detected diaSources within a difference image bound. 

230 

231 Parameters 

232 ---------- 

233 injectedCat : `astropy.table.Table` 

234 Catalog of injected sources to match to detected diaSources. 

235 diaSources : `afw.table.SourceCatalog` 

236 Catalog of difference image sources detected in ``diffIm``. 

237 

238 Returns 

239 ------- 

240 result : `lsst.pipe.base.Struct` 

241 Results struct with components. 

242 

243 - ``matchedDiaSources`` : Fakes matched to input diaSources. Has 

244 length of ``fakeCat``. (`astropy.table.Table`) 

245 """ 

246 # First match the diaSrc to the injected fakes 

247 nPossibleFakes = len(injectedCat) 

248 

249 fakeVects = self._getVectors( 

250 np.radians(injectedCat['ra']), 

251 np.radians(injectedCat['dec'])) 

252 diaSrcVects = self._getVectors( 

253 diaSources['coord_ra'], 

254 diaSources['coord_dec']) 

255 

256 diaSrcTree = cKDTree(diaSrcVects) 

257 dist, idxs = diaSrcTree.query( 

258 fakeVects, 

259 distance_upper_bound=np.radians(self.config.matchDistanceArcseconds / 3600)) 

260 # handshake matching, that is symmetrize the match by matching the 

261 # diaSrcs back to the fakes and only keeping those matches where the 

262 # same pair is returned 

263 diaSrcTreeBack = cKDTree(fakeVects) 

264 distBack, idxsBack = diaSrcTreeBack.query( 

265 diaSrcVects, 

266 distance_upper_bound=np.radians(self.config.matchDistanceArcseconds / 3600)) 

267 

268 idxsAux = np.where(np.array(idxs) < len(diaSources), idxs, -1) 

269 valid = idxsAux >= 0 

270 idxsBackMatched = np.full_like(idxsAux, -1) 

271 idxsBackMatched[valid] = idxsBack[idxsAux[valid]] 

272 idxsMatched = np.where(idxsBackMatched == np.arange(len(injectedCat)), idxs, -1) 

273 distMatched = np.where(idxsBackMatched == np.arange(len(injectedCat)), dist, np.inf) 

274 nFakesFound = np.isfinite(distMatched).sum() 

275 

276 self.log.info("Found %d out of %d possible in diaSources.", nFakesFound, nPossibleFakes) 

277 

278 # assign diaSourceId to the matched fakes 

279 diaSrcIds = diaSources['id'][np.where(np.isfinite(distMatched), idxsMatched, 0)] 

280 matchedFakes = injectedCat.copy() 

281 matchedFakes['diaSourceId'] = np.where(np.isfinite(distMatched), diaSrcIds, 0) 

282 matchedFakes['dist_diaSrc'] = np.where(np.isfinite(distMatched), 3600*np.rad2deg(distMatched), -1) 

283 return matchedFakes 

284 

285 def _getVectors(self, ras, decs): 

286 """Convert ra dec to unit vectors on the sphere. 

287 

288 Parameters 

289 ---------- 

290 ras : `numpy.ndarray`, (N,) 

291 RA coordinates in radians. 

292 decs : `numpy.ndarray`, (N,) 

293 Dec coordinates in radians. 

294 

295 Returns 

296 ------- 

297 vectors : `numpy.ndarray`, (N, 3) 

298 Vectors on the unit sphere for the given RA/DEC values. 

299 """ 

300 vectors = np.empty((len(ras), 3)) 

301 

302 vectors[:, 2] = np.sin(decs) 

303 vectors[:, 0] = np.cos(decs) * np.cos(ras) 

304 vectors[:, 1] = np.cos(decs) * np.sin(ras) 

305 

306 return vectors 

307 

308 def _addPixCoords(self, fakeCat, image): 

309 """Add pixel coordinates to the catalog of fakes. 

310 

311 Parameters 

312 ---------- 

313 fakeCat : `astropy.table.table.Table` 

314 The catalog of fake sources to be input 

315 image : `lsst.afw.image.exposure.exposure.ExposureF` 

316 The image into which the fake sources should be added 

317 Returns 

318 ------- 

319 fakeCat : `astropy.table.table.Table` 

320 """ 

321 

322 wcs = image.getWcs() 

323 

324 # Get x/y pixel coordinates for injected sources. 

325 xs, ys = wcs.skyToPixelArray( 

326 fakeCat["ra"], 

327 fakeCat["dec"], 

328 degrees=True 

329 ) 

330 fakeCat["x"] = xs 

331 fakeCat["y"] = ys 

332 

333 return fakeCat 

334 

335 def _trimFakeCat(self, fakeCat, image): 

336 """Trim the fake cat to the exact size of the input image. 

337 

338 Parameters 

339 ---------- 

340 fakeCat : `astropy.table.table.Table` 

341 The catalog of fake sources that was input 

342 image : `lsst.afw.image.exposure.exposure.ExposureF` 

343 The image into which the fake sources were added 

344 Returns 

345 ------- 

346 fakeCat : `astropy.table.table.Table` 

347 The original fakeCat trimmed to the area of the image 

348 """ 

349 

350 # fakeCat must be processed with _addPixCoords before trimming 

351 fakeCat = self._addPixCoords(fakeCat, image) 

352 

353 # Prefilter in ra/dec to avoid cases where the wcs incorrectly maps 

354 # input fakes which are really off the chip onto it. 

355 ras = fakeCat["ra"] * u.deg 

356 decs = fakeCat["dec"] * u.deg 

357 

358 isContainedRaDec = image.containsSkyCoords(ras, decs, padding=0) 

359 

360 # now use the exact pixel BBox to filter to only fakes that were inserted 

361 xs = fakeCat["x"] 

362 ys = fakeCat["y"] 

363 

364 bbox = lsstGeom.Box2D(image.getBBox()) 

365 isContainedXy = xs - self.config.trimBuffer >= bbox.minX 

366 isContainedXy &= xs + self.config.trimBuffer <= bbox.maxX 

367 isContainedXy &= ys - self.config.trimBuffer >= bbox.minY 

368 isContainedXy &= ys + self.config.trimBuffer <= bbox.maxY 

369 

370 return fakeCat[isContainedRaDec & isContainedXy] 

371 

372 def _splitVariables(self, fakeCat): 

373 """Split out the duplicated injections, that are used to generate 

374 variable sources in the fake catalog. 

375 

376 Parameters 

377 ---------- 

378 fakeCat : `astropy.table.table.Table` 

379 The catalog of fake sources that was input 

380 

381 Returns 

382 ------- 

383 initialFakeCat : `astropy.table.table.Table` 

384 Subset of the input catalog corresponding to initial sources. 

385 variableDoublesFakeCat : `astropy.table.table.Table` 

386 Subset of the input catalog corresponding to variable sources. 

387 """ 

388 if "twin_id" not in fakeCat.colnames: 

389 self.log.warning("No twin_id column found in fake catalog.") 

390 return fakeCat, None 

391 if "isVariable" not in fakeCat.colnames or "isVisitSource" not in fakeCat.colnames: 391 ↛ 392line 391 didn't jump to line 392 because the condition on line 391 was never true

392 isTwin = fakeCat["twin_id"] > 0 

393 else: 

394 isTwin = (fakeCat["twin_id"] > 0) & fakeCat["isVariable"] & ~fakeCat["isVisitSource"] 

395 

396 isTwin = (fakeCat["twin_id"] > 0) & fakeCat["isVariable"] & ~fakeCat["isVisitSource"] 

397 return fakeCat[~isTwin], fakeCat[isTwin] 

398 

399 def _add_variables_to_matched(self, matchedFakes, variableDoublesFakeCat): 

400 """Add variable sources back into the matched fakes catalog. 

401 

402 Parameters 

403 ---------- 

404 matchedFakes : `astropy.table.table.Table` 

405 Catalog of matched fakes to diaSources, corresponding to the static 

406 sources in the input fake catalog. 

407 variableDoublesFakeCat : `astropy.table.table.Table` 

408 Catalog of variable sources in the input fake catalog. 

409 

410 Returns 

411 ------- 

412 fullMatchedFakes : `astropy.table.table.Table` 

413 Catalog of matched fakes to diaSources, corresponding to both the 

414 static and variable sources in the input fake catalog. 

415 """ 

416 if variableDoublesFakeCat is None: 

417 return matchedFakes 

418 

419 # For the variable sources, we have a match to diaSources if their twins 

420 # had a match, so we fill the diaSourceId with the diaSourceId of the matched 

421 # twin if it exists and 0 otherwise, and we set the distance to -1 to 

422 # indicate that these are variable sources that were not directly matched 

423 # to diaSources. 

424 variableDoublesFakeCat = variableDoublesFakeCat.copy() 

425 variableDoublesFakeCat['diaSourceId'] = 0 

426 variableDoublesFakeCat['dist_diaSrc'] = -1 

427 

428 # Match variable sources to their twin's matched diaSource 

429 # Join on twin_id to injection_id 

430 matched = join(variableDoublesFakeCat, matchedFakes, 

431 keys_left='twin_id', keys_right='injection_id', 

432 join_type='left', table_names=['variables', 'matched'], 

433 keep_order=True) 

434 

435 # Fill diaSourceId and dist_diaSrc from matched results 

436 dia_id = np.ma.asarray(matched["diaSourceId_matched"]) 

437 dist = np.ma.asarray(matched["dist_diaSrc_matched"]) 

438 

439 variableDoublesFakeCat["diaSourceId"] = np.ma.filled(dia_id, 0).astype(np.int64) 

440 variableDoublesFakeCat["dist_diaSrc"] = np.ma.filled(dist, -1.0) 

441 

442 return vstack([matchedFakes, variableDoublesFakeCat], metadata_conflicts='silent') 

443 

444 

445class MatchInjectedToAssocDiaSourceConnections( 

446 PipelineTaskConnections, 

447 defaultTemplates={"coaddName": "deep"}, 

448 dimensions=("instrument", 

449 "visit", 

450 "detector")): 

451 

452 assocDiaSources = connTypes.Input( 

453 doc="An assocDiaSource catalog to match against fakeCat from the" 

454 "diaPipe run. Assumed to be SDMified.", 

455 name="{coaddName}Diff_assocDiaSrc", 

456 storageClass="ArrowAstropy", 

457 dimensions=("instrument", "visit", "detector"), 

458 ) 

459 matchDiaSources = connTypes.Input( 

460 doc="A catalog of those fakeCat sources that have a match in " 

461 "diaSrc. The schema is the union of the schemas for " 

462 "``fakeCat`` and ``diaSrc``.", 

463 name="{coaddName}Diff_matchDiaSrc", 

464 storageClass="ArrowAstropy", 

465 dimensions=("instrument", "visit", "detector"), 

466 ) 

467 matchAssocDiaSources = connTypes.Output( 

468 doc="A catalog of those fakeCat sources that have a match in " 

469 "associatedDiaSources. The schema is the union of the schemas for " 

470 "``fakeCat`` and ``associatedDiaSources``.", 

471 name="{coaddName}Diff_matchAssocDiaSrc", 

472 storageClass="ArrowAstropy", 

473 dimensions=("instrument", "visit", "detector"), 

474 ) 

475 

476 

477class MatchInjectedToAssocDiaSourceConfig( 

478 PipelineTaskConfig, 

479 pipelineConnections=MatchInjectedToAssocDiaSourceConnections): 

480 """Config for MatchFakesTask. 

481 """ 

482 

483 

484class MatchInjectedToAssocDiaSourceTask(PipelineTask): 

485 

486 _DefaultName = "matchInjectedToAssocDiaSource" 

487 ConfigClass = MatchInjectedToAssocDiaSourceConfig 

488 

489 def run(self, assocDiaSources, matchDiaSources): 

490 """Tag matched injected sources to associated diaSources. 

491 

492 Parameters 

493 ---------- 

494 matchDiaSources : `astropy.table.Table` 

495 Catalog of matched diaSrc to injected sources 

496 assocDiaSources : `astropy.table.Table` 

497 Catalog of associated difference image sources detected in ``diffIm``. 

498 Returns 

499 ------- 

500 result : `lsst.pipe.base.Struct` 

501 Results struct with components. 

502 

503 - ``matchAssocDiaSources`` : Fakes matched to associated diaSources. Has 

504 length of ``matchDiaSources``. (`astropy.table.Table`) 

505 """ 

506 # Match the fakes to the associated sources. For this we don't use the coordinates 

507 # but instead check for the diaSources. Since they were present in the table already 

508 matchDiaSources["diaSourceId"] = np.asarray(matchDiaSources["diaSourceId"], dtype=np.int64) 

509 assocDiaSources["diaSourceId"] = np.asarray(assocDiaSources["diaSourceId"], dtype=np.int64) 

510 

511 nPossibleFakes = len(matchDiaSources) 

512 matchDiaSources["isAssocDiaSource"] = np.isin( 

513 matchDiaSources["diaSourceId"], assocDiaSources["diaSourceId"] 

514 ) 

515 assocNFakesFound = matchDiaSources['isAssocDiaSource'].sum() 

516 self.log.info("Found %d out of %d possible in assocDiaSources."%(assocNFakesFound, nPossibleFakes)) 

517 

518 return Struct( 

519 matchAssocDiaSources=join( 

520 matchDiaSources, 

521 assocDiaSources, 

522 keys="diaSourceId", 

523 join_type="left", 

524 table_names=("ssi", "diaSrc"), 

525 uniq_col_name="{col_name}_{table_name}", 

526 ) 

527 )