Coverage for python/lsst/pipe/tasks/extended_psf.py: 23%

202 statements  

« prev     ^ index     » next       coverage.py v6.4.4, created at 2022-09-11 02:04 -0700

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"""Read preprocessed bright stars and stack them to build an extended 

23PSF model. 

24""" 

25 

26from dataclasses import dataclass 

27from typing import List 

28 

29from lsst.afw import image as afwImage 

30from lsst.afw import fits as afwFits 

31from lsst.afw import math as afwMath 

32from lsst.daf.base import PropertyList 

33from lsst.pipe import base as pipeBase 

34from lsst.pipe.tasks.assembleCoadd import AssembleCoaddTask 

35import lsst.pex.config as pexConfig 

36from lsst.geom import Extent2I 

37 

38 

39@dataclass 

40class FocalPlaneRegionExtendedPsf: 

41 """Single extended PSF over a focal plane region. 

42 

43 The focal plane region is defined through a list 

44 of detectors. 

45 

46 Parameters 

47 ---------- 

48 extended_psf_image : `lsst.afw.image.MaskedImageF` 

49 Image of the extended PSF model. 

50 detector_list : `list` [`int`] 

51 List of detector IDs that define the focal plane region over which this 

52 extended PSF model has been built (and can be used). 

53 """ 

54 extended_psf_image: afwImage.MaskedImageF 

55 detector_list: List[int] 

56 

57 

58class ExtendedPsf: 

59 """Extended PSF model. 

60 

61 Each instance may contain a default extended PSF, a set of extended PSFs 

62 that correspond to different focal plane regions, or both. At this time, 

63 focal plane regions are always defined as a subset of detectors. 

64 

65 Parameters 

66 ---------- 

67 default_extended_psf : `lsst.afw.image.MaskedImageF` 

68 Extended PSF model to be used as default (or only) extended PSF model. 

69 """ 

70 def __init__(self, default_extended_psf=None): 

71 self.default_extended_psf = default_extended_psf 

72 self.focal_plane_regions = {} 

73 self.detectors_focal_plane_regions = {} 

74 

75 def add_regional_extended_psf(self, extended_psf_image, region_name, detector_list): 

76 """Add a new focal plane region, along wit hits extended PSF, to the 

77 ExtendedPsf instance. 

78 

79 Parameters 

80 ---------- 

81 extended_psf_image : `lsst.afw.image.MaskedImageF` 

82 Extended PSF model for the region. 

83 region_name : `str` 

84 Name of the focal plane region. Will be converted to all-uppercase. 

85 detector_list : `list` [`int`] 

86 List of IDs for the detectors that define the focal plane region. 

87 """ 

88 region_name = region_name.upper() 

89 if region_name in self.focal_plane_regions: 

90 raise ValueError(f"Region name {region_name} is already used by this ExtendedPsf instance.") 

91 self.focal_plane_regions[region_name] = FocalPlaneRegionExtendedPsf( 

92 extended_psf_image=extended_psf_image, detector_list=detector_list) 

93 for det in detector_list: 

94 self.detectors_focal_plane_regions[det] = region_name 

95 

96 def __call__(self, detector=None): 

97 """Return the appropriate extended PSF. 

98 

99 If the instance contains no extended PSF defined over focal plane 

100 regions, the default extended PSF will be returned regardless of 

101 whether a detector ID was passed as argument. 

102 

103 Parameters 

104 ---------- 

105 detector : `int`, optional 

106 Detector ID. If focal plane region PSFs are defined, is used to 

107 determine which model to return. 

108 

109 Returns 

110 ------- 

111 extendedPsfImage : `lsst.afw.image.MaskedImageF` 

112 The extended PSF model. If this instance contains extended PSFs 

113 defined over focal plane regions, the extended PSF model for the 

114 region that contains ``detector`` is returned. If not, the default 

115 extended PSF is returned. 

116 """ 

117 if detector is None: 

118 if self.default_extended_psf is None: 

119 raise ValueError("No default extended PSF available; please provide detector number.") 

120 return self.default_extended_psf 

121 elif not self.focal_plane_regions: 

122 return self.default_extended_psf 

123 return self.get_regional_extended_psf(detector=detector) 

124 

125 def __len__(self): 

126 """Returns the number of extended PSF models present in the instance. 

127 

128 Note that if the instance contains both a default model and a set of 

129 focal plane region models, the length of the instance will be the 

130 number of regional models, plus one (the default). This is true even 

131 in the case where the default model is one of the focal plane 

132 region-specific models. 

133 """ 

134 n_regions = len(self.focal_plane_regions) 

135 if self.default_extended_psf is not None: 

136 n_regions += 1 

137 return n_regions 

138 

139 def get_regional_extended_psf(self, region_name=None, detector=None): 

140 """Returns the extended PSF for a focal plane region. 

141 

142 The region can be identified either by name, or through a detector ID. 

143 

144 Parameters 

145 ---------- 

146 region_name : `str` or `None`, optional 

147 Name of the region for which the extended PSF should be retrieved. 

148 Ignored if ``detector`` is provided. Must be provided if 

149 ``detector`` is None. 

150 detector : `int` or `None`, optional 

151 If provided, returns the extended PSF for the focal plane region 

152 that includes this detector. 

153 

154 Raises 

155 ------ 

156 ValueError 

157 Raised if neither ``detector`` nor ``regionName`` is provided. 

158 """ 

159 if detector is None: 

160 if region_name is None: 

161 raise ValueError("One of either a regionName or a detector number must be provided.") 

162 return self.focal_plane_regions[region_name].extended_psf_image 

163 return self.focal_plane_regions[self.detectors_focal_plane_regions[detector]].extended_psf_image 

164 

165 def write_fits(self, filename): 

166 """Write this object to a file. 

167 

168 Parameters 

169 ---------- 

170 filename : `str` 

171 Name of file to write. 

172 """ 

173 # Create primary HDU with global metadata. 

174 metadata = PropertyList() 

175 metadata["HAS_DEFAULT"] = self.default_extended_psf is not None 

176 if self.focal_plane_regions: 

177 metadata["HAS_REGIONS"] = True 

178 metadata["REGION_NAMES"] = list(self.focal_plane_regions.keys()) 

179 for region, e_psf_region in self.focal_plane_regions.items(): 

180 metadata[region] = e_psf_region.detector_list 

181 else: 

182 metadata["HAS_REGIONS"] = False 

183 fits_primary = afwFits.Fits(filename, "w") 

184 fits_primary.createEmpty() 

185 fits_primary.writeMetadata(metadata) 

186 fits_primary.closeFile() 

187 # Write default extended PSF. 

188 if self.default_extended_psf is not None: 

189 default_hdu_metadata = PropertyList() 

190 default_hdu_metadata.update({"REGION": "DEFAULT", "EXTNAME": "IMAGE"}) 

191 self.default_extended_psf.image.writeFits(filename, metadata=default_hdu_metadata, mode="a") 

192 default_hdu_metadata.update({"REGION": "DEFAULT", "EXTNAME": "MASK"}) 

193 self.default_extended_psf.mask.writeFits(filename, metadata=default_hdu_metadata, mode="a") 

194 # Write extended PSF for each focal plane region. 

195 for j, (region, e_psf_region) in enumerate(self.focal_plane_regions.items()): 

196 metadata = PropertyList() 

197 metadata.update({"REGION": region, "EXTNAME": "IMAGE"}) 

198 e_psf_region.extended_psf_image.image.writeFits(filename, metadata=metadata, mode="a") 

199 metadata.update({"REGION": region, "EXTNAME": "MASK"}) 

200 e_psf_region.extended_psf_image.mask.writeFits(filename, metadata=metadata, mode="a") 

201 

202 def writeFits(self, filename): 

203 """Alias for ``write_fits``; exists for compatibility with the Butler. 

204 """ 

205 self.write_fits(filename) 

206 

207 @classmethod 

208 def read_fits(cls, filename): 

209 """Build an instance of this class from a file. 

210 

211 Parameters 

212 ---------- 

213 filename : `str` 

214 Name of the file to read. 

215 """ 

216 # Extract info from metadata. 

217 global_metadata = afwFits.readMetadata(filename, hdu=0) 

218 has_default = global_metadata.getBool("HAS_DEFAULT") 

219 if global_metadata.getBool("HAS_REGIONS"): 

220 focal_plane_region_names = global_metadata.getArray("REGION_NAMES") 

221 else: 

222 focal_plane_region_names = [] 

223 f = afwFits.Fits(filename, "r") 

224 n_extensions = f.countHdus() 

225 extended_psf_parts = {} 

226 for j in range(1, n_extensions): 

227 md = afwFits.readMetadata(filename, hdu=j) 

228 if has_default and md["REGION"] == "DEFAULT": 

229 if md["EXTNAME"] == "IMAGE": 

230 default_image = afwImage.ImageF(filename, hdu=j) 

231 elif md["EXTNAME"] == "MASK": 

232 default_mask = afwImage.MaskX(filename, hdu=j) 

233 continue 

234 if md["EXTNAME"] == "IMAGE": 

235 extended_psf_part = afwImage.ImageF(filename, hdu=j) 

236 elif md["EXTNAME"] == "MASK": 

237 extended_psf_part = afwImage.MaskX(filename, hdu=j) 

238 extended_psf_parts.setdefault(md["REGION"], {})[md["EXTNAME"].lower()] = extended_psf_part 

239 # Handle default if present. 

240 if has_default: 

241 extended_psf = cls(afwImage.MaskedImageF(default_image, default_mask)) 

242 else: 

243 extended_psf = cls() 

244 # Ensure we recovered an extended PSF for all focal plane regions. 

245 if len(extended_psf_parts) != len(focal_plane_region_names): 

246 raise ValueError(f"Number of per-region extended PSFs read ({len(extended_psf_parts)}) does not " 

247 "match with the number of regions recorded in the metadata " 

248 f"({len(focal_plane_region_names)}).") 

249 # Generate extended PSF regions mappings. 

250 for r_name in focal_plane_region_names: 

251 extended_psf_image = afwImage.MaskedImageF(**extended_psf_parts[r_name]) 

252 detector_list = global_metadata.getArray(r_name) 

253 extended_psf.add_regional_extended_psf(extended_psf_image, r_name, detector_list) 

254 # Instantiate ExtendedPsf. 

255 return extended_psf 

256 

257 @classmethod 

258 def readFits(cls, filename): 

259 """Alias for ``readFits``; exists for compatibility with the Butler. 

260 """ 

261 return cls.read_fits(filename) 

262 

263 

264class StackBrightStarsConfig(pexConfig.Config): 

265 """Configuration parameters for StackBrightStarsTask. 

266 """ 

267 subregion_size = pexConfig.ListField( 

268 dtype=int, 

269 doc="Size, in pixels, of the subregions over which the stacking will be " 

270 "iteratively performed.", 

271 default=(100, 100) 

272 ) 

273 stacking_statistic = pexConfig.ChoiceField( 

274 dtype=str, 

275 doc="Type of statistic to use for stacking.", 

276 default="MEANCLIP", 

277 allowed={ 

278 "MEAN": "mean", 

279 "MEDIAN": "median", 

280 "MEANCLIP": "clipped mean", 

281 } 

282 ) 

283 num_sigma_clip = pexConfig.Field( 

284 dtype=float, 

285 doc="Sigma for outlier rejection; ignored if stacking_statistic != 'MEANCLIP'.", 

286 default=4 

287 ) 

288 num_iter = pexConfig.Field( 

289 dtype=int, 

290 doc="Number of iterations of outlier rejection; ignored if stackingStatistic != 'MEANCLIP'.", 

291 default=3 

292 ) 

293 bad_mask_planes = pexConfig.ListField( 

294 dtype=str, 

295 doc="Mask planes that, if set, lead to associated pixels not being included in the stacking of the " 

296 "bright star stamps.", 

297 default=('BAD', 'CR', 'CROSSTALK', 'EDGE', 'NO_DATA', 'SAT', 'SUSPECT', 'UNMASKEDNAN') 

298 ) 

299 do_mag_cut = pexConfig.Field( 

300 dtype=bool, 

301 doc="Apply magnitude cut before stacking?", 

302 default=False 

303 ) 

304 mag_limit = pexConfig.Field( 

305 dtype=float, 

306 doc="Magnitude limit, in Gaia G; all stars brighter than this value will be stacked", 

307 default=18 

308 ) 

309 

310 

311class StackBrightStarsTask(pipeBase.Task): 

312 """Stack bright stars together to build an extended PSF model. 

313 """ 

314 ConfigClass = StackBrightStarsConfig 

315 _DefaultName = "stack_bright_stars" 

316 

317 def _set_up_stacking(self, example_stamp): 

318 """Configure stacking statistic and control from config fields. 

319 """ 

320 stats_control = afwMath.StatisticsControl() 

321 stats_control.setNumSigmaClip(self.config.num_sigma_clip) 

322 stats_control.setNumIter(self.config.num_iter) 

323 if bad_masks := self.config.bad_mask_planes: 

324 and_mask = example_stamp.mask.getPlaneBitMask(bad_masks[0]) 

325 for bm in bad_masks[1:]: 

326 and_mask = and_mask | example_stamp.mask.getPlaneBitMask(bm) 

327 stats_control.setAndMask(and_mask) 

328 stats_flags = afwMath.stringToStatisticsProperty(self.config.stacking_statistic) 

329 return stats_control, stats_flags 

330 

331 def run(self, bss_ref_list, region_name=None): 

332 """Read input bright star stamps and stack them together. 

333 

334 The stacking is done iteratively over smaller areas of the final model 

335 image to allow for a great number of bright star stamps to be used. 

336 

337 Parameters 

338 ---------- 

339 bss_ref_list : `list` of 

340 `lsst.daf.butler._deferredDatasetHandle.DeferredDatasetHandle` 

341 List of available bright star stamps data references. 

342 region_name : `str`, optional 

343 Name of the focal plane region, if applicable. Only used for 

344 logging purposes, when running over multiple such regions 

345 (typically from `MeasureExtendedPsfTask`) 

346 """ 

347 if region_name: 

348 region_message = f' for region "{region_name}".' 

349 else: 

350 region_message = '' 

351 self.log.info('Building extended PSF from stamps extracted from %d detector images%s', 

352 len(bss_ref_list), region_message) 

353 # read in example set of full stamps 

354 example_bss = bss_ref_list[0].get(datasetType="brightStarStamps", immediate=True) 

355 example_stamp = example_bss[0].stamp_im 

356 # create model image 

357 ext_psf = afwImage.MaskedImageF(example_stamp.getBBox()) 

358 # divide model image into smaller subregions 

359 subregion_size = Extent2I(*self.config.subregion_size) 

360 sub_bboxes = AssembleCoaddTask._subBBoxIter(ext_psf.getBBox(), subregion_size) 

361 # compute approximate number of subregions 

362 n_subregions = int(ext_psf.getDimensions()[0]/subregion_size[0] + 1)*int( 

363 ext_psf.getDimensions()[1]/subregion_size[1] + 1) 

364 self.log.info("Stacking will performed iteratively over approximately %d " 

365 "smaller areas of the final model image.", n_subregions) 

366 # set up stacking statistic 

367 stats_control, stats_flags = self._set_up_stacking(example_stamp) 

368 # perform stacking 

369 for jbbox, bbox in enumerate(sub_bboxes): 

370 all_stars = None 

371 for bss_ref in bss_ref_list: 

372 read_stars = bss_ref.get(datasetType="brightStarStamps", parameters={'bbox': bbox}) 

373 if self.config.do_mag_cut: 

374 read_stars = read_stars.selectByMag(magMax=self.config.mag_limit) 

375 if all_stars: 

376 all_stars.extend(read_stars) 

377 else: 

378 all_stars = read_stars 

379 # TODO: DM-27371 add weights to bright stars for stacking 

380 coadd_sub_bbox = afwMath.statisticsStack(all_stars.getMaskedImages(), stats_flags, stats_control) 

381 ext_psf.assign(coadd_sub_bbox, bbox) 

382 return ext_psf 

383 

384 

385class MeasureExtendedPsfConnections(pipeBase.PipelineTaskConnections, 

386 dimensions=("band", "instrument")): 

387 input_brightStarStamps = pipeBase.connectionTypes.Input( 

388 doc="Input list of bright star collections to be stacked.", 

389 name="brightStarStamps", 

390 storageClass="BrightStarStamps", 

391 dimensions=("visit", "detector"), 

392 deferLoad=True, 

393 multiple=True 

394 ) 

395 extended_psf = pipeBase.connectionTypes.Output( 

396 doc="Extended PSF model built by stacking bright stars.", 

397 name="extended_psf", 

398 storageClass="ExtendedPsf", 

399 dimensions=("band",), 

400 ) 

401 

402 

403class MeasureExtendedPsfConfig(pipeBase.PipelineTaskConfig, 

404 pipelineConnections=MeasureExtendedPsfConnections): 

405 """Configuration parameters for MeasureExtendedPsfTask. 

406 """ 

407 stack_bright_stars = pexConfig.ConfigurableField( 

408 target=StackBrightStarsTask, 

409 doc="Stack selected bright stars", 

410 ) 

411 detectors_focal_plane_regions = pexConfig.DictField( 

412 keytype=int, 

413 itemtype=str, 

414 doc="Mapping from detector IDs to focal plane region names. If empty, a constant " 

415 "extended PSF model is built from all selected bright stars.", 

416 default={} 

417 ) 

418 

419 

420class MeasureExtendedPsfTask(pipeBase.Task): 

421 """Build and save extended PSF model. 

422 

423 The model is built by stacking bright star stamps, extracted and 

424 preprocessed by 

425 `lsst.pipe.tasks.processBrightStars.ProcessBrightStarsTask`. 

426 If a mapping from detector IDs to focal plane regions is provided, 

427 a different extended PSF model will be built for each focal plane 

428 region. If not, a single, constant extended PSF model is built using 

429 all available data. 

430 """ 

431 ConfigClass = MeasureExtendedPsfConfig 

432 _DefaultName = "measureExtendedPsf" 

433 

434 def __init__(self, initInputs=None, *args, **kwargs): 

435 pipeBase.Task.__init__(self, *args, **kwargs) 

436 self.makeSubtask("stack_bright_stars") 

437 self.focal_plane_regions = {region: [] for region in 

438 set(self.config.detectors_focal_plane_regions.values())} 

439 for det, region in self.config.detectors_focal_plane_regions.items(): 

440 self.focal_plane_regions[region].append(det) 

441 # make no assumption on what detector IDs should be, but if we come 

442 # across one where there are processed bright stars, but no 

443 # corresponding focal plane region, make sure we keep track of 

444 # it (eg to raise a warning only once) 

445 self.regionless_dets = [] 

446 

447 def select_detector_refs(self, ref_list): 

448 """Split available sets of bright star stamps according to focal plane 

449 regions. 

450 

451 Parameters 

452 ---------- 

453 ref_list : `list` of 

454 `lsst.daf.butler._deferredDatasetHandle.DeferredDatasetHandle` 

455 List of available bright star stamps data references. 

456 """ 

457 region_ref_list = {region: [] for region in self.focal_plane_regions.keys()} 

458 for dataset_handle in ref_list: 

459 det_id = dataset_handle.ref.dataId["detector"] 

460 if det_id in self.regionless_dets: 

461 continue 

462 try: 

463 region_name = self.config.detectors_focal_plane_regions[det_id] 

464 except KeyError: 

465 self.log.warning('Bright stars were available for detector %d, but it was missing ' 

466 'from the "detectors_focal_plane_regions" config field, so they will not ' 

467 'be used to build any of the extended PSF models', det_id) 

468 self.regionless_dets.append(det_id) 

469 continue 

470 region_ref_list[region_name].append(dataset_handle) 

471 return region_ref_list 

472 

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

474 input_data = butlerQC.get(inputRefs) 

475 bss_ref_list = input_data['input_brightStarStamps'] 

476 # Handle default case of a single region with empty detector list 

477 if not self.config.detectors_focal_plane_regions: 

478 self.log.info("No detector groups were provided to MeasureExtendedPsfTask; computing a single, " 

479 "constant extended PSF model over all available observations.") 

480 output_e_psf = ExtendedPsf(self.stack_bright_stars.run(bss_ref_list)) 

481 else: 

482 output_e_psf = ExtendedPsf() 

483 region_ref_list = self.select_detector_refs(bss_ref_list) 

484 for region_name, ref_list in region_ref_list.items(): 

485 if not ref_list: 

486 # no valid references found 

487 self.log.warning('No valid brightStarStamps reference found for region "%s"; ' 

488 'skipping it.', region_name) 

489 continue 

490 ext_psf = self.stack_bright_stars.run(ref_list, region_name) 

491 output_e_psf.add_regional_extended_psf(ext_psf, region_name, 

492 self.focal_plane_regions[region_name]) 

493 output = pipeBase.Struct(extended_psf=output_e_psf) 

494 butlerQC.put(output, outputRefs)