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