Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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.CmdLineTask): 

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

313 """ 

314 ConfigClass = StackBrightStarsConfig 

315 _DefaultName = "stack_bright_stars" 

316 

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

318 pipeBase.CmdLineTask.__init__(self, *args, **kwargs) 

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 log_message = f'Building extended PSF from stamps extracted from {len(bss_ref_list)} detector images' 

351 if region_name: 

352 log_message += f' for region "{region_name}".' 

353 self.log.info(log_message) 

354 # read in example set of full stamps 

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

356 example_stamp = example_bss[0].stamp_im 

357 # create model image 

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

359 # divide model image into smaller subregions 

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

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

362 # compute approximate number of subregions 

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

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

365 self.log.info(f"Stacking will performed iteratively over approximately {n_subregions} " 

366 "smaller areas of the final model image.") 

367 # set up stacking statistic 

368 stats_control, stats_flags = self._set_up_stacking(example_stamp) 

369 # perform stacking 

370 for jbbox, bbox in enumerate(sub_bboxes): 

371 all_stars = None 

372 for bss_ref in bss_ref_list: 

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

374 if self.config.do_mag_cut: 

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

376 if all_stars: 

377 all_stars.extend(read_stars) 

378 else: 

379 all_stars = read_stars 

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

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

382 ext_psf.assign(coadd_sub_bbox, bbox) 

383 return ext_psf 

384 

385 

386class MeasureExtendedPsfConnections(pipeBase.PipelineTaskConnections, 

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

388 input_brightStarStamps = pipeBase.connectionTypes.Input( 

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

390 name="brightStarStamps", 

391 storageClass="BrightStarStamps", 

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

393 deferLoad=True, 

394 multiple=True 

395 ) 

396 extended_psf = pipeBase.connectionTypes.Output( 

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

398 name="extended_psf", 

399 storageClass="ExtendedPsf", 

400 dimensions=("band",), 

401 ) 

402 

403 

404class MeasureExtendedPsfConfig(pipeBase.PipelineTaskConfig, 

405 pipelineConnections=MeasureExtendedPsfConnections): 

406 """Configuration parameters for MeasureExtendedPsfTask. 

407 """ 

408 stack_bright_stars = pexConfig.ConfigurableField( 

409 target=StackBrightStarsTask, 

410 doc="Stack selected bright stars", 

411 ) 

412 detectors_focal_plane_regions = pexConfig.DictField( 

413 keytype=int, 

414 itemtype=str, 

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

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

417 default={} 

418 ) 

419 

420 

421class MeasureExtendedPsfTask(pipeBase.CmdLineTask): 

422 """Build and save extended PSF model. 

423 

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

425 preprocessed by 

426 `lsst.pipe.tasks.processBrightStars.ProcessBrightStarsTask`. 

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

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

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

430 all available data. 

431 """ 

432 ConfigClass = MeasureExtendedPsfConfig 

433 _DefaultName = "measureExtendedPsf" 

434 

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

436 pipeBase.CmdLineTask.__init__(self, *args, **kwargs) 

437 self.makeSubtask("stack_bright_stars") 

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

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

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

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

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

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

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

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

446 self.regionless_dets = [] 

447 

448 def select_detector_refs(self, ref_list): 

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

450 regions. 

451 

452 Parameters 

453 ---------- 

454 ref_list : `list` of 

455 `lsst.daf.butler._deferredDatasetHandle.DeferredDatasetHandle` 

456 List of available bright star stamps data references. 

457 """ 

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

459 for dataset_handle in ref_list: 

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

461 if det_id in self.regionless_dets: 

462 continue 

463 try: 

464 region_name = self.config.detectors_focal_plane_regions[det_id] 

465 except KeyError: 

466 self.log.warn(f'Bright stars were available for detector {det_id}, but it was missing ' 

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

468 'be used to build any of the extended PSF models') 

469 self.regionless_dets.append(det_id) 

470 continue 

471 region_ref_list[region_name].append(dataset_handle) 

472 return region_ref_list 

473 

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

475 input_data = butlerQC.get(inputRefs) 

476 bss_ref_list = input_data['input_brightStarStamps'] 

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

478 if not self.config.detectors_focal_plane_regions: 

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

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

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

482 else: 

483 output_e_psf = ExtendedPsf() 

484 region_ref_list = self.select_detector_refs(bss_ref_list) 

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

486 if not ref_list: 

487 # no valid references found 

488 self.log.warn(f'No valid brightStarStamps reference found for region "{region_name}"; ' 

489 'skipping it.') 

490 continue 

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

492 output_e_psf.add_regional_extended_psf(ext_psf, region_name, 

493 self.focal_plane_regions[region_name]) 

494 output = pipeBase.Struct(extended_psf=output_e_psf) 

495 butlerQC.put(output, outputRefs)