lsst.pipe.drivers  14.0-9-gb9632f1+1
coaddDriver.py
Go to the documentation of this file.
1 from __future__ import absolute_import, division, print_function
2 
3 from builtins import zip
4 from builtins import map
5 
6 import lsst.afw.image as afwImage
7 import lsst.afw.geom as afwGeom
8 from lsst.afw.fits.fitsLib import FitsError
9 from lsst.ctrl.pool.parallel import BatchPoolTask
10 from lsst.ctrl.pool.pool import Pool, abortOnError, NODE
11 from lsst.geom import convexHull
12 from lsst.pex.config import Config, Field, ConfigurableField
13 from lsst.pipe.base import Struct, ArgumentParser
14 from lsst.pipe.tasks.coaddBase import CoaddTaskRunner
15 from lsst.pipe.tasks.makeCoaddTempExp import MakeCoaddTempExpTask
16 from lsst.pipe.tasks.multiBand import DetectCoaddSourcesTask
17 from lsst.pipe.tasks.selectImages import WcsSelectImagesTask
18 from lsst.pipe.tasks.assembleCoadd import SafeClipAssembleCoaddTask
19 from lsst.pipe.drivers.utils import getDataRef, NullSelectImagesTask, TractDataIdContainer
20 
21 
22 class CoaddDriverConfig(Config):
23  coaddName = Field(dtype=str, default="deep", doc="Name for coadd")
24  select = ConfigurableField(
25  target=WcsSelectImagesTask, doc="Select images to process")
26  makeCoaddTempExp = ConfigurableField(
27  target=MakeCoaddTempExpTask, doc="Warp images to sky")
28  doBackgroundReference = Field(
29  dtype=bool, default=False, doc="Build background reference?")
30  backgroundReference = ConfigurableField(
31  target=NullSelectImagesTask, doc="Build background reference")
32  assembleCoadd = ConfigurableField(
33  target=SafeClipAssembleCoaddTask, doc="Assemble warps into coadd")
34  doDetection = Field(dtype=bool, default=True,
35  doc="Run detection on the coaddition product")
36  detectCoaddSources = ConfigurableField(
37  target=DetectCoaddSourcesTask, doc="Detect sources on coadd")
38 
39  def setDefaults(self):
40  self.makeCoaddTempExp.select.retarget(NullSelectImagesTask)
41  self.assembleCoadd.select.retarget(NullSelectImagesTask)
42  self.assembleCoadd.doWrite = False
43  self.assembleCoadd.doMatchBackgrounds = False
44  self.makeCoaddTempExp.bgSubtracted = True
45  self.assembleCoadd.badMaskPlanes = [
46  'BAD', 'EDGE', 'SAT', 'INTRP', 'NO_DATA']
47 
48  def validate(self):
49  if self.makeCoaddTempExp.coaddName != self.coaddName:
50  raise RuntimeError(
51  "makeCoaddTempExp.coaddName and coaddName don't match")
52  if self.assembleCoadd.coaddName != self.coaddName:
53  raise RuntimeError(
54  "assembleCoadd.coaddName and coaddName don't match")
55 
56 
58 
59  def __init__(self, TaskClass, parsedCmd, doReturnResults=False):
60  CoaddTaskRunner.__init__(self, TaskClass, parsedCmd, doReturnResults)
61  self.reuse = parsedCmd.reuse
62 
63  def makeTask(self, parsedCmd=None, args=None):
64  return self.TaskClass(config=self.config, log=self.log, reuse=self.reuse)
65 
66  @staticmethod
67  def getTargetList(parsedCmd, **kwargs):
68  """!Get bare butler into Task
69 
70  @param parsedCmd results of parsing command input
71  """
72  kwargs["butler"] = parsedCmd.butler
73  kwargs["selectIdList"] = [
74  ref.dataId for ref in parsedCmd.selectId.refList]
75  return [(parsedCmd.id.refList, kwargs), ]
76 
77 
79  ConfigClass = CoaddDriverConfig
80  _DefaultName = "coaddDriver"
81  RunnerClass = CoaddDriverTaskRunner
82 
83  def __init__(self, reuse=tuple(), **kwargs):
84  BatchPoolTask.__init__(self, **kwargs)
85  self.reuse = reuse
86  self.makeSubtask("select")
87  self.makeSubtask("makeCoaddTempExp", reuse=("makeCoaddTempExp" in self.reuse))
88  self.makeSubtask("backgroundReference")
89  self.makeSubtask("assembleCoadd")
90  self.makeSubtask("detectCoaddSources")
91 
92  @classmethod
93  def _makeArgumentParser(cls, **kwargs):
94  """!Build argument parser
95 
96  Selection references are not cheap (reads Wcs), so are generated
97  only if we're not doing a batch submission.
98  """
99  parser = ArgumentParser(name=cls._DefaultName)
100  parser.add_id_argument("--id", "deepCoadd", help="data ID, e.g. --id tract=12345 patch=1,2",
101  ContainerClass=TractDataIdContainer)
102  parser.add_id_argument(
103  "--selectId", "calexp", help="data ID, e.g. --selectId visit=6789 ccd=0..9")
104  parser.addReuseOption(["makeCoaddTempExp", "assembleCoadd", "detectCoaddSources"])
105  return parser
106 
107  @classmethod
108  def batchWallTime(cls, time, parsedCmd, numCores):
109  """!
110  Return walltime request for batch job
111 
112  @param time: Requested time per iteration
113  @param parsedCmd: Results of argument parsing
114  @param numCores: Number of cores
115  @return float walltime request length
116  """
117  numTargets = len(parsedCmd.selectId.refList)
118  return time*numTargets/float(numCores)
119 
120  @abortOnError
121  def run(self, tractPatchRefList, butler, selectIdList=[]):
122  """!Determine which tracts are non-empty before processing
123 
124  @param tractPatchRefList: List of tracts and patches to include in the coaddition
125  @param butler: butler reference object
126  @param selectIdList: List of data Ids (i.e. visit, ccd) to consider when making the coadd
127  @return list of references to sel.runTract function evaluation for each tractPatchRefList member
128  """
129  pool = Pool("tracts")
130  pool.storeSet(butler=butler, skymap=butler.get(
131  self.config.coaddName + "Coadd_skyMap"))
132  tractIdList = []
133  for patchRefList in tractPatchRefList:
134  tractSet = set([patchRef.dataId["tract"]
135  for patchRef in patchRefList])
136  assert len(tractSet) == 1
137  tractIdList.append(tractSet.pop())
138 
139  selectDataList = [data for data in pool.mapNoBalance(self.readSelection, selectIdList) if
140  data is not None]
141  nonEmptyList = pool.mapNoBalance(
142  self.checkTract, tractIdList, selectDataList)
143  tractPatchRefList = [patchRefList for patchRefList, nonEmpty in
144  zip(tractPatchRefList, nonEmptyList) if nonEmpty]
145  self.log.info("Non-empty tracts (%d): %s" % (len(tractPatchRefList),
146  [patchRefList[0].dataId["tract"] for patchRefList in
147  tractPatchRefList]))
148 
149  # Install the dataRef in the selectDataList
150  for data in selectDataList:
151  data.dataRef = getDataRef(butler, data.dataId, "calexp")
152 
153  # Process the non-empty tracts
154  return [self.runTract(patchRefList, butler, selectDataList) for patchRefList in tractPatchRefList]
155 
156  @abortOnError
157  def runTract(self, patchRefList, butler, selectDataList=[]):
158  """!Run stacking on a tract
159 
160  This method only runs on the master node.
161 
162  @param patchRefList: List of patch data references for tract
163  @param butler: Data butler
164  @param selectDataList: List of SelectStruct for inputs
165  """
166  pool = Pool("stacker")
167  pool.cacheClear()
168  pool.storeSet(butler=butler, warpType=self.config.coaddName + "Coadd_directWarp",
169  coaddType=self.config.coaddName + "Coadd")
170  patchIdList = [patchRef.dataId for patchRef in patchRefList]
171 
172  selectedData = pool.map(self.warp, patchIdList, selectDataList)
173  if self.config.doBackgroundReference:
174  self.backgroundReference.run(patchRefList, selectDataList)
175 
176  def refNamer(patchRef):
177  return tuple(map(int, patchRef.dataId["patch"].split(",")))
178 
179  lookup = dict(zip(map(refNamer, patchRefList), selectedData))
180  coaddData = [Struct(patchId=patchRef.dataId, selectDataList=lookup[refNamer(patchRef)]) for
181  patchRef in patchRefList]
182  pool.map(self.coadd, coaddData)
183 
184  def readSelection(self, cache, selectId):
185  """!Read Wcs of selected inputs
186 
187  This method only runs on slave nodes.
188  This method is similar to SelectDataIdContainer.makeDataRefList,
189  creating a Struct like a SelectStruct, except with a dataId instead
190  of a dataRef (to ease MPI).
191 
192  @param cache: Pool cache
193  @param selectId: Data identifier for selected input
194  @return a SelectStruct with a dataId instead of dataRef
195  """
196  try:
197  ref = getDataRef(cache.butler, selectId, "calexp")
198  self.log.info("Reading Wcs from %s" % (selectId,))
199  md = ref.get("calexp_md", immediate=True)
200  wcs = afwImage.makeWcs(md)
201  data = Struct(dataId=selectId, wcs=wcs, bbox=afwImage.bboxFromMetadata(md))
202  except FitsError:
203  self.log.warn("Unable to construct Wcs from %s" % (selectId,))
204  return None
205  return data
206 
207  def checkTract(self, cache, tractId, selectIdList):
208  """!Check whether a tract has any overlapping inputs
209 
210  This method only runs on slave nodes.
211 
212  @param cache: Pool cache
213  @param tractId: Data identifier for tract
214  @param selectDataList: List of selection data
215  @return whether tract has any overlapping inputs
216  """
217  def makePolygon(wcs, bbox):
218  """Return a polygon for the image, given Wcs and bounding box"""
219  return convexHull([wcs.pixelToSky(afwGeom.Point2D(coord)).getVector() for
220  coord in bbox.getCorners()])
221 
222  skymap = cache.skymap
223  tract = skymap[tractId]
224  tractWcs = tract.getWcs()
225  tractPoly = makePolygon(tractWcs, tract.getBBox())
226 
227  for selectData in selectIdList:
228  if not hasattr(selectData, "poly"):
229  selectData.poly = makePolygon(selectData.wcs, selectData.bbox)
230  if tractPoly.intersects(selectData.poly):
231  return True
232  return False
233 
234  def warp(self, cache, patchId, selectDataList):
235  """!Warp all images for a patch
236 
237  Only slave nodes execute this method.
238 
239  Because only one argument may be passed, it is expected to
240  contain multiple elements, which are:
241 
242  @param patchRef: data reference for patch
243  @param selectDataList: List of SelectStruct for inputs
244  @return selectDataList with non-overlapping elements removed
245  """
246  patchRef = getDataRef(cache.butler, patchId, cache.coaddType)
247  selectDataList = self.selectExposures(patchRef, selectDataList)
248  with self.logOperation("warping %s" % (patchRef.dataId,), catch=True):
249  self.makeCoaddTempExp.run(patchRef, selectDataList)
250  return selectDataList
251 
252  def coadd(self, cache, data):
253  """!Construct coadd for a patch and measure
254 
255  Only slave nodes execute this method.
256 
257  Because only one argument may be passed, it is expected to
258  contain multiple elements, which are:
259 
260  @param patchRef: data reference for patch
261  @param selectDataList: List of SelectStruct for inputs
262  """
263  patchRef = getDataRef(cache.butler, data.patchId, cache.coaddType)
264  selectDataList = data.selectDataList
265  coadd = None
266 
267  # We skip the assembleCoadd step if either the *Coadd dataset exists
268  # or we aren't configured to write it, we're supposed to reuse
269  # detectCoaddSources outputs too, and those outputs already exist.
270  canSkipDetection = (
271  "detectCoaddSources" in self.reuse and
272  patchRef.datasetExists(self.detectCoaddSources.config.coaddName+"Coadd_det", write=True)
273  )
274  if "assembleCoadd" in self.reuse:
275  if patchRef.datasetExists(cache.coaddType, write=True):
276  self.log.info("%s: Skipping assembleCoadd for %s; outputs already exist." %
277  (NODE, patchRef.dataId))
278  coadd = patchRef.get(cache.coaddType, immediate=True)
279  elif not self.config.assembleCoadd.doWrite and self.config.doDetection and canSkipDetection:
280  self.log.info(
281  "%s: Skipping assembleCoadd and detectCoaddSources for %s; outputs already exist." %
282  (NODE, patchRef.dataId)
283  )
284  return
285  if coadd is None:
286  with self.logOperation("coadding %s" % (patchRef.dataId,), catch=True):
287  coaddResults = self.assembleCoadd.run(patchRef, selectDataList)
288  if coaddResults is not None:
289  coadd = coaddResults.coaddExposure
290  canSkipDetection = False # can't skip it because coadd may have changed
291  if coadd is None:
292  return
293 
294  # The section of code below determines if the detection task should be
295  # run. If detection is run, then the products are written out as
296  # deepCoadd_calexp. If detection is not run, then the outputs of the
297  # assemble task are written out as deepCoadd.
298  if self.config.doDetection:
299  if canSkipDetection:
300  self.log.info("%s: Skipping detectCoaddSources for %s; outputs already exist." %
301  (NODE, patchRef.dataId))
302  return
303  with self.logOperation("detection on {}".format(patchRef.dataId),
304  catch=True):
305  idFactory = self.detectCoaddSources.makeIdFactory(patchRef)
306  # This includes background subtraction, so do it before writing
307  # the coadd
308  detResults = self.detectCoaddSources.runDetection(coadd, idFactory)
309  self.detectCoaddSources.write(coadd, detResults, patchRef)
310  else:
311  patchRef.put(coadd, self.assembleCoadd.config.coaddName+"Coadd")
312 
313  def selectExposures(self, patchRef, selectDataList):
314  """!Select exposures to operate upon, via the SelectImagesTask
315 
316  This is very similar to CoaddBaseTask.selectExposures, except we return
317  a list of SelectStruct (same as the input), so we can plug the results into
318  future uses of SelectImagesTask.
319 
320  @param patchRef data reference to a particular patch
321  @param selectDataList list of references to specific data products (i.e. visit, ccd)
322  @return filtered list of SelectStruct
323  """
324  def key(dataRef):
325  return tuple(dataRef.dataId[k] for k in sorted(dataRef.dataId))
326  inputs = dict((key(select.dataRef), select)
327  for select in selectDataList)
328  skyMap = patchRef.get(self.config.coaddName + "Coadd_skyMap")
329  tract = skyMap[patchRef.dataId["tract"]]
330  patch = tract[(tuple(int(i)
331  for i in patchRef.dataId["patch"].split(",")))]
332  bbox = patch.getOuterBBox()
333  wcs = tract.getWcs()
334  cornerPosList = afwGeom.Box2D(bbox).getCorners()
335  coordList = [wcs.pixelToSky(pos) for pos in cornerPosList]
336  dataRefList = self.select.runDataRef(
337  patchRef, coordList, selectDataList=selectDataList).dataRefList
338  return [inputs[key(dataRef)] for dataRef in dataRefList]
339 
340  def writeMetadata(self, dataRef):
341  pass
def batchWallTime(cls, time, parsedCmd, numCores)
Return walltime request for batch job.
Definition: coaddDriver.py:108
def run(self, tractPatchRefList, butler, selectIdList=[])
Determine which tracts are non-empty before processing.
Definition: coaddDriver.py:121
def selectExposures(self, patchRef, selectDataList)
Select exposures to operate upon, via the SelectImagesTask.
Definition: coaddDriver.py:313
def runTract(self, patchRefList, butler, selectDataList=[])
Run stacking on a tract.
Definition: coaddDriver.py:157
def makeTask(self, parsedCmd=None, args=None)
Definition: coaddDriver.py:63
def getDataRef(butler, dataId, datasetType="raw")
Definition: utils.py:17
def __init__(self, reuse=tuple(), kwargs)
Definition: coaddDriver.py:83
def coadd(self, cache, data)
Construct coadd for a patch and measure.
Definition: coaddDriver.py:252
def warp(self, cache, patchId, selectDataList)
Warp all images for a patch.
Definition: coaddDriver.py:234
def logOperation(self, operation, catch=False, trace=True)
def getTargetList(parsedCmd, kwargs)
Get bare butler into Task.
Definition: coaddDriver.py:67
def readSelection(self, cache, selectId)
Read Wcs of selected inputs.
Definition: coaddDriver.py:184
def __init__(self, TaskClass, parsedCmd, doReturnResults=False)
Definition: coaddDriver.py:59
def checkTract(self, cache, tractId, selectIdList)
Check whether a tract has any overlapping inputs.
Definition: coaddDriver.py:207