lsst.pipe.drivers  14.0-4-g61c0bd4+4
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 
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 
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, bbox=afwImage.bboxFromMetadata(md))
195  except FitsError:
196  self.log.warn("Unable to construct Wcs from %s" % (selectId,))
197  return None
198  return data
199 
200  def checkTract(self, cache, tractId, selectIdList):
201  """!Check whether a tract has any overlapping inputs
202 
203  This method only runs on slave nodes.
204 
205  @param cache: Pool cache
206  @param tractId: Data identifier for tract
207  @param selectDataList: List of selection data
208  @return whether tract has any overlapping inputs
209  """
210  def makePolygon(wcs, bbox):
211  """Return a polygon for the image, given Wcs and bounding box"""
212  return convexHull([wcs.pixelToSky(afwGeom.Point2D(coord)).getVector() for
213  coord in bbox.getCorners()])
214 
215  skymap = cache.skymap
216  tract = skymap[tractId]
217  tractWcs = tract.getWcs()
218  tractPoly = makePolygon(tractWcs, tract.getBBox())
219 
220  for selectData in selectIdList:
221  if not hasattr(selectData, "poly"):
222  selectData.poly = makePolygon(selectData.wcs, selectData.bbox)
223  if tractPoly.intersects(selectData.poly):
224  return True
225  return False
226 
227  def warp(self, cache, patchId, selectDataList):
228  """!Warp all images for a patch
229 
230  Only slave nodes execute this method.
231 
232  Because only one argument may be passed, it is expected to
233  contain multiple elements, which are:
234 
235  @param patchRef: data reference for patch
236  @param selectDataList: List of SelectStruct for inputs
237  @return selectDataList with non-overlapping elements removed
238  """
239  patchRef = getDataRef(cache.butler, patchId, cache.coaddType)
240  selectDataList = self.selectExposures(patchRef, selectDataList)
241  with self.logOperation("warping %s" % (patchRef.dataId,), catch=True):
242  self.makeCoaddTempExp.run(patchRef, selectDataList)
243  return selectDataList
244 
245  def coadd(self, cache, data):
246  """!Construct coadd for a patch and measure
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  """
256  patchRef = getDataRef(cache.butler, data.patchId, cache.coaddType)
257  selectDataList = data.selectDataList
258  coadd = None
259  with self.logOperation("coadding %s" % (patchRef.dataId,), catch=True):
260  if self.config.doOverwriteCoadd or not patchRef.datasetExists(cache.coaddType):
261  coaddResults = self.assembleCoadd.run(patchRef, selectDataList)
262  if coaddResults is not None:
263  coadd = coaddResults.coaddExposure
264  elif patchRef.datasetExists(cache.coaddType):
265  self.log.info("%s: Reading coadd %s" % (NODE, patchRef.dataId))
266  coadd = patchRef.get(cache.coaddType, immediate=True)
267 
268  if coadd is None:
269  return
270 
271  # The section of code below determines if the detection task should be
272  # run. If detection is run, then the products are written out as
273  # deepCoadd_calexp. If detection is not run, then the outputs of the
274  # assemble task are written out as deepCoadd.
275  if self.config.doDetection:
276  with self.logOperation("detection on {}".format(patchRef.dataId),
277  catch=True):
278  idFactory = self.detectCoaddSources.makeIdFactory(patchRef)
279  # This includes background subtraction, so do it before writing
280  # the coadd
281  detResults = self.detectCoaddSources.runDetection(coadd, idFactory)
282  self.detectCoaddSources.write(coadd, detResults, patchRef)
283  else:
284  patchRef.put(coadd, self.assembleCoadd.config.coaddName+"Coadd")
285 
286  def selectExposures(self, patchRef, selectDataList):
287  """!Select exposures to operate upon, via the SelectImagesTask
288 
289  This is very similar to CoaddBaseTask.selectExposures, except we return
290  a list of SelectStruct (same as the input), so we can plug the results into
291  future uses of SelectImagesTask.
292 
293  @param patchRef data reference to a particular patch
294  @param selectDataList list of references to specific data products (i.e. visit, ccd)
295  @return filtered list of SelectStruct
296  """
297  def key(dataRef):
298  return tuple(dataRef.dataId[k] for k in sorted(dataRef.dataId))
299  inputs = dict((key(select.dataRef), select)
300  for select in selectDataList)
301  skyMap = patchRef.get(self.config.coaddName + "Coadd_skyMap")
302  tract = skyMap[patchRef.dataId["tract"]]
303  patch = tract[(tuple(int(i)
304  for i in patchRef.dataId["patch"].split(",")))]
305  bbox = patch.getOuterBBox()
306  wcs = tract.getWcs()
307  cornerPosList = afwGeom.Box2D(bbox).getCorners()
308  coordList = [wcs.pixelToSky(pos) for pos in cornerPosList]
309  dataRefList = self.select.runDataRef(
310  patchRef, coordList, selectDataList=selectDataList).dataRefList
311  return [inputs[key(dataRef)] for dataRef in dataRefList]
312 
313  def writeMetadata(self, dataRef):
314  pass
def batchWallTime(cls, time, parsedCmd, numCores)
Return walltime request for batch job.
Definition: coaddDriver.py:101
def run(self, tractPatchRefList, butler, selectIdList=[])
Determine which tracts are non-empty before processing.
Definition: coaddDriver.py:114
def selectExposures(self, patchRef, selectDataList)
Select exposures to operate upon, via the SelectImagesTask.
Definition: coaddDriver.py:286
def runTract(self, patchRefList, butler, selectDataList=[])
Run stacking on a tract.
Definition: coaddDriver.py:150
def getDataRef(butler, dataId, datasetType="raw")
Definition: utils.py:17
def coadd(self, cache, data)
Construct coadd for a patch and measure.
Definition: coaddDriver.py:245
def warp(self, cache, patchId, selectDataList)
Warp all images for a patch.
Definition: coaddDriver.py:227
def logOperation(self, operation, catch=False, trace=True)
def getTargetList(parsedCmd, kwargs)
Get bare butler into Task.
Definition: coaddDriver.py:62
def readSelection(self, cache, selectId)
Read Wcs of selected inputs.
Definition: coaddDriver.py:177
def checkTract(self, cache, tractId, selectIdList)
Check whether a tract has any overlapping inputs.
Definition: coaddDriver.py:200