Reorganize workflow steps (#943)

* Reorganize workflow steps.

- Move batch control script out of toast submodule

- Fix det indexing in toast readout filter application

- Extract noise model selection into helper functions

- Extract common preprocess workflow setup / application
  into helper functions

- If simulating signal from scanning from a map, cleanup
  any custom pointing model that is used.

* Fix entry for so_batch_control in pyproject.toml

* Fix typo

Author: tskisner

pyproject.toml

@@ -62,14 +62,14 @@ so_hardware_sim = "sotodlib.scripts.hardware_sim:main"
 so_hardware_plot = "sotodlib.scripts.hardware_plot:main"
 so_hardware_trim = "sotodlib.scripts.hardware_trim:main"
 so_hardware_info = "sotodlib.scripts.hardware_info:main"
+so_batch_control = "sotodlib.scripts.so_batch_control:main"
 so-metadata = "sotodlib.core.metadata.cli:main"
 so-site-pipeline = "sotodlib.site_pipeline.cli:main"
 so-data-package = "sotodlib.io.imprinter_cli:main"
 toast_so_sim = "sotodlib.toast.scripts.so_sim:cli"
 toast_so_map = "sotodlib.toast.scripts.so_map:cli"
 toast_so_transfer = "sotodlib.toast.scripts.so_transfer:cli"
 toast_so_convert = "sotodlib.toast.scripts.so_convert:cli"
-toast_so_batch_control = "sotodlib.toast.scripts.so_batch_control:cli"
 get_wafer_offset = "sotodlib.toast.scripts.get_wafer_offset:main"
 
 [tool.versioneer]

sotodlib/toast/scripts/so_batch_control.py renamed to sotodlib/scripts/so_batch_control.py

@@ -12,28 +12,10 @@
 """
 
 import argparse
-import datetime
 import os
 import re
-import sys
 import time
 
-import numpy as np
-
-from astropy import units as u
-
-# Import sotodlib.toast first, since that sets default object names
-# to use in toast.
-import sotodlib.toast as sotoast
-
-import toast
-import toast.ops
-from toast.mpi import MPI, Comm
-from toast.observation import default_values as defaults
-
-from .. import ops as so_ops
-from .. import workflows as wrk
-
 
 def load_obs_file(path):
     """Load the observation file into a list.
@@ -152,9 +134,6 @@ def find_obs(all_obs, n_job_obs, out_root, ignore_running=False, timeout_hours=2
 
 
 def main():
-    # Get optional MPI parameters
-    comm, procs, rank = toast.get_world()
-
     parser = argparse.ArgumentParser(description="Get observation IDs for a job")
     parser.add_argument(
         "--out_root",
@@ -220,42 +199,32 @@ def main():
 
     args = parser.parse_args()
 
-    if rank == 0:
-        if args.get_batch is not None:
-            # We are getting the next batch of observations
-            all_obs = load_obs_file(args.observations)
-            batch_obs = find_obs(all_obs, args.get_batch, args.out_root)
-            if args.batch_list:
-                batch_str = ",".join(batch_obs)
-                print(f"{batch_str}", flush=True)
-            else:
-                for obs in batch_obs:
-                    print(f"{obs}", flush=True)
-        elif args.get_state is not None:
-            state = get_obs_state(args.out_root, args.get_state)
-            print(f"{state}", flush=True)
-        elif args.clear_state is not None:
-            clear_obs_state(args.out_root, args.clear_state)
-        elif args.set_state_done is not None:
-            set_obs_state(args.out_root, args.set_state_done, "done")
-        elif args.set_state_running is not None:
-            set_obs_state(args.out_root, args.set_state_running, "running")
-        elif args.cleanup:
-            all_obs = load_obs_file(args.observations)
-            for obs in all_obs:
-                state = get_obs_state(args.out_root, obs)
-                if state is not None and state == "running":
-                    clear_obs_state(args.out_root, obs)
-
-    if comm is not None:
-        comm.barrier()
-
-
-def cli():
-    world, procs, rank = toast.mpi.get_world()
-    with toast.mpi.exception_guard(comm=world):
-        main()
+    if args.get_batch is not None:
+        # We are getting the next batch of observations
+        all_obs = load_obs_file(args.observations)
+        batch_obs = find_obs(all_obs, args.get_batch, args.out_root)
+        if args.batch_list:
+            batch_str = ",".join(batch_obs)
+            print(f"{batch_str}", flush=True)
+        else:
+            for obs in batch_obs:
+                print(f"{obs}", flush=True)
+    elif args.get_state is not None:
+        state = get_obs_state(args.out_root, args.get_state)
+        print(f"{state}", flush=True)
+    elif args.clear_state is not None:
+        clear_obs_state(args.out_root, args.clear_state)
+    elif args.set_state_done is not None:
+        set_obs_state(args.out_root, args.set_state_done, "done")
+    elif args.set_state_running is not None:
+        set_obs_state(args.out_root, args.set_state_running, "running")
+    elif args.cleanup:
+        all_obs = load_obs_file(args.observations)
+        for obs in all_obs:
+            state = get_obs_state(args.out_root, obs)
+            if state is not None and state == "running":
+                clear_obs_state(args.out_root, obs)
 
 
 if __name__ == "__main__":
-    cli()
+    main()

sotodlib/toast/ops/readout_filter.py

@@ -65,15 +65,16 @@ def _exec(self, data, detectors=None, **kwargs):
 
             # Get the rows of the focalplane table containing these dets
             det_table = ob.telescope.focalplane.detector_data
-            fp_rows = np.array(
-                [x for x, y in enumerate(det_table["name"]) if y in local_set_dets]
-            )
+            det_to_row = {
+                y: x for x, y in enumerate(det_table["name"]) if y in local_set_dets
+            }
+            fp_rows = np.array([det_to_row[x] for x in local_dets])
 
             # Get the set of all stream IDs
             all_wafers = set(det_table[self.wafer_key][fp_rows])
 
-            # The IIR filter parameters will either be in a single, top-level dictionary
-            # or they are organized per-UFM.
+            # The IIR filter parameters will either be in a single,
+            # top-level dictionary or they are organized per-UFM.
             if (
                 "per_stream" in ob[self.iir_params] and
                 ob[self.iir_params]["per_stream"]

sotodlib/toast/ops/splits.py

@@ -503,12 +503,12 @@ def _exec(self, data, detectors=None, **kwargs):
             "write_rcond",
             "write_solver_products",
             "save_cleaned",
-            "reset_pix_dist",
         ]
 
         # Possible traits we want to enable
         mapmaker_enable_traits = [
             "keep_final_products",
+            "reset_pix_dist",
         ]
 
         if hasattr(self.mapmaker, "map_binning"):

sotodlib/toast/scripts/so_map.py

@@ -58,25 +58,28 @@
 
 def reduce_data(job, otherargs, runargs, data):
     log = toast.utils.Logger.get()
+    job_ops = job.operators
+    job_tmpl = job.templates
 
-    wrk.simple_jumpcorrect(job, otherargs, runargs, data)
-    wrk.simple_deglitch(job, otherargs, runargs, data)
+    # Preprocess data to get flags / cuts
+    wrk.preprocess(job, otherargs, runargs, data)
 
-    wrk.flag_diff_noise_outliers(job, otherargs, runargs, data)
-    wrk.flag_noise_outliers(job, otherargs, runargs, data)
-    wrk.deconvolve_detector_timeconstant(job, otherargs, runargs, data)
+    # Now apply preprocessing to good detectors
     wrk.raw_statistics(job, otherargs, runargs, data)
 
-    wrk.filter_hwpss(job, otherargs, runargs, data)
-    wrk.filter_common_mode(job, otherargs, runargs, data)
+    wrk.diff_noise_estimation(job, otherargs, runargs, data)
+    wrk.noise_estimation(job, otherargs, runargs, data)
+
+    data = wrk.demodulate(job, otherargs, runargs, data)
+
     wrk.filter_ground(job, otherargs, runargs, data)
+    wrk.filter_common_mode(job, otherargs, runargs, data)
     wrk.filter_poly1d(job, otherargs, runargs, data)
     wrk.filter_poly2d(job, otherargs, runargs, data)
+
     wrk.diff_noise_estimation(job, otherargs, runargs, data)
     wrk.noise_estimation(job, otherargs, runargs, data)
 
-    data = wrk.demodulate(job, otherargs, runargs, data)
-
     wrk.processing_mask(job, otherargs, runargs, data)
     wrk.flag_sso(job, otherargs, runargs, data)
     wrk.hn_map(job, otherargs, runargs, data)
@@ -152,6 +155,19 @@ def main():
         action="store_true",
         help="Map each interval separately.",
     )
+    parser.add_argument(
+        "--preprocess_copy",
+        required=False,
+        default=False,
+        action="store_true",
+        help="Perform preprocessing on a copy of the data.",
+    )
+    parser.add_argument(
+        "--log_config",
+        required=False,
+        default=None,
+        help="Dump out config log to this yaml file",
+    )
 
     # The operators and templates we want to configure from the command line
     # or a parameter file.
@@ -160,23 +176,18 @@ def main():
     templates = list()
 
     wrk.setup_load_or_simulate_observing(parser, operators)
+    wrk.setup_preprocess(parser, operators)
 
-    wrk.setup_simple_jumpcorrect(operators)
-    wrk.setup_simple_deglitch(operators)
-
-    wrk.setup_flag_diff_noise_outliers(operators)
-    wrk.setup_flag_noise_outliers(operators)
-    wrk.setup_deconvolve_detector_timeconstant(operators)
     wrk.setup_raw_statistics(operators)
 
-    wrk.setup_filter_hwpss(operators)
-    wrk.setup_filter_common_mode(operators)
+    wrk.setup_demodulate(operators)
+
     wrk.setup_filter_ground(operators)
+    wrk.setup_filter_common_mode(operators)
     wrk.setup_filter_poly1d(operators)
     wrk.setup_filter_poly2d(operators)
-    wrk.setup_noise_estimation(operators)
 
-    wrk.setup_demodulate(operators)
+    wrk.setup_noise_estimation(operators)
 
     wrk.setup_processing_mask(operators)
     wrk.setup_flag_sso(operators)
@@ -201,8 +212,8 @@ def main():
             os.makedirs(otherargs.out_dir, exist_ok=True)
 
     # Log the config that was actually used at runtime.
-    outlog = os.path.join(otherargs.out_dir, "config_log.toml")
-    toast.config.dump_toml(outlog, config, comm=comm)
+    if otherargs.log_config is not None:
+        toast.config.dump_yaml(otherargs.log_config, config, comm=comm)
 
     # If this is a dry run, exit
     if otherargs.dry_run:

sotodlib/toast/scripts/so_sim.py

@@ -89,24 +89,23 @@ def simulate_data(job, otherargs, runargs, data):
 def reduce_data(job, otherargs, runargs, data):
     log = toast.utils.Logger.get()
 
-    wrk.simple_jumpcorrect(job, otherargs, runargs, data)
-    wrk.simple_deglitch(job, otherargs, runargs, data)
+    wrk.preprocess(job, otherargs, runargs, data)
 
-    wrk.flag_diff_noise_outliers(job, otherargs, runargs, data)
-    wrk.flag_noise_outliers(job, otherargs, runargs, data)
-    wrk.deconvolve_detector_timeconstant(job, otherargs, runargs, data)
     wrk.raw_statistics(job, otherargs, runargs, data)
 
-    wrk.filter_hwpss(job, otherargs, runargs, data)
-    wrk.filter_common_mode(job, otherargs, runargs, data)
+    wrk.diff_noise_estimation(job, otherargs, runargs, data)
+    wrk.noise_estimation(job, otherargs, runargs, data)
+
+    data = wrk.demodulate(job, otherargs, runargs, data)
+
     wrk.filter_ground(job, otherargs, runargs, data)
+    wrk.filter_common_mode(job, otherargs, runargs, data)
     wrk.filter_poly1d(job, otherargs, runargs, data)
     wrk.filter_poly2d(job, otherargs, runargs, data)
+
     wrk.diff_noise_estimation(job, otherargs, runargs, data)
     wrk.noise_estimation(job, otherargs, runargs, data)
 
-    data = wrk.demodulate(job, otherargs, runargs, data)
-
     wrk.processing_mask(job, otherargs, runargs, data)
     wrk.flag_sso(job, otherargs, runargs, data)
     wrk.hn_map(job, otherargs, runargs, data)
@@ -193,6 +192,19 @@ def main():
         action="store_true",
         help="Zero out detector data loaded from disk",
     )
+    parser.add_argument(
+        "--preprocess_copy",
+        required=False,
+        default=False,
+        action="store_true",
+        help="Perform preprocessing on a copy of the data.",
+    )
+    parser.add_argument(
+        "--log_config",
+        required=False,
+        default=None,
+        help="Dump out config log to this yaml file",
+    )
 
     # The operators and templates we want to configure from the command line
     # or a parameter file.
@@ -201,6 +213,7 @@ def main():
     templates = list()
 
     wrk.setup_load_or_simulate_observing(parser, operators)
+    wrk.setup_preprocess(parser, operators)
 
     wrk.setup_simulate_atmosphere_signal(operators)
     wrk.setup_simulate_sky_map_signal(operators)
@@ -220,15 +233,8 @@ def main():
     wrk.setup_simulate_readout_effects(operators)
     wrk.setup_save_data_hdf5(operators)
 
-    wrk.setup_simple_jumpcorrect(operators)
-    wrk.setup_simple_deglitch(operators)
-
-    wrk.setup_flag_diff_noise_outliers(operators)
-    wrk.setup_flag_noise_outliers(operators)
-    wrk.setup_deconvolve_detector_timeconstant(operators)
     wrk.setup_raw_statistics(operators)
 
-    wrk.setup_filter_hwpss(operators)
     wrk.setup_filter_common_mode(operators)
     wrk.setup_filter_ground(operators)
     wrk.setup_filter_poly1d(operators)
@@ -260,8 +266,13 @@ def main():
             os.makedirs(otherargs.out_dir, exist_ok=True)
 
     # Log the config that was actually used at runtime.
-    outlog = os.path.join(otherargs.out_dir, "config_log.toml")
-    toast.config.dump_toml(outlog, config, comm=comm)
+    if otherargs.log_config is not None:
+        # User wants to specify the location of the yaml config log
+        toast.config.dump_yaml(otherargs.log_config, config, comm=comm)
+    else:
+        # Log a toml file in the output directory
+        outlog = os.path.join(otherargs.out_dir, "config_log.toml")
+        toast.config.dump_toml(outlog, config, comm=comm)
 
     # If this is a dry run, exit
     if otherargs.dry_run: