Mss003 tuning

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)
@@ -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)

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)