Common mode filter

src/toast/pipeline_tools/__init__.py

@@ -23,8 +23,10 @@
 from .filters import (
     add_polyfilter_args,
     add_polyfilter2D_args,
+    add_common_mode_filter_args,
     apply_polyfilter,
     apply_polyfilter2D,
+    apply_common_mode_filter,
     add_groundfilter_args,
     apply_groundfilter,
 )

src/toast/pipeline_tools/filters.py

@@ -9,14 +9,53 @@
 from toast.timing import function_timer, Timer
 from toast.utils import Logger, Environment
 
-from ..tod import OpPolyFilter, OpPolyFilter2D
+from ..tod import OpPolyFilter, OpPolyFilter2D, OpCommonModeFilter
 from ..todmap import OpGroundFilter
 
 #
 # Polynomial filter
 #
 
 
+def add_common_mode_filter_args(parser):
+    """Add the common mode filter arguments to argparser"""
+    parser.add_argument(
+        "--common-mode-filter",
+        required=False,
+        default=False,
+        action="store_true",
+        help="Apply common mode filter",
+        dest="apply_common_mode_filter",
+    )
+    parser.add_argument(
+        "--no-common-mode-filter",
+        required=False,
+        action="store_false",
+        help="Do not apply common mode filter",
+        dest="apply_common_mode_filter",
+    )
+    parser.set_defaults(apply_common_mode_filter=False)
+
+    parser.add_argument(
+        "--common-mode-filter-key",
+        required=False,
+        help="Focalplane key ('telescope', 'band', 'tube', 'wafer') "
+        "to match in common mode filter",
+    )
+    # Common flag mask may already be added
+    try:
+        parser.add_argument(
+            "--common-flag-mask",
+            required=False,
+            default=1,
+            type=np.uint8,
+            help="Common flag mask",
+        )
+    except argparse.ArgumentError:
+        pass
+    return
+
+
 def add_polyfilter2D_args(parser):
     """Add the polynomial filter arguments to argparser"""
     parser.add_argument(
@@ -97,6 +136,27 @@ def add_polyfilter_args(parser):
     return
 
 
+@function_timer
+def apply_common_mode_filter(args, comm, data, cache_name=None, verbose=True):
+    """Apply the common mode filter to data under `cache_name`."""
+    if not args.apply_common_mode_filter:
+        return
+    log = Logger.get()
+    timer = Timer()
+    timer.start()
+    if comm.world_rank == 0 and verbose:
+        log.info("Common mode filtering signal")
+    commonfilter = OpCommonModeFilter(
+        name=cache_name,
+        common_flag_mask=args.common_flag_mask,
+        focalplane_key=args.common_mode_filter_key,
+    )
+    commonfilter.exec(data)
+    if comm.world_rank == 0 and verbose:
+        timer.report_clear("Common mode filtering")
+    return
+
+
 @function_timer
 def apply_polyfilter2D(args, comm, data, cache_name=None, verbose=True):
     """Apply the 2D polynomial filter to data under `cache_name`."""

src/toast/tests/ops_filterbin.py

@@ -84,6 +84,25 @@ def combine_observation_matrix(rootname):
     return
 
 
+def make_input_map(lmax, nside, fname):
+    cl = np.ones([4, lmax + 1])
+    cl[:, 0:2] = 0
+    fwhm = np.radians(10)
+    inmap = hp.synfast(
+        cl,
+        nside,
+        lmax=lmax,
+        mmax=lmax,
+        pol=True,
+        pixwin=True,
+        fwhm=np.radians(30),
+        verbose=False,
+    )
+    inmap = hp.reorder(inmap, r2n=True)
+    hp.write_map(fname, inmap, overwrite=True, nest=True)
+    return
+
+
 class OpFilterBinTest(MPITestCase):
     def setUp(self):
         fixture_name = os.path.splitext(os.path.basename(__file__))[0]
@@ -99,8 +118,6 @@ def setUp(self):
                     os.remove(fname)
                 except OSError:
                     pass
-        if self.comm is not None:
-            self.comm.barrier()
 
         self.nobs = 1
         self.data = create_distdata(self.comm, obs_per_group=self.nobs)
@@ -217,28 +234,14 @@ def read_boresight_az(self, local_start=0, n=None):
 
         # Synthesize an input map
         self.lmax = 2 * self.sim_nside
-        self.cl = np.ones([4, self.lmax + 1])
-        self.cl[:, 0:2] = 0
-        # self.cl[1:] = 0  # DEBUG
-        fwhm = np.radians(10)
-        self.inmap = hp.synfast(
-            self.cl,
-            self.sim_nside,
-            lmax=self.lmax,
-            mmax=self.lmax,
-            pol=True,
-            pixwin=True,
-            fwhm=np.radians(30),
-            verbose=False,
-        )
-        self.inmap = hp.reorder(self.inmap, r2n=True)
-        self.inmapfile = os.path.join(self.outdir, "input_map.fits")
-        hp.write_map(self.inmapfile, self.inmap, overwrite=True, nest=True)
 
         return
 
     def test_filterbin(self):
 
+        inmapfile = os.path.join(self.outdir, "input_map1.fits")
+        make_input_map(self.lmax, self.sim_nside, inmapfile)
+
         name = "testtod2"
         init = OpCacheInit(name=name, init_val=0)
 
@@ -259,7 +262,7 @@ def test_filterbin(self):
 
         # Scan the signal from a map
         distmap = DistPixels(self.data, nnz=self.nnz, dtype=np.float32)
-        distmap.read_healpix_fits(self.inmapfile)
+        distmap.read_healpix_fits(inmapfile)
 
         scansim = OpSimScan(input_map=distmap, out=name)
         scansim.exec(self.data)
@@ -302,18 +305,18 @@ def test_filterbin(self):
             fname = os.path.join(self.outdir, outprefix + "filtered.fits.gz")
             outmap = hp.read_map(fname, None, nest=True)
 
-            inmap = hp.read_map(self.inmapfile, None, nest=True)
+            inmap = hp.read_map(inmapfile, None, nest=True)
             outmap_test = obs_matrix.dot(inmap.ravel()).reshape([self.nnz, -1])
 
             np.testing.assert_array_almost_equal(outmap, outmap_test)
 
-        if self.comm is not None:
-            self.comm.Barrier()
-
         return
 
     def test_filterbin_deproject(self):
 
+        inmapfile = os.path.join(self.outdir, "input_map2.fits")
+        make_input_map(self.lmax, self.sim_nside, inmapfile)
+
         name = "testtod2"
         init = OpCacheInit(name=name, init_val=0)
 
@@ -335,7 +338,7 @@ def test_filterbin_deproject(self):
 
         # Scan the signal from a map
         distmap = DistPixels(self.data, nnz=self.nnz, dtype=np.float32)
-        distmap.read_healpix_fits(self.inmapfile)
+        distmap.read_healpix_fits(inmapfile)
 
         scansim = OpSimScan(input_map=distmap, out=name)
         scansim.exec(self.data)
@@ -392,10 +395,8 @@ def test_filterbin_deproject(self):
         # Make a deprojection template map
         dpmap_file = os.path.join(self.outdir, "deprojection_map.fits")
         if self.rank == 0:
-            tmap = hp.read_map(self.inmapfile, nest=True)
+            tmap = hp.read_map(inmapfile, nest=True)
             hp.write_map(dpmap_file, tmap, nest=True)
-        if self.comm is not None:
-            self.comm.Barrier()
 
         # Run FilterBin
 
@@ -439,7 +440,7 @@ def test_filterbin_deproject(self):
             fname = os.path.join(self.outdir, outprefix + "filtered.fits.gz")
             outmap = hp.read_map(fname, None, nest=True)
 
-            inmap = hp.read_map(self.inmapfile, None, nest=True)
+            inmap = hp.read_map(inmapfile, None, nest=True)
             outmap_test = obs_matrix.dot(inmap.ravel()).reshape([self.nnz, -1])
 
             # The gain fluctuations can change the overall calibration
@@ -449,7 +450,4 @@ def test_filterbin_deproject(self):
 
             np.testing.assert_array_almost_equal(outmap, outmap_test)
 
-        if self.comm is not None:
-            self.comm.Barrier()
-
         return

src/toast/tests/ops_polyfilter.py

@@ -8,7 +8,14 @@
 
 import numpy as np
 
-from ..tod import OpPolyFilter, AnalyticNoise, OpSimNoise, Interval
+from ..tod import (
+    OpPolyFilter,
+    OpPolyFilter2D,
+    OpCommonModeFilter,
+    AnalyticNoise,
+    OpSimNoise,
+    Interval,
+)
 from ..todmap import TODHpixSpiral
 
 from ._helpers import create_outdir, create_distdata, boresight_focalplane
@@ -91,7 +98,17 @@ def setUp(self):
         self.data.obs[0]["noise"] = nse
         self.data.obs[0]["intervals"] = intervals
 
-    def test_filter(self):
+        # Add a focalplane to the observation
+        focalplane = {}
+        for idet, (det, quat) in enumerate(dquat.items()):
+            focalplane[det] = {
+                "quat" : quat,
+                "tube" : f"{idet // 4:03}",
+                "wafer" : f"{idet // 2:03}",
+            }
+        self.data.obs[0]["focalplane"] = focalplane
+
+    def test_1D_filter(self):
         # generate timestreams
         op = OpSimNoise()
         op.exec(self.data)
@@ -130,7 +147,85 @@ def test_filter(self):
                 old = orms[det]
                 if np.abs(rms / old) > 1e-6:
                     raise RuntimeError(
-                        "det {} old rms = {}, new rms = {}" "".format(det, old, rms)
+                        "det {} old rms = {}, new rms = {}".format(det, old, rms)
+                    )
+                del y
+        return
+
+    def test_2D_filter(self):
+        # generate timestreams
+        op = OpSimNoise()
+        op.exec(self.data)
+
+        # Replace the noise with time stamps (common for all detectors)
+        old_rms = []
+        for ob in self.data.obs:
+            tod = ob["tod"]
+            orms = {}
+            t = tod.read_times()
+            for det in tod.local_dets:
+                cachename = "noise_{}".format(det)
+                y = tod.cache.reference(cachename)
+                y[:] = t
+                orms[det] = np.std(y)
+                del y
+            old_rms.append(orms)
+
+        # Filter timestreams
+        op = OpPolyFilter2D(name="noise", order=self.order)
+        op.exec(self.data)
+
+        # Ensure all timestreams are zeroed out by the filter.
+
+        for ob, orms in zip(self.data.obs, old_rms):
+            tod = ob["tod"]
+            for det in tod.local_dets:
+                cachename = "noise_{}".format(det)
+                y = tod.cache.reference(cachename)
+                rms = np.std(y)
+                old = orms[det]
+                if np.abs(rms / old) > 1e-6:
+                    raise RuntimeError(
+                        "det {} old rms = {}, new rms = {}".format(det, old, rms)
+                    )
+                del y
+        return
+
+    def test_common_filter(self):
+        # generate timestreams
+        op = OpSimNoise()
+        op.exec(self.data)
+
+        # Replace the noise with time stamps (common for all detectors)
+        old_rms = []
+        for ob in self.data.obs:
+            tod = ob["tod"]
+            orms = {}
+            t = tod.read_times()
+            for det in tod.local_dets:
+                cachename = "noise_{}".format(det)
+                y = tod.cache.reference(cachename)
+                y[:] = t
+                orms[det] = np.std(y)
+                del y
+            old_rms.append(orms)
+
+        # Filter timestreams
+        op = OpCommonModeFilter(name="noise", focalplane_key="wafer")
+        op.exec(self.data)
+
+        # Ensure all timestreams are zeroed out by the filter.
+
+        for ob, orms in zip(self.data.obs, old_rms):
+            tod = ob["tod"]
+            for det in tod.local_dets:
+                cachename = "noise_{}".format(det)
+                y = tod.cache.reference(cachename)
+                rms = np.std(y)
+                old = orms[det]
+                if np.abs(rms / old) > 1e-6:
+                    raise RuntimeError(
+                        "det {} old rms = {}, new rms = {}".format(det, old, rms)
                     )
                 del y
         return

src/toast/tod/__init__.py

@@ -35,7 +35,7 @@
 
 from .noise import Noise
 
-from .polyfilter import OpPolyFilter, OpPolyFilter2D
+from .polyfilter import OpPolyFilter, OpPolyFilter2D, OpCommonModeFilter
 
 from .gainscrambler import OpGainScrambler
 from .applygain import OpApplyGain, write_calibration_file