Move PerTractCcdDataIdContainer to forcedPhotCcd.py

Currently, forcedPhotCcd.py is the only place where the
PerTractCcdDataIdContainer is being used, so it seems appropriate
just to include it there (as opposed to having in its own file).
The only other place currently planned to make use of it is in
meas_mosaic, which is not yet implemented.  If a more suitable
location for this container becomes apparent (e.g. in the process
of getting meas_mosaic running with the stack), it can be moved
accordingly then.

python/lsst/meas/base/forcedPhotCcd.py

@@ -20,21 +20,124 @@
 # the GNU General Public License along with this program.  If not,
 # see <https://www.lsstcorp.org/LegalNotices/>.
 #
+import collections
 
 import lsst.pex.config
+import lsst.pex.exceptions
+import lsst.pex.logging
 import lsst.pipe.base
 import lsst.afw.image
 import lsst.afw.table
+from lsst.geom import convexHull
 
 from .forcedPhotImage import ProcessImageForcedTask, ProcessImageForcedConfig
-from .dataIds import PerTractCcdDataIdContainer
 
 try:
     from lsst.meas.mosaic import applyMosaicResults
 except ImportError:
     applyMosaicResults = None
 
-__all__ = ("ForcedPhotCcdConfig", "ForcedPhotCcdTask")
+__all__ = ("PerTractCcdDataIdContainer", "ForcedPhotCcdConfig", "ForcedPhotCcdTask")
+
+
+class PerTractCcdDataIdContainer(lsst.pipe.base.DataIdContainer):
+    """A version of lsst.pipe.base.DataIdContainer that combines raw data IDs with a tract.
+
+    Required because we need to add "tract" to the raw data ID keys (defined as whatever we
+    use for 'src') when no tract is provided (so that the user is not required to know
+    which tracts are spanned by the raw data ID).
+
+    This IdContainer assumes that a calexp is being measured using the detection information,
+    a set of reference catalogs, from the set of coadds which intersect with the calexp.
+    It needs the calexp id (e.g. visit, raft, sensor), but is also uses the tract to decide
+    what set of coadds to use.  The references from the tract whose patches intersect with
+    the calexp are used.
+    """
+    def _addDataRef(self, namespace, dataId, tract):
+        """Construct a dataRef based on dataId, but with an added tract key"""
+        forcedDataId = dataId.copy()
+        forcedDataId['tract'] = tract
+        dataRef = namespace.butler.dataRef(datasetType=self.datasetType, dataId=forcedDataId)
+        self.refList.append(dataRef)
+
+    def makeDataRefList(self, namespace):
+        """Make self.refList from self.idList
+        """
+        if self.datasetType is None:
+            raise RuntimeError("Must call setDatasetType first")
+        log = lsst.pex.logging.Log.getDefaultLog()
+        skymap = None
+        visitTract = {}  # Set of tracts for each visit
+        visitRefs = {}   # List of data references for each visit
+        for dataId in self.idList:
+            if "tract" not in dataId:
+                # Discover which tracts the data overlaps
+                log.info("Reading WCS for components of dataId=%s to determine tracts" % (dict(dataId),))
+                if skymap is None:
+                    skymap = namespace.butler.get(namespace.config.coaddName + "Coadd_skyMap")
+
+                for ref in namespace.butler.subset("calexp", dataId=dataId):
+                    if not ref.datasetExists("calexp"):
+                        continue
+
+                    visit = ref.dataId["visit"]
+                    visitRefs = collections.defaultdict(list)
+                    visitRefs[visit].append(ref)
+
+                    md = ref.get("calexp_md", immediate=True)
+                    wcs = lsst.afw.image.makeWcs(md)
+                    box = lsst.afw.geom.Box2D(lsst.afw.geom.Point2D(0, 0),
+                                              lsst.afw.geom.Point2D(md.get("NAXIS1"), md.get("NAXIS2")))
+                    # Going with just the nearest tract.  Since we're throwing all tracts for the visit
+                    # together, this shouldn't be a problem unless the tracts are much smaller than a CCD.
+                    tract = skymap.findTract(wcs.pixelToSky(box.getCenter()))
+                    if overlapsTract(tract, wcs, box):
+                        visitTract = collections.defaultdict(set)
+                        visitTract[visit].add(tract.getId())
+            else:
+                tract = dataId.pop("tract")
+                # making a DataRef for src fills out any missing keys and allows us to iterate
+                for ref in namespace.butler.subset("src", dataId=dataId):
+                    self._addDataRef(namespace, ref.dataId, tract)
+
+        # Ensure all components of a visit are kept together by putting them all in the same set of tracts
+        for visit, tractSet in visitTract.iteritems():
+            for ref in visitRefs[visit]:
+                for tract in tractSet:
+                    self._addDataRef(namespace, ref.dataId, tract)
+        if visitTract:
+            tractCounter = collections.Counter()
+            for tractSet in visitTract.itervalues():
+                tractCounter.update(tractSet)
+            log.info("Number of visits for each tract: %s" % (dict(tractCounter),))
+
+def overlapsTract(tract, imageWcs, imageBox):
+    """Return whether the image (specified by Wcs and bounding box) overlaps the tract
+
+    @param tract: TractInfo specifying a tract
+    @param imageWcs: Wcs for image
+    @param imageBox: Bounding box for image
+    @return bool
+    """
+    tractWcs = tract.getWcs()
+    tractCorners = [tractWcs.pixelToSky(lsst.afw.geom.Point2D(coord)).getVector() for
+                    coord in tract.getBBox().getCorners()]
+    tractPoly = convexHull(tractCorners)
+
+    try:
+        imageCorners = [imageWcs.pixelToSky(lsst.afw.geom.Point2D(pix)) for pix in imageBox.getCorners()]
+    except lsst.pex.exceptions.LsstCppException, e:
+        # Protecting ourselves from awful Wcs solutions in input images
+        if (not isinstance(e.message, lsst.pex.exceptions.DomainErrorException) and
+            not isinstance(e.message, lsst.pex.exceptions.RuntimeErrorException)):
+            raise
+        return False
+
+    imagePoly = convexHull([coord.getVector() for coord in imageCorners])
+    if imagePoly is None:
+        return False
+    return tractPoly.intersects(imagePoly) # "intersects" also covers "contains" or "is contained by"
+
 
 class ForcedPhotCcdConfig(ProcessImageForcedConfig):
     doApplyUberCal = lsst.pex.config.Field(