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

151 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-26 00:48 -0700

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: 388 ↛ 392line 388 didn't jump to line 392 because the condition on line 388 was always true

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

390 return fakeCat, None 

391 

392 isVariable = fakeCat["twin_id"] > 0 

393 

394 return fakeCat[~isVariable], fakeCat[isVariable] 

395 

396 def _add_variables_to_matched(self, matchedFakes, variableDoublesFakeCat): 

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

398 

399 Parameters 

400 ---------- 

401 matchedFakes : `astropy.table.table.Table` 

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

403 sources in the input fake catalog. 

404 variableDoublesFakeCat : `astropy.table.table.Table` 

405 Catalog of variable sources in the input fake catalog. 

406 

407 Returns 

408 ------- 

409 fullMatchedFakes : `astropy.table.table.Table` 

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

411 static and variable sources in the input fake catalog. 

412 """ 

413 if variableDoublesFakeCat is None: 413 ↛ 421line 413 didn't jump to line 421 because the condition on line 413 was always true

414 return matchedFakes 

415 

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

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

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

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

420 # to diaSources. 

421 variableDoublesFakeCat = variableDoublesFakeCat.copy() 

422 variableDoublesFakeCat['diaSourceId'] = 0 

423 variableDoublesFakeCat['dist_diaSrc'] = -1 

424 

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

426 # Join on twin_id to injection_id 

427 matched = join(variableDoublesFakeCat, matchedFakes, 

428 keys_left='twin_id', keys_right='injection_id', 

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

430 keep_order=True) 

431 

432 # Fill diaSourceId and dist_diaSrc from matched results 

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

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

435 

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

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

438 

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

440 

441 

442class MatchInjectedToAssocDiaSourceConnections( 

443 PipelineTaskConnections, 

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

445 dimensions=("instrument", 

446 "visit", 

447 "detector")): 

448 

449 assocDiaSources = connTypes.Input( 

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

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

452 name="{coaddName}Diff_assocDiaSrc", 

453 storageClass="ArrowAstropy", 

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

455 ) 

456 matchDiaSources = connTypes.Input( 

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

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

459 "``fakeCat`` and ``diaSrc``.", 

460 name="{coaddName}Diff_matchDiaSrc", 

461 storageClass="ArrowAstropy", 

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

463 ) 

464 matchAssocDiaSources = connTypes.Output( 

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

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

467 "``fakeCat`` and ``associatedDiaSources``.", 

468 name="{coaddName}Diff_matchAssocDiaSrc", 

469 storageClass="ArrowAstropy", 

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

471 ) 

472 

473 

474class MatchInjectedToAssocDiaSourceConfig( 

475 PipelineTaskConfig, 

476 pipelineConnections=MatchInjectedToAssocDiaSourceConnections): 

477 """Config for MatchFakesTask. 

478 """ 

479 

480 

481class MatchInjectedToAssocDiaSourceTask(PipelineTask): 

482 

483 _DefaultName = "matchInjectedToAssocDiaSource" 

484 ConfigClass = MatchInjectedToAssocDiaSourceConfig 

485 

486 def run(self, assocDiaSources, matchDiaSources): 

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

488 

489 Parameters 

490 ---------- 

491 matchDiaSources : `astropy.table.Table` 

492 Catalog of matched diaSrc to injected sources 

493 assocDiaSources : `astropy.table.Table` 

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

495 Returns 

496 ------- 

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

498 Results struct with components. 

499 

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

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

502 """ 

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

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

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

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

507 

508 nPossibleFakes = len(matchDiaSources) 

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

510 matchDiaSources["diaSourceId"], assocDiaSources["diaSourceId"] 

511 ) 

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

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

514 

515 return Struct( 

516 matchAssocDiaSources=join( 

517 matchDiaSources, 

518 assocDiaSources, 

519 keys="diaSourceId", 

520 join_type="left", 

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

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

523 ) 

524 )