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