Coverage for python/lsst/pipe/tasks/diffractionSpikeMask.py: 25%

154 statements  

« prev     ^ index     » next       coverage.py v7.11.0, created at 2025-11-04 02:26 -0800

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 

22 

23__all__ = ["DiffractionSpikeMaskConfig", "DiffractionSpikeMaskTask"] 

24 

25 

26import math 

27import numpy as np 

28import astropy.units as u 

29 

30import lsst.afw.geom as afwGeom 

31from lsst.afw.image import abMagErrFromFluxErr 

32import lsst.afw.table as afwTable 

33import lsst.geom 

34from lsst.pex.config import Config, ConfigField, Field 

35from lsst.pipe.base import Struct, Task 

36from lsst.meas.algorithms import getRefFluxField, LoadReferenceObjectsConfig 

37import lsst.sphgeom 

38from .colorterms import ColortermLibrary 

39 

40 

41class DiffractionSpikeMaskConfig(Config): 

42 """Config for BrightStarMaskTask. 

43 """ 

44 refObjLoader = ConfigField(dtype=LoadReferenceObjectsConfig, 

45 doc="Configuration of reference object loader") 

46 applyColorTerms = Field( 

47 dtype=bool, 

48 default=False, 

49 doc=("Apply photometric color terms to reference stars?\n" 

50 "`True`: attempt to apply color terms; fail if color term data is " 

51 "not available for the specified reference catalog and filter.\n" 

52 "`False`: do not apply color terms."), 

53 optional=True, 

54 ) 

55 colorterms = ConfigField( 

56 dtype=ColortermLibrary, 

57 doc="Library of photometric reference catalog name: color term dict" 

58 " (see also applyColorTerms).", 

59 ) 

60 photoCatName = Field( 

61 dtype=str, 

62 optional=True, 

63 doc=("Name of photometric reference catalog; used to select a color" 

64 " term dict in colorterms. See also applyColorTerms."), 

65 ) 

66 raKey = Field( 

67 dtype=str, 

68 default="coord_ra", 

69 doc="RA column name in the reference catalog.", 

70 ) 

71 decKey = Field( 

72 dtype=str, 

73 default="coord_dec", 

74 doc="Declination column name in the reference catalog.", 

75 ) 

76 angleMargin = Field( 

77 dtype=float, 

78 default=60., 

79 doc="Margin outside the exposure bounding box to include bright " 

80 "sources. In arcseconds.", 

81 ) 

82 magnitudeThreshold = Field( 

83 dtype=float, 

84 default=15, 

85 doc="Threshold magnitude for treating a star from the reference catalog" 

86 " as bright.", 

87 ) 

88 diffractAngle = Field( 

89 dtype=float, 

90 default=45, 

91 doc="Angle in degrees of the location of diffraction spikes with " 

92 "respect to camera at 0 rotation angle.", 

93 ) 

94 spikeAspectRatio = Field( 

95 dtype=float, 

96 default=10, 

97 doc="Ratio of the length of a diffraction spike to it's width in the" 

98 " core of the star.", 

99 ) 

100 magSlope = Field( 

101 dtype=float, 

102 default=-0.12, 

103 doc="Slope of the fit for the log(spike length) as a function of" 

104 " magnitude.", 

105 ) 

106 magOffset = Field( 

107 dtype=float, 

108 default=3.8, 

109 doc="Intercept of the fit for the log(spike length) as a function of" 

110 " magnitude.", 

111 ) 

112 fallbackMagnitude = Field( 

113 dtype=float, 

114 default=12., 

115 doc="Default magnitude to use for sources in the reference catalog" 

116 " with missing magnitudes that land in regions of saturated pixels.", 

117 ) 

118 fallbackFluxField = Field( 

119 dtype=str, 

120 default="phot_g_mean", 

121 doc="Fallback flux field in the reference catalog to use for sources" 

122 " that don't have measurements in the science image's band.", 

123 ) 

124 spikeMask = Field( 

125 dtype=str, 

126 default="SPIKE", 

127 doc="Name of the mask plane indicating likely contamination from" 

128 " a diffraction spike.", 

129 ) 

130 saturatedMaskPlane = Field( 

131 dtype=str, 

132 default="SAT", 

133 doc="Name of the mask plane indicating a saturated pixel.", 

134 ) 

135 

136 

137class DiffractionSpikeMaskTask(Task): 

138 """Load a reference catalog to identify bright stars that are likely to be 

139 saturated and have visible diffraction spikes that need to be masked. 

140 

141 Attributes 

142 ---------- 

143 angles : `numpy.ndarray` 

144 Expected angles of diffraction spikes for bright sources. 

145 refObjLoader : `lsst.meas.algorithms.ReferenceObjectLoader` 

146 An instance of a reference object loader. 

147 """ 

148 ConfigClass = DiffractionSpikeMaskConfig 

149 _DefaultName = "diffractionSpikeMask" 

150 

151 def __init__(self, refObjLoader=None, **kwargs): 

152 Task.__init__(self, **kwargs) 

153 self.refObjLoader = refObjLoader 

154 

155 def setRefObjLoader(self, refObjLoader): 

156 """Set the reference object loader for the task. 

157 

158 Parameters 

159 ---------- 

160 refObjLoader : `lsst.meas.algorithms.ReferenceObjectLoader` 

161 An instance of a reference object loader. 

162 """ 

163 self.refObjLoader = refObjLoader 

164 

165 def run(self, exposure): 

166 """Load reference objects and mask bright stars on an exposure. 

167 

168 Parameters 

169 ---------- 

170 exposure : `lsst.afw.image.Exposure` 

171 Science exposure to set the SPIKE plane. Will be modified in place. 

172 

173 Returns 

174 ------- 

175 spikeCat : `lsst.afw.table.SourceCatalog` 

176 The entries from the reference catalog selected as stars with 

177 diffraction spikes. 

178 """ 

179 if self.refObjLoader is None: 

180 raise RuntimeError("Running diffraction spike mask task with no refObjLoader set in" 

181 " __init__ or setRefObjLoader") 

182 # First set the angles of the diffraction spikes from the exposure 

183 # metadata. 

184 self.set_diffraction_angle(exposure) 

185 filterLabel = exposure.getFilter() 

186 region = getRegion(exposure, lsst.sphgeom.Angle.fromDegrees(self.config.angleMargin/3600.)) 

187 refCat = self.refObjLoader.loadRegion(region, filterLabel.bandLabel).refCat 

188 

189 # Mask any sources with known or estimated magnitudes for the current 

190 # filter, including sources off of the image which may have diffraction 

191 # spikes that extend onto the image. 

192 magnitudes = self.extractMagnitudes(refCat, filterLabel).refMag 

193 radii = self.calculateReferenceRadius(magnitudes) 

194 bright = (magnitudes < self.config.magnitudeThreshold) 

195 

196 nBright = np.count_nonzero(bright) 

197 

198 mask = exposure.mask 

199 mask.addMaskPlane(self.config.spikeMask) 

200 if nBright > 0: 

201 xvals, yvals = exposure.wcs.skyToPixelArray(refCat[bright][self.config.raKey], 

202 refCat[bright][self.config.decKey]) 

203 spikeCandidates = self.selectSources(xvals, yvals, mask) 

204 nSpike = len(spikeCandidates) 

205 else: 

206 nSpike = 0 

207 if nSpike > 0: 

208 self.log.info("Calculating mask for %d stars brighter than magnitude %f", nSpike, 

209 self.config.magnitudeThreshold) 

210 self.maskSources(xvals[spikeCandidates], 

211 yvals[spikeCandidates], 

212 radii[bright][spikeCandidates], 

213 mask) 

214 else: 

215 self.log.info("No bright stars found in the reference catalog; not masking diffraction spikes.") 

216 return afwTable.SourceCatalog(refCat.schema) 

217 

218 return refCat[bright][spikeCandidates].copy(deep=True) 

219 

220 def selectSources(self, xvals, yvals, mask): 

221 """Select saturated sources, and bright sources that are off the image. 

222 

223 Parameters 

224 ---------- 

225 xvals, yvals : `numpy.ndarray` 

226 Array of x- and y-values of bright sources to mask. 

227 mask : `lsst.afw.image.Mask` 

228 The mask plane of the image to set the BRIGHT mask plane. 

229 

230 Returns 

231 ------- 

232 candidates : `numpy.ndarray` 

233 Array of boolean flags indicating whether the given coordinates 

234 should have a diffraction spike mask calculated. 

235 """ 

236 candidates = np.zeros(len(xvals), dtype=bool) 

237 saturatedBitMask = mask.getPlaneBitMask(self.config.saturatedMaskPlane) 

238 bbox = mask.getBBox() 

239 

240 for i, (x, y) in enumerate(zip(xvals, yvals)): 

241 srcBox = lsst.geom.Box2I.makeCenteredBox(lsst.geom.Point2D(x, y), lsst.geom.Extent2I(3, 3)) 

242 if bbox.contains(srcBox): 

243 candidates[i] = np.all((mask[srcBox].array & saturatedBitMask) > 0) 

244 else: 

245 candidates[i] = True 

246 return candidates 

247 

248 def maskSources(self, xvals, yvals, radii, mask): 

249 """Apply the SPIKE mask for a given set of coordinates. The mask plane 

250 will be modified in place. 

251 

252 Parameters 

253 ---------- 

254 xvals, yvals : `numpy.ndarray` 

255 Array of x- and y-values of bright sources to mask. 

256 radii : `numpy.ndarray` 

257 Array of radius values for each bright source. 

258 mask : `lsst.afw.image.Mask` 

259 The mask plane of the image to set the BRIGHT mask plane. 

260 """ 

261 bbox = mask.getBBox() 

262 for x, y, r in zip(xvals, yvals, radii): 

263 maskSingle = self.makeSingleMask(x, y, r) 

264 singleBBox = maskSingle.getBBox() 

265 if bbox.overlaps(singleBBox): 

266 singleBBox = singleBBox.clippedTo(bbox) 

267 mask[singleBBox] |= maskSingle[singleBBox] 

268 

269 def makeSingleMask(self, x, y, r): 

270 """Create a mask plane centered on a single source with the BRIGHT mask 

271 set. This mask does not have to be fully contained in the bounding box 

272 of the mask of the science image. 

273 

274 Parameters 

275 ---------- 

276 x, y : `float` 

277 Coordinates of the source to be masked. 

278 r : `float` 

279 Expected length of a diffraction spike for a source with a 

280 magnitude 

281 

282 Returns 

283 ------- 

284 mask : `lsst.afw.image.Mask` 

285 A mask plane centered on the single source being modeled. 

286 """ 

287 polygons = [] 

288 for angle in self.angles: 

289 # Model the diffraction spike as an isosceles triangle with a long 

290 # side along the spike and a short base. For efficiency, model the 

291 # diffraction spike in the opposite direction at the same time, 

292 # so the overall shape is a narrow diamond with equal length sides. 

293 xLong = math.cos(np.deg2rad(angle))*r 

294 yLong = math.sin(np.deg2rad(angle))*r 

295 xShort = -math.sin(np.deg2rad(angle))*r/self.config.spikeAspectRatio 

296 yShort = math.cos(np.deg2rad(angle))*r/self.config.spikeAspectRatio 

297 

298 corners = [lsst.geom.Point2D(x + xLong, y + yLong), 

299 lsst.geom.Point2D(x + xShort, y + yShort), 

300 lsst.geom.Point2D(x - xLong, y - yLong), 

301 lsst.geom.Point2D(x - xShort, y - yShort)] 

302 polygons.append(afwGeom.Polygon(corners)) 

303 # Combine all of the polygons into a single region 

304 polygon = polygons[0] 

305 for poly in polygons[1:]: 

306 polygon = polygon.unionSingle(poly) 

307 bbox = lsst.geom.Box2I(polygon.getBBox()) 

308 polyImage = polygon.createImage(bbox).array 

309 mask = lsst.afw.image.Mask(bbox) 

310 mask.array[polyImage > 0] = mask.getPlaneBitMask(self.config.spikeMask) 

311 return mask 

312 

313 def set_diffraction_angle(self, exposure): 

314 """Calculate the angle of diffration spikes on the image given the 

315 camera rotation. 

316 

317 Parameters 

318 ---------- 

319 exposure : `lsst.afw.image.Exposure` 

320 The exposure to calculate the expected angle of diffraction spikes. 

321 """ 

322 angle = (exposure.visitInfo.boresightParAngle.asDegrees() 

323 - exposure.visitInfo.boresightRotAngle.asDegrees() 

324 + self.config.diffractAngle) 

325 self.angles = np.array([angle, angle + 90]) % 360 

326 

327 def calculateReferenceRadius(self, magnitudes): 

328 """Calculate the size of the region to mask for each bright star. 

329 

330 Parameters 

331 ---------- 

332 magnitudes : `numpy.ndarray` 

333 Magnitudes of the bright stars. 

334 

335 Returns 

336 ------- 

337 radii : `numpy.ndarray` 

338 Array of radius values for the given magnitudes. 

339 """ 

340 # In pixels 

341 radii = [10**(self.config.magSlope*magnitude + self.config.magOffset) for magnitude in magnitudes] 

342 return np.array(radii) 

343 

344 def extractMagnitudes(self, refCat, filterLabel): 

345 """Extract magnitude and magnitude error arrays from the given catalog. 

346 

347 Parameters 

348 ---------- 

349 refCat : `lsst.afw.table.SourceCatalog` 

350 The input reference catalog. 

351 filterLabel : `str` 

352 Label of filter being calibrated. 

353 

354 Returns 

355 ------- 

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

357 Results as a struct with attributes: 

358 

359 ``refMag`` 

360 Reference magnitude (`np.array`). 

361 ``refMagErr`` 

362 Reference magnitude error (`np.array`). 

363 ``refFluxFieldList`` 

364 A list of field names of the reference catalog used for fluxes (1 or 2 strings) (`list`). 

365 """ 

366 refSchema = refCat.schema 

367 

368 if self.config.applyColorTerms: 

369 self.log.info("Applying color terms for filter=%r, config.photoCatName=%s", 

370 filterLabel.physicalLabel, self.config.photoCatName) 

371 colorterm = self.config.colorterms.getColorterm(filterLabel.physicalLabel, 

372 self.config.photoCatName, 

373 doRaise=True) 

374 

375 refMagArr, refMagErrArr = colorterm.getCorrectedMagnitudes(refCat) 

376 fluxFieldList = [getRefFluxField(refSchema, filt) for filt in (colorterm.primary, 

377 colorterm.secondary)] 

378 else: 

379 self.log.info("Not applying color terms.") 

380 colorterm = None 

381 

382 fluxFieldList = [getRefFluxField(refSchema, filterLabel.bandLabel)] 

383 fluxField = getRefFluxField(refSchema, filterLabel.bandLabel) 

384 fluxKey = refSchema.find(fluxField).key 

385 refFluxArr = np.array(refCat.get(fluxKey)) 

386 

387 try: 

388 fluxErrKey = refSchema.find(fluxField + "Err").key 

389 refFluxErrArr = np.array(refCat.get(fluxErrKey)) 

390 except KeyError: 

391 # Reference catalogue may not have flux uncertainties; HACK DM-2308 

392 self.log.warning("Reference catalog does not have flux uncertainties for %s;" 

393 " using sqrt(flux).", fluxField) 

394 refFluxErrArr = np.sqrt(refFluxArr) 

395 

396 refMagArr = u.Quantity(refFluxArr, u.nJy).to_value(u.ABmag) 

397 # HACK convert to Jy until we have a replacement for this (DM-16903) 

398 refMagErrArr = abMagErrFromFluxErr(refFluxErrArr*1e-9, refFluxArr*1e-9) 

399 # Not all sources in the reference catalog will have flux measurements 

400 # in the current band. Since we are only masking sources use a fallback 

401 # flux measurement in these cases to get an approximate magnitude to 

402 # check for bright sources and to set the size of the mask. 

403 badMagnitudes = np.isnan(refMagArr) 

404 if np.count_nonzero(badMagnitudes): 

405 fluxName = f"{self.config.fallbackFluxField}_flux" 

406 # If the configured fallback flux field is missing, skip it. 

407 try: 

408 fallbackFluxKey = refSchema.find(fluxName).key 

409 except KeyError: 

410 self.log.warning(f"Field {fluxName} not found in photometric reference catalog." 

411 " Not masking bright stars with missing magnitudes.") 

412 else: 

413 fallbackFluxArr = np.array(refCat.get(fallbackFluxKey)) 

414 fallbackMagArr = u.Quantity(fallbackFluxArr[badMagnitudes], u.nJy).to_value(u.ABmag) 

415 refMagArr[badMagnitudes] = fallbackMagArr 

416 

417 return Struct( 

418 refMag=refMagArr, 

419 refMagErr=refMagErrArr, 

420 refFluxFieldList=fluxFieldList, 

421 ) 

422 

423 

424def getRegion(exposure, margin=None): 

425 """Calculate an enveloping region for an exposure. 

426 

427 Parameters 

428 ---------- 

429 exposure : `lsst.afw.image.Exposure` 

430 Exposure object with calibrated WCS. 

431 

432 Returns 

433 ------- 

434 region : `lsst.sphgeom.Region` 

435 Region enveloping an exposure. 

436 """ 

437 # Bounding box needs to be a `Box2D` not a `Box2I` for `wcs.pixelToSky()` 

438 bbox = lsst.geom.Box2D(exposure.getBBox()) 

439 wcs = exposure.getWcs() 

440 

441 region = lsst.sphgeom.ConvexPolygon([pp.getVector() for pp in wcs.pixelToSky(bbox.getCorners())]) 

442 if margin is not None: 

443 # This is an ad-hoc, approximate implementation. It should be good 

444 # enough for catalog loading, but is not a general-purpose solution. 

445 center = lsst.geom.SpherePoint(region.getCentroid()) 

446 corners = [lsst.geom.SpherePoint(c) for c in region.getVertices()] 

447 # Approximate the region as a Euclidian square 

448 # geom.Angle(sphgeom.Angle) converter not pybind-wrapped??? 

449 diagonal_margin = lsst.geom.Angle(margin.asRadians() * math.sqrt(2.0)) 

450 padded = [c.offset(center.bearingTo(c), diagonal_margin) for c in corners] 

451 return lsst.sphgeom.ConvexPolygon.convexHull([c.getVector() for c in padded]) 

452 

453 return region 

454 

455 

456def computePsfWidthFromMoments(psf, angle=0.): 

457 """Calculate the width of an elliptical PSF along a given direction. 

458 

459 Parameters 

460 ---------- 

461 psf : `lsst.afw.detection.Psf` 

462 The point spread function of the image. 

463 angle : `float`, optional 

464 Rotation CCW from the +x axis to calculate the width of the PSF along. 

465 

466 Returns 

467 ------- 

468 fwhm : `float` 

469 Full width at half maximum of the fitted shape of the PSF along the 

470 given `angle`. In pixels. 

471 """ 

472 psfShape = psf.computeShape(psf.getAveragePosition()) 

473 c = np.cos(np.deg2rad(angle)) 

474 s = np.sin(np.deg2rad(angle)) 

475 sigma2 = c*c*psfShape.getIxx() + 2*c*s*psfShape.getIxy() + s*s*psfShape.getIyy() 

476 sigma = np.sqrt(max(sigma2, 0.0)) # rms width in pixels 

477 

478 # 5) optional: Gaussian-equivalent FWHM along angle 

479 fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0)) * sigma 

480 return fwhm