diff --git a/scripts/gcorr/README b/scripts/gcorr/README index 5de8c60e..ee465ac4 100644 --- a/scripts/gcorr/README +++ b/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: diff --git a/scripts/gcorr/iterate_gcorr.py b/scripts/gcorr/iterate_gcorr.py index 903dc03d..b64aad54 100644 --- a/scripts/gcorr/iterate_gcorr.py +++ b/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) diff --git a/xfaster/gcorr_tools.py b/xfaster/gcorr_tools.py index 22b5d4a8..6ee809da 100644 --- a/xfaster/gcorr_tools.py +++ b/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