lsst.pipe.drivers  14.0-10-gd86e9bf+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 
78 def unpickle(factory, args, kwargs):
79  """Unpickle something by calling a factory"""
80  return factory(*args, **kwargs)
81 
82 
84  ConfigClass = CoaddDriverConfig
85  _DefaultName = "coaddDriver"
86  RunnerClass = CoaddDriverTaskRunner
87 
88  def __init__(self, reuse=tuple(), **kwargs):
89  BatchPoolTask.__init__(self, **kwargs)
90  self.reuse = reuse
91  self.makeSubtask("select")
92  self.makeSubtask("makeCoaddTempExp", reuse=("makeCoaddTempExp" in self.reuse))
93  self.makeSubtask("backgroundReference")
94  self.makeSubtask("assembleCoadd")
95  self.makeSubtask("detectCoaddSources")
96 
97  def __reduce__(self):
98  """Pickler"""
99  return unpickle, (self.__class__, [], dict(config=self.config, name=self._name,
100  parentTask=self._parentTask, log=self.log,
101  reuse=self.reuse))
102 
103  @classmethod
104  def _makeArgumentParser(cls, **kwargs):
105  """!Build argument parser
106 
107  Selection references are not cheap (reads Wcs), so are generated
108  only if we're not doing a batch submission.
109  """
110  parser = ArgumentParser(name=cls._DefaultName)
111  parser.add_id_argument("--id", "deepCoadd", help="data ID, e.g. --id tract=12345 patch=1,2",
112  ContainerClass=TractDataIdContainer)
113  parser.add_id_argument(
114  "--selectId", "calexp", help="data ID, e.g. --selectId visit=6789 ccd=0..9")
115  parser.addReuseOption(["makeCoaddTempExp", "assembleCoadd", "detectCoaddSources"])
116  return parser
117 
118  @classmethod
119  def batchWallTime(cls, time, parsedCmd, numCores):
120  """!
121  Return walltime request for batch job
122 
123  @param time: Requested time per iteration
124  @param parsedCmd: Results of argument parsing
125  @param numCores: Number of cores
126  @return float walltime request length
127  """
128  numTargets = len(parsedCmd.selectId.refList)
129  return time*numTargets/float(numCores)
130 
131  @abortOnError
132  def run(self, tractPatchRefList, butler, selectIdList=[]):
133  """!Determine which tracts are non-empty before processing
134 
135  @param tractPatchRefList: List of tracts and patches to include in the coaddition
136  @param butler: butler reference object
137  @param selectIdList: List of data Ids (i.e. visit, ccd) to consider when making the coadd
138  @return list of references to sel.runTract function evaluation for each tractPatchRefList member
139  """
140  pool = Pool("tracts")
141  pool.storeSet(butler=butler, skymap=butler.get(
142  self.config.coaddName + "Coadd_skyMap"))
143  tractIdList = []
144  for patchRefList in tractPatchRefList:
145  tractSet = set([patchRef.dataId["tract"]
146  for patchRef in patchRefList])
147  assert len(tractSet) == 1
148  tractIdList.append(tractSet.pop())
149 
150  selectDataList = [data for data in pool.mapNoBalance(self.readSelection, selectIdList) if
151  data is not None]
152  nonEmptyList = pool.mapNoBalance(
153  self.checkTract, tractIdList, selectDataList)
154  tractPatchRefList = [patchRefList for patchRefList, nonEmpty in
155  zip(tractPatchRefList, nonEmptyList) if nonEmpty]
156  self.log.info("Non-empty tracts (%d): %s" % (len(tractPatchRefList),
157  [patchRefList[0].dataId["tract"] for patchRefList in
158  tractPatchRefList]))
159 
160  # Install the dataRef in the selectDataList
161  for data in selectDataList:
162  data.dataRef = getDataRef(butler, data.dataId, "calexp")
163 
164  # Process the non-empty tracts
165  return [self.runTract(patchRefList, butler, selectDataList) for patchRefList in tractPatchRefList]
166 
167  @abortOnError
168  def runTract(self, patchRefList, butler, selectDataList=[]):
169  """!Run stacking on a tract
170 
171  This method only runs on the master node.
172 
173  @param patchRefList: List of patch data references for tract
174  @param butler: Data butler
175  @param selectDataList: List of SelectStruct for inputs
176  """
177  pool = Pool("stacker")
178  pool.cacheClear()
179  pool.storeSet(butler=butler, warpType=self.config.coaddName + "Coadd_directWarp",
180  coaddType=self.config.coaddName + "Coadd")
181  patchIdList = [patchRef.dataId for patchRef in patchRefList]
182 
183  selectedData = pool.map(self.warp, patchIdList, selectDataList)
184  if self.config.doBackgroundReference:
185  self.backgroundReference.run(patchRefList, selectDataList)
186 
187  def refNamer(patchRef):
188  return tuple(map(int, patchRef.dataId["patch"].split(",")))
189 
190  lookup = dict(zip(map(refNamer, patchRefList), selectedData))
191  coaddData = [Struct(patchId=patchRef.dataId, selectDataList=lookup[refNamer(patchRef)]) for
192  patchRef in patchRefList]
193  pool.map(self.coadd, coaddData)
194 
195  def readSelection(self, cache, selectId):
196  """!Read Wcs of selected inputs
197 
198  This method only runs on slave nodes.
199  This method is similar to SelectDataIdContainer.makeDataRefList,
200  creating a Struct like a SelectStruct, except with a dataId instead
201  of a dataRef (to ease MPI).
202 
203  @param cache: Pool cache
204  @param selectId: Data identifier for selected input
205  @return a SelectStruct with a dataId instead of dataRef
206  """
207  try:
208  ref = getDataRef(cache.butler, selectId, "calexp")
209  self.log.info("Reading Wcs from %s" % (selectId,))
210  md = ref.get("calexp_md", immediate=True)
211  wcs = afwImage.makeWcs(md)
212  data = Struct(dataId=selectId, wcs=wcs, bbox=afwImage.bboxFromMetadata(md))
213  except FitsError:
214  self.log.warn("Unable to construct Wcs from %s" % (selectId,))
215  return None
216  return data
217 
218  def checkTract(self, cache, tractId, selectIdList):
219  """!Check whether a tract has any overlapping inputs
220 
221  This method only runs on slave nodes.
222 
223  @param cache: Pool cache
224  @param tractId: Data identifier for tract
225  @param selectDataList: List of selection data
226  @return whether tract has any overlapping inputs
227  """
228  def makePolygon(wcs, bbox):
229  """Return a polygon for the image, given Wcs and bounding box"""
230  return convexHull([wcs.pixelToSky(afwGeom.Point2D(coord)).getVector() for
231  coord in bbox.getCorners()])
232 
233  skymap = cache.skymap
234  tract = skymap[tractId]
235  tractWcs = tract.getWcs()
236  tractPoly = makePolygon(tractWcs, tract.getBBox())
237 
238  for selectData in selectIdList:
239  if not hasattr(selectData, "poly"):
240  selectData.poly = makePolygon(selectData.wcs, selectData.bbox)
241  if tractPoly.intersects(selectData.poly):
242  return True
243  return False
244 
245  def warp(self, cache, patchId, selectDataList):
246  """!Warp all images for a patch
247 
248  Only slave nodes execute this method.
249 
250  Because only one argument may be passed, it is expected to
251  contain multiple elements, which are:
252 
253  @param patchRef: data reference for patch
254  @param selectDataList: List of SelectStruct for inputs
255  @return selectDataList with non-overlapping elements removed
256  """
257  patchRef = getDataRef(cache.butler, patchId, cache.coaddType)
258  selectDataList = self.selectExposures(patchRef, selectDataList)
259  with self.logOperation("warping %s" % (patchRef.dataId,), catch=True):
260  self.makeCoaddTempExp.run(patchRef, selectDataList)
261  return selectDataList
262 
263  def coadd(self, cache, data):
264  """!Construct coadd for a patch and measure
265 
266  Only slave nodes execute this method.
267 
268  Because only one argument may be passed, it is expected to
269  contain multiple elements, which are:
270 
271  @param patchRef: data reference for patch
272  @param selectDataList: List of SelectStruct for inputs
273  """
274  patchRef = getDataRef(cache.butler, data.patchId, cache.coaddType)
275  selectDataList = data.selectDataList
276  coadd = None
277 
278  # We skip the assembleCoadd step if either the *Coadd dataset exists
279  # or we aren't configured to write it, we're supposed to reuse
280  # detectCoaddSources outputs too, and those outputs already exist.
281  canSkipDetection = (
282  "detectCoaddSources" in self.reuse and
283  patchRef.datasetExists(self.detectCoaddSources.config.coaddName+"Coadd_det", write=True)
284  )
285  if "assembleCoadd" in self.reuse:
286  if patchRef.datasetExists(cache.coaddType, write=True):
287  self.log.info("%s: Skipping assembleCoadd for %s; outputs already exist." %
288  (NODE, patchRef.dataId))
289  coadd = patchRef.get(cache.coaddType, immediate=True)
290  elif not self.config.assembleCoadd.doWrite and self.config.doDetection and canSkipDetection:
291  self.log.info(
292  "%s: Skipping assembleCoadd and detectCoaddSources for %s; outputs already exist." %
293  (NODE, patchRef.dataId)
294  )
295  return
296  if coadd is None:
297  with self.logOperation("coadding %s" % (patchRef.dataId,), catch=True):
298  coaddResults = self.assembleCoadd.run(patchRef, selectDataList)
299  if coaddResults is not None:
300  coadd = coaddResults.coaddExposure
301  canSkipDetection = False # can't skip it because coadd may have changed
302  if coadd is None:
303  return
304 
305  # The section of code below determines if the detection task should be
306  # run. If detection is run, then the products are written out as
307  # deepCoadd_calexp. If detection is not run, then the outputs of the
308  # assemble task are written out as deepCoadd.
309  if self.config.doDetection:
310  if canSkipDetection:
311  self.log.info("%s: Skipping detectCoaddSources for %s; outputs already exist." %
312  (NODE, patchRef.dataId))
313  return
314  with self.logOperation("detection on {}".format(patchRef.dataId),
315  catch=True):
316  idFactory = self.detectCoaddSources.makeIdFactory(patchRef)
317  # This includes background subtraction, so do it before writing
318  # the coadd
319  detResults = self.detectCoaddSources.runDetection(coadd, idFactory)
320  self.detectCoaddSources.write(coadd, detResults, patchRef)
321  else:
322  patchRef.put(coadd, self.assembleCoadd.config.coaddName+"Coadd")
323 
324  def selectExposures(self, patchRef, selectDataList):
325  """!Select exposures to operate upon, via the SelectImagesTask
326 
327  This is very similar to CoaddBaseTask.selectExposures, except we return
328  a list of SelectStruct (same as the input), so we can plug the results into
329  future uses of SelectImagesTask.
330 
331  @param patchRef data reference to a particular patch
332  @param selectDataList list of references to specific data products (i.e. visit, ccd)
333  @return filtered list of SelectStruct
334  """
335  def key(dataRef):
336  return tuple(dataRef.dataId[k] for k in sorted(dataRef.dataId))
337  inputs = dict((key(select.dataRef), select)
338  for select in selectDataList)
339  skyMap = patchRef.get(self.config.coaddName + "Coadd_skyMap")
340  tract = skyMap[patchRef.dataId["tract"]]
341  patch = tract[(tuple(int(i)
342  for i in patchRef.dataId["patch"].split(",")))]
343  bbox = patch.getOuterBBox()
344  wcs = tract.getWcs()
345  cornerPosList = afwGeom.Box2D(bbox).getCorners()
346  coordList = [wcs.pixelToSky(pos) for pos in cornerPosList]
347  dataRefList = self.select.runDataRef(
348  patchRef, coordList, selectDataList=selectDataList).dataRefList
349  return [inputs[key(dataRef)] for dataRef in dataRefList]
350 
351  def writeMetadata(self, dataRef):
352  pass
def batchWallTime(cls, time, parsedCmd, numCores)
Return walltime request for batch job.
Definition: coaddDriver.py:119
def unpickle(factory, args, kwargs)
Definition: coaddDriver.py:78
def run(self, tractPatchRefList, butler, selectIdList=[])
Determine which tracts are non-empty before processing.
Definition: coaddDriver.py:132
def selectExposures(self, patchRef, selectDataList)
Select exposures to operate upon, via the SelectImagesTask.
Definition: coaddDriver.py:324
def runTract(self, patchRefList, butler, selectDataList=[])
Run stacking on a tract.
Definition: coaddDriver.py:168
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:88
def coadd(self, cache, data)
Construct coadd for a patch and measure.
Definition: coaddDriver.py:263
def warp(self, cache, patchId, selectDataList)
Warp all images for a patch.
Definition: coaddDriver.py:245
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:195
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:218