Fix parallelization problems in noise estimation (#728)

* Fix parallelization problems in noise estimation * Small fixes for self.view=None * Some small fixes, discovered with unit tests and batch jobs (#730) with 8 or 16 processes per group: * When duplicating an observation, also duplicate per-detector flags. * When redistributing an observation, reset per-detector flags before setting. * Make the common mode removal an option (default True) prior to noise estimation. This is useful if the data has already had the common mode removed prior to this operator. * When high-pass filtering data that is distributed by time slices there may be some chunks that are fully flagged. This should not be an error. Instead, the flagged timestream is set to zero. * When communicating overlaps, ensure that the span used (lagmax + half_average) is less than the number of samples. * Ignore flagged detectors in poly2d fit --------- Co-authored-by: Theodore Kisner <[email protected]> Co-authored-by: Ted Kisner <[email protected]>
hpc4cmb · Jan 26, 2024 · f62013b · f62013b
1 parent 59c9bba
commit f62013b
Show file tree

Hide file tree

Showing 9 changed files with 515 additions and 318 deletions.
diff --git a/src/libtoast/include/toast/fod_psd.hpp b/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
diff --git a/src/libtoast/src/toast_fod_psd.cpp b/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;

diff --git a/src/toast/_libtoast/fod_psd.cpp b/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.

diff --git a/src/toast/observation.py b/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):

diff --git a/src/toast/observation_dist.py b/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,
+    )