Add y_log and y_min arguments to hist()

nnfwtbn/plot.py

 
 import numpy as np
 import matplotlib.pyplot as plt
-from matplotlib.ticker import AutoMinorLocator
+from matplotlib.ticker import AutoMinorLocator, LogLocator
 import pandas as pd
 import seaborn as sns
 
@@ -41,7 +41,7 @@ class HistogramFactory:
 
 def hist(dataframe, variable, bins, stacks, selection=None,
          range=None, blind=None, axes=None, figure=None,
-         weight=None, **kwds):
+         weight=None, y_log=False, y_min=None, **kwds):
     """
     Creates a histogram of stacked processes. The first argument is the
     dataframe to operate on. The 'variable' argument defines the x-axis. The
@@ -73,6 +73,11 @@ def hist(dataframe, variable, bins, stacks, selection=None,
     weight if omitted. The argument can be a string name of a column or a
     variable object.
 
+    If the y_log argument is set to True, the y axis will be logarithmic. The
+    axis is enlarged on a logarithmic scale to make room for the ATLAS labels.
+    The optional y_min argument can be used to set the lower limit of the y
+    axis. The default is 0 for linear scale, and 1 for logarithmic scale.
+
     Any other keyword argument is used to define the style. If the keyword
     argument is a list, it's length must equal the number of stacks. If the
     list item is a list, it is cycled for every stack member. If the list
@@ -181,8 +186,22 @@ def hist(dataframe, variable, bins, stacks, selection=None,
                     **process_kwds)
             bottom += n 
 
+    # Configure x-axis
     axes.set_xlim((bins.min(), bins.max()))
-    axes.set_ylim((0, axes.get_ylim()[1] * 1.6))
+
+    # Configure y-axis
+    _, y_max = axes.get_ylim()
+
+    if y_log:
+        y_min = y_min if y_min is not None else 1
+        y_max *= 10**1.6
+        axes.set_yscale('log')
+    else:
+        y_min = y_min if y_min is not None else 0
+        y_max *= 1.6
+
+    axes.set_ylim((y_min, y_max))
+
     axes.legend(frameon=False, loc=1)
 
     if variable.unit is not None:
@@ -205,7 +224,8 @@ def hist(dataframe, variable, bins, stacks, selection=None,
     axes.tick_params("x", which="both", top=True)
     axes.tick_params("y", which="both", right=True)
     axes.xaxis.set_minor_locator(AutoMinorLocator())
-    axes.yaxis.set_minor_locator(AutoMinorLocator())
+    if not y_log:
+        axes.yaxis.set_minor_locator(AutoMinorLocator())
 
     if ATLAS is not None:
         axes.text(0.04, 0.89, "ATLAS", transform=axes.transAxes,