Source code for freeforestml.process


from freeforestml.cut import Cut
import freeforestml.error as err

[docs]class Process: """ This class represents a physics process to be selected during training and plotting. The class stores the cuts to select the process' events from a dataframe, its style and human-readable name for plotting. """ DEFAULT_RANGE_VAR = 'fpid'
[docs] def __init__(self, label, selection=None, range=None, range_var=None): r""" Returns a new process object. The process has a human-readable name (potentially using latex) and a selection cut. The selection argument can be a cut object or any callable. Stacking of processes is handled by the plotting method. >>> Process("Top", lambda d: d.is_top) <Process 'Top': (func)> >>> Process("VBF", lambda d: d.is_VBFH) <Process 'VBF': (func)> The optional argument range accepts a two-value tuple and is a shortcut to defined a selection cut accepting events whose 'range_var' is between (including boundaries) the given values. The range_var can be a string naming a column in the dataframe or a Variable object. >>> Process("Z\\rightarrow\\ell\\ell", range=(-599, -500)) <Process 'Z\\rightarrow\\ell\\ell': [-599, -500]> If the range_var argument is omitted, the value of Process.DEFAULT_RANGE_VAR is used, this defaults to 'fpid'. A process behaves like a cut in many ways. For example, the call() and idx_array methods are identical. """ ####################################################### # Selection if (selection is None) and (range is None): raise err.InvalidProcessSelection("Exactly one of selection or " "range can be used.") elif (selection is None) and (range is not None): # Range is the only given method if not isinstance(range, tuple) or len(range) != 2: raise err.InvalidProcessSelection("Range argument must be a " "tuple of two numbers.") if range_var is None: # Use default range_var range_var = Process.DEFAULT_RANGE_VAR self.range = range self.range_var = range_var 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 if not isinstance(selection, Cut): # Wrap callable in selection selection = Cut(selection) self.selection = selection self.range = None else: raise err.InvalidProcessSelection("Only one of selection or range" " can be used.") ####################################################### # Label self.label = label
[docs] def __call__(self, dataframe): """ Returns a dataframe containing only the events of this process. """ return self.selection(dataframe)
[docs] def idx_array(self, dataframe): """ Returns the index array of the given dataframe which selects all events of this process. """ return self.selection.idx_array(dataframe)
[docs] def __repr__(self): """ Returns a string representation of the process. """ if self.range is None: return "<Process %s: (func)>" % repr(self.label) else: return "<Process %s: [%g, %g]>" % (repr(self.label), *self.range)