from nnfwtbn.cut import Cut
import nnfwtbn.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)