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