Number of contour levels and log/lin scale configurable for plot_NN_vs_var_2D

plotting.py

@@ -5,6 +5,7 @@ import math
 
 import matplotlib.pyplot as plt
 import matplotlib.colors
+from matplotlib.ticker import LogFormatter
 import numpy as np
 
 import meme
@@ -43,7 +44,9 @@ def plot_NN_vs_var_2D(plotname, means,
                       var2_index, var2_range,
                       var1_label=None,
                       var2_label=None,
-                      contourdistance=0.1):
+                      logscale=False,
+                      ncontours=20,
+                      black_contourlines=False):
 
     print("Creating varied events (2d)")
     # example: vary var1 vs var2
@@ -73,9 +76,19 @@ def plot_NN_vs_var_2D(plotname, means,
     zmin = np.min(scores)
     zmax = np.max(scores)
 
-    # TODO: find out on how to set (in a reasonable way) the contour levels and z-axis ticks
-    pcm = ax.contourf(sequence1, sequence2, scores, norm=matplotlib.colors.LogNorm(vmin=zmin, vmax=zmax))
-    cbar = fig.colorbar(pcm, ax=ax, extend='max')
+    if logscale:
+        lvls = np.logspace(math.log10(zmin), math.log10(zmax), ncontours)
+        pcm = ax.contourf(sequence1, sequence2, scores, levels=lvls, norm=matplotlib.colors.LogNorm(vmin=zmin, vmax=zmax))
+        if black_contourlines:
+            ax.contour(sequence1, sequence2, scores, levels=lvls, colors="k", linewidths=1)
+        l_f = LogFormatter(10, labelOnlyBase=False, minor_thresholds=(np.inf, np.inf))
+        cbar = fig.colorbar(pcm, ax=ax, extend='max', ticks=lvls, format=l_f)
+    else:
+        pcm = ax.contourf(sequence1, sequence2, scores, ncontours, norm=matplotlib.colors.Normalize(vmin=0, vmax=1))
+        if black_contourlines:
+            ax.contour(sequence1, sequence2, scores, ncontours, colors="k", linewidths=1)
+        cbar = fig.colorbar(pcm, ax=ax, extend='max')
+
     cbar.set_label("NN output")
     if var1_label is not None:
         ax.set_xlabel(var1_label)
@@ -114,3 +127,4 @@ if __name__ == "__main__":
                       var1_index=c.branches.index("met"), var1_range=(0, 1000, 10),
                       var2_index=c.branches.index("mt"), var2_range=(0, 500, 10),
                       var1_label="met [GeV]", var2_label="mt [GeV]")
+