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