Fix bug in writing out FITS maps without MPI

hpc4cmb · Oct 14, 2019 · 66aef04 · 66aef04
1 parent 58ec4e8
commit 66aef04
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/src/toast/map/pixels.py b/src/toast/map/pixels.py
@@ -641,12 +641,16 @@ def write_healpix_fits(self, path, comm_bytes=None):
 
         # Determine which processes "own" each submap.
 
-        owners = np.zeros(nsubmap, dtype=np.int32)
-        allowners = None
         if self._comm is None:
-            allowners = owners
+            NO_OWNER = 1
+            allowners = np.zeros(nsubmap, dtype=np.int32)
+            allowners.fill(NO_OWNER)
+            for m in self._local:
+                allowners[m] = rank
         else:
-            owners.fill(self._comm.size)
+            NO_OWNER = self._comm.size
+            owners = np.zeros(nsubmap, dtype=np.int32)
+            owners.fill(NO_OWNER)
             for m in self._local:
                 owners[m] = self._comm.rank
             allowners = np.zeros_like(owners)
@@ -677,9 +681,11 @@ def write_healpix_fits(self, path, comm_bytes=None):
             while submap_off < nsubmap:
                 if submap_off + ncomm > nsubmap:
                     ncomm = nsubmap - submap_off
-                if np.sum(allowners[submap_off : submap_off + ncomm]) != ncomm:
+                if np.any(allowners[submap_off : submap_off + ncomm] != NO_OWNER):
                     # at least one submap has some hits
                     for c in range(ncomm):
+                        if allowners[submap_off + c] == NO_OWNER:
+                            continue
                         dview[c, :, :] = self.data[self._glob2loc[submap_off + c], :, :]
                     # copy into FITS buffers
                     for c in range(ncomm):
@@ -709,8 +715,7 @@ def write_healpix_fits(self, path, comm_bytes=None):
                 if submap_off + ncomm > nsubmap:
                     ncomm = nsubmap - submap_off
                 if (
-                    np.sum(allowners[submap_off : submap_off + ncomm])
-                    != ncomm * self._comm.size
+                    np.any(allowners[submap_off : submap_off + ncomm] != NO_OWNER)
                 ):
                     # at least one submap has some hits.  reduce.
                     for c in range(ncomm):