Modify mainly source_flags so that we can use multiple sources and taua

Author: Yoshinori Sueno

sotodlib/coords/planets.py

@@ -39,6 +39,17 @@
                ('QSO_J2253+1608', 343.4952422, 16.14301323),
                ('galactic_center', -93.5833, -29.0078)]
 
+def get_source_list_fromstr(target_source):
+    """Get a source_list from SOURCE_LIST by name, or raise ValueError if not found.
+    """
+    for isource in SOURCE_LIST:
+        if isinstance(isource, str):
+            isource_name = isource
+        elif isinstance(isource, tuple):
+            isource_name = isource[0]
+        if isource_name.lower() == target_source.lower():
+            return isource
+    raise ValueError(f'Source "{target_source}" not found in {SOURCE_LIST}.')
 
 class SlowSource:
     """Class to track the time-dependent position of a slow-moving source,
@@ -119,7 +130,11 @@ def get_scan_q(tod, planet, boresight_offset=None, refq=None):
     el = np.median(tod.boresight.el[::10])
     az = np.median(tod.boresight.az[::10])
     t = (tod.timestamps[0] + tod.timestamps[-1]) / 2
-    if isinstance(planet, str):
+    if isinstance(planet, (list, tuple)):
+        planet_name, ra, dec = planet
+        planet = SlowSource(t, float(ra) * coords.DEG,
+                                        float(dec) * coords.DEG)
+    else:
         planet = SlowSource.for_named_source(planet, t)
 
     def scan_q_model(t, az, el, planet):

sotodlib/core/flagman.py

@@ -278,6 +278,61 @@ def has_all_cut(flag):
 def count_cuts(flag):
     return np.array([len(x.ranges()) for x in flag], dtype='int')
 
+def has_ratio_cuts(flag, ratio):
+    """Determine if the ratio of flag samples to total samples in each flag 
+    exceeds a given threshold.
+
+    Args:
+      flag (RangesMatrix): so3g.proj.RangesMatrix
+      ratio (float or int): The threshold ratio above which a flag is 
+        considered to have too many cuts.
+    """
+
+    if not isinstance(ratio, (float, int)):
+        raise ValueError('Ratio must be float or int')
+    len_samps = flag.shape[-1]
+    return np.array([np.sum(np.diff(x.ranges()))/len_samps > ratio for x in flag])
+
+def flag_cut_select(flags, cut, kind):
+    """Determine which detectors to select based on flag conditions.
+
+    Args:
+        flags (RangesMatrix): An instance of so3g.proj.RangesMatrix indicating flagged time ranges.
+        cut (bool): If True, returns detectors to be excluded (cut); if False, returns detectors to be selected.
+        kind (str or float): One of the following:
+            - 'any': Select/cut detectors with any flagged samples.
+            - 'all': Select/cut detectors with all samples flagged.
+            - float: A threshold ratio (0.0–1.0); selects/cuts detectors whose flagged ratio exceeds the threshold.
+
+    Returns:
+        np.ndarray: Boolean array indicating which detectors to **keep** (True) or **drop** (False),
+                    depending on `cut` and `kind`.
+
+    Examples:
+        1. cut=True, kind='any' → Select detectors with **no** True flags (e.g., for Moon cut).
+        2. cut=False, kind='any' → Select detectors with **any** True flags (e.g., for planet selection).
+        3. cut=True, kind=0.4 → Select detectors with <40% of True flags.
+    """
+    if cut:
+        if kind == 'any':
+            return ~has_any_cuts(flags)
+        elif kind == 'all':
+            return ~has_all_cut(flags)
+        elif isinstance(kind, float):
+            return ~has_ratio_cuts(flags, ratio=kind)
+        else:
+            raise ValueError("kind must be 'any', 'all', or a float between 0.0 and 1.0")
+    else:
+        if kind == 'any':
+            return has_any_cuts(flags)
+        elif kind == 'all':
+            return has_all_cut(flags)
+        elif isinstance(kind, float):
+            return has_ratio_cuts(flags, ratio=kind)
+        else:
+            raise ValueError("kind must be 'any', 'all', or a float between 0.0 and 1.0")
+
+
 def sparse_to_ranges_matrix(arr, buffer=0, close_gaps=0, val=True):
     """Convert a csr sparse array into a ranges matrix
     
@@ -299,3 +354,23 @@ def sparse_to_ranges_matrix(arr, buffer=0, close_gaps=0, val=True):
         x[i].close_gaps(close_gaps)
     return x
 
+def find_common_edge_idx(flags):
+    """Find the common valid range across multiple RangesMatrix objects.
+
+    Args:
+        flags (RangesMatrix): An instance of so3g.proj.RangesMatrix indicating flagged time ranges.
+    Returns:
+        min_idx, max_idx: minmum and maximum indices that has False flag across all detectros.
+    """
+    max_val = max(arr.complement().ranges()[:, 1].max() for arr in flags)
+    common_mask = np.ones(max_val, dtype=bool)
+    for arr in flags:
+        mask = np.zeros(max_val, dtype=bool)
+        for start, end in arr.complement().ranges():
+            mask[start:end] = True
+        common_mask &= mask
+
+    valid_indices = np.where(common_mask)[0]
+    if len(valid_indices) == 0:
+        raise ValueError("No common valid range found across all flags.")
+    return valid_indices[0], valid_indices[-1]

sotodlib/preprocess/processes.py

@@ -12,7 +12,7 @@
 import sotodlib.coords.planets as planets
 
 from sotodlib.core.flagman import (has_any_cuts, has_all_cut,
-                                   count_cuts,
+                                   count_cuts, flag_cut_select,
                                    sparse_to_ranges_matrix)
 
 from .pcore import _Preprocess, _FracFlaggedMixIn
@@ -245,7 +245,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}_glitch_signal_diff'))
             plot_flag_stats(aman, proc_aman[self.glitch_name], flag_type='glitches', filename=filename.replace('{name}', f'{ufm}_glitch_stats'))
 
-
 class FixJumps(_Preprocess):
     """
     Repairs the jump heights given a set of jump flags and heights.
@@ -1281,6 +1280,7 @@ class SourceFlags(_Preprocess):
      Example config block::
 
         - name : "source_flags"
+          source_flags_name: "my_source_flags"
           calc:
             mask: {'shape': 'circle',
                    'xyr': [0, 0, 1.]}
@@ -1290,80 +1290,113 @@ class SourceFlags(_Preprocess):
             distance: 0 # max distance of footprint from source in degrees
           save: True
           select: True # optional
+            select_source: 'jupiter' # list of str or str. If not provided, all sources from center_on are selected.
+            cut: True # True performs cut based on flags
+            kind: 'any' # 'any', 'all', or 'cut'
+            Examples:
+                1. cut=True, kind='any' → Select detectors with **no** True flags (e.g., for Moon cut).
+                2. cut=False, kind='any' → Select detectors with **any** True flags (e.g., for planet selection).
+                3. cut=True, kind=0.4 → Select detectors with <40% of True flags.
 
     .. autofunction:: sotodlib.tod_ops.flags.get_source_flags
     """
     name = "source_flags"
+    def __init__(self, step_cfgs):
+        self.source_flags_name = step_cfgs.get('source_flags_name', 'source_flags')
+        super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
-        source_list = np.atleast_1d(self.calc_cfgs.get('center_on', 'planet'))
-        if source_list == ['planet']:
-            from sotodlib.coords.planets import SOURCE_LIST
+        from sotodlib.coords.planets import SOURCE_LIST
+        source_name = np.atleast_1d(self.calc_cfgs.get('center_on', 'planet'))
+        if 'planet' in source_name:
             source_list = [x for x in aman.tags if x in SOURCE_LIST]
             if len(source_list) == 0:
                 raise ValueError("No tags match source list")
+        else:
+            source_list = [planets.get_source_list_fromstr(isource) for isource in source_name]
 
         # find if source is within footprint + distance
         positions = planets.get_nearby_sources(tod=aman, source_list=source_list,
                                                distance=self.calc_cfgs.get('distance', 0))
-
         source_aman = core.AxisManager(aman.dets, aman.samps)
         for p in positions:
+            center_on = planets.get_source_list_fromstr(p[0])
             source_flags = tod_ops.flags.get_source_flags(aman,
                                                           merge=self.calc_cfgs.get('merge', False),
                                                           overwrite=self.calc_cfgs.get('overwrite', True),
                                                           source_flags_name=self.calc_cfgs.get('source_flags_name', None),
                                                           mask=self.calc_cfgs.get('mask', None),
-                                                          center_on=p[0],
+                                                          center_on=center_on,
                                                           res=self.calc_cfgs.get('res', None),
                                                           max_pix=self.calc_cfgs.get('max_pix', None))
 
             source_aman.wrap(p[0], source_flags, [(0, 'dets'), (1, 'samps')])
 
+            # if inv_flag is set, add inverse of source flags
+            if self.calc_cfgs.get('inv_flag'):
+                source_aman.wrap(p[0]+ '_inv',
+                                RangesMatrix.from_mask(~source_flags.mask()),
+                                [(0, 'dets'), (1, 'samps')])
+
         # add sources that were not nearby from source list
-        for source in source_list:
+        for source in source_name:
             if source not in source_aman._fields:
                 source_aman.wrap(source, RangesMatrix.zeros([aman.dets.count, aman.samps.count]),
                                  [(0, 'dets'), (1, 'samps')])
+                
+                if self.calc_cfgs.get('inv_flag'):
+                    source_aman.wrap(source + '_inv',
+                                    RangesMatrix.ones([aman.dets.count, aman.samps.count]),
+                                    [(0, 'dets'), (1, 'samps')])
 
         self.save(proc_aman, source_aman)
 
     def save(self, proc_aman, source_aman):
         if self.save_cfgs is None:
             return
         if self.save_cfgs:
-            proc_aman.wrap("source_flags", source_aman)
+            proc_aman.wrap(self.source_flags_name, source_aman)
 
     def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             source_flags = meta.preprocess.source_flags
         else:
-            source_flags = proc_aman.source_flags
+            source_flags = proc_aman[self.source_flags_name]
 
-        source_list = np.atleast_1d(self.calc_cfgs.get('center_on', 'planet'))
-        if source_list == ['planet']:
+        select_list = np.atleast_1d(self.select_cfgs.get("select_source", self.calc_cfgs.get('center_on', 'planet')))
+        if 'planet' in select_list:
             from sotodlib.coords.planets import SOURCE_LIST
-            source_list = [x for x in aman.tags if x in SOURCE_LIST]
-            if len(source_list) == 0:
+            select_list = [x for x in meta.tags if x in SOURCE_LIST]
+            if len(select_list) == 0:
                 raise ValueError("No tags match source list")
 
+        cuts = self.select_cfgs.get("cut", True) # default of True is for backward compatibility
+        if isinstance(cuts, bool):
+            cuts = [cuts]*len(select_list)
+        elif len(cuts) != len(select_list):
+            raise ValueError("Length of cuts must match length of select_source, or just bool")
+        
+        kinds = self.select_cfgs.get("kind", 'all') # default of 'all' is for backward compatibility
+        if isinstance(kinds, (str, float)):
+            kinds = [kinds]*len(select_list)
+        elif len(kinds) != len(select_list):
+            raise ValueError("Length of kinds must match length of select_source, or just str")
+
         keep_all = np.ones(meta.dets.count, dtype=bool)
 
-        for source in source_list:
+        for source, kind, cut in zip(select_list, kinds, cuts):
             if source in source_flags._fields:
-                keep = ~has_all_cut(source_flags[source])
-                if in_place:
-                    meta.restrict("dets", meta.dets.vals[keep])
-                    source_flags.restrict("dets", source_flags.dets.vals[keep])
-                else:
-                    keep_all &= keep
+                keep_all &= flag_cut_select(source_flags[source], cut, kind)
+
         if in_place:
+            meta.restrict("dets", meta.dets.vals[keep_all])
+            source_flags.restrict("dets", source_flags.dets.vals[keep_all])
             return meta
         else:
             return keep_all
-
+    
 class HWPAngleModel(_Preprocess):
     """Apply hwp angle model to the TOD.
 
@@ -1689,7 +1722,7 @@ def __init__(self, step_cfgs):
 
         super().__init__(step_cfgs)
 
-    def process(self, aman, proc_aman):
+    def process(self, aman, proc_aman, sim=False):
         n_modes = self.process_cfgs.get('n_modes')
         signal = aman.get(self.signal)
         flags = aman.flags.get(self.process_cfgs.get('source_flags'))
@@ -2146,10 +2179,33 @@ class CorrectIIRParams(_Preprocess):
     """
     name = "correct_iir_params"
 
-    def process(self, aman, proc_aman):
+    def process(self, aman, proc_aman, sim=False):
         from sotodlib.obs_ops import correct_iir_params
         correct_iir_params(aman)
 
+class TrimFlagEdge(_Preprocess):
+    """Trim edge until given flags of all detectors are False
+    To find first and last sample id that has False (i.e., no flags applied) for all detectors.
+    This is for avoiding glitchfill problem for data whose edge has flags of True.
+
+    Example config block::
+    
+        - name: "trim_flag_edge"
+          process:
+            flags: "pca_exclude"
+    
+    .. autofunction:: sotodlib.core.flagman.find_common_edge_idx
+    """
+    name = 'trim_flag_edge'
+
+    def process(self, aman, proc_aman, sim=False):
+        flags = aman.flags.get(self.process_cfgs.get('flags'))
+        trimst, trimen = core.flagman.find_common_edge_idx(flags)
+        aman.restrict('samps', (aman.samps.offset + trimst,
+                                aman.samps.offset + trimen))
+        proc_aman.restrict('samps', (proc_aman.samps.offset + trimst,
+                                     proc_aman.samps.offset + trimen))
+
 _Preprocess.register(SplitFlags)
 _Preprocess.register(SubtractT2P)
 _Preprocess.register(EstimateT2P)
@@ -2193,3 +2249,4 @@ def process(self, aman, proc_aman):
 _Preprocess.register(PointingModel)  
 _Preprocess.register(BadSubscanFlags)
 _Preprocess.register(CorrectIIRParams)
+_Preprocess.register(TrimFlagEdge)