Coverage for python/lsst/pipe/tasks/extended_psf.py: 23%
202 statements
« prev ^ index » next coverage.py v6.4.4, created at 2022-08-18 12:37 -0700
« prev ^ index » next coverage.py v6.4.4, created at 2022-08-18 12:37 -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"""
26from dataclasses import dataclass
27from typing import List
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
39@dataclass
40class FocalPlaneRegionExtendedPsf:
41 """Single extended PSF over a focal plane region.
43 The focal plane region is defined through a list
44 of detectors.
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]
58class ExtendedPsf:
59 """Extended PSF model.
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.
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 = {}
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.
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
96 def __call__(self, detector=None):
97 """Return the appropriate extended PSF.
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.
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.
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)
125 def __len__(self):
126 """Returns the number of extended PSF models present in the instance.
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
139 def get_regional_extended_psf(self, region_name=None, detector=None):
140 """Returns the extended PSF for a focal plane region.
142 The region can be identified either by name, or through a detector ID.
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.
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
165 def write_fits(self, filename):
166 """Write this object to a file.
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")
202 def writeFits(self, filename):
203 """Alias for ``write_fits``; exists for compatibility with the Butler.
204 """
205 self.write_fits(filename)
207 @classmethod
208 def read_fits(cls, filename):
209 """Build an instance of this class from a file.
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
257 @classmethod
258 def readFits(cls, filename):
259 """Alias for ``readFits``; exists for compatibility with the Butler.
260 """
261 return cls.read_fits(filename)
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 )
311class StackBrightStarsTask(pipeBase.Task):
312 """Stack bright stars together to build an extended PSF model.
313 """
314 ConfigClass = StackBrightStarsConfig
315 _DefaultName = "stack_bright_stars"
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
331 def run(self, bss_ref_list, region_name=None):
332 """Read input bright star stamps and stack them together.
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.
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
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 )
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 )
420class MeasureExtendedPsfTask(pipeBase.Task):
421 """Build and save extended PSF model.
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"
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 = []
447 def select_detector_refs(self, ref_list):
448 """Split available sets of bright star stamps according to focal plane
449 regions.
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
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)