From 61e69f40c28f8f7faffa875e495b28cd31aee606 Mon Sep 17 00:00:00 2001
From: Sasha Rahlin <arahlin@fnal.gov>
Date: Tue, 5 Sep 2023 18:03:13 -0500
Subject: [PATCH] option to collate maps into separate coadds, refactor for
 cleanup

---
 maps/bin/spt3g-coadd-maps | 142 ++++++++++++++++++++++++--------------
 1 file changed, 92 insertions(+), 50 deletions(-)

diff --git a/maps/bin/spt3g-coadd-maps b/maps/bin/spt3g-coadd-maps
index b8f05e71..d38ad1eb 100644
--- a/maps/bin/spt3g-coadd-maps
+++ b/maps/bin/spt3g-coadd-maps
@@ -1,62 +1,104 @@
 #!/usr/bin/env python
 
-import argparse as ap
-import os, glob
 from spt3g import core, maps
 
-P = ap.ArgumentParser(
-    description="Coadd map frames from input files into a single output map frame"
-)
-P.add_argument("inputs", nargs="+", help="Input g3 files.  May be glob-able strings")
-P.add_argument("-o", "--output", default="map_coadd.g3", help="Output g3 file")
-P.add_argument(
-    "-m",
-    "--map-ids",
-    nargs="+",
-    help="Id's of map frames to include in the coadd.  If not set, all map frames are included.",
-)
-P.add_argument("-i", "--output-map-id", help="Id for output coadd map frame")
-P.add_argument(
-    "-w",
-    "--weighted",
-    action="store_true",
-    help="Ensure that weights are applied before coadding.",
-)
-
-args = P.parse_args()
-
-args.inputs = sum([glob.glob(x) for x in args.inputs], [])
-args.inputs = [x for x in args.inputs if os.path.splitext(x)[-1].startswith(".g3")]
-
-pipe = core.G3Pipeline()
-
-pipe.Add(core.G3Reader, filename=args.inputs)
-
-# drop metadata frames
-pipe.Add(lambda fr: fr.type == core.G3FrameType.Map)
-
-# build coadded map with consistently handled weights
-if args.weighted:
-    pipe.Add(maps.ApplyWeights)
-pipe.Add(
-    maps.CoaddMaps,
-    map_ids=args.map_ids,
-    output_map_id=args.output_map_id,
-    ensure_weighted_maps=False,
-)
-pipe.Add(lambda fr: "Id" not in fr or fr["Id"] == args.output_map_id)
-
 
 # record input parameters
-def RecordInputs(fr):
+def RecordInputs(fr, coadd_id=None, map_ids=None, inputs=None):
     if fr.type != core.G3FrameType.Map:
         return
-    fr["InputFiles"] = core.G3VectorString(args.inputs)
-    fr["InputMapIds"] = core.G3VectorString(args.map_ids)
+    if coadd_id and "Id" in fr and fr["Id"] != coadd_id:
+        return
+    fr["InputFiles"] = core.G3VectorString(inputs)
+    fr["InputMapIds"] = core.G3VectorString(map_ids)
+
+
+def coadd_maps(
+    inputs,
+    output="map_coadd.g3",
+    map_ids=None,
+    collate=False,
+    output_map_id=None,
+    weighted=False,
+):
+    pipe = core.G3Pipeline()
+    pipe.Add(core.G3Reader, filename=inputs)
+
+    # drop metadata frames
+    pipe.Add(lambda fr: fr.type == core.G3FrameType.Map)
+
+    # build coadded map with consistently handled weights
+    if weighted:
+        pipe.Add(maps.ApplyWeights)
+
+    if collate:
+        # coadd each map_id into a separate output frame
+        for mid in map_ids:
+            cid = output_map_id + mid if output_map_id else None
+            pipe.Add(
+                maps.CoaddMaps,
+                map_ids=[mid],
+                output_map_id=cid,
+                ensure_weighted_maps=False,
+            )
+            pipe.Add(RecordInputs, coadd_id=cid, map_ids=[mid], inputs=inputs)
+        pipe.Add(lambda fr: "Id" not in fr or fr["Id"].startswith(output_map_id))
+
+    else:
+        # coadd all map_id's into one output frame
+        pipe.Add(
+            maps.CoaddMaps,
+            map_ids=map_ids,
+            output_map_id=output_map_id,
+            ensure_weighted_maps=False,
+        )
+        pipe.Add(lambda fr: "Id" not in fr or fr["Id"] == output_map_id)
+        pipe.Add(
+            RecordInputs,
+            coadd_id=output_map_id,
+            map_ids=map_ids,
+            inputs=inputs,
+        )
+
+    pipe.Add(core.G3Writer, filename=output)
+    pipe.Run()
+
+
+if __name__ == "__main__":
+    import argparse as ap
+    import os, glob
 
+    P = ap.ArgumentParser(
+        description="Coadd map frames from input files into a single output map frame"
+    )
+    P.add_argument(
+        "inputs", nargs="+", help="Input g3 files.  May be glob-able strings"
+    )
+    P.add_argument("-o", "--output", default="map_coadd.g3", help="Output g3 file")
+    P.add_argument(
+        "-m",
+        "--map-ids",
+        nargs="+",
+        help="Id's of map frames to include in the coadd.  If not set, all map frames are included.",
+    )
+    P.add_argument(
+        "-c",
+        "--collate",
+        action="store_true",
+        help="Coadd each input map Id into a separate output map frame.  In this case, the "
+        "output map Id is treated as a prefix, with the input map Id appended to the string.",
+    )
+    P.add_argument("-i", "--output-map-id", help="Id for output coadd map frame")
+    P.add_argument(
+        "-w",
+        "--weighted",
+        action="store_true",
+        help="Ensure that weights are applied before coadding.",
+    )
 
-pipe.Add(RecordInputs)
+    args = P.parse_args()
 
-pipe.Add(core.G3Writer, filename=args.output)
+    args.inputs = sum([glob.glob(x) for x in args.inputs], [])
+    args.inputs = [x for x in args.inputs if os.path.splitext(x)[-1].startswith(".g3")]
 
-pipe.Run()
+    coadd_maps(**vars(args))