From ecb57fa46d9e01cd81855f2f84682b036d95d1b5 Mon Sep 17 00:00:00 2001
From: Jesse Sandberg <js2746@belenus.ccib.rutgers.edu>
Date: Thu, 14 Nov 2024 15:25:27 -0500
Subject: [PATCH] colorbar now works

---
 python/plotting.py | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/python/plotting.py b/python/plotting.py
index e1aab4f..a7454a9 100644
--- a/python/plotting.py
+++ b/python/plotting.py
@@ -149,13 +149,35 @@ def create_heatmap_figure_and_axes(lipids, cmap, v_vals, figwidth, figheight, he
             ax = plot_helices(helices[0], False, ax, 50)
         else:
             ax = plot_helices(helices[1], False, ax, 50)
+    return fig, fig.axes
+
+
+def make_colorbar(fig, v_vals, cmap):
+    """
+    Generate the colorbar. Must be done after plot_heatmap.
+
+    Parameters
+    ----------
+    fig : fig
+        The figure object from matplotlib.
+    v_vals : list or tuple
+        min, mid, and max values for the colorbar.
+    cmap : colormap
+        matplotlib colormap object.
+
+    Returns
+    -------
+    fig
+
+    """
+    vmin, vmid, vmax = v_vals
     fig.subplots_adjust(right=0.8)
     cbar_ax = fig.add_axes([0.21, 1, 0.5, 0.02])
     sm = mpl.cm.ScalarMappable(cmap=cmap)
     cbar = fig.colorbar(sm, cax=cbar_ax, orientation="horizontal")
     cbar.set_ticks(np.linspace(0, 1, 5))
     cbar.ax.set_xticklabels([vmin, (vmin + vmid) / 2, vmid, (vmid + vmax) / 2, vmax])
-    return fig, fig.axes
+    return fig
 
 
 def bin_prep(bin_info):