From 2e7004eecd3068c2fbf48f103fcf529192d1c699 Mon Sep 17 00:00:00 2001
From: Daniel Wysocki <daniel.wysocki@ligo.org>
Date: Fri, 4 Sep 2020 20:18:03 -0500
Subject: [PATCH] Updating auto-correlation plots

---
 src/pop_models/sample_extraction/cli.py       | 18 ++++++------
 .../sample_extraction/integrated_acorr.py     | 28 ++++++++++++++++---
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/src/pop_models/sample_extraction/cli.py b/src/pop_models/sample_extraction/cli.py
index dee5570..c1a73e1 100644
--- a/src/pop_models/sample_extraction/cli.py
+++ b/src/pop_models/sample_extraction/cli.py
@@ -124,14 +124,14 @@ def main(raw_args=None):
             )
         )
 #        delta_n_acorr = numpy.tile(n_acorr, (f.n_walkers, f.n_dim))
-        print(
-            "Auto-correlation time is {dn} for the worst parameter."
-            .format(dn=n_acorr)
-        )
-        print(
-            "Auto-correlation time for all parameters are",
-            ", ".join([str(x) for x in acorr_times[n_acorr]]),
-        )
+        # print(
+        #     "Auto-correlation time is {dn} for the worst parameter."
+        #     .format(dn=n_acorr)
+        # )
+        # print(
+        #     "Auto-correlation time for all parameters are",
+        #     ", ".join([str(x) for x in acorr_times[n_acorr]]),
+        # )
         # Take the most auto-correlated walker and parameter throughout.
         n_thinning = numpy.max(n_acorr)
         print("Thinning is", n_thinning)
@@ -140,6 +140,8 @@ def main(raw_args=None):
         integrated_acorr.plot_integrated_auto_correlation(
             cli_args.acorr_plot,
             acorr_times, None,
+            param_names=f.variable_names,
+            safety_factor=cli_args.acorr_safety_factor,
         )
     elif cli_args.fixed_thinning is not None:
         n_thinning = cli_args.fixed_thinning
diff --git a/src/pop_models/sample_extraction/integrated_acorr.py b/src/pop_models/sample_extraction/integrated_acorr.py
index a18d0ce..05b2d2f 100644
--- a/src/pop_models/sample_extraction/integrated_acorr.py
+++ b/src/pop_models/sample_extraction/integrated_acorr.py
@@ -222,7 +222,11 @@ def plot_auto_correlation(
     fig.savefig(filename)
 
 
-def plot_integrated_auto_correlation(filename, acorr_times, acorr_time_errs):
+def plot_integrated_auto_correlation(
+        filename,
+        acorr_times, acorr_time_errs,
+        param_names=None, safety_factor=None,
+    ):
     import itertools
 
     import numpy
@@ -233,14 +237,31 @@ def plot_integrated_auto_correlation(filename, acorr_times, acorr_time_errs):
     K_final = K_final_plus_one - 1
     Ks = numpy.arange(K_final_plus_one)
 
-    fig, ax = plt.subplots(figsize=(8,4))
+    fig, ax = plt.subplots(figsize=(8,4), constrained_layout=True)
 
+
+    # Plot K = m*tau_K line
+    if safety_factor is not None:
+        ax.plot(
+            Ks, Ks/safety_factor,
+            color="black", linestyle="dashed",
+            label="$K = m \\, \\tau_K$",
+        )
+
+    # Plot autocorrelation curves.
     for d, color in zip(range(n_dim), itertools.cycle(color_cycle)):
+        # Determine label.
+        if param_names is None:
+            label = "Param #{i}".format(i=d+1)
+        else:
+            label = param_names[d]
+
+        # Plot autocorrelation curve.
         ac = acorr_times[...,d]
         ax.plot(
             Ks, ac,
             color=color,
-            label="Param #{i}".format(i=d+1),
+            label=label,
         )
 
         if acorr_time_errs is not None:
@@ -265,5 +286,4 @@ def plot_integrated_auto_correlation(filename, acorr_times, acorr_time_errs):
 
     ax.legend(loc="best")
 
-    fig.tight_layout()
     fig.savefig(filename)
-- 
GitLab