Fix parallelization problems in noise estimation

src/libtoast/include/toast/fod_psd.hpp

@@ -1,5 +1,5 @@
 
-// Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+// Copyright (c) 2015-2024 by the parties listed in the AUTHORS file.
 // All rights reserved.  Use of this source code is governed by
 // a BSD-style license that can be found in the LICENSE file.
 
@@ -11,11 +11,12 @@
 
 namespace toast {
 void fod_autosums(int64_t n, const double * x, const uint8_t * good,
-                  int64_t lagmax, double * sums, int64_t * hits);
+                  int64_t lagmax, double * sums, int64_t * hits,
+		  int64_t all_sums);
 
 void fod_crosssums(int64_t n, const double * x, const double * y,
                    const uint8_t * good, int64_t lagmax, double * sums,
-                   int64_t * hits);
+                   int64_t * hits, int64_t all_sums, int64_t symmetric);
 }
 
 #endif // ifndef TOAST_FOD_PSD_HPP

src/libtoast/src/toast_fod_psd.cpp

@@ -1,5 +1,5 @@
 
-// Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+// Copyright (c) 2015-2024 by the parties listed in the AUTHORS file.
 // All rights reserved.  Use of this source code is governed by
 // a BSD-style license that can be found in the LICENSE file.
 
@@ -10,7 +10,8 @@
 
 
 void toast::fod_autosums(int64_t n, double const * x, uint8_t const * good,
-                         int64_t lagmax, double * sums, int64_t * hits) {
+                         int64_t lagmax, double * sums, int64_t * hits,
+			 int64_t all_sums) {
     toast::AlignedVector <double> xgood(n);
     toast::AlignedVector <uint8_t> gd(n);
 
@@ -24,27 +25,30 @@ void toast::fod_autosums(int64_t n, double const * x, uint8_t const * good,
         }
     }
 
-    #pragma \
-    omp parallel for default(none) shared(n, gd, lagmax, xgood, sums, hits) schedule(static, 100)
+#pragma								\
+  omp parallel for default(none)				\
+  shared(n, gd, lagmax, xgood, sums, hits, all_sums)		\
+  schedule(static, 100)
     for (int64_t lag = 0; lag < lagmax; ++lag) {
         int64_t j = lag;
         double lagsum = 0.0;
         int64_t hitsum = 0;
-        for (int64_t i = 0; i < (n - lag); ++i) {
+	int64_t imax = all_sums ? n - lag : n - lagmax;
+        for (int64_t i = 0; i < imax; ++i) {
             lagsum += xgood[i] * xgood[j];
             hitsum += gd[i] * gd[j];
             j++;
         }
-        sums[lag] = lagsum;
-        hits[lag] = hitsum;
+        sums[lag] += lagsum;
+        hits[lag] += hitsum;
     }
 
     return;
 }
 
 void toast::fod_crosssums(int64_t n, double const * x, double const * y,
                           uint8_t const * good, int64_t lagmax, double * sums,
-                          int64_t * hits) {
+                          int64_t * hits, int64_t all_sums, int64_t symmetric) {
     toast::AlignedVector <double> xgood(n);
     toast::AlignedVector <double> ygood(n);
     toast::AlignedVector <uint8_t> gd(n);
@@ -61,23 +65,28 @@ void toast::fod_crosssums(int64_t n, double const * x, double const * y,
         }
     }
 
-    #pragma \
-    omp parallel for default(none) shared(n, gd, lagmax, xgood, ygood, sums, hits) schedule(static, 100)
+#pragma									\
+  omp parallel for default(none)					\
+  shared(n, gd, lagmax, xgood, ygood, sums, hits, all_sums, symmetric)	\
+  schedule(static, 100)
     for (int64_t lag = 0; lag < lagmax; ++lag) {
         int64_t i, j;
         double lagsum = 0.0;
         int64_t hitsum = 0;
-        for (i = 0, j = lag; i < (n - lag); ++i, ++j) {
+	int64_t imax = all_sums ? n - lag : n - lagmax;
+        for (i = 0, j = lag; i < imax; ++i, ++j) {
             lagsum += xgood[i] * ygood[j];
             hitsum += gd[i] * gd[j];
         }
-
-        // Use symmetry to double the statistics
-        for (i = 0, j = lag; i < (n - lag); ++i, ++j) {
-            lagsum += xgood[j] * ygood[i];
-        }
-        sums[lag] = lagsum;
-        hits[lag] = 2 * hitsum;
+	if (symmetric && lag != 0) {
+	  // Use symmetry to double the statistics
+	  for (i = 0, j = lag; i < imax; ++i, ++j) {
+              lagsum += xgood[j] * ygood[i];
+	  }
+	  hitsum *= 2;
+	}
+        sums[lag] += lagsum;
+        hits[lag] += hitsum;
     }
 
     return;

src/toast/_libtoast/fod_psd.cpp

@@ -1,5 +1,5 @@
 
-// Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+// Copyright (c) 2015-2024 by the parties listed in the AUTHORS file.
 // All rights reserved.  Use of this source code is governed by
 // a BSD-style license that can be found in the LICENSE file.
 
@@ -8,7 +8,8 @@
 
 void init_fod_psd(py::module & m) {
     m.def("fod_crosssums", [](py::buffer x, py::buffer y, py::buffer good,
-                              int64_t lagmax, py::buffer sums, py::buffer hits) {
+                              int64_t lagmax, py::buffer sums, py::buffer hits,
+			      int64_t all_sums, int64_t symmetric) {
               pybuffer_check_1D <double> (x);
               pybuffer_check_1D <double> (y);
               pybuffer_check_1D <uint8_t> (good);
@@ -42,11 +43,12 @@ void init_fod_psd(py::module & m) {
               uint8_t * rawgood = reinterpret_cast <uint8_t *> (info_good.ptr);
               double * rawsums = reinterpret_cast <double *> (info_sums.ptr);
               int64_t * rawhits = reinterpret_cast <int64_t *> (info_hits.ptr);
-              toast::fod_crosssums(n, rawx, rawy, rawgood, lagmax, rawsums, rawhits);
+              toast::fod_crosssums(n, rawx, rawy, rawgood, lagmax,
+				   rawsums, rawhits, all_sums, symmetric);
               return;
           }, py::arg("x"), py::arg("y"), py::arg("good"), py::arg("lagmax"),
-          py::arg("sums"), py::arg(
-              "hits"), R"(
+          py::arg("sums"), py::arg("hits"), py::arg("all_sums"),
+	  py::arg("symmetric"), R"(
         Accumulate the time domain covariance between two vectors.
 
         Args:
@@ -56,14 +58,17 @@ void init_fod_psd(py::module & m) {
             lagmax (int): The maximum sample distance to consider.
             sums (array like, float64): The vector of sums^2 to accumulate for each lag.
             hits (array_like, int64):  The vector of hits to accumulate for each lag.
+            all_sums (int):  If non-zero, evaluate lags < `lagmax` even in the last
+                `lagmax` samples of `x` and `y`.
 
         Returns:
             None.
 
     )");
 
     m.def("fod_autosums", [](py::buffer x, py::buffer good, int64_t lagmax,
-                             py::buffer sums, py::buffer hits) {
+                             py::buffer sums, py::buffer hits,
+			     int64_t all_sums) {
               pybuffer_check_1D <double> (x);
               pybuffer_check_1D <uint8_t> (good);
               pybuffer_check_1D <double> (sums);
@@ -93,11 +98,11 @@ void init_fod_psd(py::module & m) {
               uint8_t * rawgood = reinterpret_cast <uint8_t *> (info_good.ptr);
               double * rawsums = reinterpret_cast <double *> (info_sums.ptr);
               int64_t * rawhits = reinterpret_cast <int64_t *> (info_hits.ptr);
-              toast::fod_autosums(n, rawx, rawgood, lagmax, rawsums, rawhits);
+              toast::fod_autosums(n, rawx, rawgood, lagmax,
+				  rawsums, rawhits, all_sums);
               return;
           }, py::arg("x"), py::arg("good"), py::arg("lagmax"),
-          py::arg("sums"), py::arg(
-              "hits"), R"(
+          py::arg("sums"), py::arg("hits"), py::arg("all_sums"), R"(
         Accumulate the time domain covariance.
 
         Args:
@@ -106,6 +111,8 @@ void init_fod_psd(py::module & m) {
             lagmax (int): The maximum sample distance to consider.
             sums (array like, float64): The vector of sums^2 to accumulate for each lag.
             hits (array_like, int64):  The vector of hits to accumulate for each lag.
+            all_sums (int):  If non-zero, evaluate lags < `lagmax` even in the last
+                `lagmax` samples of `x`.
 
         Returns:
             None.

src/toast/observation.py

@@ -656,6 +656,7 @@ def duplicate(
             sample_sets=self.all_sample_sets,
             process_rows=self.dist.process_rows,
         )
+        new_obs.set_local_detector_flags(self.local_detector_flags)
         for k, v in self._internal.items():
             if meta is None or k in meta:
                 new_obs[k] = copy.deepcopy(v)
@@ -720,6 +721,7 @@ def redistribute(
         times=None,
         override_sample_sets=False,
         override_detector_sets=False,
+        return_global_intervals=False,
     ):
         """Take the currently allocated observation and redistribute in place.
 
@@ -738,9 +740,11 @@ def redistribute(
                 existing sample set boundaries in the redistributed data.
             override_detector_sets (False, None or list):  If not False, override
                 existing detector set boundaries in the redistributed data.
+            return_global_intervals (bool):  Return a list of global intervals for
+                reference
 
         Returns:
-            None
+            None or global_intervals
 
         """
         log = Logger.get()
@@ -781,15 +785,16 @@ def redistribute(
         )
 
         # Do the actual redistribution
-        new_shr_manager, new_det_manager, new_intervals_manager = redistribute_data(
-            self.dist,
-            new_dist,
-            self.shared,
-            self.detdata,
-            self.intervals,
-            times=times,
-            dbg=self.name,
-        )
+        new_shr_manager, new_det_manager, new_intervals_manager, global_intervals = \
+            redistribute_data(
+                self.dist,
+                new_dist,
+                self.shared,
+                self.detdata,
+                self.intervals,
+                times=times,
+                dbg=self.name,
+            )
 
         # Redistribute any metadata objects that support it.
         for k, v in self._internal.items():
@@ -813,10 +818,17 @@ def redistribute(
         self.intervals = new_intervals_manager
 
         # Restore detector flags for our new local detectors
+        self._detflags = {x: int(0) for x in self.dist.dets[self.dist.comm.group_rank]}
         self.set_local_detector_flags(
             {x: all_det_flags[x] for x in self.local_detectors}
         )
 
+        if return_global_intervals:
+            global_intervals = self.dist.comm.comm_group.bcast(global_intervals)
+            return global_intervals
+        else:
+            return
+
     # Accelerator use
 
     def accel_create(self, names):

src/toast/observation_dist.py

@@ -917,4 +917,9 @@ def redistribute_data(
             glb = global_intervals[field]
         new_intervals_manager.create(field, glb, new_shared_manager[times], fromrank=0)
 
-    return new_shared_manager, new_detdata_manager, new_intervals_manager
+    return (
+        new_shared_manager,
+        new_detdata_manager,
+        new_intervals_manager,
+        global_intervals,
+    )