Junna preproc

sotodlib/preprocess/processes.py

@@ -376,7 +376,6 @@ def plot(self, aman, proc_aman, filename):
                              plot_ds_factor=self.plot_cfgs.get("plot_ds_factor", 50), filename=filename.replace('{name}', f'{ufm}_jump_signal_diff'))
             plot_flag_stats(aman, proc_aman[name], flag_type='jumps', filename=filename.replace('{name}', f'{ufm}_jumps_stats'))
 
-
 class PSDCalc(_Preprocess):
     """ Calculate the PSD of the data and add it to the Preprocessing AxisManager under the
     "psd" field.
@@ -455,7 +454,6 @@ def plot(self, aman, proc_aman, filename):
             plot_psd(aman, signal=attrgetter(f"{self.wrap}.Pxx")(proc_aman),
                      xx=attrgetter(f"{self.wrap}.freqs")(proc_aman), filename=filename, **self.plot_cfgs)
 
-
 class GetStats(_Preprocess):
     """
     Get basic statistics from a TOD or its power spectrum.
@@ -1735,6 +1733,7 @@ def process(self, aman, proc_aman, sim=False):
         field = self.process_cfgs.get("noise_fit_array", None)
         if field:
             noise_fit = proc_aman[field]
+            noise_fit = noise_fit.restrict('dets', aman.dets.vals, in_place=False)
         else:
             noise_fit = None
 
@@ -1753,6 +1752,7 @@ def process(self, aman, proc_aman, sim=False):
                 noise_fit=noise_fit,
                 filter_params=spec.get("filter_params")
             )
+
             ffun = getattr(tod_ops.filters, fname)
             filters.append(ffun(**params))
 
@@ -1954,16 +1954,21 @@ class PCAFilter(_Preprocess):
 
     def __init__(self, step_cfgs):
         self.signal = step_cfgs.get('signal', 'signal')
+        self.model_signal = step_cfgs.get('model_signal', None)
 
         super().__init__(step_cfgs)
 
     def process(self, aman, proc_aman, sim=False):
         n_modes = self.process_cfgs.get('n_modes')
         signal = aman.get(self.signal)
+        if self.model_signal is not None:
+            model_signal = aman.get(self.model_signal)
+        else:
+            model_signal = signal.copy()
         if aman.dets.count < n_modes:
             raise ValueError(f'The number of pca modes {n_modes} is '
                              f'larger than the number of detectors {aman.dets.count}.')
-        model = tod_ops.pca.get_pca_model(aman, signal=signal, n_modes=n_modes)
+        model = tod_ops.pca.get_pca_model(aman, signal=model_signal, n_modes=n_modes)
         _ = tod_ops.pca.add_model(aman, model, signal=signal, scale=-1)
         return aman, proc_aman
 
@@ -1974,28 +1979,45 @@ class GetCommonMode(_Preprocess):
     example config file entry::
 
       - name: "get_common_mode"
+        noise_fit: True
+        f_max: 2.0
+        wrap_name: "common_demodQ"
         calc:
             signal: "signal"
             method: "median"
-            wrap: "signal_commonmode"
+
         save: True
 
+    If ``noise_fit`` is True, the 1/f noise fit parameters of the common mode
+    is wrapped together.
     .. autofunction:: sotodlib.tod_ops.pca.get_common_mode
     """
     name = 'get_common_mode'
+    def __init__(self, step_cfgs):
+        self.noise_fit = step_cfgs.get('noise_fit', False)
+        self.f_max = step_cfgs.get('f_max', None)
+        self.wrap_name = step_cfgs.get('wrap_name', 'common_mode')
+
+        super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
         common_mode = tod_ops.pca.get_common_mode(aman, **self.calc_cfgs)
-        common_aman = core.AxisManager(aman.samps)
-        common_aman.wrap(self.calc_cfgs['wrap'], common_mode, [(0, 'samps')])
+        if self.noise_fit:
+            common_aman = tod_ops.fft_ops.get_common_noise_params(aman, signal=common_mode,
+                                                                  f_max=self.f_max)
+            samps = core.OffsetAxis('samps', aman.samps.count)
+            common_aman.wrap(self.wrap_name, common_mode, [(0, samps)])
+        else:
+            common_aman = core.AxisManager(aman.samps)
+            common_aman.wrap(self.wrap_name, common_mode, [(0, 'samps')])
         self.save(proc_aman, common_aman)
         return aman, proc_aman
 
     def save(self, proc_aman, common_aman):
         if self.save_cfgs is None:
             return
         if self.save_cfgs:
-            proc_aman.wrap(self.calc_cfgs['wrap'], common_aman)
+            proc_aman.wrap(self.wrap_name, common_aman)
 
 class FilterForSources(_Preprocess):
     """
@@ -2134,6 +2156,7 @@ class EstimateT2P(_Preprocess):
      Example config block::
 
         - name : "estimate_t2p"
+          fit_in_freq : False
           calc:
             T_sig_name: 'dsT'
             Q_sig_name: 'demodQ'
@@ -2151,8 +2174,16 @@ class EstimateT2P(_Preprocess):
     """
     name = "estimate_t2p"
 
+    def __init__(self, step_cfgs):
+        self.fit_in_freq = step_cfgs.get('fit_in_freq', False)
+
+        super().__init__(step_cfgs)
+
     def calc_and_save(self, aman, proc_aman):
-        t2p_aman = tod_ops.t2pleakage.get_t2p_coeffs(aman, **self.calc_cfgs)
+        if self.fit_in_freq:
+            t2p_aman = tod_ops.t2pleakage.get_t2p_coeffs_in_freq(aman, **self.calc_cfgs)
+        else:
+            t2p_aman = tod_ops.t2pleakage.get_t2p_coeffs(aman, **self.calc_cfgs)
         self.save(proc_aman, t2p_aman)
 
         return aman, proc_aman
@@ -2178,8 +2209,10 @@ class SubtractT2P(_Preprocess):
     name = "subtract_t2p"
 
     def process(self, aman, proc_aman, sim=False):
-        tod_ops.t2pleakage.subtract_t2p(aman, proc_aman['t2p'],
-                                        **self.process_cfgs)
+        if 't2p' not in proc_aman._fields:
+            raise ValueError("No t2p field found in proc_aman. Please run estimate_t2p first.")
+            return meta, proc_aman
+        tod_ops.t2pleakage.subtract_t2p(aman, proc_aman['t2p'])
         return aman, proc_aman
 
 class SplitFlags(_Preprocess):
@@ -2389,6 +2422,35 @@ def process(self, aman, proc_aman, sim=False):
                                                       self.signal_name_U, merge=True)
         return aman, proc_aman
 
+class ScanFreqCut(_Preprocess):
+    """Apply high-pass cut at the scan frequency.
+
+    Example config block::
+
+        - name : 'scan_freq_cut'
+          signal_name_T: 'dsT'
+          signal_name_Q: 'demodQ'
+          signal_name_U: 'demodU'
+          process: True
+
+    """
+    name = "scan_freq_cut"
+
+    def __init__(self, step_cfgs):
+        self.signal_name_T = step_cfgs.get('signal_name_T', 'dsT')
+        self.signal_name_Q = step_cfgs.get('signal_name_Q', 'demodQ')
+        self.signal_name_U = step_cfgs.get('signal_name_U', 'demodU')
+        super().__init__(step_cfgs)
+
+    def process(self, aman, proc_aman, sim=False):
+        scan_freq = tod_ops.utils.get_scan_freq(aman)
+        hpf_cfg = {'type': 'sine2', 'cutoff': scan_freq, 'trans_width': scan_freq/10}
+        filt = tod_ops.get_hpf(hpf_cfg)
+        aman[self.signal_name_T] = tod_ops.fourier_filter(aman, filt, signal_name=self.signal_name_T, detrend='linear')
+        aman[self.signal_name_Q] = tod_ops.fourier_filter(aman, filt, signal_name=self.signal_name_Q, detrend='linear')
+        aman[self.signal_name_U] = tod_ops.fourier_filter(aman, filt, signal_name=self.signal_name_U, detrend='linear')
+        return aman, proc_aman
+
 class FocalplaneNanFlags(_Preprocess):
     """Find additional detectors which have nans 
        in their focal plane coordinates.
@@ -2692,7 +2754,8 @@ def process(self, aman, proc_aman, sim=False):
 _Preprocess.register(CombineFlags)
 _Preprocess.register(RotateFocalPlane)
 _Preprocess.register(RotateQU)
-_Preprocess.register(SubtractQUCommonMode) 
+_Preprocess.register(SubtractQUCommonMode)
+_Preprocess.register(ScanFreqCut)
 _Preprocess.register(FocalplaneNanFlags) 
 _Preprocess.register(PointingModel)  
 _Preprocess.register(BadSubscanFlags)

sotodlib/tod_ops/__init__.py

@@ -1,6 +1,6 @@
 from .detrend import detrend_tod
 from .fft_ops import rfft
-from .filters import fourier_filter, fft_trim
+from .filters import fourier_filter, fft_trim, get_hpf
 from .gapfill import \
     get_gap_fill, get_gap_fill_single, \
     get_gap_model, get_gap_model_single
@@ -12,3 +12,4 @@
 from .azss import get_azss
 from .t2pleakage import get_t2p_coeffs, subtract_t2p
 from . import deproject
+from . import utils