Coverage for python/lsst/pipe/tasks/matchFakes.py: 41%

109 statements  

« prev     ^ index     » next       coverage.py v6.4.1, created at 2022-06-05 12:15 +0000

1# This file is part of pipe_tasks. 

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 astropy.units as u 

23import numpy as np 

24import pandas as pd 

25from scipy.spatial import cKDTree 

26 

27from lsst.geom import Box2D 

28import lsst.pex.config as pexConfig 

29from lsst.pipe.base import PipelineTask, PipelineTaskConnections, Struct 

30import lsst.pipe.base.connectionTypes as connTypes 

31from lsst.skymap import BaseSkyMap 

32 

33from lsst.pipe.tasks.insertFakes import InsertFakesConfig 

34 

35__all__ = ["MatchFakesTask", 

36 "MatchFakesConfig", 

37 "MatchVariableFakesConfig", 

38 "MatchVariableFakesTask"] 

39 

40 

41class MatchFakesConnections(PipelineTaskConnections, 

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

43 "fakesType": "fakes_"}, 

44 dimensions=("instrument", 

45 "visit", 

46 "detector")): 

47 skyMap = connTypes.Input( 

48 doc="Input definition of geometry/bbox and projection/wcs for " 

49 "template exposures. Needed to test which tract to generate ", 

50 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

51 dimensions=("skymap",), 

52 storageClass="SkyMap", 

53 ) 

54 fakeCats = connTypes.Input( 

55 doc="Catalog of fake sources inserted into an image.", 

56 name="{fakesType}fakeSourceCat", 

57 storageClass="DataFrame", 

58 dimensions=("tract", "skymap"), 

59 deferLoad=True, 

60 multiple=True 

61 ) 

62 diffIm = connTypes.Input( 

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

64 name="{fakesType}{coaddName}Diff_differenceExp", 

65 storageClass="ExposureF", 

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

67 ) 

68 associatedDiaSources = connTypes.Input( 

69 doc="A DiaSource catalog to match against fakeCat. Assumed " 

70 "to be SDMified.", 

71 name="{fakesType}{coaddName}Diff_assocDiaSrc", 

72 storageClass="DataFrame", 

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

74 ) 

75 matchedDiaSources = connTypes.Output( 

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

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

78 "``fakeCat`` and ``associatedDiaSources``.", 

79 name="{fakesType}{coaddName}Diff_matchDiaSrc", 

80 storageClass="DataFrame", 

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

82 ) 

83 

84 

85class MatchFakesConfig( 

86 InsertFakesConfig, 

87 pipelineConnections=MatchFakesConnections): 

88 """Config for MatchFakesTask. 

89 """ 

90 matchDistanceArcseconds = pexConfig.RangeField( 

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

92 dtype=float, 

93 default=0.5, 

94 min=0, 

95 max=10, 

96 ) 

97 

98 doMatchVisit = pexConfig.Field( 

99 dtype=bool, 

100 default=False, 

101 doc="Match visit to trim the fakeCat" 

102 ) 

103 

104 trimBuffer = pexConfig.Field( 

105 doc="Size of the pixel buffer surrounding the image. Only those fake sources with a centroid" 

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

107 dtype=int, 

108 default=100, 

109 ) 

110 

111 

112class MatchFakesTask(PipelineTask): 

113 """Match a pre-existing catalog of fakes to a catalog of detections on 

114 a difference image. 

115 

116 This task is generally for injected sources that cannot be easily 

117 identified by their footprints such as in the case of detector sources 

118 post image differencing. 

119 """ 

120 

121 _DefaultName = "matchFakes" 

122 ConfigClass = MatchFakesConfig 

123 

124 def run(self, fakeCats, skyMap, diffIm, associatedDiaSources): 

125 """Compose fakes into a single catalog and match fakes to detected 

126 diaSources within a difference image bound. 

127 

128 Parameters 

129 ---------- 

130 fakeCats : `pandas.DataFrame` 

131 List of catalog of fakes to match to detected diaSources. 

132 skyMap : `lsst.skymap.SkyMap` 

133 SkyMap defining the tracts and patches the fakes are stored over. 

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

135 Difference image where ``associatedDiaSources`` were detected. 

136 associatedDiaSources : `pandas.DataFrame` 

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

138 

139 Returns 

140 ------- 

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

142 Results struct with components. 

143 

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

145 length of ``fakeCat``. (`pandas.DataFrame`) 

146 """ 

147 fakeCat = self.composeFakeCat(fakeCats, skyMap) 

148 

149 if self.config.doMatchVisit: 

150 fakeCat = self.getVisitMatchedFakeCat(fakeCat, diffIm) 

151 

152 return self._processFakes(fakeCat, diffIm, associatedDiaSources) 

153 

154 def _processFakes(self, fakeCat, diffIm, associatedDiaSources): 

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

156 

157 Parameters 

158 ---------- 

159 fakeCat : `pandas.DataFrame` 

160 Catalog of fakes to match to detected diaSources. 

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

162 Difference image where ``associatedDiaSources`` were detected. 

163 associatedDiaSources : `pandas.DataFrame` 

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

165 

166 Returns 

167 ------- 

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

169 Results struct with components. 

170 

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

172 length of ``fakeCat``. (`pandas.DataFrame`) 

173 """ 

174 trimmedFakes = self._trimFakeCat(fakeCat, diffIm) 

175 nPossibleFakes = len(trimmedFakes) 

176 

177 fakeVects = self._getVectors(trimmedFakes[self.config.ra_col], 

178 trimmedFakes[self.config.dec_col]) 

179 diaSrcVects = self._getVectors( 

180 np.radians(associatedDiaSources.loc[:, "ra"]), 

181 np.radians(associatedDiaSources.loc[:, "decl"])) 

182 

183 diaSrcTree = cKDTree(diaSrcVects) 

184 dist, idxs = diaSrcTree.query( 

185 fakeVects, 

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

187 nFakesFound = np.isfinite(dist).sum() 

188 

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

190 diaSrcIds = associatedDiaSources.iloc[np.where(np.isfinite(dist), idxs, 0)]["diaSourceId"].to_numpy() 

191 matchedFakes = trimmedFakes.assign(diaSourceId=np.where(np.isfinite(dist), diaSrcIds, 0)) 

192 

193 return Struct( 

194 matchedDiaSources=matchedFakes.merge( 

195 associatedDiaSources.reset_index(drop=True), on="diaSourceId", how="left") 

196 ) 

197 

198 def composeFakeCat(self, fakeCats, skyMap): 

199 """Concatenate the fakeCats from tracts that may cover the exposure. 

200 

201 Parameters 

202 ---------- 

203 fakeCats : `list` of `lst.daf.butler.DeferredDatasetHandle` 

204 Set of fake cats to concatenate. 

205 skyMap : `lsst.skymap.SkyMap` 

206 SkyMap defining the geometry of the tracts and patches. 

207 

208 Returns 

209 ------- 

210 combinedFakeCat : `pandas.DataFrame` 

211 All fakes that cover the inner polygon of the tracts in this 

212 quantum. 

213 """ 

214 if len(fakeCats) == 1: 

215 return fakeCats[0].get( 

216 datasetType=self.config.connections.fakeCats) 

217 outputCat = [] 

218 for fakeCatRef in fakeCats: 

219 cat = fakeCatRef.get( 

220 datasetType=self.config.connections.fakeCats) 

221 tractId = fakeCatRef.dataId["tract"] 

222 # Make sure all data is within the inner part of the tract. 

223 outputCat.append(cat[ 

224 skyMap.findTractIdArray(cat[self.config.ra_col], 

225 cat[self.config.dec_col], 

226 degrees=False) 

227 == tractId]) 

228 

229 return pd.concat(outputCat) 

230 

231 def getVisitMatchedFakeCat(self, fakeCat, exposure): 

232 """Trim the fakeCat to select particular visit 

233 

234 Parameters 

235 ---------- 

236 fakeCat : `pandas.core.frame.DataFrame` 

237 The catalog of fake sources to add to the exposure 

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

239 The exposure to add the fake sources to 

240 

241 Returns 

242 ------- 

243 movingFakeCat : `pandas.DataFrame` 

244 All fakes that belong to the visit 

245 """ 

246 selected = exposure.getInfo().getVisitInfo().getId() == fakeCat["visit"] 

247 

248 return fakeCat[selected] 

249 

250 def _addPixCoords(self, fakeCat, image): 

251 

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

253 

254 Parameters 

255 ---------- 

256 fakeCat : `pandas.core.frame.DataFrame` 

257 The catalog of fake sources to be input 

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

259 The image into which the fake sources should be added 

260 Returns 

261 ------- 

262 fakeCat : `pandas.core.frame.DataFrame` 

263 """ 

264 wcs = image.getWcs() 

265 ras = fakeCat[self.config.ra_col].values 

266 decs = fakeCat[self.config.dec_col].values 

267 xs, ys = wcs.skyToPixelArray(ras, decs) 

268 fakeCat["x"] = xs 

269 fakeCat["y"] = ys 

270 

271 return fakeCat 

272 

273 def _trimFakeCat(self, fakeCat, image): 

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

275 

276 Parameters 

277 ---------- 

278 fakeCat : `pandas.core.frame.DataFrame` 

279 The catalog of fake sources that was input 

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

281 The image into which the fake sources were added 

282 Returns 

283 ------- 

284 fakeCat : `pandas.core.frame.DataFrame` 

285 The original fakeCat trimmed to the area of the image 

286 """ 

287 

288 # fakeCat must be processed with _addPixCoords before trimming 

289 if ('x' not in fakeCat.columns) or ('y' not in fakeCat.columns): 

290 fakeCat = self._addPixCoords(fakeCat, image) 

291 

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

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

294 ras = fakeCat[self.config.ra_col].values * u.rad 

295 decs = fakeCat[self.config.dec_col].values * u.rad 

296 

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

298 

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

300 xs = fakeCat["x"].values 

301 ys = fakeCat["y"].values 

302 

303 bbox = Box2D(image.getBBox()) 

304 isContainedXy = xs >= bbox.minX 

305 isContainedXy &= xs <= bbox.maxX 

306 isContainedXy &= ys >= bbox.minY 

307 isContainedXy &= ys <= bbox.maxY 

308 

309 return fakeCat[isContainedRaDec & isContainedXy] 

310 

311 def _getVectors(self, ras, decs): 

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

313 

314 Parameters 

315 ---------- 

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

317 RA coordinates in radians. 

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

319 Dec coordinates in radians. 

320 

321 Returns 

322 ------- 

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

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

325 """ 

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

327 

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

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

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

331 

332 return vectors 

333 

334 

335class MatchVariableFakesConnections(MatchFakesConnections): 

336 ccdVisitFakeMagnitudes = connTypes.Input( 

337 doc="Catalog of fakes with magnitudes scattered for this ccdVisit.", 

338 name="{fakesType}ccdVisitFakeMagnitudes", 

339 storageClass="DataFrame", 

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

341 ) 

342 

343 

344class MatchVariableFakesConfig(MatchFakesConfig, 

345 pipelineConnections=MatchVariableFakesConnections): 

346 """Config for MatchFakesTask. 

347 """ 

348 pass 

349 

350 

351class MatchVariableFakesTask(MatchFakesTask): 

352 """Match injected fakes to their detected sources in the catalog and 

353 compute their expected brightness in a difference image assuming perfect 

354 subtraction. 

355 

356 This task is generally for injected sources that cannot be easily 

357 identified by their footprints such as in the case of detector sources 

358 post image differencing. 

359 """ 

360 _DefaultName = "matchVariableFakes" 

361 ConfigClass = MatchVariableFakesConfig 

362 

363 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

364 inputs = butlerQC.get(inputRefs) 

365 inputs["band"] = butlerQC.quantum.dataId["band"] 

366 

367 outputs = self.run(**inputs) 

368 butlerQC.put(outputs, outputRefs) 

369 

370 def run(self, fakeCats, ccdVisitFakeMagnitudes, skyMap, diffIm, associatedDiaSources, band): 

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

372 

373 Parameters 

374 ---------- 

375 fakeCat : `pandas.DataFrame` 

376 Catalog of fakes to match to detected diaSources. 

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

378 Difference image where ``associatedDiaSources`` were detected in. 

379 associatedDiaSources : `pandas.DataFrame` 

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

381 

382 Returns 

383 ------- 

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

385 Results struct with components. 

386 

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

388 length of ``fakeCat``. (`pandas.DataFrame`) 

389 """ 

390 fakeCat = self.composeFakeCat(fakeCats, skyMap) 

391 self.computeExpectedDiffMag(fakeCat, ccdVisitFakeMagnitudes, band) 

392 return self._processFakes(fakeCat, diffIm, associatedDiaSources) 

393 

394 def computeExpectedDiffMag(self, fakeCat, ccdVisitFakeMagnitudes, band): 

395 """Compute the magnitude expected in the difference image for this 

396 detector/visit. Modify fakeCat in place. 

397 

398 Negative magnitudes indicate that the source should be detected as 

399 a negative source. 

400 

401 Parameters 

402 ---------- 

403 fakeCat : `pandas.DataFrame` 

404 Catalog of fake sources. 

405 ccdVisitFakeMagnitudes : `pandas.DataFrame` 

406 Magnitudes for variable sources in this specific ccdVisit. 

407 band : `str` 

408 Band that this ccdVisit was observed in. 

409 """ 

410 magName = self.config.mag_col % band 

411 magnitudes = fakeCat[magName].to_numpy() 

412 visitMags = ccdVisitFakeMagnitudes["variableMag"].to_numpy() 

413 diffFlux = (visitMags * u.ABmag).to_value(u.nJy) - (magnitudes * u.ABmag).to_value(u.nJy) 

414 diffMag = np.where(diffFlux > 0, 

415 (diffFlux * u.nJy).to_value(u.ABmag), 

416 -(-diffFlux * u.nJy).to_value(u.ABmag)) 

417 

418 noVisit = ~fakeCat["isVisitSource"] 

419 noTemplate = ~fakeCat["isTemplateSource"] 

420 both = np.logical_and(fakeCat["isVisitSource"], 

421 fakeCat["isTemplateSource"]) 

422 

423 fakeCat.loc[noVisit, magName] = -magnitudes[noVisit] 

424 fakeCat.loc[noTemplate, magName] = visitMags[noTemplate] 

425 fakeCat.loc[both, magName] = diffMag[both]