azss correction per subscan init

sotodlib/preprocess/processes.py

@@ -1664,6 +1664,51 @@ def calc_and_save(self, aman, proc_aman):
         from sotodlib.coords import pointing_model
         if self.calc_cfgs:
             pointing_model.apply_pointing_model(aman)
+
+class EstimateAzSSsubscan(_Preprocess):
+    """
+    .. autofunction:: sotodlib.tod_ops.azss.get_azss_lr
+    """
+    name = "estimate_azss_lr"
+
+    def calc_and_save(self, aman, proc_aman):
+        calc_aman_left, calc_aman_right, model_left, model_right = tod_ops.azss.get_azss_lr(aman, **self.calc_cfgs)
+        self.save(proc_aman, calc_aman_left, "left", "stats")
+        self.save(proc_aman, calc_aman_right, "right", "stats")
+        self.save(proc_aman, model_left, "left", "model")
+        self.save(proc_aman, model_right, "right", "model")
+
+    def save(self, proc_aman, calc_aman, leftright, modelstats):
+        if self.save_cfgs is None:
+            return
+        if self.save_cfgs:
+            if modelstats == "stats":
+                _azss_stats_dict = {
+                    f'{leftright}_{self.calc_cfgs["signal"]}': calc_aman
+                    }
+                azss_stats = AxisManager(calc_aman.dets, calc_aman.bin_az_samps)
+                azss_stats.wrap('binned_az', calc_aman['binned_az'], [(0, 'bin_az_samps')])
+                for key, val in _azss_stats_dict.items():
+                    azss_stats.wrap('binned_signal_'+key, val['binned_signal'], [(0, 'dets'), (1, 'bin_az_samps')])
+                    azss_stats.wrap('binned_signal_sigma_'+key, val['binned_signal_sigma'], [(0, 'dets'), (1, 'bin_az_samps')])
+                    azss_stats.wrap('uniform_binned_signal_sigma_'+key, val['uniform_binned_signal_sigma'], [(0, 'dets')])
+                    azss_stats.wrap('binned_model_'+key, val['binned_model'], [(0, 'dets'), (1, 'bin_az_samps')])
+                    azss_stats.wrap('redchi2s_'+key, val['redchi2s'], [(0, 'dets')])
+                proc_aman.wrap(self.calc_cfgs["azss_stats_name"]+"_"+leftright, azss_stats)
+            elif modelstats == "model":
+                azss_model = core.AxisManager(proc_aman.dets, proc_aman.samps)
+                azss_model.wrap('azss_model', calc_aman)
+                proc_aman.wrap(self.calc_cfgs["azss_model_name"]+"_"+leftright, azss_model) #calc_aman)
+
+class SubtractAzSSsubscan(_Preprocess):
+    """
+    .. autofunction:: sotodlib.tod_ops.azss.subtract_azss_lr
+    """
+    name = "subtract_azss_lr"
+
+    def process(self, aman, proc_aman):
+        tod_ops.azss.subtract_azss_lr(aman, **self.process_cfgs)
+
 
 _Preprocess.register(SplitFlags)
 _Preprocess.register(SubtractT2P)

sotodlib/tod_ops/azss.py

@@ -325,4 +325,140 @@ def subtract_azss(aman, signal='signal', azss_template_name='azss_model',
                   [(0, 'dets'), (1, 'samps')])
 
     if remove_template:
-        aman.move(azss_template_name, None)
+        aman.move(azss_template_name, None)
+
+
+def get_azss_lr(aman, signal='signal', bins=100, 
+                apodize_edges_sec=100, apodize_flags_sec=1,
+                method='fit', max_mode=10,
+                merge_stats=True, azss_stats_name='azss_stats',
+                merge_model=True, azss_model_name='azss_model'):
+    """Compute the azimuthal scan synchronous signal (AZSS) for left and right scans.
+    
+       Parameters:
+       ----------
+        aman : object
+            The data manager object containing the scan data and flags.
+        signal : str, optional
+            The name of the signal to process. Default is 'signal'.
+        bins : int, optional
+            Number of bins for the histogram. Default is 100.
+        apodize_edges_sec : int, optional
+            Number of seconds to apodize at the edges of the scan. Default is 100.
+        apodize_flags_sec : int, optional
+            Number of seconds to apodize around flagged regions. Default is 1.
+        bad_flags : array-like, optional
+            Additional flags to consider as bad. Default is None.
+        method : str, optional
+            Method to use for fitting the AZSS. Default is 'fit'.
+        max_mode : int, optional
+            Maximum mode to consider for the fit. Default is 10.
+        merge_stats : bool, optional
+            Whether to merge the computed statistics into the data manager. Default is True.
+        azss_stats_name : str, optional
+            Name to use for the merged statistics in the data manager. Default is 'azss_stats'.
+        merge_model : bool, optional
+            Whether to merge the computed model into the data manager. Default is True.
+        azss_model_name : str, optional
+            Name to use for the merged model in the data manager. Default is 'azss_model'.
+
+        Returns:
+        ----------
+        tuple
+            A tuple containing the left and right AZSS statistics and models:
+            (azss_stats_left, azss_stats_right, azss_model_left, azss_model_right).
+        """
+
+    mask_for_valid_left_scans = aman.flags.right_scan
+    mask_for_valid_right_scans = aman.flags.left_scan
+
+    if "bad_subscan" in aman.flags:
+        if "valid_right_scans" not in aman: 
+            aman.flags.reduce(flags=['bad_subscan', 'left_scan'], 
+                                method='union', wrap=True, 
+                                new_flag='valid_left_scans', remove_reduced=False)
+            aman.flags.reduce(flags=['bad_subscan', 'right_scan'], 
+                                method='union', wrap=True, 
+                                new_flag='valid_right_scans', remove_reduced=False)
+        mask_for_valid_left_scans = aman.flags.valid_left_scans
+        mask_for_valid_right_scans = aman.flags.valid_right_scans
+
+    apodize_edges_samps = int(apodize_edges_sec * np.round(1/np.median(np.diff(aman.timestamps))))
+    apodize_flags_samps = int(apodize_flags_sec * np.round(1/np.median(np.diff(aman.timestamps))))
+
+    azss_stats_left, azss_model_left = get_azss(aman, signal=signal, merge_stats=False, merge_model=False,
+                                                flags=mask_for_valid_left_scans, bins=bins, apodize_flags=True,
+                                                method=method, max_mode=max_mode, apodize_flags_samps=apodize_flags_samps,
+                                                apodize_edges_samps=apodize_edges_samps)
+
+    azss_stats_right, azss_model_right = get_azss(aman, signal=signal, merge_stats=False, merge_model=False,
+                                                  flags=mask_for_valid_right_scans, bins=bins, apodize_flags=True,
+                                                  method=method, max_mode=max_mode, apodize_flags_samps=apodize_flags_samps,
+                                                  apodize_edges_samps=apodize_edges_samps)
+    aman.flags.move('valid_left_scans', None)
+    aman.flags.move('valid_right_scans', None)
+
+    # If AZSS is not modeled for all AZ set zero to nan
+    if np.any(np.isnan(azss_model_left)) or np.any(np.isnan(azss_model_right)):
+        azss_model_left[np.isnan(azss_model_left)] = 0
+        azss_model_right[np.isnan(azss_model_right)] = 0
+
+    if merge_stats:
+        aman.wrap(azss_stats_name+"_left", azss_stats_left)
+        aman.wrap(azss_stats_name+"_right", azss_stats_right)
+    if merge_model:
+        aman.wrap(azss_model_name+"_left", azss_model_left, [(0, 'dets'), (1, 'samps')])
+        aman.wrap(azss_model_name+"_right", azss_model_right, [(0, 'dets'), (1, 'samps')])
+
+    return azss_stats_left, azss_stats_right, azss_model_left, azss_model_right
+
+
+def subtract_azss_lr(aman, signal='signal', merge_stats=True, merge_model=True,
+                     azss_stats_name='azss_stats', azss_model_name='azss_model', 
+                     bins=100, apodize_edges_sec=100, apodize_flags_sec=1,
+                     method='fit', max_mode=10,): 
+    """
+    Subtracts the azimuth synchronous signal (AZSS) from the given signal in the aman object, 
+    for both left and right scans.
+
+    Parameters:
+    ----------
+    aman : object
+        The data object containing the signal and flags.
+    signal : str, optional
+        The name of the signal to be processed. Default is 'signal'.
+    merge_stats : bool, optional
+        If True, merges statistics after subtraction into aman. Default is True.
+    azss_model_name : str, optional
+        The base name of the AzSS model to be subtracted. Default is 'azss_model'.
+    bins : int, optional
+        The number of bins to use for AzSS calculation. Default is 100.
+
+    Returns:
+    ----------
+    None
+    """
+    left_flags = aman.flags.left_scan.mask()
+    right_flags = aman.flags.right_scan.mask()
+
+    aman[signal][:, left_flags] -= aman[f"{azss_model_name}{signal[-1]}_left"][:, left_flags]
+    aman[signal][:, right_flags] -= aman[f"{azss_model_name}{signal[-1]}_right"][:, right_flags]
+
+    if merge_stats or merge_model:
+        apodize_edges_samps = int(apodize_edges_sec * np.round(1/np.median(np.diff(aman.timestamps))))
+        apodize_flags_samps = int(apodize_flags_sec * np.round(1/np.median(np.diff(aman.timestamps))))
+
+        azss_stats_left, azss_stats_right, \
+            azss_model_left, azss_model_right = \
+                get_azss_lr(aman, signal=signal, merge_stats=False, merge_model=False, bins=bins,
+                            apodize_edges_sec=apodize_edges_samps, apodize_flags_sec=apodize_flags_samps,
+                            method=method, max_mode=max_mode)
+        if merge_stats:
+            aman.wrap(azss_stats_name+"_left_after", azss_stats_left)
+            aman.wrap(azss_stats_name+"_right_after", azss_stats_right)
+        if merge_model:
+            aman.wrap(azss_model_name+"_left_after", azss_model_left, [(0, 'dets'), (1, 'samps')])
+            aman.wrap(azss_model_name+"_right_after", azss_model_right, [(0, 'dets'), (1, 'samps')])
+
+        del azss_model_left, azss_model_right, azss_stats_left, azss_stats_right
+    return