re-entrant options to rerun iterations until convergence

scripts/gcorr/README

@@ -42,8 +42,10 @@ Gcorr run procedure:
       python iterate_gcorr.py --gcorr-config path-to-my-gcorr-config.ini
 
 3. Run iterate_gcorr.py until convergence is reached. In practice, you will run
-   the command above and wait for it to finish. Then look at the
-   correction-to-the-correction that it both prints and plots (it should
+   the command above and wait for it to finish. If you include the `--max-iters`
+   option with a non-zero value, the code will try to determine whether
+   convergence has been reached and stop on its own.  Otherwise, you can look at
+   the correction-to-the-correction that it both prints and plots (it should
    converge to 1s for TT, EE, BB), and if it hasn't converged, up+enter (redo)
    the same command you just did.  In much more detail, here's what the code
    does:

scripts/gcorr/iterate_gcorr.py

@@ -4,6 +4,7 @@
 """
 import os
 import argparse as ap
+import subprocess as sp
 from xfaster import gcorr_tools as gt
 from xfaster.batch_tools import batch_sub
 from matplotlib import use
@@ -15,11 +16,10 @@
     "--gcorr-config", required=True, help="The config file for gcorr computation"
 )
 P.add_argument(
-    "-n",
-    "--no-submit",
-    dest="submit",
-    action="store_false",
-    help="Don't submit jobs; run serially on current session",
+    "-s",
+    "--submit",
+    action="store_true",
+    help="Submit jobs, instead of running serially on current session",
 )
 P.add_argument(
     "--force-restart",
@@ -55,9 +55,34 @@
     action="store_true",
     help="Compute and update gcorr files for the current iteration",
 )
+P.add_argument(
+    "-c",
+    "--converge-criteria",
+    type=float,
+    default=0.01,
+    help="Maximum fractional change in gcorr that indicates "
+    "convergence and stops iteration",
+)
+P.add_argument(
+    "--max-iters",
+    default=0,
+    type=int,
+    help="Maximum number of iterations to run.  If 0, run once and exit.",
+)
+P.add_argument(
+    "--submit-next",
+    action="store_true",
+    default=None,
+    help="Submit jobs for the next iteration (rather than running them locally). "
+    "This option is used internally by the script, do not set this option manually.",
+)
 
 args = P.parse_args()
 
+if args.submit_next is None:
+    args.submit_next = args.submit
+
+# load configuration
 g_cfg = gt.get_gcorr_config(args.gcorr_config)
 
 tags = g_cfg["gcorr_opts"]["map_tags"]
@@ -86,30 +111,43 @@
     gcorr_fit_hist=args.gcorr_fit_hist,
     allow_extreme=args.allow_extreme,
     keep_iters=args.keep_iters,
+    converge_criteria=args.converge_criteria,
+    max_iters=args.max_iters,
 )
 
 # build command for this script
 if args.submit:
     run_opts.update(submit=True, **g_cfg["submit_opts"])
 
-    cmd = [
-        "python",
-        os.path.abspath(__file__),
-        "--gcorr-config",
-        os.path.abspath(args.gcorr_config),
-    ]
-    for k in ["allow_extreme", "gcorr_fit_hist", "keep_iters", "force_restart"]:
-        if getattr(args, k):
-            cmd += ["--{}".format(k.replace("_", "-"))]
+cmd = [
+    "python",
+    os.path.abspath(__file__),
+    "--gcorr-config",
+    os.path.abspath(args.gcorr_config),
+]
+for k in [
+    "allow_extreme",
+    "gcorr_fit_hist",
+    "keep_iters",
+    "force_restart",
+    "submit_next",
+]:
+    if getattr(args, k):
+        cmd += ["--{}".format(k.replace("_", "-"))]
+for k in ["converge_criteria", "max_iters"]:
+    cmd += ["--{}".format(k.replace("_", "-")), str(getattr(args, k))]
 
 # run
 for tag in tags:
     if not args.analyze_only:
         jobs = gt.xfaster_gcorr(output_tag=tag, **run_opts)
 
+    tag_cmd = cmd + ["-t", tag]
+
     if args.submit:
-        batch_sub(
-            cmd + ["-n", "-a", "-t", tag], dep_afterok=jobs, **g_cfg["submit_opts"]
-        )
+        # submit analysis job
+        batch_sub(tag_cmd + ["-a"], dep_afterok=jobs, **g_cfg["submit_opts"])
     else:
-        gt.process_gcorr(output_tag=tag, **gcorr_opts)
+        if not gt.process_gcorr(output_tag=tag, **gcorr_opts):
+            # run again if not converged or reached max_iters
+            sp.check_call(tag_cmd + ["-s"] * args.submit_next)

xfaster/gcorr_tools.py

@@ -449,17 +449,23 @@ def plot_gcorr(output_root="xfaster_gcal", output_tag=None):
     import matplotlib.pyplot as plt
     from matplotlib import colormaps
 
-    fig, axs = plt.subplots(2, 3, sharex=True, sharey="row")
+    fig = axs = None
     colors = colormaps["viridis"].resampled(len(files)).colors
-    axs[1, 1].set_xlabel("Bin Number")
-    axs[0, 0].set_ylabel("{} G Total".format(output_tag))
-    axs[1, 0].set_ylabel("{} G Correction".format(output_tag))
 
     for f, c in zip(files, colors):
         gdata = pt.load_and_parse(f)["gcorr"]
         cdata = pt.load_and_parse(f.replace("total", "corr"))["gcorr"]
 
+        if fig is None:
+            pol = "ee" in gdata
+            fig, axs = plt.subplots(2, 3 if pol else 1, sharex=True, sharey="row")
+            axs[1, 1].set_xlabel("Bin Number")
+            axs[0, 0].set_ylabel("{} G Total".format(output_tag))
+            axs[1, 0].set_ylabel("{} G Correction".format(output_tag))
+
         for spec, ax1 in zip(["tt", "ee", "bb"], axs.T):
+            if spec not in gdata:
+                continue
             ax1[0].plot(gdata[spec], color=c)
             ax1[0].set_title(spec.upper())
             ax1[1].plot(cdata[spec], color=c)
@@ -479,6 +485,8 @@ def process_gcorr(
     gcorr_fit_hist=False,
     allow_extreme=False,
     keep_iters=False,
+    converge_criteria=0.01,
+    max_iters=0,
 ):
     """
     Run gcorr analysis on a complete ensemble of sim bandpowers. This function
@@ -512,6 +520,16 @@ def process_gcorr(
     keep_iters : bool
         If True, store the transfer function and bandpowers outputs from each
         iteration in a separate sub-directory, rather than deleting these files.
+    converge_criteria : float
+        Maximum fractional change in gcorr that indicates convergence and stops
+        iteration.
+    max_iters : int
+        Maximum number of iterations to run.  If 0, run once and return.
+
+    Returns
+    -------
+    retval : bool
+        True if convergence or iteration index is reached, False otherwise.
     """
     iternum = get_next_iter(output_root, output_tag)
 
@@ -568,3 +586,17 @@ def process_gcorr(
 
     print("{} iter {} gcorr correction: {}".format(output_tag, iternum, out["gcorr"]))
     print("{} iter {} total gcorr: {}".format(output_tag, iternum, gcorr["gcorr"]))
+
+    gcc = out["gcorr"]
+    v = np.array([gcc[spec][1:] for spec in ["tt", "ee", "bb"] if spec in gcc])
+    gmax = np.max(np.abs(v - 1))
+    print("Maximum fractional gcorr correction (among TT/EE/BB bins 1+):", gmax)
+    if gmax <= converge_criteria:
+        print("{} gcorr converged after {} iterations".format(output_tag, iternum))
+        return True
+
+    if iternum >= max_iters:
+        print("{} gcorr not converged after {} iterations".format(output_tag, iternum))
+        return True
+
+    return False