lsst.pipe.tasks  21.0.0-89-g04f5858e+73638ba61f
extended_psf.py
Go to the documentation of this file.
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
23 PSF model.
24 """
25 
26 from dataclasses import dataclass
27 from typing import List
28 
29 from lsst.afw import image as afwImage
30 from lsst.afw import fits as afwFits
31 from lsst.afw import math as afwMath
32 from lsst.daf.base import PropertyList
33 from lsst.pipe import base as pipeBase
34 from lsst.pipe.tasks.assembleCoadd import AssembleCoaddTask
35 import lsst.pex.config as pexConfig
36 from lsst.geom import Extent2I
37 
38 
39 @dataclass
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 
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_psfdefault_extended_psf = default_extended_psf
72  self.focal_plane_regionsfocal_plane_regions = {}
73  self.detectors_focal_plane_regionsdetectors_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_regionsfocal_plane_regions:
90  raise ValueError(f"Region name {region_name} is already used by this ExtendedPsf instance.")
91  self.focal_plane_regionsfocal_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_regionsdetectors_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_psfdefault_extended_psf is None:
119  raise ValueError("No default extended PSF available; please provide detector number.")
120  return self.default_extended_psfdefault_extended_psf
121  elif not self.focal_plane_regionsfocal_plane_regions:
122  return self.default_extended_psfdefault_extended_psf
123  return self.get_regional_extended_psfget_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_regionsfocal_plane_regions)
135  if self.default_extended_psfdefault_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_regionsfocal_plane_regions[region_name].extended_psf_image
163  return self.focal_plane_regionsfocal_plane_regions[self.detectors_focal_plane_regionsdetectors_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_psfdefault_extended_psf is not None
176  if self.focal_plane_regionsfocal_plane_regions:
177  metadata["HAS_REGIONS"] = True
178  metadata["REGION_NAMES"] = list(self.focal_plane_regionsfocal_plane_regions.keys())
179  for region, e_psf_region in self.focal_plane_regionsfocal_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_psfdefault_extended_psf is not None:
189  default_hdu_metadata = PropertyList()
190  default_hdu_metadata.update({"REGION": "DEFAULT", "EXTNAME": "IMAGE"})
191  self.default_extended_psfdefault_extended_psf.image.writeFits(filename, metadata=default_hdu_metadata, mode="a")
192  default_hdu_metadata.update({"REGION": "DEFAULT", "EXTNAME": "MASK"})
193  self.default_extended_psfdefault_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_regionsfocal_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_fitswrite_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_fitsread_fits(filename)
262 
263 
264 class 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 
311 class 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_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 
386 class 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 
404 class 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 
421 class 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_regionsfocal_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_regionsfocal_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_detsregionless_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_regionsfocal_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_detsregionless_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_detsregionless_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_refsselect_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_regionsfocal_plane_regions[region_name])
494  output = pipeBase.Struct(extended_psf=output_e_psf)
495  butlerQC.put(output, outputRefs)
def __call__(self, detector=None)
Definition: extended_psf.py:96
def get_regional_extended_psf(self, region_name=None, detector=None)
def add_regional_extended_psf(self, extended_psf_image, region_name, detector_list)
Definition: extended_psf.py:75
def __init__(self, default_extended_psf=None)
Definition: extended_psf.py:70
def __init__(self, initInputs=None, *args, **kwargs)
def runQuantum(self, butlerQC, inputRefs, outputRefs)
def run(self, bss_ref_list, region_name=None)
def __init__(self, initInputs=None, *args, **kwargs)
def run(self, skyInfo, tempExpRefList, imageScalerList, weightList, altMaskList=None, mask=None, supplementaryData=None)