TOAST mlmapmaker: Add option to write components of div

Add option to write individual components of `div`.
Reduce peak memory by writing the per-pixel covariance matrix div one
component at a time. So instead of MPI send/receive 3 x 3 x npix matrix,
only send/receive npix matrix.

Author: Bai-Chiang

sotodlib/toast/ops/mlmapmaker.py

@@ -217,7 +217,12 @@ class MLMapmaker(Operator):
         help="Truncate TOD to an easily factorizable length to ensure efficient FFT.",
     )
 
-    write_div = Bool(True, help="Write out the noise weight map")
+    write_div = Unicode(
+        "all",
+        allow_none=True,
+        help="Components (must be 'T', 'QU', 'TQU', 'all', None, or '')"
+    )
+
     write_hits= Bool(True, help="Write out the hitcount map")
 
     write_rhs = Bool(
@@ -231,7 +236,7 @@ class MLMapmaker(Operator):
     )
 
     @traitlets.validate("comps")
-    def _check_mode(self, proposal):
+    def _check_comps(self, proposal):
         check = proposal["value"]
         if check not in ["T", "QU", "TQU"]:
             raise traitlets.TraitError("Invalid comps (must be 'T', 'QU' or 'TQU')")
@@ -259,7 +264,7 @@ def _check_det_flag_mask(self, proposal):
         return check
 
     @traitlets.validate("dtype_map")
-    def _check_det_flag_mask(self, proposal):
+    def _check_dtype_map(self, proposal):
         check = proposal["value"]
         if check not in ["float", "float64"]:
             raise traitlets.TraitError(
@@ -306,6 +311,13 @@ def _check_nmat_type(self, proposal):
             raise traitlets.TraitError(msg)
         return check
 
+    @traitlets.validate("write_div")
+    def _check_write_div(self, proposal):
+        check = proposal["value"]
+        if check not in ["T", "QU", "TQU", None, 'all', '']:
+            raise traitlets.TraitError("Invalid write_div (must be 'T', 'QU', 'TQU', 'all', None or '')")
+        return check
+
     def __init__(self, **kwargs):
         self.shape = None
         self.wcs = None
@@ -513,20 +525,22 @@ def _init_mapmaker(
                 fname = signal_map.write(prefix, "rhs", signal_map.rhs)
                 log.info_rank(f"Wrote rhs to {fname}", comm=comm)
 
-        if self.write_div:
-            fname = f"{prefix}sky_div.fits"
-            if self.skip_existing and os.path.isfile(fname):
-                log.info_rank(f"Skipping existing div in {fname}", comm=comm)
-            else:
-                # FIXME : only writing the TT variance to avoid integer overflow in communication
-                fname = signal_map.write(prefix, "div", signal_map.div)
-                # fname = signal_map.write(prefix, "div", signal_map.div[0, 0])
-                log.info_rank(f"Wrote div to {fname}", comm=comm)
+        if self.write_div is not None:
+            # Write each covariance element seperately, to reduce peak memory.
+            for i,ci in enumerate(self.comps):
+                for j,cj in enumerate(self.comps):
+                    if ci in self.write_div and cj in self.write_div:
+                        fname = f"{prefix}sky_div{ci}{cj}.fits"
+                        if self.skip_existing and os.path.isfile(fname):
+                            log.info_rank(f"Skipping existing div{ci}{cj} in {fname}", comm=comm)
+                        else:
+                            fname = signal_map.write(prefix, f"div{ci}{cj}", signal_map.div[i, j])
+                            log.info_rank(f"Wrote div{ci}{cj} to {fname}", comm=comm)
 
         if self.write_hits:
             fname = f"{prefix}sky_hits.fits"
             if self.skip_existing and os.path.isfile(fname):
-                log.info_rank(f"Skipping existing div in {fname}", comm=comm)
+                log.info_rank(f"Skipping existing hits in {fname}", comm=comm)
             else:
                 fname = signal_map.write(prefix, "hits", signal_map.hits)
                 log.info_rank(f"Wrote hits to {fname}", comm=comm)
@@ -610,6 +624,17 @@ def _exec(self, data, detectors=None, **kwargs):
         gcomm = data.comm.comm_group
         timer.start()
 
+        if self.write_div == 'all':
+            self.write_div = self.comps
+        elif self.write_div == '':
+            self.write_div = None
+        elif self.write_div is not None:
+            # Make sure all components in self.write_div is in self.comps
+            for i in self.write_div:
+                if i not in self.comps:
+                    msg = f"Component '{i}' in write_div={self.write_div} not present in comps={self.comps}"
+                    raise RuntimeError(msg)
+
         if data.comm.group_size != 1:
             raise RuntimeError(
                 "The ML mapmaker requires the TOAST process group size to be exactly one."

tests/test_mapmaker_pointing.py

@@ -139,7 +139,7 @@ def test_mapmaker_pointing(self):
             truncate_tod=False,
             write_hits=True,
             write_rhs=False,
-            write_div=False,
+            write_div="all",
             write_bin=True,
             deslope=False,
         )