Restructure unit tests. Compute global coverage on the sky

when autoscaling.  Several unrelated small fixes.

Author: tskisner

src/toast/ops/noise_estimation.py

@@ -128,7 +128,8 @@ class NoiseEstim(Operator):
     )
 
     output_dir = Unicode(
-        ".",
+        None,
+        allow_none=True,
         help="If specified, write output data products to this directory",
     )
 
@@ -1075,21 +1076,21 @@ def process_noise_estimate(
             )
             log.debug_rank("Discard outliers", timer=timer)
 
-            self.save_psds(
-                binfreq, all_psds, all_times, det1, det2, fsample, fileroot, all_cov
-            )
-
-            if nbad > 0:
+            if self.output_dir is not None:
                 self.save_psds(
-                    binfreq,
-                    good_psds,
-                    good_times,
-                    det1,
-                    det2,
-                    fsample,
-                    fileroot + "_good",
-                    good_cov,
+                    binfreq, all_psds, all_times, det1, det2, fsample, fileroot, all_cov
                 )
+                if nbad > 0:
+                    self.save_psds(
+                        binfreq,
+                        good_psds,
+                        good_times,
+                        det1,
+                        det2,
+                        fsample,
+                        fileroot + "_good",
+                        good_cov,
+                    )
 
             final_freqs = binfreq
             final_psd = np.mean(np.array(good_psds), axis=0)

src/toast/ops/noise_model.py

@@ -5,7 +5,7 @@
 import numpy as np
 import traitlets
 from astropy import units as u
-from scipy.optimize import Bounds, least_squares
+from scipy.optimize import least_squares
 
 from ..noise import Noise
 from ..noise_sim import AnalyticNoise
@@ -105,8 +105,12 @@ def _exec(self, data, detectors=None, **kwargs):
                 freqs = in_model.freq(det)
                 in_psd = in_model.psd(det)
                 fitted, result = self._fit_psd(freqs, in_psd, params)
-                if result.success:
+                if result.success and len(result.x) == 3:
+                    # This was a good fit
                     params = result.x
+                else:
+                    msg = f"FitNoiseModel observation {ob.name}, det {det} failed.  Params = {result.x}"
+                    log.warning(msg)
                 nse_freqs[det] = freqs
                 nse_psds[det] = fitted
             if ob.comm.comm_group is not None:
@@ -156,12 +160,18 @@ def _evaluate_model(self, freqs, fmin, net, fknee, alpha):
         return psd
 
     def _fit_fun(self, x, *args, **kwargs):
-        net = x[0]
-        fknee = x[1]
-        alpha = x[2]
         freqs = kwargs["freqs"]
         data = kwargs["data"]
         fmin = kwargs["fmin"]
+        if "net" in kwargs:
+            # We are fixing the NET value
+            net = kwargs["net"]
+            fknee = x[0]
+            alpha = x[1]
+        else:
+            net = x[0]
+            fknee = x[1]
+            alpha = x[2]
         current = self._evaluate_model(freqs, fmin, net, fknee, alpha)
         # We weight the residual so that the high-frequency values specifying
         # the white noise plateau / NET are more important.  Also the lowest
@@ -186,6 +196,9 @@ def _fit_psd(self, freqs, data, guess=None):
         result = least_squares(
             self._fit_fun,
             x_0,
+            xtol=1.0e-8,
+            gtol=1.0e-8,
+            ftol=1.0e-8,
             kwargs={"freqs": raw_freqs, "data": raw_data, "fmin": raw_fmin},
         )
         fit_data = data
@@ -196,6 +209,31 @@ def _fit_psd(self, freqs, data, guess=None):
                 )
                 * psd_unit
             )
+        # else:
+        #     # Try fixing the NET based on the last few elements
+        #     fixed_net = np.sqrt(np.mean(raw_data[-5:]))
+        #     x_0 = np.array([0.1, 1.0])
+        #     result = least_squares(
+        #         self._fit_fun,
+        #         x_0,
+        #         xtol=1.0e-8,
+        #         gtol=1.0e-8,
+        #         ftol=1.0e-8,
+        #         kwargs={
+        #             "freqs": raw_freqs,
+        #             "data": raw_data,
+        #             "fmin": raw_fmin,
+        #             "net": fixed_net,
+        #         },
+        #     )
+        #     if result.success:
+        #         fit_data = (
+        #             self._evaluate_model(
+        #                 raw_freqs, raw_fmin, fixed_net, result.x[0], result.x[1]
+        #             )
+        #             * psd_unit
+        #         )
+
         return fit_data, result
 
     def _finalize(self, data, **kwargs):

src/toast/ops/pixels_wcs.py

@@ -282,6 +282,7 @@ def _set_wcs(self, proj, center, bounds, dims, res):
         self.pix_ra = self.wcs_shape[0]
         self.pix_dec = self.wcs_shape[1]
         self._n_pix = self.pix_ra * self.pix_dec
+        dbgrank = MPI.COMM_WORLD.rank
         self._n_pix_submap = self._n_pix // self.submaps
         if self._n_pix_submap * self.submaps < self._n_pix:
             self._n_pix_submap += 1
@@ -317,6 +318,19 @@ def _exec(self, data, detectors=None, **kwargs):
                 lonmax = max(lonmax, lnmax)
                 latmin = min(latmin, ltmin)
                 latmax = max(latmax, ltmax)
+            if ob.comm.comm_group_rank is not None:
+                # Synchronize between groups
+                if ob.comm.group_rank == 0:
+                    lonmin = ob.comm.comm_group_rank.allreduce(lonmin, op=MPI.MIN)
+                    latmin = ob.comm.comm_group_rank.allreduce(latmin, op=MPI.MIN)
+                    lonmax = ob.comm.comm_group_rank.allreduce(lonmax, op=MPI.MAX)
+                    latmax = ob.comm.comm_group_rank.allreduce(latmax, op=MPI.MAX)
+            # Broadcast result within the group
+            if ob.comm.comm_group is not None:
+                lonmin = ob.comm.comm_group.bcast(lonmin, root=0)
+                lonmax = ob.comm.comm_group.bcast(lonmax, root=0)
+                latmin = ob.comm.comm_group.bcast(latmin, root=0)
+                latmax = ob.comm.comm_group.bcast(latmax, root=0)
             new_bounds = (
                 (lonmax.to(u.degree), latmin.to(u.degree)),
                 (lonmin.to(u.degree), latmax.to(u.degree)),
@@ -494,8 +508,8 @@ def _get_scan_range(self, obs):
                 lon.append(phi)
                 lat.append(np.pi / 2 - theta)
             else:
-                lon.append(phi - center_lonlat[:, 0])
-                lat.append((np.pi / 2 - theta) - center_lonlat[:, 1])
+                lon.append(phi - center_lonlat[rank::ntask, 0])
+                lat.append((np.pi / 2 - theta) - center_lonlat[rank::ntask, 1])
         lon = np.unwrap(np.hstack(lon))
         lat = np.hstack(lat)
 

src/toast/pixels.py

@@ -316,6 +316,7 @@ def alltoallv_info(self):
             - The locations in the receive buffer of each submap.
 
         """
+        log = Logger.get()
         if self._alltoallv_info is not None:
             # Already computed
             return self._alltoallv_info
@@ -390,6 +391,17 @@ def alltoallv_info(self):
             recv_locations,
         )
 
+        msg_rank = 0
+        if self._comm is not None:
+            msg_rank = self._comm.rank
+        msg = f"alltoallv_info[{msg_rank}]:\n"
+        msg += f"  send_counts={send_counts} "
+        msg += f"send_displ={send_displ}\n"
+        msg += f"  recv_counts={recv_counts} "
+        msg += f"recv_displ={recv_displ} "
+        msg += f"recv_locations={recv_locations}"
+        log.verbose(msg)
+
         return self._alltoallv_info
 
 
@@ -744,6 +756,17 @@ def setup_alltoallv(self):
             for sm, locs in recv_locations.items():
                 self._recv_locations[sm] = scale * np.array(locs, dtype=np.int32)
 
+            msg_rank = 0
+            if self._dist.comm is not None:
+                msg_rank = self._dist.comm.rank
+            msg = f"setup_alltoallv[{msg_rank}]:\n"
+            msg += f"  send_counts={self._send_counts} "
+            msg += f"send_displ={self._send_displ}\n"
+            msg += f"  recv_counts={self._recv_counts} "
+            msg += f"recv_displ={self._recv_displ} "
+            msg += f"recv_locations={self._recv_locations}"
+            log.verbose(msg)
+
             # Allocate a persistent single-submap buffer
             self._reduce_buf_raw = self.storage_class.zeros(self._n_submap_value)
             self.reduce_buf = self._reduce_buf_raw.array()
@@ -774,6 +797,9 @@ def setup_alltoallv(self):
                         buf_check_fail = True
 
                     # Allocate a persistent receive buffer
+                    msg = f"{msg_rank}:  allocate receive buffer of "
+                    msg += f"{recv_buf_size} elements"
+                    log.verbose(msg)
                     self._receive_raw = self.storage_class.zeros(recv_buf_size)
                     self.receive = self._receive_raw.array()
             except:
@@ -796,6 +822,7 @@ def forward_alltoallv(self):
             None.
 
         """
+        log = Logger.get()
         gt = GlobalTimers.get()
         self.setup_alltoallv()
 

src/toast/tests/_helpers.py

@@ -100,6 +100,52 @@ def create_space_telescope(group_size, sample_rate=10.0 * u.Hz, pixel_per_proces
     site = SpaceSite("L2")
     return Telescope("test", focalplane=fp, site=site)
 
+def create_boresight_telescope(group_size, sample_rate=10.0 * u.Hz):
+    """Create a fake telescope with one boresight detector per process."""
+    nullquat = np.array([0, 0, 0, 1], dtype=np.float64)
+    n_det = group_size
+    det_names = [f"d{x:03d}" for x in range(n_det)]
+    pol_ang = np.array([(2*np.pi*x/n_det) for x in range(n_det)])
+
+    det_table = QTable(
+        [
+            Column(name="name", data=det_names),
+            Column(name="quat", data=[nullquat for x in range(n_det)]),
+            Column(name="pol_leakage", length=n_det, unit=None),
+            Column(name="psi_pol", data=pol_ang, unit=u.rad),
+            Column(name="fwhm", length=n_det, unit=u.arcmin),
+            Column(name="psd_fmin", length=n_det, unit=u.Hz),
+            Column(name="psd_fknee", length=n_det, unit=u.Hz),
+            Column(name="psd_alpha", length=n_det, unit=None),
+            Column(name="psd_net", length=n_det, unit=(u.K * np.sqrt(1.0 * u.second))),
+            Column(name="bandcenter", length=n_det, unit=u.GHz),
+            Column(name="bandwidth", length=n_det, unit=u.GHz),
+            Column(
+                name="pixel", data=[0 for x in range(n_det)]
+            ),
+        ]
+    )
+
+    fwhm = 5.0 * u.arcmin
+
+    for idet in range(len(det_table)):
+        det_table[idet]["pol_leakage"] = 0.0
+        det_table[idet]["fwhm"] = fwhm
+        det_table[idet]["bandcenter"] = 150.0 * u.GHz
+        det_table[idet]["bandwidth"] = 20.0 * u.GHz
+        det_table[idet]["psd_fmin"] = 1.0e-5 * u.Hz
+        det_table[idet]["psd_fknee"] = 0.05 * u.Hz
+        det_table[idet]["psd_alpha"] = 1.0
+        det_table[idet]["psd_net"] = 100 * (u.K * np.sqrt(1.0 * u.second))
+
+    fp = Focalplane(
+        detector_data=det_table,
+        sample_rate=sample_rate,
+        field_of_view=1.1 * (2 * fwhm),
+    )
+
+    site = SpaceSite("L2")
+    return Telescope("test", focalplane=fp, site=site)
 
 def create_ground_telescope(
     group_size, sample_rate=10.0 * u.Hz, pixel_per_process=1, fknee=None