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