Er/bundler updates 260112

pipeline/bundling/bundle_maps.py

@@ -35,6 +35,13 @@ def main(args, parallelizor=None):
         else:
             raise ValueError(f"patch {patch} not recognized.")
 
+    atomic_list = None
+    if args.atomic_list is not None:
+        if '.npz' in args.atomic_list:
+            atomic_list = np.load(args.atomic_list)["atomic_list"]
+        else:
+            atomic_list = np.load(args.atomic_list)
+
     if os.path.isfile(args.bundle_db) and not args.overwrite:
         print(f"Loading from {args.bundle_db}.")
         bundle_coordinator = BundleCoordinator.from_dbfile(
@@ -46,7 +53,7 @@ def main(args, parallelizor=None):
         bundle_coordinator = BundleCoordinator(
             args.atomic_db, n_bundles=args.n_bundles,
             seed=args.seed, null_props=args.inter_obs_props,
-            query_restrict=query_restrict
+            query_restrict=query_restrict, atomic_list=atomic_list
         )
         bundle_coordinator.save_db(args.bundle_db)
 
@@ -61,25 +68,16 @@ def main(args, parallelizor=None):
     out_dir = args.output_dir
     os.makedirs(out_dir, exist_ok=True)
 
-    atomic_list = None
-    if args.atomic_list is not None:
-        if '.npz' in args.atomic_list:
-            atomic_list = np.load(args.atomic_list)["atomic_list"]
-        else:
-            atomic_list = np.load(args.atomic_list)
-
     car_map_template = args.car_map_template
 
     bundler = Bundler(
-        atomic_db=args.atomic_db,
         bundle_db=args.bundle_db,
         freq_channel=args.freq_channel,
         wafer=args.wafer,
         pix_type=args.pix_type,
         atomic_list=atomic_list,
         car_map_template=car_map_template,
-        telescope=args.tel,
-        query_restrict=query_restrict
+        telescope=args.tel
     )
 
     bundle_ids = range(args.n_bundles)
@@ -91,9 +89,9 @@ def main(args, parallelizor=None):
 
         bundled_map, weights_map, hits_map, fnames = bundler.bundle(
             bundle_id,
+            args.map_dir,
             split_label=split_intra_obs,
             null_prop_val=split_inter_obs,
-            map_dir=args.map_dir,
             abscal=args.abscal,
             parallelizor=parallelizor
         )

pipeline/bundling/bundling_utils.py

@@ -9,7 +9,7 @@
 from typing import Optional
 from dataclasses import dataclass
 import yaml
-
+import pandas as pd
 
 def _check_pix_type(pix_type):
     """
@@ -530,6 +530,63 @@ def make_full(template, split_pair, nbundles, pix_type, do_hits=True,
         return out
 
 
+def filter_by_atomic_list(arr, atomic_list, obs_id_only=False, return_index=False):
+    """ Filter an array by atomic_list
+
+    Parameters
+    ----------
+    arr: Array
+         First column should be obs_ids. Next columns should be wafer, freq_channel, if also filtering by them.
+         May also be a DataFrame; see filter_by_atomic_list_df
+    atomic_list: Array
+        Array of [(obs_id, wafer, freq_channel), ...]
+    obs_id_only: bool
+        If True filter with only obs_id. Else also use wafer, freq_channel.
+    return_index: bool
+        If True also return indices of selected rows.
+
+    Returns
+    -------
+    arr: Array
+         Filtered array
+    ind: Array [optional]
+         Indices of selected rows
+    """
+    if type(arr) is pd.DataFrame:
+        return filter_by_atomic_list_df(arr, atomic_list, obs_id_only, return_index)
+    if atomic_list is None:
+        if return_index:
+            return arr, slice(None)
+        else:
+            return arr
+    else:
+        atomic_list = np.asarray(atomic_list)
+        if obs_id_only:
+            arr_query = arr if arr.ndim == 1 else arr[:,0]
+            atomic_list_query = atomic_list if atomic_list.ndim == 1 else atomic_list[:,0]
+            ind = np.isin(arr_query, atomic_list_query)
+        else:
+            tags1 = [' '.join(line[:3]) for line in arr]
+            tags2 = [' '.join(line) for line in atomic_list]
+            ind = np.isin(tags1, tags2)
+        if return_index:
+            return arr[ind], ind
+        else:
+            return arr[ind]
+
+def filter_by_atomic_list_df(df, atomic_list, obs_id_only=False, return_index=False, reindex=True):
+    """Filter a DataFrame by atomic_list; see filter_by_atomic_list"""
+    if 'wafer' in df.columns and 'freq_channel' in df.columns:
+        arr = np.array([df.obs_id, df.wafer, df.freq_channel]).T
+    else:
+        arr = df.obs_id.to_numpy()
+    _, ind = filter_by_atomic_list(arr, atomic_list, obs_id_only, return_index=True)
+    if reindex:
+        return (df[ind]).reset_index(drop=True)
+    else:
+        return df[ind]
+
+
 @dataclass
 class Cfg:
     """