Merge branch '23-implement-stack-class'

nnfwtbn/stack.py

+
+import numpy as np
+
+from nnfwtbn.variable import Variable
+from nnfwtbn.cut import Cut
+import nnfwtbn.error as err
+
+def not_none_or_default(value, default):
+    """
+    Returns the value if it is not None. Otherwise returns the default.
+    """
+    if value is not None:
+        return value
+    return default
+
+class Stack:
+    """
+    This class represents a collection of Prcesses drawn as a stack in
+    histograms created with hist(). The Stack class stores information about
+    the plotting style (e.g. markersize, linestyle), the histogram type
+    (step, stepfilled, points), the color wheel, and the method to compute
+    the total uncertainty of the stack.
+
+    A stack is not tied to a specific plot. It can be reused for plot with
+    different binning, different variables or different selections.
+    """
+
+    def __init__(self, *processes, histtype='stepfilled',
+                 data_uncertainty=False, **aux):
+        """
+        Creates a new stack and sets its default properties. If a process is
+        added (via add_process()) to the stack without specifying a custom style, the
+        defaults are used.
+
+        The object is initialized with the processes passed to the method.
+        """
+        self.processes = list(processes)
+        self.histtypes = [None for _ in processes]
+        self.data_uncertainties = [None for _ in processes]
+        self.aux = [{} for _ in processes]
+
+        self.default_aux = aux
+        self.default_histtype = histtype
+        self.default_data_uncertainty = data_uncertainty
+
+    def add_process(self, process, histtype=None, data_uncertainty=None,
+            **aux):
+        """
+        Adds a new process to the stack. Arguments passed to this method the
+        precedence over the default passed to the constructor.
+        
+        The process argument must be a Process object with information about
+        the selection and the label in the legend. The histtype argument can
+        take the value 'step', 'stepfilled', 'line', 'points'. This argument
+        controls the type of the histogram. If the data_uncertainty is set to
+        True, then the get_total_uncertainty() will return sqrt(get_total()).
+        This method is useful when plotting Asimov data. If the option is
+        False, the weights are used to compute the uncertainty.
+
+        Additional keyword arguments are stored internally for the plotting
+        method to be forwarded to matplotlib.
+        """
+        self.processes.append(process)
+        self.histtypes.append(histtype)
+        self.data_uncertainties.append(data_uncertainty)
+        self.aux.append(aux)
+
+    def get_hist(self, df, i, bins, variable, weight):
+        """
+        Returns the yields per bin for the i-th process in the stack. The bins
+        argument specifies the bin edges.
+        """
+        process = self.processes[i]
+        df = process(df)
+        total, _ = np.histogram(variable(df), bins=bins, weights=weight(df))
+        return total
+
+    def get_total(self, df, bins, variable, weight):
+        """
+        Returns the sum of yields per bin of all processes. The bins argument
+        specifies the bin edges.
+        """
+        total = 0
+        for i in range(len(self.processes)):
+            total += self.get_hist(df, i, bins, variable, weight)
+
+        return total
+
+    def get_total_uncertainty(self, df, bins, variable, weight):
+        """
+        Returns the uncertainty of the total yield per bin. The bins argument
+        specifies the bin edges.
+        """
+        weight_2 = Variable("$w^2$", lambda d: weight(d)**2)
+        uncertainty_2 = 0
+
+        for i in range(len(self.processes)):
+            if self.is_data_uncertainty(i):
+                uncertainty_2 += self.get_hist(df, i, bins, variable, weight)
+            else:
+                uncertainty_2 += self.get_hist(df, i, bins, variable,
+                                               weight=weight_2)
+
+        return np.sqrt(uncertainty_2)
+
+
+    def get_histtype(self, i):
+        """
+        Return the histtype of process i.
+        """
+        return not_none_or_default(self.histtypes[i], self.default_histtype)
+
+    def get_aux(self, i):
+        """
+        Returns the auxiliary keyword arguments. The returned dict is a mix of
+        the default keyword arguments updated by the ones used when adding a
+        process.
+        """
+        return dict(self.default_aux, **self.aux[i])
+
+    def is_data_uncertainty(self, i):
+        """
+        Returns True if process i uses data_uncertainties.
+        """
+        return not_none_or_default(self.data_uncertainties[i],
+                                   self.default_data_uncertainty)
+
+    def __len__(self):
+        """
+        Returns the number of processes in this stack.
+        """
+        return len(self.processes)
+
+class McStack(Stack):
+    """
+    Short-hand class for a Stack with only Monte-Carlo-like processes.
+    """
+
+    def __init__(self, *args, histtype='stepfilled', **kwds):
+        if 'data_uncertainty' in kwds:
+            raise TypeError("Cannot pass 'data_uncertainty' to McStack. "
+                            "Use generic Stack instead.")
+
+        kwds['data_uncertainty'] = False
+        super().__init__(*args, **kwds)
+
+    def add_process(self, *args, **kwds):
+        if 'data_uncertainty' in kwds:
+            raise TypeError("Cannot pass 'data_uncertainty' to McStack. "
+                            "Use generic Stack instead.")
+
+        super().add_process(*args, histtype=histtype, **kwds)
+
+class DataStack(Stack):
+    """
+    Short-hand class for a Stack with only data-like processes.
+    """
+
+    def __init__(self, *args, histtype='points', **kwds):
+        if 'data_uncertainty' in kwds:
+            raise TypeError("Cannot pass 'data_uncertainty' to DataStack. "
+                            "Use generic Stack instead.")
+
+        kwds['data_uncertainty'] = True
+        super().__init__(*args, **kwds)
+
+    def add_process(self, *args, **kwds):
+        if 'data_uncertainty' in kwds:
+            raise TypeError("Cannot pass 'data_uncertainty' to DataStack. "
+                            "Use generic Stack instead.")
+
+        super().add_process(*args, histtype=histtype, **kwds)