Implement basic histogram plotting

nnfwtbn/plot.py

 
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.ticker import AutoMinorLocator
+
+from nnfwtbn.process import Process
+from nnfwtbn.cut import Cut
+from nnfwtbn.variable import Variable
+
+ATLAS = "Work in Progress"
+INFO = "$\sqrt{s} = 13\,\mathrm{TeV}$, $36.1\,\mathrm{fb}^{-1}$\nSome selection"
+
 class HistogramFactory:
     """
     Short-cut to create multiple histogram with the same set of processes or
@@ -23,8 +34,21 @@ class HistogramFactory:
         """
 
 
+def type_to_histtype(type):
+    """
+    Returns the matplotlib histogram type for a given process plotting type.
+
+    >>> type_to_histtype("fill")
+    'stepfilled'
+    >>> type_to_histtype("line")
+    'step'
+    """
+    type_map = {"fill": "stepfilled", "line": "step"} 
+    return type_map[type]
+
 def hist(dataframe, variable, bins, *stacks, data=None, selection=None,
-         range=None, color=None, blind=None, axes=None, figure=None):
+         range=None, color=None, blind=None, axes=None, figure=None,
+         weight=None):
     """
     Creates a histogram of stacked processes. The first argument is the
     dataframe to operate on. The 'variable' argument defines the x-axis. The
@@ -51,9 +75,102 @@ def hist(dataframe, variable, bins, *stacks, data=None, selection=None,
     argument can be a list of processes to blind. By default blinding is
     applied to data. Use an empty list to disable blinding.
 
-    If the axes argument and/or figure arguments are omitted, this method
-    creates a new axes/figure.
+    If the figure argument is omitted, this method creates a new axes and
+    figure. If axes only is omitted, the method creates a new axes from the
+    figure.
 
     The method returns (figure, axes) which were used during plotting. This
     might be identical to the figure and axes arguments.
+
+    The weight is used to weight the entries. Entries have unit
+    weight if omitted. The argument can be a string name of a column or a
+    variable object.
     """
+    # Wrap column string by variable
+    if isinstance(variable, str):
+        variable = Variable(variable, variable)
+
+    if weight is None:
+        weight = Variable("unity", lambda d: variable(d) * 0 + 1)
+    elif isinstance(weight, str):
+        weight = Variable(weight, weight)
+
+    # Handle axes, figure
+    if figure is None:
+        figure, axes = plt.subplots()
+    elif axes is None:
+        axes = figure.subplots()
+
+    # Handle selection
+    if selection is None:
+        selection = Cut(lambda d: variable(d) * 0 == 0)
+    elif not isinstance(selection, Cut):
+        selection = Cut(selection)
+
+    # Handle range/bins
+    equidistant_bins = False
+    if range is not None:
+        # Build bins
+        if not isinstance(bins, int):
+            raise err.InvalidBins("When range is given, bins must be int.")
+        if not isinstance(range, tuple)  or len(range) != 2:
+            raise err.InvalidProcessSelection("Range argument must be a "
+                                              "tuple of two numbers.")
+        bins = np.linspace(range[0], range[1], bins + 1)
+        equidistant_bins = True
+
+    # Handle stack
+    for stack in stacks:
+        if isinstance(stack, Process):
+            # Wrap single process
+            stack = [stack]
+
+        bottom = np.zeros(len(bins) - 1)
+        for process in stack:
+            sel = selection & process.selection
+
+            n, _, _ = axes.hist(variable(dataframe[sel(dataframe)]),
+                                bins=bins, range=range,
+                                bottom=bottom,
+                                label=process.label,
+                                histtype=type_to_histtype(process.type),
+                                weights=weight(dataframe[sel(dataframe)]))
+            bottom += n 
+
+    axes.set_xlim((bins.min(), bins.max()))
+    axes.set_ylim((0, axes.get_ylim()[1] * 1.4))
+    axes.legend(frameon=False)
+
+    if variable.unit is not None:
+        axes.set_xlabel("%s in %s" % (variable.name, variable.unit))
+    else:
+        axes.set_xlabel(variable.name)
+
+    if equidistant_bins:
+        axes.set_ylabel("Events / %g %s" % (bins[1] - bins[0], variable.unit))
+    else:
+        axes.set_ylabel("Events / bin")
+    
+
+
+    axes.tick_params("both", which="both", direction="in")
+    axes.tick_params("both", which="major", length=6)
+    axes.tick_params("both", which="minor", length=3)
+    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 ATLAS is not None:
+        axes.text(0.04, 0.89, "ATLAS", transform=axes.transAxes,
+                  fontdict={"size": 18, "style": "italic", "weight": "bold"})
+        if isinstance(ATLAS, str):
+            axes.text(0.25, 0.89, ATLAS, transform=axes.transAxes,
+                      fontdict={"size": 12, })
+        if isinstance(INFO, str):
+            axes.text(0.04, 0.86, INFO, transform=axes.transAxes,
+                      fontdict={"size": 12, }, verticalalignment='top')
+
+
+
+    return (figure, axes)

nnfwtbn/process.py

@@ -60,7 +60,8 @@ class Process:
 
             self.range = range
             self.range_var = range_var
-            self.selection = None
+            self.selection = Cut(lambda d: (range[0] <= d[range_var])
+                                           & (d[range_var] <= range[1]))
 
         elif (selection is not None) and (range is None):
             # Selection is the only given method