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

207 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-11 10:21 +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 

22"""Read preprocessed bright stars and stack to build an extended PSF model.""" 

23 

24__all__ = [ 

25 "FocalPlaneRegionExtendedPsf", 

26 "ExtendedPsf", 

27 "StackBrightStarsConfig", 

28 "StackBrightStarsTask", 

29 "MeasureExtendedPsfConfig", 

30 "MeasureExtendedPsfTask", 

31 "DetectorsInRegion", 

32] 

33 

34from dataclasses import dataclass 

35 

36from lsst.afw.fits import Fits, readMetadata 

37from lsst.afw.image import ImageF, MaskedImageF, MaskX 

38from lsst.afw.math import StatisticsControl, statisticsStack, stringToStatisticsProperty 

39from lsst.daf.base import PropertyList 

40from lsst.geom import Extent2I 

41from lsst.pex.config import ChoiceField, Config, ConfigDictField, ConfigurableField, Field, ListField 

42from lsst.pipe.base import PipelineTaskConfig, PipelineTaskConnections, Struct, Task 

43from lsst.pipe.base.connectionTypes import Input, Output 

44from lsst.pipe.tasks.coaddBase import subBBoxIter 

45 

46 

47def find_region_for_detector(detector_id, detectors_focal_plane_regions): 

48 """Find the focal plane region that contains a given detector. 

49 

50 Parameters 

51 ---------- 

52 detector_id : `int` 

53 The detector ID. 

54 

55 detectors_focal_plane_regions : 

56 `dict` [`str`, `lsst.pipe.tasks.extended_psf.DetectorsInRegion`] 

57 A dictionary containing focal plane region names as keys, and the 

58 corresponding detector IDs encoded within the values. 

59 

60 Returns 

61 ------- 

62 key: `str` 

63 The name of the region to which the given detector belongs. 

64 

65 Raises 

66 ------ 

67 KeyError 

68 Raised if the given detector is not included in any focal plane region. 

69 """ 

70 for region_id, detectors_in_region in detectors_focal_plane_regions.items(): 

71 if detector_id in detectors_in_region.detectors: 

72 return region_id 

73 raise KeyError( 

74 "Detector %d is not included in any focal plane region.", 

75 detector_id, 

76 ) 

77 

78 

79class DetectorsInRegion(Config): 

80 """Provides a list of detectors that define a region.""" 

81 

82 detectors = ListField[int]( 

83 doc="A list containing the detectors IDs.", 

84 default=[], 

85 ) 

86 

87 

88@dataclass 

89class FocalPlaneRegionExtendedPsf: 

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

91 

92 The focal plane region is defined through a list of detectors. 

93 

94 Parameters 

95 ---------- 

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

97 Image of the extended PSF model. 

98 region_detectors : `lsst.pipe.tasks.extended_psf.DetectorsInRegion` 

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

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

101 """ 

102 

103 extended_psf_image: MaskedImageF 

104 region_detectors: DetectorsInRegion 

105 

106 

107class ExtendedPsf: 

108 """Extended PSF model. 

109 

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

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

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

113 

114 Parameters 

115 ---------- 

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

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

118 """ 

119 

120 def __init__(self, default_extended_psf=None): 

121 self.default_extended_psf = default_extended_psf 

122 self.focal_plane_regions = {} 

123 self.detectors_focal_plane_regions = {} 

124 

125 def add_regional_extended_psf(self, extended_psf_image, region_name, region_detectors): 

126 """Add a new focal plane region, along with its extended PSF, to the 

127 ExtendedPsf instance. 

128 

129 Parameters 

130 ---------- 

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

132 Extended PSF model for the region. 

133 region_name : `str` 

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

135 region_detectors : `lsst.pipe.tasks.extended_psf.DetectorsInRegion` 

136 List of detector IDs for the detectors that define a region on the 

137 focal plane. 

138 """ 

139 region_name = region_name.upper() 

140 if region_name in self.focal_plane_regions: 

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

142 self.focal_plane_regions[region_name] = FocalPlaneRegionExtendedPsf( 

143 extended_psf_image=extended_psf_image, region_detectors=region_detectors 

144 ) 

145 self.detectors_focal_plane_regions[region_name] = region_detectors 

146 

147 def __call__(self, detector=None): 

148 """Return the appropriate extended PSF. 

149 

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

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

152 whether a detector ID was passed as argument. 

153 

154 Parameters 

155 ---------- 

156 detector : `int`, optional 

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

158 determine which model to return. 

159 

160 Returns 

161 ------- 

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

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

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

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

166 extended PSF is returned. 

167 """ 

168 if detector is None: 

169 if self.default_extended_psf is None: 

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

171 return self.default_extended_psf 

172 elif not self.focal_plane_regions: 

173 return self.default_extended_psf 

174 return self.get_extended_psf(region_name=detector) 

175 

176 def __len__(self): 

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

178 

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

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

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

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

183 region-specific models. 

184 """ 

185 n_regions = len(self.focal_plane_regions) 

186 if self.default_extended_psf is not None: 

187 n_regions += 1 

188 return n_regions 

189 

190 def get_extended_psf(self, region_name): 

191 """Returns the extended PSF for a focal plane region or detector. 

192 

193 This method takes either a region name or a detector ID as input. If 

194 the input is a `str` type, it is assumed to be the region name and if 

195 the input is a `int` type it is assumed to be the detector ID. 

196 

197 Parameters 

198 ---------- 

199 region_name : `str` or `int` 

200 Name of the region (str) or detector (int) for which the extended 

201 PSF should be retrieved. 

202 

203 Returns 

204 ------- 

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

206 The extended PSF model for the requested region or detector. 

207 

208 Raises 

209 ------ 

210 ValueError 

211 Raised if the input is not in the correct type. 

212 """ 

213 if isinstance(region_name, str): 

214 return self.focal_plane_regions[region_name].extended_psf_image 

215 elif isinstance(region_name, int): 

216 region_name = find_region_for_detector(region_name, self.detectors_focal_plane_regions) 

217 return self.focal_plane_regions[region_name].extended_psf_image 

218 else: 

219 raise ValueError("A region name with `str` type or detector number with `int` must be provided") 

220 

221 def write_fits(self, filename): 

222 """Write this object to a file. 

223 

224 Parameters 

225 ---------- 

226 filename : `str` 

227 Name of file to write. 

228 """ 

229 # Create primary HDU with global metadata. 

230 metadata = PropertyList() 

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

232 if self.focal_plane_regions: 

233 metadata["HAS_REGIONS"] = True 

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

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

236 metadata[region] = e_psf_region.region_detectors.detectors 

237 else: 

238 metadata["HAS_REGIONS"] = False 

239 fits_primary = Fits(filename, "w") 

240 fits_primary.createEmpty() 

241 fits_primary.writeMetadata(metadata) 

242 fits_primary.closeFile() 

243 # Write default extended PSF. 

244 if self.default_extended_psf is not None: 

245 default_hdu_metadata = PropertyList() 

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

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

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

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

250 # Write extended PSF for each focal plane region. 

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

252 metadata = PropertyList() 

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

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

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

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

257 

258 def writeFits(self, filename): 

259 """Alias for ``write_fits``; for compatibility with the Butler.""" 

260 self.write_fits(filename) 

261 

262 @classmethod 

263 def read_fits(cls, filename): 

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

265 

266 Parameters 

267 ---------- 

268 filename : `str` 

269 Name of the file to read. 

270 """ 

271 # Extract info from metadata. 

272 global_metadata = readMetadata(filename, hdu=0) 

273 has_default = global_metadata.getBool("HAS_DEFAULT") 

274 if global_metadata.getBool("HAS_REGIONS"): 

275 focal_plane_region_names = global_metadata.getArray("REGION_NAMES") 

276 else: 

277 focal_plane_region_names = [] 

278 f = Fits(filename, "r") 

279 n_extensions = f.countHdus() 

280 extended_psf_parts = {} 

281 for j in range(1, n_extensions): 

282 md = readMetadata(filename, hdu=j) 

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

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

285 default_image = ImageF(filename, hdu=j) 

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

287 default_mask = MaskX(filename, hdu=j) 

288 continue 

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

290 extended_psf_part = ImageF(filename, hdu=j) 

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

292 extended_psf_part = MaskX(filename, hdu=j) 

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

294 # Handle default if present. 

295 if has_default: 

296 extended_psf = cls(MaskedImageF(default_image, default_mask)) 

297 else: 

298 extended_psf = cls() 

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

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

301 raise ValueError( 

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

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

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

305 ) 

306 # Generate extended PSF regions mappings. 

307 for r_name in focal_plane_region_names: 

308 extended_psf_image = MaskedImageF(**extended_psf_parts[r_name]) 

309 region_detectors = DetectorsInRegion() 

310 region_detectors.detectors = global_metadata.getArray(r_name) 

311 extended_psf.add_regional_extended_psf(extended_psf_image, r_name, region_detectors) 

312 # Instantiate ExtendedPsf. 

313 return extended_psf 

314 

315 @classmethod 

316 def readFits(cls, filename): 

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

318 return cls.read_fits(filename) 

319 

320 

321class StackBrightStarsConfig(Config): 

322 """Configuration parameters for StackBrightStarsTask.""" 

323 

324 subregion_size = ListField[int]( 

325 doc="Size, in pixels, of the subregions over which the stacking will be " "iteratively performed.", 

326 default=(100, 100), 

327 ) 

328 stacking_statistic = ChoiceField[str]( 

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

330 default="MEANCLIP", 

331 allowed={ 

332 "MEAN": "mean", 

333 "MEDIAN": "median", 

334 "MEANCLIP": "clipped mean", 

335 }, 

336 ) 

337 num_sigma_clip = Field[float]( 

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

339 default=4, 

340 ) 

341 num_iter = Field[int]( 

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

343 default=3, 

344 ) 

345 bad_mask_planes = ListField[str]( 

346 doc="Mask planes that define pixels to be excluded from the stacking of the bright star stamps.", 

347 default=("BAD", "CR", "CROSSTALK", "EDGE", "NO_DATA", "SAT", "SUSPECT", "UNMASKEDNAN"), 

348 ) 

349 do_mag_cut = Field[bool]( 

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

351 default=False, 

352 ) 

353 mag_limit = Field[float]( 

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

355 default=18, 

356 ) 

357 

358 

359class StackBrightStarsTask(Task): 

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

361 

362 ConfigClass = StackBrightStarsConfig 

363 _DefaultName = "stack_bright_stars" 

364 

365 def _set_up_stacking(self, example_stamp): 

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

367 stats_control = StatisticsControl( 

368 numSigmaClip=self.config.num_sigma_clip, 

369 numIter=self.config.num_iter, 

370 ) 

371 if bad_masks := self.config.bad_mask_planes: 

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

373 for bm in bad_masks[1:]: 

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

375 stats_control.setAndMask(and_mask) 

376 stats_flags = stringToStatisticsProperty(self.config.stacking_statistic) 

377 return stats_control, stats_flags 

378 

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

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

381 

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

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

384 

385 Parameters 

386 ---------- 

387 bss_ref_list : `list` of 

388 `lsst.daf.butler._deferredDatasetHandle.DeferredDatasetHandle` 

389 List of available bright star stamps data references. 

390 region_name : `str`, optional 

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

392 logging purposes, when running over multiple such regions 

393 (typically from `MeasureExtendedPsfTask`) 

394 """ 

395 if region_name: 

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

397 else: 

398 region_message = "." 

399 self.log.info( 

400 "Building extended PSF from stamps extracted from %d detector images%s", 

401 len(bss_ref_list), 

402 region_message, 

403 ) 

404 # read in example set of full stamps 

405 example_bss = bss_ref_list[0].get() 

406 example_stamp = example_bss[0].stamp_im 

407 # create model image 

408 ext_psf = MaskedImageF(example_stamp.getBBox()) 

409 # divide model image into smaller subregions 

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

411 sub_bboxes = subBBoxIter(ext_psf.getBBox(), subregion_size) 

412 # compute approximate number of subregions 

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

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

415 ) 

416 self.log.info( 

417 "Stacking performed iteratively over approximately %d smaller areas of the final model image.", 

418 n_subregions, 

419 ) 

420 # set up stacking statistic 

421 stats_control, stats_flags = self._set_up_stacking(example_stamp) 

422 # perform stacking 

423 for jbbox, bbox in enumerate(sub_bboxes): 

424 all_stars = None 

425 for bss_ref in bss_ref_list: 

426 read_stars = bss_ref.get(parameters={"bbox": bbox}) 

427 if self.config.do_mag_cut: 

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

429 if all_stars: 

430 all_stars.extend(read_stars) 

431 else: 

432 all_stars = read_stars 

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

434 coadd_sub_bbox = statisticsStack(all_stars.getMaskedImages(), stats_flags, stats_control) 

435 ext_psf.assign(coadd_sub_bbox, bbox) 

436 return ext_psf 

437 

438 

439class MeasureExtendedPsfConnections(PipelineTaskConnections, dimensions=("band", "instrument")): 

440 input_brightStarStamps = Input( 

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

442 name="brightStarStamps", 

443 storageClass="BrightStarStamps", 

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

445 deferLoad=True, 

446 multiple=True, 

447 ) 

448 extended_psf = Output( 

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

450 name="extended_psf", 

451 storageClass="ExtendedPsf", 

452 dimensions=("band",), 

453 ) 

454 

455 

456class MeasureExtendedPsfConfig(PipelineTaskConfig, pipelineConnections=MeasureExtendedPsfConnections): 

457 """Configuration parameters for MeasureExtendedPsfTask.""" 

458 

459 stack_bright_stars = ConfigurableField( 

460 target=StackBrightStarsTask, 

461 doc="Stack selected bright stars", 

462 ) 

463 detectors_focal_plane_regions = ConfigDictField( 

464 keytype=str, 

465 itemtype=DetectorsInRegion, 

466 doc=( 

467 "Mapping from focal plane region names to detector IDs. " 

468 "If empty, a constant extended PSF model is built from all selected bright stars. " 

469 "It's possible for a single detector to be included in multiple regions if so desired." 

470 ), 

471 default={}, 

472 ) 

473 

474 

475class MeasureExtendedPsfTask(Task): 

476 """Build and save extended PSF model. 

477 

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

479 preprocessed by 

480 `lsst.pipe.tasks.processBrightStars.ProcessBrightStarsTask`. 

481 

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

483 different extended PSF model will be built for each focal plane region. If 

484 not, a single constant extended PSF model is built with all available data. 

485 """ 

486 

487 ConfigClass = MeasureExtendedPsfConfig 

488 _DefaultName = "measureExtendedPsf" 

489 

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

491 super().__init__(*args, **kwargs) 

492 self.makeSubtask("stack_bright_stars") 

493 self.detectors_focal_plane_regions = self.config.detectors_focal_plane_regions 

494 self.regionless_dets = [] 

495 

496 def select_detector_refs(self, ref_list): 

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

498 regions. 

499 

500 Parameters 

501 ---------- 

502 ref_list : `list` of 

503 `lsst.daf.butler._deferredDatasetHandle.DeferredDatasetHandle` 

504 List of available bright star stamps data references. 

505 """ 

506 region_ref_list = {region: [] for region in self.detectors_focal_plane_regions.keys()} 

507 for dataset_handle in ref_list: 

508 detector_id = dataset_handle.ref.dataId["detector"] 

509 if detector_id in self.regionless_dets: 

510 continue 

511 try: 

512 region_name = find_region_for_detector(detector_id, self.detectors_focal_plane_regions) 

513 except KeyError: 

514 self.log.warning( 

515 "Bright stars were available for detector %d, but it was missing from the %s config " 

516 "field, so they will not be used to build any of the extended PSF models.", 

517 detector_id, 

518 "'detectors_focal_plane_regions'", 

519 ) 

520 self.regionless_dets.append(detector_id) 

521 continue 

522 region_ref_list[region_name].append(dataset_handle) 

523 return region_ref_list 

524 

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

526 input_data = butlerQC.get(inputRefs) 

527 bss_ref_list = input_data["input_brightStarStamps"] 

528 if not self.config.detectors_focal_plane_regions: 

529 self.log.info( 

530 "No detector groups were provided to MeasureExtendedPsfTask; computing a single, " 

531 "constant extended PSF model over all available observations." 

532 ) 

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

534 else: 

535 output_e_psf = ExtendedPsf() 

536 region_ref_list = self.select_detector_refs(bss_ref_list) 

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

538 if not ref_list: 

539 # no valid references found 

540 self.log.warning( 

541 "No valid brightStarStamps reference found for region '%s'; skipping it.", 

542 region_name, 

543 ) 

544 continue 

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

546 output_e_psf.add_regional_extended_psf( 

547 ext_psf, region_name, self.detectors_focal_plane_regions[region_name] 

548 ) 

549 output = Struct(extended_psf=output_e_psf) 

550 butlerQC.put(output, outputRefs)