From 917765d513dc63144fb5b420be777f428f48840c Mon Sep 17 00:00:00 2001
From: Sasha Rahlin <arahlin@fnal.gov>
Date: Wed, 23 Mar 2022 11:33:18 -0500
Subject: [PATCH] add ReprojectMaps pipeline module to simplify reprojection of
 entire map frames

---
 maps/python/map_modules.py | 65 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/maps/python/map_modules.py b/maps/python/map_modules.py
index 437aa637..0aa8e028 100644
--- a/maps/python/map_modules.py
+++ b/maps/python/map_modules.py
@@ -14,6 +14,7 @@
     "InjectMaps",
     "ReplicateMaps",
     "CoaddMaps",
+    "ReprojectMaps",
 ]
 
 
@@ -530,3 +531,67 @@ def __call__(self, frame):
 
         if not input_weighted:
             RemoveWeights(frame)
+
+
+@core.indexmod
+class ReprojectMaps(object):
+    """
+    Reproject a map frame into a different projection.  Original data are
+    dropped and replaced by reprojected maps in the input frames.
+
+    Arguments
+    ---------
+    map_stub : G3SkyMap object
+        A stub (empty) sky map object to be used to construct the output maps.
+    rebin : int
+        If supplied and >1, subdivide the output pixel by n x n with each
+        sub-pixel taking on the input map values at pixel center (with interp or
+        nearest neighbor). The output pixel takes on the average of the
+        sub-pixel values.  In the case that the input map has higher resolution
+        than the output map (and that the input map is not low-pass filtered to
+        remove information above the Nyquist freq. of the output map pixel),
+        this reduces aliasing compared with direct sampling. But there would
+        still be aliased power from the input map from freq above the ouput map
+        pixel's Nyquist.
+    interp : bool
+        If True, use bilinear interpolation to extract values from the input
+        map.  Otherwise, the nearest-neighbor value is used.
+
+    Notes
+    -----
+    Maps can be rotated between Equatorial and Galactic coordinates, and/or
+    change polarization convention between COSMO and IAU, by setting the
+    appropriate attributes of the input and output maps.  Attributes not defined
+    in the output map are assumed to be that of the input map.
+    """
+
+    def __init__(self, map_stub=None, rebin=1, interp=False):
+        assert map_stub is not None, "map_stub argument required"
+        self.stub = map_stub
+        self.rebin = rebin
+        self.interp = interp
+
+    def __call__(self, frame):
+
+        if frame.type != core.G3FrameType.Map:
+            return
+
+        for key in ["T", "Q", "U", "Wpol", "Wunpol"]:
+
+            if key not in frame:
+                continue
+
+            m = frame.pop(key)
+
+            if key in "TQU":
+                mnew = self.stub.clone(False)
+                maps.reproj_map(m, mnew, rebin=self.rebin, interp=self.interp)
+
+            elif key in ["Wpol", "Wunpol"]:
+                mnew = core.G3SkyMapWeights(self.stub, key == "Wpol")
+                for wkey in mnew.keys():
+                    maps.reproj_map(m[wkey], mnew[wkey], rebin=self.rebin, interp=self.interp)
+
+            frame[key] = mnew
+
+        return frame