split plot_profile_2D into 2 functions

plotting.py

@@ -303,15 +303,10 @@ def plot_cond_avg_actmax_2D(plotname, model, layer, neuron, ranges,
     plot_hist_2D(plotname, xedges, yedges, hist, zlabel="Neuron output", **kwargs)
 
 
-def plot_profile_2D(plotname, valsx, valsy, scores,
-                    nbinsx, xmin, xmax,
-                    nbinsy, ymin, ymax,
-                    metric=np.mean,
-                    weights=None,
-                    **kwargs):
-
-    kwargs["zlabel"] = kwargs.get("zlabel", "Profile")
-
+def get_profile_2D(valsx, valsy, scores,
+                   nbinsx, xmin, xmax,
+                   nbinsy, ymin, ymax,
+                   metric=np.mean, weights=None):
     xedges = np.linspace(xmin, xmax, nbinsx)
     yedges = np.linspace(ymin, ymax, nbinsy)
 
@@ -337,6 +332,25 @@ def plot_profile_2D(plotname, valsx, valsy, scores,
     hist = np.array(hist)
     hist = hist.T # had a list of columns - needs to be list of rows
 
+    return hist, xedges, yedges
+
+
+def plot_profile_2D(plotname, valsx, valsy, scores,
+                    nbinsx, xmin, xmax,
+                    nbinsy, ymin, ymax,
+                    metric=np.mean,
+                    weights=None,
+                    **kwargs):
+
+    kwargs["zlabel"] = kwargs.get("zlabel", "Profile")
+
+    hist, xedges, yedges = get_profile_2D(
+        valsx, valsy, scores,
+        nbinsx, xmin, xmax,
+        nbinsy, ymin, ymax,
+        metric=metric, weights=weights
+    )
+
     plot_hist_2D(plotname, xedges, yedges, hist, **kwargs)