Rotate weights with flatten_pol if possible, warn if weights are miss…

…ing (#73)

maps/include/maps/G3SkyMap.h

@@ -317,6 +317,12 @@ class G3SkyMapWeights : public G3FrameObject {
 			TT->IsCompatible(*UU));
 	}
 
+	bool IsCompatible(const G3SkyMap & other) const {
+		if (!TT)
+			return false;
+		return TT->IsCompatible(other);
+	}
+
 	MuellerMatrix operator [] (int i) {
 		return (!IsPolarized()) ? MuellerMatrix((*TT)[i]) :
 		    MuellerMatrix((*TT)[i], (*TQ)[i], (*TU)[i], (*QQ)[i],

maps/include/maps/maputils.h

@@ -40,7 +40,7 @@ void ReprojMap(G3SkyMapConstPtr in_map, G3SkyMapPtr out_map, int rebin=1, bool i
 
 // Flatten or unflatten Q and U polarization maps using the polarization gradient across the map.
 // The h parameter controls the pixel width over which a gradient is computed.
-void FlattenPol(FlatSkyMapPtr Q, FlatSkyMapPtr U, double h=0.001, bool invert=false);
+void FlattenPol(FlatSkyMapPtr Q, FlatSkyMapPtr U, G3SkyMapWeightsPtr W=NULL, double h=0.001, bool invert=false);
 
 // Compute map moments up to fourth order (mean, var, skewness, kurtosis) of the input map,
 // optionally excluding any pixels that are zero in the input mask, or zero/nan/inf in the input map

maps/python/map_modules.py

@@ -121,7 +121,13 @@ def FlattenPol(frame, invert=False):
         return
 
     qmap, umap = frame.pop("Q"), frame.pop("U")
-    maps.flatten_pol(qmap, umap, invert=invert)
+
+    if "Wpol" in frame:
+        wmap = frame.pop("Wpol")
+        maps.flatten_pol(qmap, umap, wmap, invert=invert)
+        frame["Wpol"] = wmap
+    else:
+        maps.flatten_pol(qmap, umap, invert=invert)
 
     frame["Q"] = qmap
     frame["U"] = umap

maps/src/G3SkyMap.cxx

@@ -1057,6 +1057,9 @@ PYBINDINGS("maps") {
 	      "True if all components are set, False if only the TT component is set")
 	    .add_property("congruent", &G3SkyMapWeights::IsCongruent,
 	      "True if all components are internally compatible with each other")
+	    .def("compatible", &G3SkyMapWeights::IsCompatible,
+	      "Returns true if the input argument is a map with matching dimensions "
+	      "and boundaries on the sky.")
 	    .def("rebin", &G3SkyMapWeights::Rebin, (bp::arg("scale")),
 	      "Rebin the weights into larger pixels by summing scale-x-scale blocks "
 	      "of pixels together.  Returns a new weights object.  Map dimensions "

maps/src/maputils.cxx

@@ -232,15 +232,24 @@ G3SkyMapPtr GetRaDecMask(G3SkyMapConstPtr m, double ra_left, double ra_right,
 }
 
 
-void FlattenPol(FlatSkyMapPtr Q, FlatSkyMapPtr U, double h, bool invert)
+void FlattenPol(FlatSkyMapPtr Q, FlatSkyMapPtr U, G3SkyMapWeightsPtr W, double h, bool invert)
 {
 	if (U->GetPolConv() == G3SkyMap::ConvNone)
 		log_warn("Missing pol_conv attribute for flatten_pol, assuming "
 			 "U.pol_conv is set to IAU. This will raise an error "
 			 "in the future.");
 
+	if (!W)
+		log_warn("Missing weights for flatten_pol");
+
 	g3_assert(Q->IsCompatible(*U));
 	g3_assert(Q->IsPolFlat() == U->IsPolFlat());
+	FlatSkyMapPtr flatptr;
+	if (!!W) {
+		g3_assert(W->IsCompatible(*Q));
+		flatptr = boost::dynamic_pointer_cast<FlatSkyMap>(W->TQ);
+		g3_assert(flatptr->IsPolFlat() == Q->IsPolFlat());
+	}
 
 	if (Q->IsPolFlat() && !invert)
 		return;
@@ -264,10 +273,41 @@ void FlattenPol(FlatSkyMapPtr Q, FlatSkyMapPtr U, double h, bool invert)
 
 		(*Q)[i.first] = cr * q - sr * u;
 		(*U)[i.first] = sr * q + cr * u;
+
+		if (!W)
+			continue;
+
+		MuellerMatrix w = (*W)[i.first];
+
+		double sr2 = 2.0 * sr * cr;
+		double cr2 = 1.0 - 2.0 * sr * sr;
+		double ws = (w.qq + w.uu) / 2.0;
+		double wd = (w.qq - w.uu) / 2.0;
+		double delta = wd * cr2 - w.qu * sr2;
+		double tq = w.tq;
+		double tu = w.tu;
+		w.tq = cr * tq - sr * tu;
+		w.tu = sr * tq + cr * tu;
+		w.qq = ws + delta;
+		w.uu = ws - delta;
+		w.qu = wd * sr2 + w.qu * cr2;
 	}
 
 	Q->SetFlatPol(!invert);
 	U->SetFlatPol(!invert);
+
+	if (!!W) {
+		flatptr = boost::dynamic_pointer_cast<FlatSkyMap>(W->TQ);
+		flatptr->SetFlatPol(!invert);
+		flatptr = boost::dynamic_pointer_cast<FlatSkyMap>(W->TU);
+		flatptr->SetFlatPol(!invert);
+		flatptr = boost::dynamic_pointer_cast<FlatSkyMap>(W->QQ);
+		flatptr->SetFlatPol(!invert);
+		flatptr = boost::dynamic_pointer_cast<FlatSkyMap>(W->QU);
+		flatptr->SetFlatPol(!invert);
+		flatptr = boost::dynamic_pointer_cast<FlatSkyMap>(W->UU);
+		flatptr->SetFlatPol(!invert);
+	}
 }
 
 
@@ -608,15 +648,17 @@ void maputils_pybindings(void){
 		"Returns a mask that is nonzero for any pixels within the given ra and dec ranges");
 
 	bp::def("flatten_pol", FlattenPol,
-		(bp::arg("Q"), bp::arg("U"), bp::arg("h")=0.001, bp::arg("invert")=false),
+		(bp::arg("Q"), bp::arg("U"), bp::arg("W")=G3SkyMapWeightsPtr(),
+		 bp::arg("h")=0.001, bp::arg("invert")=false),
 		"For maps defined on the sphere the direction of the polarization angle is "
 		"is defined relative to the direction of North.  When making maps we follow "
 		"this definition.\n\nFor any flat sky estimators, the polarization angle is "
 		"defined relative to the vertical axis.  For some map projections the "
 		"direction of north is not the same as the vertical axis.  This function "
 		"applies a rotation to the Q and U values to switch the curved sky Q/U "
 		"definition to the flat sky Q/U definition.\n\nIf for whatever reason you "
-		"want to reverse the process set the invert argument to True.");
+		"want to reverse the process set the invert argument to True. Also applies "
+		"the appropriate rotation to the Q and u elements of the associated weights.");
 
 	bp::def("reproj_map", ReprojMap,
 		(bp::arg("in_map"), bp::arg("out_map"), bp::arg("rebin")=1, bp::arg("interp")=false),

maps/tests/weights_operators.py

@@ -2,12 +2,12 @@
 
 import numpy as np
 from spt3g import core
-from spt3g.maps import G3SkyMapWeights, FlatSkyMap, MapPolType, MapPolConv
-from spt3g.maps import get_mask_map, remove_weights, apply_weights
+from spt3g.maps import G3SkyMapWeights, FlatSkyMap, MapPolType, MapPolConv, MapProjection
+from spt3g.maps import get_mask_map, remove_weights, remove_weights_t, apply_weights, apply_weights_t, flatten_pol
 
 for pol in [True, False]:
     # allocation
-    m = FlatSkyMap(500, 20, core.G3Units.arcmin, pol_conv=MapPolConv.IAU)
+    m = FlatSkyMap(500, 20, core.G3Units.arcmin, pol_conv=MapPolConv.IAU, proj=MapProjection.Proj5)
     mt = m.clone(False)
     mt.pol_type = MapPolType.T
     if pol:
@@ -39,9 +39,9 @@
         mt[15] = vec
         assert(np.allclose(mt[15], vec))
 
-    mat = np.array([[2, 0.1, 0.1],
-                    [0.1, 0.5, 0.05],
-                    [0.1, 0.05, 0.5]]) if pol else 2.
+    mat = np.array([[2, 0.2, 0.1],
+                    [0.2, 0.5, 0.05],
+                    [0.1, 0.05, 0.25]]) if pol else 2.
     mw[15] = mat
     assert(mw.npix_allocated == 1)
     assert(np.allclose(mw[15], mat))
@@ -63,8 +63,9 @@
 
     m[15] = 10
     mt *= m
-    mq *= m
-    mu *= m
+    if pol:
+        mq *= m
+        mu *= m
     mw *= m
     assert(np.allclose(mw[15], mat * 5))
     assert(mw.npix_allocated == 1)
@@ -111,22 +112,55 @@
     assert(not mt.sparse)
     assert(mt.npix_allocated == mt.size)
 
+    # compactify maps back to sparse
     mt.compact(zero_nans=True)
     assert(mt.npix_allocated == 1)
     assert(np.allclose(mt[15], np.atleast_1d(vec * 10)[0]))
     mask = get_mask_map(mt)
     assert(mask[15] == 1)
     assert(mask.npix_allocated == 1)
 
-    # compactify maps back to sparse
     if pol:
         mq.compact(zero_nans=True)
         mu.compact(zero_nans=True)
 
     # check memory-efficient weights functions
-    remove_weights(mt, mq, mu, mw, zero_nans=True)
+    if pol:
+        remove_weights(mt, mq, mu, mw, zero_nans=True)
+    else:
+        remove_weights_t(mt, mw, zero_nans=True)
     assert(mt.npix_allocated == 1)
 
-    apply_weights(mt, mq, mu, mw)
+    if pol:
+        apply_weights(mt, mq, mu, mw)
+        assert(np.allclose([mt[15], mq[15], mu[15]], vec * 10))
+    else:
+        apply_weights_t(mt, mw)
+        assert(np.allclose(mt[15], vec * 10))
     assert(mt.npix_allocated == 1)
-    assert(mt[15] == np.atleast_1d(vec * 10)[0])
+
+    if pol:
+        # check flatten_pol
+        mq2 = mq.copy()
+        mu2 = mu.copy()
+        mw2 = mw.copy()
+        assert(not mq2.flat_pol)
+        flatten_pol(mq2, mu2, mw2)
+        assert(mq2.flat_pol)
+        assert(not np.allclose(mq2, mq))
+        assert(not np.allclose(mu2, mu))
+        assert(not np.allclose(mw2.TQ, mw.TQ))
+        assert(not np.allclose(mw2.TU, mw.TU))
+        assert(not np.allclose(mw2.QQ, mw.QQ))
+        assert(not np.allclose(mw2.QU, mw.QU))
+        assert(not np.allclose(mw2.UU, mw.UU))
+
+        flatten_pol(mq2, mu2, mw2, invert=True)
+        assert(not mq2.flat_pol)
+        assert(np.allclose(mq2, mq))
+        assert(np.allclose(mu2, mu))
+        assert(np.allclose(mw2.TQ, mw.TQ))
+        assert(np.allclose(mw2.TU, mw.TU))
+        assert(np.allclose(mw2.QQ, mw.QQ))
+        assert(np.allclose(mw2.QU, mw.QU))
+        assert(np.allclose(mw2.UU, mw.UU))