From 10dd57843cb946c2a6c60651f6b3dad372334e48 Mon Sep 17 00:00:00 2001
From: Reijo Keskitalo <reijo.keskitalo@gmail.com>
Date: Mon, 23 Dec 2024 08:52:05 -0800
Subject: [PATCH 1/3] Add PWV cuts

---
 sotodlib/toast/workflows/sim_observe.py | 35 +++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/sotodlib/toast/workflows/sim_observe.py b/sotodlib/toast/workflows/sim_observe.py
index 5ba820eb2..ae5d8b203 100644
--- a/sotodlib/toast/workflows/sim_observe.py
+++ b/sotodlib/toast/workflows/sim_observe.py
@@ -98,6 +98,13 @@ def setup_simulate_observing(parser, operators):
         help="Realization index",
         type=int,
     )
+    parser.add_argument(
+        "--pwv_limit",
+        required=False,
+        type=float,
+        help="If set, discard observations with simulated PWV "
+        "higher than the limit [mm]",
+    )
 
     operators.append(
         toast.ops.SimGround(
@@ -230,6 +237,34 @@ def simulate_observing(job, otherargs, runargs, comm):
     job_ops.mem_count.prefix = "After Scan Simulation"
     job_ops.mem_count.apply(data)
 
+    if otherargs.pwv_limit is not None:
+        iobs = 0
+        ngood = 0
+        nbad = 0
+        while iobs < len(data.obs):
+            pwv = data.obs[iobs].telescope.site.weather.pwv.to_value(u.mm)
+            if pwv <= otherargs.pwv_limit:
+                ngood += 1
+                iobs += 1
+            else:
+                nbad += 1
+                del data.obs[iobs]
+                if len(data.obs) == 0:
+                    msg = (
+                        f"PWV limit = {otherargs.pwv_limit} mm rejected all "
+                        f"{nbad} observations assigned to this process"
+                    )
+                    raise RuntimeError(msg)
+        if toast_comm.comm_group_rank is not None:
+            nbad = toast_comm.comm_group_rank.allreduce(nbad)
+            ngood = toast_comm.comm_group_rank.allreduce(ngood)
+        log.info_rank(
+            f"  Discarded {nbad} / {ngood + nbad} observations "
+            f"with PWV > {otherargs.pwv_limit} mm in",
+            comm=comm,
+            timer=timer,
+        )
+
     # Apply LAT co-rotation
     if job_ops.corotate_lat.enabled:
         log.info_rank("  Running simulated LAT corotation...", comm=comm)

From 77265de65dbbc8162803397e7ec20358e4656989 Mon Sep 17 00:00:00 2001
From: Reijo Keskitalo <reijo.keskitalo@gmail.com>
Date: Mon, 23 Dec 2024 08:52:33 -0800
Subject: [PATCH 2/3] TOAST now produces lists with detector slices but
 axismanager requires 2D arrays

---
 sotodlib/toast/ops/mlmapmaker.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sotodlib/toast/ops/mlmapmaker.py b/sotodlib/toast/ops/mlmapmaker.py
index cd23a4197..1f44d168e 100644
--- a/sotodlib/toast/ops/mlmapmaker.py
+++ b/sotodlib/toast/ops/mlmapmaker.py
@@ -419,7 +419,7 @@ def _wrap_obs(self, ob, dets, passinfo):
         axobs.wrap("timestamps", ob.shared[self.times], axis_map=[(0, axsamps)])
         axobs.wrap(
             "signal",
-            ob.detdata[self.det_data][dets, :],
+            np.vstack(ob.detdata[self.det_data][dets, :]),
             axis_map=[(0, axdets), (1, axsamps)],
         )
         axobs.wrap("boresight", axbore)

From 6d3df42beea084d0c0db41e63e625c5938cc1cc7 Mon Sep 17 00:00:00 2001
From: Reijo Keskitalo <reijo.keskitalo@gmail.com>
Date: Mon, 23 Dec 2024 10:01:04 -0800
Subject: [PATCH 3/3] Put the add_obs call into a function to profile it
 separately

---
 sotodlib/toast/ops/mlmapmaker.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/sotodlib/toast/ops/mlmapmaker.py b/sotodlib/toast/ops/mlmapmaker.py
index 1f44d168e..d995f1378 100644
--- a/sotodlib/toast/ops/mlmapmaker.py
+++ b/sotodlib/toast/ops/mlmapmaker.py
@@ -453,6 +453,21 @@ def _wrap_obs(self, ob, dets, passinfo):
 
         return axobs
 
+    @function_timer
+    def _add_obs(self, mapmaker, axobs, nmat, name):
+        """ Add data to the mapmaker
+
+        Singled out for more granular timing
+        """
+        mapmaker.add_obs(
+            name,
+            axobs,
+            deslope=self.deslope,
+            noise_model=nmat,
+            signal_estimate=None,
+        )
+        return
+
     @function_timer
     def _init_mapmaker(
             self, mapmaker, signal_map, mapmaker_prev, x_prev, comm, gcomm, prefix,
@@ -655,13 +670,7 @@ def _exec(self, data, detectors=None, **kwargs):
                 nmat, nmat_file = self._load_noise_model(ob, npass, ipass, gcomm)
 
                 axobs = self._wrap_obs(ob, dets, passinfo)
-                mapmaker.add_obs(
-                    ob.name,
-                    axobs,
-                    deslope=self.deslope,
-                    noise_model=nmat,
-                    signal_estimate=None,
-                )
+                self._add_obs(mapmaker, axobs, nmat, ob.name)
                 del axobs
 
                 self._save_noise_model(mapmaker, nmat, nmat_file, gcomm)