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

203 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-14 03:16 -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"""Read preprocessed bright stars and stack them to build an extended 

23PSF model. 

24""" 

25 

26__all__ = ["FocalPlaneRegionExtendedPsf", "ExtendedPsf", "StackBrightStarsConfig", 

27 "StackBrightStarsTask", "MeasureExtendedPsfConfig", "MeasureExtendedPsfTask"] 

28 

29from dataclasses import dataclass 

30from typing import List 

31 

32from lsst.afw import image as afwImage 

33from lsst.afw import fits as afwFits 

34from lsst.afw import math as afwMath 

35from lsst.daf.base import PropertyList 

36from lsst.pipe import base as pipeBase 

37from lsst.pipe.tasks.assembleCoadd import AssembleCoaddTask 

38import lsst.pex.config as pexConfig 

39from lsst.geom import Extent2I 

40 

41 

42@dataclass 

43class FocalPlaneRegionExtendedPsf: 

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

45 

46 The focal plane region is defined through a list 

47 of detectors. 

48 

49 Parameters 

50 ---------- 

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

52 Image of the extended PSF model. 

53 detector_list : `list` [`int`] 

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

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

56 """ 

57 extended_psf_image: afwImage.MaskedImageF 

58 detector_list: List[int] 

59 

60 

61class ExtendedPsf: 

62 """Extended PSF model. 

63 

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

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

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

67 

68 Parameters 

69 ---------- 

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

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

72 """ 

73 def __init__(self, default_extended_psf=None): 

74 self.default_extended_psf = default_extended_psf 

75 self.focal_plane_regions = {} 

76 self.detectors_focal_plane_regions = {} 

77 

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

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

80 ExtendedPsf instance. 

81 

82 Parameters 

83 ---------- 

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

85 Extended PSF model for the region. 

86 region_name : `str` 

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

88 detector_list : `list` [`int`] 

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

90 """ 

91 region_name = region_name.upper() 

92 if region_name in self.focal_plane_regions: 

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

94 self.focal_plane_regions[region_name] = FocalPlaneRegionExtendedPsf( 

95 extended_psf_image=extended_psf_image, detector_list=detector_list) 

96 for det in detector_list: 

97 self.detectors_focal_plane_regions[det] = region_name 

98 

99 def __call__(self, detector=None): 

100 """Return the appropriate extended PSF. 

101 

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

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

104 whether a detector ID was passed as argument. 

105 

106 Parameters 

107 ---------- 

108 detector : `int`, optional 

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

110 determine which model to return. 

111 

112 Returns 

113 ------- 

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

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

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

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

118 extended PSF is returned. 

119 """ 

120 if detector is None: 

121 if self.default_extended_psf is None: 

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

123 return self.default_extended_psf 

124 elif not self.focal_plane_regions: 

125 return self.default_extended_psf 

126 return self.get_regional_extended_psf(detector=detector) 

127 

128 def __len__(self): 

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

130 

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

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

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

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

135 region-specific models. 

136 """ 

137 n_regions = len(self.focal_plane_regions) 

138 if self.default_extended_psf is not None: 

139 n_regions += 1 

140 return n_regions 

141 

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

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

144 

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

146 

147 Parameters 

148 ---------- 

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

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

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

152 ``detector`` is None. 

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

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

155 that includes this detector. 

156 

157 Raises 

158 ------ 

159 ValueError 

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

161 """ 

162 if detector is None: 

163 if region_name is None: 

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

165 return self.focal_plane_regions[region_name].extended_psf_image 

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

167 

168 def write_fits(self, filename): 

169 """Write this object to a file. 

170 

171 Parameters 

172 ---------- 

173 filename : `str` 

174 Name of file to write. 

175 """ 

176 # Create primary HDU with global metadata. 

177 metadata = PropertyList() 

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

179 if self.focal_plane_regions: 

180 metadata["HAS_REGIONS"] = True 

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

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

183 metadata[region] = e_psf_region.detector_list 

184 else: 

185 metadata["HAS_REGIONS"] = False 

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

187 fits_primary.createEmpty() 

188 fits_primary.writeMetadata(metadata) 

189 fits_primary.closeFile() 

190 # Write default extended PSF. 

191 if self.default_extended_psf is not None: 

192 default_hdu_metadata = PropertyList() 

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

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

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

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

197 # Write extended PSF for each focal plane region. 

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

199 metadata = PropertyList() 

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

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

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

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

204 

205 def writeFits(self, filename): 

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

207 """ 

208 self.write_fits(filename) 

209 

210 @classmethod 

211 def read_fits(cls, filename): 

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

213 

214 Parameters 

215 ---------- 

216 filename : `str` 

217 Name of the file to read. 

218 """ 

219 # Extract info from metadata. 

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

221 has_default = global_metadata.getBool("HAS_DEFAULT") 

222 if global_metadata.getBool("HAS_REGIONS"): 

223 focal_plane_region_names = global_metadata.getArray("REGION_NAMES") 

224 else: 

225 focal_plane_region_names = [] 

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

227 n_extensions = f.countHdus() 

228 extended_psf_parts = {} 

229 for j in range(1, n_extensions): 

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

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

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

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

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

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

236 continue 

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

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

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

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

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

242 # Handle default if present. 

243 if has_default: 

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

245 else: 

246 extended_psf = cls() 

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

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

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

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

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

252 # Generate extended PSF regions mappings. 

253 for r_name in focal_plane_region_names: 

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

255 detector_list = global_metadata.getArray(r_name) 

256 extended_psf.add_regional_extended_psf(extended_psf_image, r_name, detector_list) 

257 # Instantiate ExtendedPsf. 

258 return extended_psf 

259 

260 @classmethod 

261 def readFits(cls, filename): 

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

263 """ 

264 return cls.read_fits(filename) 

265 

266 

267class StackBrightStarsConfig(pexConfig.Config): 

268 """Configuration parameters for StackBrightStarsTask. 

269 """ 

270 subregion_size = pexConfig.ListField( 

271 dtype=int, 

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

273 "iteratively performed.", 

274 default=(100, 100) 

275 ) 

276 stacking_statistic = pexConfig.ChoiceField( 

277 dtype=str, 

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

279 default="MEANCLIP", 

280 allowed={ 

281 "MEAN": "mean", 

282 "MEDIAN": "median", 

283 "MEANCLIP": "clipped mean", 

284 } 

285 ) 

286 num_sigma_clip = pexConfig.Field( 

287 dtype=float, 

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

289 default=4 

290 ) 

291 num_iter = pexConfig.Field( 

292 dtype=int, 

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

294 default=3 

295 ) 

296 bad_mask_planes = pexConfig.ListField( 

297 dtype=str, 

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

299 "bright star stamps.", 

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

301 ) 

302 do_mag_cut = pexConfig.Field( 

303 dtype=bool, 

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

305 default=False 

306 ) 

307 mag_limit = pexConfig.Field( 

308 dtype=float, 

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

310 default=18 

311 ) 

312 

313 

314class StackBrightStarsTask(pipeBase.Task): 

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

316 """ 

317 ConfigClass = StackBrightStarsConfig 

318 _DefaultName = "stack_bright_stars" 

319 

320 def _set_up_stacking(self, example_stamp): 

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

322 """ 

323 stats_control = afwMath.StatisticsControl() 

324 stats_control.setNumSigmaClip(self.config.num_sigma_clip) 

325 stats_control.setNumIter(self.config.num_iter) 

326 if bad_masks := self.config.bad_mask_planes: 

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

328 for bm in bad_masks[1:]: 

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

330 stats_control.setAndMask(and_mask) 

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

332 return stats_control, stats_flags 

333 

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

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

336 

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

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

339 

340 Parameters 

341 ---------- 

342 bss_ref_list : `list` of 

343 `lsst.daf.butler._deferredDatasetHandle.DeferredDatasetHandle` 

344 List of available bright star stamps data references. 

345 region_name : `str`, optional 

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

347 logging purposes, when running over multiple such regions 

348 (typically from `MeasureExtendedPsfTask`) 

349 """ 

350 if region_name: 

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

352 else: 

353 region_message = '' 

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

355 len(bss_ref_list), region_message) 

356 # read in example set of full stamps 

357 example_bss = bss_ref_list[0].get() 

358 example_stamp = example_bss[0].stamp_im 

359 # create model image 

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

361 # divide model image into smaller subregions 

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

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

364 # compute approximate number of subregions 

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

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

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

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

369 # set up stacking statistic 

370 stats_control, stats_flags = self._set_up_stacking(example_stamp) 

371 # perform stacking 

372 for jbbox, bbox in enumerate(sub_bboxes): 

373 all_stars = None 

374 for bss_ref in bss_ref_list: 

375 read_stars = bss_ref.get(parameters={'bbox': bbox}) 

376 if self.config.do_mag_cut: 

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

378 if all_stars: 

379 all_stars.extend(read_stars) 

380 else: 

381 all_stars = read_stars 

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

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

384 ext_psf.assign(coadd_sub_bbox, bbox) 

385 return ext_psf 

386 

387 

388class MeasureExtendedPsfConnections(pipeBase.PipelineTaskConnections, 

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

390 input_brightStarStamps = pipeBase.connectionTypes.Input( 

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

392 name="brightStarStamps", 

393 storageClass="BrightStarStamps", 

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

395 deferLoad=True, 

396 multiple=True 

397 ) 

398 extended_psf = pipeBase.connectionTypes.Output( 

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

400 name="extended_psf", 

401 storageClass="ExtendedPsf", 

402 dimensions=("band",), 

403 ) 

404 

405 

406class MeasureExtendedPsfConfig(pipeBase.PipelineTaskConfig, 

407 pipelineConnections=MeasureExtendedPsfConnections): 

408 """Configuration parameters for MeasureExtendedPsfTask. 

409 """ 

410 stack_bright_stars = pexConfig.ConfigurableField( 

411 target=StackBrightStarsTask, 

412 doc="Stack selected bright stars", 

413 ) 

414 detectors_focal_plane_regions = pexConfig.DictField( 

415 keytype=int, 

416 itemtype=str, 

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

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

419 default={} 

420 ) 

421 

422 

423class MeasureExtendedPsfTask(pipeBase.Task): 

424 """Build and save extended PSF model. 

425 

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

427 preprocessed by 

428 `lsst.pipe.tasks.processBrightStars.ProcessBrightStarsTask`. 

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

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

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

432 all available data. 

433 """ 

434 ConfigClass = MeasureExtendedPsfConfig 

435 _DefaultName = "measureExtendedPsf" 

436 

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

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

439 self.makeSubtask("stack_bright_stars") 

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

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

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

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

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

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

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

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

448 self.regionless_dets = [] 

449 

450 def select_detector_refs(self, ref_list): 

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

452 regions. 

453 

454 Parameters 

455 ---------- 

456 ref_list : `list` of 

457 `lsst.daf.butler._deferredDatasetHandle.DeferredDatasetHandle` 

458 List of available bright star stamps data references. 

459 """ 

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

461 for dataset_handle in ref_list: 

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

463 if det_id in self.regionless_dets: 

464 continue 

465 try: 

466 region_name = self.config.detectors_focal_plane_regions[det_id] 

467 except KeyError: 

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

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

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

471 self.regionless_dets.append(det_id) 

472 continue 

473 region_ref_list[region_name].append(dataset_handle) 

474 return region_ref_list 

475 

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

477 input_data = butlerQC.get(inputRefs) 

478 bss_ref_list = input_data['input_brightStarStamps'] 

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

480 if not self.config.detectors_focal_plane_regions: 

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

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

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

484 else: 

485 output_e_psf = ExtendedPsf() 

486 region_ref_list = self.select_detector_refs(bss_ref_list) 

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

488 if not ref_list: 

489 # no valid references found 

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

491 'skipping it.', region_name) 

492 continue 

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

494 output_e_psf.add_regional_extended_psf(ext_psf, region_name, 

495 self.focal_plane_regions[region_name]) 

496 output = pipeBase.Struct(extended_psf=output_e_psf) 

497 butlerQC.put(output, outputRefs)