Coverage for python/lsst/pipe/tasks/skyCorrection.py: 21%
173 statements
« prev ^ index » next coverage.py v6.4.4, created at 2022-09-15 03:34 -0700
« prev ^ index » next coverage.py v6.4.4, created at 2022-09-15 03:34 -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__all__ = ["SkyCorrectionTask", "SkyCorrectionConfig"]
23import numpy as np
25import lsst.afw.math as afwMath
26import lsst.afw.image as afwImage
27import lsst.pipe.base as pipeBase
29from lsst.afw.cameraGeom.utils import makeImageFromCamera
30from lsst.daf.butler import DimensionGraph
31from lsst.pex.config import Config, Field, ConfigurableField, ConfigField
32import lsst.pipe.base.connectionTypes as cT
34from .background import (SkyMeasurementTask, FocalPlaneBackground,
35 FocalPlaneBackgroundConfig, MaskObjectsTask)
38__all__ = ["SkyCorrectionConfig", "SkyCorrectionTask"]
41def reorderAndPadList(inputList, inputKeys, outputKeys, padWith=None):
42 """Match the order of one list to another, padding if necessary
44 Parameters
45 ----------
46 inputList : list
47 List to be reordered and padded. Elements can be any type.
48 inputKeys : iterable
49 Iterable of values to be compared with outputKeys.
50 Length must match `inputList`
51 outputKeys : iterable
52 Iterable of values to be compared with inputKeys.
53 padWith :
54 Any value to be inserted where inputKey not in outputKeys
56 Returns
57 -------
58 list
59 Copy of inputList reordered per outputKeys and padded with `padWith`
60 so that the length matches length of outputKeys.
61 """
62 outputList = []
63 for d in outputKeys:
64 if d in inputKeys:
65 outputList.append(inputList[inputKeys.index(d)])
66 else:
67 outputList.append(padWith)
68 return outputList
71def _makeCameraImage(camera, exposures, binning):
72 """Make and write an image of an entire focal plane
74 Parameters
75 ----------
76 camera : `lsst.afw.cameraGeom.Camera`
77 Camera description.
78 exposures : `dict` mapping detector ID to `lsst.afw.image.Exposure`
79 CCD exposures, binned by `binning`.
80 binning : `int`
81 Binning size that has been applied to images.
82 """
83 class ImageSource:
84 """Source of images for makeImageFromCamera"""
85 def __init__(self, exposures):
86 """Constructor
88 Parameters
89 ----------
90 exposures : `dict` mapping detector ID to `lsst.afw.image.Exposure`
91 CCD exposures, already binned.
92 """
93 self.isTrimmed = True
94 self.exposures = exposures
95 self.background = np.nan
97 def getCcdImage(self, detector, imageFactory, binSize):
98 """Provide image of CCD to makeImageFromCamera"""
99 detId = detector.getId()
100 if detId not in self.exposures:
101 dims = detector.getBBox().getDimensions()/binSize
102 image = imageFactory(*[int(xx) for xx in dims])
103 image.set(self.background)
104 else:
105 image = self.exposures[detector.getId()]
106 if hasattr(image, "getMaskedImage"):
107 image = image.getMaskedImage()
108 if hasattr(image, "getMask"):
109 mask = image.getMask()
110 isBad = mask.getArray() & mask.getPlaneBitMask("NO_DATA") > 0
111 image = image.clone()
112 image.getImage().getArray()[isBad] = self.background
113 if hasattr(image, "getImage"):
114 image = image.getImage()
116 image = afwMath.rotateImageBy90(image, detector.getOrientation().getNQuarter())
118 return image, detector
120 image = makeImageFromCamera(
121 camera,
122 imageSource=ImageSource(exposures),
123 imageFactory=afwImage.ImageF,
124 binSize=binning
125 )
126 return image
129def makeCameraImage(camera, exposures, filename=None, binning=8):
130 """Make and write an image of an entire focal plane
132 Parameters
133 ----------
134 camera : `lsst.afw.cameraGeom.Camera`
135 Camera description.
136 exposures : `list` of `tuple` of `int` and `lsst.afw.image.Exposure`
137 List of detector ID and CCD exposures (binned by `binning`).
138 filename : `str`, optional
139 Output filename.
140 binning : `int`
141 Binning size that has been applied to images.
142 """
143 image = _makeCameraImage(camera, dict(exp for exp in exposures if exp is not None), binning)
144 if filename is not None:
145 image.writeFits(filename)
146 return image
149def _skyLookup(datasetType, registry, quantumDataId, collections):
150 """Lookup function to identify sky frames
152 Parameters
153 ----------
154 datasetType : `lsst.daf.butler.DatasetType`
155 Dataset to lookup.
156 registry : `lsst.daf.butler.Registry`
157 Butler registry to query.
158 quantumDataId : `lsst.daf.butler.DataCoordinate`
159 Data id to transform to find sky frames.
160 The ``detector`` entry will be stripped.
161 collections : `lsst.daf.butler.CollectionSearch`
162 Collections to search through.
164 Returns
165 -------
166 results : `list` [`lsst.daf.butler.DatasetRef`]
167 List of datasets that will be used as sky calibration frames
168 """
169 newDataId = quantumDataId.subset(DimensionGraph(registry.dimensions, names=["instrument", "visit"]))
170 skyFrames = []
171 for dataId in registry.queryDataIds(["visit", "detector"], dataId=newDataId).expanded():
172 skyFrame = registry.findDataset(datasetType, dataId, collections=collections,
173 timespan=dataId.timespan)
174 skyFrames.append(skyFrame)
176 return skyFrames
179class SkyCorrectionConnections(pipeBase.PipelineTaskConnections, dimensions=("instrument", "visit")):
180 rawLinker = cT.Input(
181 doc="Raw data to provide exp-visit linkage to connect calExp inputs to camera/sky calibs.",
182 name="raw",
183 multiple=True,
184 deferLoad=True,
185 storageClass="Exposure",
186 dimensions=["instrument", "exposure", "detector"],
187 )
188 calExpArray = cT.Input(
189 doc="Input exposures to process",
190 name="calexp",
191 multiple=True,
192 storageClass="ExposureF",
193 dimensions=["instrument", "visit", "detector"],
194 )
195 calBkgArray = cT.Input(
196 doc="Input background files to use",
197 multiple=True,
198 name="calexpBackground",
199 storageClass="Background",
200 dimensions=["instrument", "visit", "detector"],
201 )
202 camera = cT.PrerequisiteInput(
203 doc="Input camera to use.",
204 name="camera",
205 storageClass="Camera",
206 dimensions=["instrument"],
207 isCalibration=True,
208 )
209 skyCalibs = cT.PrerequisiteInput(
210 doc="Input sky calibrations to use.",
211 name="sky",
212 multiple=True,
213 storageClass="ExposureF",
214 dimensions=["instrument", "physical_filter", "detector"],
215 isCalibration=True,
216 lookupFunction=_skyLookup,
217 )
218 calExpCamera = cT.Output(
219 doc="Output camera image.",
220 name='calexp_camera',
221 storageClass="ImageF",
222 dimensions=["instrument", "visit"],
223 )
224 skyCorr = cT.Output(
225 doc="Output sky corrected images.",
226 name='skyCorr',
227 multiple=True,
228 storageClass="Background",
229 dimensions=["instrument", "visit", "detector"],
230 )
233class SkyCorrectionConfig(pipeBase.PipelineTaskConfig, pipelineConnections=SkyCorrectionConnections):
234 """Configuration for SkyCorrectionTask"""
235 bgModel = ConfigField(dtype=FocalPlaneBackgroundConfig, doc="Background model")
236 bgModel2 = ConfigField(dtype=FocalPlaneBackgroundConfig, doc="2nd Background model")
237 sky = ConfigurableField(target=SkyMeasurementTask, doc="Sky measurement")
238 maskObjects = ConfigurableField(target=MaskObjectsTask, doc="Mask Objects")
239 doMaskObjects = Field(dtype=bool, default=True, doc="Mask objects to find good sky?")
240 doBgModel = Field(dtype=bool, default=True, doc="Do background model subtraction?")
241 doBgModel2 = Field(dtype=bool, default=True, doc="Do cleanup background model subtraction?")
242 doSky = Field(dtype=bool, default=True, doc="Do sky frame subtraction?")
243 binning = Field(dtype=int, default=8, doc="Binning factor for constructing focal-plane images")
244 calexpType = Field(dtype=str, default="calexp",
245 doc="Should be set to fakes_calexp if you want to process calexps with fakes in.")
247 def setDefaults(self):
248 Config.setDefaults(self)
249 self.bgModel2.doSmooth = True
250 self.bgModel2.minFrac = 0.5
251 self.bgModel2.xSize = 256
252 self.bgModel2.ySize = 256
253 self.bgModel2.smoothScale = 1.0
256class SkyCorrectionTask(pipeBase.PipelineTask):
257 """Correct sky over entire focal plane"""
258 ConfigClass = SkyCorrectionConfig
259 _DefaultName = "skyCorr"
261 def runQuantum(self, butlerQC, inputRefs, outputRefs):
263 # Reorder the skyCalibs, calBkgArray, and calExpArray inputRefs and the
264 # skyCorr outputRef sorted by detector id to ensure reproducibility.
265 detectorOrder = [ref.dataId['detector'] for ref in inputRefs.calExpArray]
266 detectorOrder.sort()
267 inputRefs.calExpArray = reorderAndPadList(inputRefs.calExpArray,
268 [ref.dataId['detector'] for ref in inputRefs.calExpArray],
269 detectorOrder)
270 inputRefs.skyCalibs = reorderAndPadList(inputRefs.skyCalibs,
271 [ref.dataId['detector'] for ref in inputRefs.skyCalibs],
272 detectorOrder)
273 inputRefs.calBkgArray = reorderAndPadList(inputRefs.calBkgArray,
274 [ref.dataId['detector'] for ref in inputRefs.calBkgArray],
275 detectorOrder)
276 outputRefs.skyCorr = reorderAndPadList(outputRefs.skyCorr,
277 [ref.dataId['detector'] for ref in outputRefs.skyCorr],
278 detectorOrder)
279 inputs = butlerQC.get(inputRefs)
280 inputs.pop("rawLinker", None)
281 outputs = self.run(**inputs)
282 butlerQC.put(outputs, outputRefs)
284 def __init__(self, *args, **kwargs):
285 super().__init__(**kwargs)
287 self.makeSubtask("sky")
288 self.makeSubtask("maskObjects")
290 def focalPlaneBackgroundRun(self, camera, cacheExposures, idList, config):
291 """Perform full focal-plane background subtraction
293 This method runs on the master node.
295 Parameters
296 ----------
297 camera : `lsst.afw.cameraGeom.Camera`
298 Camera description.
299 cacheExposures : `list` of `lsst.afw.image.Exposures`
300 List of loaded and processed input calExp.
301 idList : `list` of `int`
302 List of detector ids to iterate over.
303 config : `lsst.pipe.drivers.background.FocalPlaneBackgroundConfig`
304 Configuration to use for background subtraction.
306 Returns
307 -------
308 exposures : `list` of `lsst.afw.image.Image`
309 List of binned images, for creating focal plane image.
310 newCacheBgList : `list` of `lsst.afwMath.backgroundList`
311 Background lists generated.
312 cacheBgModel : `FocalPlaneBackground`
313 Full focal plane background model.
314 """
315 bgModel = FocalPlaneBackground.fromCamera(config, camera)
316 data = [pipeBase.Struct(id=id, bgModel=bgModel.clone()) for id in idList]
318 bgModelList = []
319 for nodeData, cacheExp in zip(data, cacheExposures):
320 nodeData.bgModel.addCcd(cacheExp)
321 bgModelList.append(nodeData.bgModel)
323 for ii, bg in enumerate(bgModelList):
324 self.log.info("Background %d: %d pixels", ii, bg._numbers.getArray().sum())
325 bgModel.merge(bg)
327 exposures = []
328 newCacheBgList = []
329 cacheBgModel = []
330 for cacheExp in cacheExposures:
331 nodeExp, nodeBgModel, nodeBgList = self.subtractModelRun(cacheExp, bgModel)
332 exposures.append(afwMath.binImage(nodeExp.getMaskedImage(), self.config.binning))
333 cacheBgModel.append(nodeBgModel)
334 newCacheBgList.append(nodeBgList)
336 return exposures, newCacheBgList, cacheBgModel
338 def run(self, calExpArray, calBkgArray, skyCalibs, camera):
339 """Performa sky correction on an exposure.
341 Parameters
342 ----------
343 calExpArray : `list` of `lsst.afw.image.Exposure`
344 Array of detector input calExp images for the exposure to
345 process.
346 calBkgArray : `list` of `lsst.afw.math.BackgroundList`
347 Array of detector input background lists matching the
348 calExps to process.
349 skyCalibs : `list` of `lsst.afw.image.Exposure`
350 Array of SKY calibrations for the input detectors to be
351 processed.
352 camera : `lsst.afw.cameraGeom.Camera`
353 Camera matching the input data to process.
355 Returns
356 -------
357 results : `pipeBase.Struct` containing
358 calExpCamera : `lsst.afw.image.Exposure`
359 Full camera image of the sky-corrected data.
360 skyCorr : `list` of `lsst.afw.math.BackgroundList`
361 Detector-level sky-corrected background lists.
362 """
363 # To allow SkyCorrectionTask to run in the Gen3 butler
364 # environment, a new run() method was added that performs the
365 # same operations in a serial environment (pipetask processing
366 # does not support MPI processing as of 2019-05-03). Methods
367 # used in runDataRef() are used as appropriate in run(), but
368 # some have been rewritten in serial form. Please ensure that
369 # any updates to runDataRef() or the methods it calls with
370 # pool.mapToPrevious() are duplicated in run() and its
371 # methods.
372 #
373 # Variable names here should match those in runDataRef() as
374 # closely as possible. Variables matching data stored in the
375 # pool cache have a prefix indicating this. Variables that
376 # would be local to an MPI processing client have a prefix
377 # "node".
378 idList = [exp.getDetector().getId() for exp in calExpArray]
380 # Construct arrays that match the cache in self.runDataRef() after
381 # self.loadImage() is map/reduced.
382 cacheExposures = []
383 cacheBgList = []
384 exposures = []
385 for calExp, calBgModel in zip(calExpArray, calBkgArray):
386 nodeExp, nodeBgList = self.loadImageRun(calExp, calBgModel)
387 cacheExposures.append(nodeExp)
388 cacheBgList.append(nodeBgList)
389 exposures.append(afwMath.binImage(nodeExp.getMaskedImage(), self.config.binning))
391 if self.config.doBgModel:
392 # Generate focal plane background, updating backgrounds in the "cache".
393 exposures, newCacheBgList, cacheBgModel = self.focalPlaneBackgroundRun(
394 camera, cacheExposures, idList, self.config.bgModel
395 )
396 for cacheBg, newBg in zip(cacheBgList, newCacheBgList):
397 cacheBg.append(newBg)
399 if self.config.doSky:
400 # Measure the sky frame scale on all inputs. Results in
401 # values equal to self.measureSkyFrame() and
402 # self.sky.solveScales() in runDataRef().
403 cacheSky = []
404 measScales = []
405 for cacheExp, skyCalib in zip(cacheExposures, skyCalibs):
406 skyExp = self.sky.exposureToBackground(skyCalib)
407 cacheSky.append(skyExp)
408 scale = self.sky.measureScale(cacheExp.getMaskedImage(), skyExp)
409 measScales.append(scale)
411 scale = self.sky.solveScales(measScales)
412 self.log.info("Sky frame scale: %s" % (scale, ))
414 # Subtract sky frame, as in self.subtractSkyFrame(), with
415 # appropriate scale from the "cache".
416 exposures = []
417 newBgList = []
418 for cacheExp, nodeSky, nodeBgList in zip(cacheExposures, cacheSky, cacheBgList):
419 self.sky.subtractSkyFrame(cacheExp.getMaskedImage(), nodeSky, scale, nodeBgList)
420 exposures.append(afwMath.binImage(cacheExp.getMaskedImage(), self.config.binning))
422 if self.config.doBgModel2:
423 # As above, generate a focal plane background model and
424 # update the cache models.
425 exposures, newBgList, cacheBgModel = self.focalPlaneBackgroundRun(
426 camera, cacheExposures, idList, self.config.bgModel2
427 )
428 for cacheBg, newBg in zip(cacheBgList, newBgList):
429 cacheBg.append(newBg)
431 # Generate camera-level image of calexp and return it along
432 # with the list of sky corrected background models.
433 image = makeCameraImage(camera, zip(idList, exposures))
435 return pipeBase.Struct(
436 calExpCamera=image,
437 skyCorr=cacheBgList,
438 )
440 def loadImageRun(self, calExp, calExpBkg):
441 """Serial implementation of self.loadImage() for Gen3.
443 Load and restore background to calExp and calExpBkg.
445 Parameters
446 ----------
447 calExp : `lsst.afw.image.Exposure`
448 Detector level calExp image to process.
449 calExpBkg : `lsst.afw.math.BackgroundList`
450 Detector level background list associated with the calExp.
452 Returns
453 -------
454 calExp : `lsst.afw.image.Exposure`
455 Background restored calExp.
456 bgList : `lsst.afw.math.BackgroundList`
457 New background list containing the restoration background.
458 """
459 image = calExp.getMaskedImage()
461 for bgOld in calExpBkg:
462 statsImage = bgOld[0].getStatsImage()
463 statsImage *= -1
465 image -= calExpBkg.getImage()
466 bgList = afwMath.BackgroundList()
467 for bgData in calExpBkg:
468 bgList.append(bgData)
470 if self.config.doMaskObjects:
471 self.maskObjects.findObjects(calExp)
473 return (calExp, bgList)
475 def subtractModelRun(self, exposure, bgModel):
476 """Serial implementation of self.subtractModel() for Gen3.
478 Load and restore background to calExp and calExpBkg.
480 Parameters
481 ----------
482 exposure : `lsst.afw.image.Exposure`
483 Exposure to subtract the background model from.
484 bgModel : `lsst.pipe.drivers.background.FocalPlaneBackground`
485 Full camera level background model.
487 Returns
488 -------
489 exposure : `lsst.afw.image.Exposure`
490 Background subtracted input exposure.
491 bgModelCcd : `lsst.afw.math.BackgroundList`
492 Detector level realization of the full background model.
493 bgModelMaskedImage : `lsst.afw.image.MaskedImage`
494 Background model from the bgModelCcd realization.
495 """
496 image = exposure.getMaskedImage()
497 detector = exposure.getDetector()
498 bbox = image.getBBox()
499 bgModelCcd = bgModel.toCcdBackground(detector, bbox)
500 image -= bgModelCcd.getImage()
502 return (exposure, bgModelCcd, bgModelCcd[0])