From 7933e2606add5a69de1b307040b28701dd455503 Mon Sep 17 00:00:00 2001
From: David Dotson <dotsdl@gmail.com>
Date: Tue, 18 May 2021 19:41:37 -0700
Subject: [PATCH] Hotfixes to plotting, etc.

---
 fah_xchem/analysis/plots.py      | 47 ++++++++++++++++----------------
 fah_xchem/analysis/structures.py |  2 +-
 2 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/fah_xchem/analysis/plots.py b/fah_xchem/analysis/plots.py
index 24bebb14..1001b8e4 100644
--- a/fah_xchem/analysis/plots.py
+++ b/fah_xchem/analysis/plots.py
@@ -106,7 +106,7 @@ def _filter_inclusive(
 
 def plot_relative_distribution(
     relative_delta_fs: List[float], min_delta_f: float = -30, max_delta_f: float = 30
-) -> None:
+) -> plt.Figure:
     """
     Plot the distribution of relative free energies
 
@@ -124,18 +124,20 @@ def plot_relative_distribution(
     )
     valid_relative_delta_fs_kcal = valid_relative_delta_fs * KT_KCALMOL
 
-    sns.displot(
-        valid_relative_delta_fs_kcal,
-        kind="kde",
-        rug=True,
-        color="hotpink",
-        fill=True,
-        rug_kws=dict(alpha=0.5),
-        label=f"$N={len(relative_delta_fs)}$",
-    )
+    fgrid = sns.displot(
+            valid_relative_delta_fs_kcal,
+            kind="kde",
+            rug=True,
+            color="hotpink",
+            fill=True,
+            rug_kws=dict(alpha=0.5),
+            label=f"$N={len(relative_delta_fs)}$",
+        )
     plt.xlabel(r"Relative free energy to reference fragment / kcal mol$^{-1}$")
     plt.legend()
 
+    return fgrid.fig
+
 
 def plot_convergence(
     complex_gens: List[int],
@@ -322,7 +324,7 @@ def plot_cumulative_distribution(
     cmap: str = "PiYG",
     n_bins: int = 100,
     markers_kcal: List[float] = [-6, -5, -4, -3, -2, -1, 0, 1, 2],
-) -> None:
+) -> plt.Figure:
     """
     Plot cumulative distribution of ligand affinities
 
@@ -357,7 +359,8 @@ def plot_cumulative_distribution(
     x_span = X.max() - X.min()
     C = [cm(((X.max() - x) / x_span)) for x in X]
 
-    plt.bar(X[:-1], Y, color=C, width=X[1] - X[0], edgecolor="k")
+    fig, ax  = plt.subplots()
+    ax.bar(X[:-1], Y, color=C, width=X[1] - X[0], edgecolor="k")
 
     for marker_kcal in markers_kcal:
         n_below = (relative_delta_fs_kcal < marker_kcal).astype(int).sum()
@@ -373,6 +376,8 @@ def plot_cumulative_distribution(
     plt.xlabel(r"Relative free energy to reference fragment / kcal mol$^{-1}$")
     plt.ylabel("Cumulative $N$ ligands")
 
+    return fig
+
 
 def _bootstrap(
     gens: List[GenAnalysis],
@@ -730,17 +735,13 @@ def generate_plots(
     # Summary plots
 
     # we always regenerate these, since they concern all data
-    with save_summary_plot(
-        name="relative_fe_dist",
-    ):
-        plot_relative_distribution(binding_delta_fs)
-        plt.title("Relative free energy")
-
-    with save_summary_plot(
-        name="cumulative_fe_dist",
-    ):
-        plot_cumulative_distribution(binding_delta_fs)
-        plt.title("Cumulative distribution")
+    fig = plot_relative_distribution(binding_delta_fs)
+    plt.title("Relative free energy")
+    save_summary_plot(name="relative_fe_dist", fig=fig)
+
+    fig = plot_cumulative_distribution(binding_delta_fs)
+    plt.title("Cumulative distribution")
+    save_summary_plot(name="cumulative_fe_dist", fig=fig)
         
     with _save_table_pdf(path=output_dir, name="poor_complex_convergence_fe_table"):
         plot_poor_convergence_fe_table(series.transformations)
diff --git a/fah_xchem/analysis/structures.py b/fah_xchem/analysis/structures.py
index 18ac8b26..232ac2b4 100644
--- a/fah_xchem/analysis/structures.py
+++ b/fah_xchem/analysis/structures.py
@@ -350,8 +350,8 @@ def generate_representative_snapshot(
     None
     """
     # create output directory if not present
-    os.makedirs(os.path.join(output_dir, f"RUN{run_id}"), exist_ok=True)
     run_id = transformation.transformation.run_id
+    os.makedirs(os.path.join(output_dir, f"RUN{run_id}"), exist_ok=True)
 
     if (
         max_binding_free_energy is not None