Pull out azss_model generation fn.

sotodlib/preprocess/processes.py

@@ -804,37 +804,74 @@ def process(self, aman, proc_aman):
                                     aman.samps.offset + aman.samps.count - trim))
 
 
-class EstimateAzSS(_Preprocess):
+class AzSS(_Preprocess):
     """Estimates Azimuth Synchronous Signal (AzSS) by binning signal by azimuth of boresight.
     All process confgis go to `get_azss`. If `method` is 'interpolate', no fitting applied 
     and binned signal is directly used as AzSS model. If `method` is 'fit', Legendre polynominal
     fitting will be applied and used as AzSS model.
 
     Example configuration block::
 
-      - name: "estimate_azss"
+      - name: "azss"
+        azss_stats_name: 'azss_statsQ'
+        proc_aman_turnaround_info: 'turnaround_flags'
         calc:
           signal: 'demodQ'
-          azss_stats_name: 'azss_statsQ'
-          range: [-1.57079, 7.85398]
+          frange: [-1.57079, 7.85398]
           bins: 1080
-          merge_stats: False
-          merge_model: False
+          left_right: True
         save: True
+        process:
+            subtract: True
 
     .. autofunction:: sotodlib.tod_ops.azss.get_azss
     """
-    name = "estimate_azss"
+    name = "azss"
+    def __init__(self, step_cfgs):
+        self.azss_stats_name = step_cfgs.get('azss_stats_name', 'azss_stats')
+        self.proc_aman_turnaround_info = step_cfgs.get('proc_aman_turnaround_info', None)
+
+        super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
-        calc_aman, _ = tod_ops.azss.get_azss(aman, **self.calc_cfgs)
-        self.save(proc_aman, calc_aman)
+        # If process is run then just wrap info from process step
+        if self.process_cfgs:
+            self.save(proc_aman, aman[self.azss_stats_name])
+        else:
+            if self.proc_aman_turnaround_info:
+                _f = attergetter(self.proc_aman_turnaround_info)
+                turnaround_info = _f(proc_aman)
+            else:
+                turnaround_info = None
+            azss_stats, _ = tod_ops.azss.get_azss(aman, turnaround_info=turnaround_info,
+                                                  azss_stats_name=self.azss_stats_name,
+                                                  merge_stats = False, merge_model=False,
+                                                  **self.calc_cfgs)
+            self.save(proc_aman, azss_stats)
 
     def save(self, proc_aman, azss_stats):
         if self.save_cfgs is None:
             return
         if self.save_cfgs:
-            proc_aman.wrap(self.calc_cfgs["azss_stats_name"], azss_stats)
+            proc_aman.wrap(self.azss_stats_name, azss_stats)
+
+    def process(self, aman, proc_aman):
+        subtract = self.process_cfgs.get('subtract', False)
+        if self.proc_aman_turnaround_info:
+            _f = attergetter(self.proc_aman_turnaround_info)
+            turnaround_info = _f(proc_aman)
+        else:
+            turnaround_info = None
+        if self.azss_stats_name in proc_aman:
+            if subtract:
+                tod_ops.azss.subtract_azss(aman, proc_aman[self.azss_stats_name],
+                                           signal = self.calc_cfgs.get('signal'),
+                                           in_place=True)
+        else:
+            tod_ops.azss.get_azss(aman, azss_stats_name=self.azss_stats_name,
+                                  turnaround_info=turnaround_info,
+                                  merge_stats = True, merge_model=False,
+                                  subtract_in_place=subtract, **self.calc_cfgs)
 
 class GlitchFill(_Preprocess):
     """Fill glitches. All process configs go to `fill_glitches`.
@@ -1641,7 +1678,7 @@ def process(self, aman, proc_aman):
 _Preprocess.register(SubtractHWPSS)
 _Preprocess.register(Apodize)
 _Preprocess.register(Demodulate)
-_Preprocess.register(EstimateAzSS)
+_Preprocess.register(AzSS)
 _Preprocess.register(GlitchFill)
 _Preprocess.register(FlagTurnarounds)
 _Preprocess.register(SubPolyf)