import numpy as np
# the class
[docs]
class LogicalCut:
"""
A class for the application of logical cuts to given data.
:param initial_condition: An initial condition can be applied, e.g. ph > 0.01, where ph is an 1D array. The number
of events is equal to the length of the array.
:type initial_condition: 1D bool array or None
"""
def __init__(self, initial_condition=None):
if initial_condition is not None:
if len(initial_condition.shape) != 1:
raise KeyError('initial_condition needs to be a 1-dimensional array!')
if initial_condition.dtype != 'bool':
raise KeyError('initial_condition needs to be numpy array with dtype bool!')
self.all_idx = np.arange(len(initial_condition))
else:
self.all_idx = None
self.cut_flag = initial_condition
def __len__(self):
return len(self.get_idx())
[docs]
def add_condition(self, condition):
"""
Add and condition that is linked with and AND statement to all other conditions.
:param condition: A condition to be applied to all events, e.g. ph > 0.01.
:type condition: 1D bool array
"""
if len(condition.shape) != 1:
raise KeyError('condition needs to be a 1-dimensional array!')
if condition.dtype != 'bool':
raise KeyError('condition needs to be numpy array with dtype bool!')
if self.cut_flag is not None:
self.cut_flag = np.logical_and(self.cut_flag, condition)
else:
self.cut_flag = condition
if self.all_idx is None:
self.all_idx = np.arange(len(condition))
[docs]
def force_true(self, idx):
"""
Force certain events to true.
:param idx: All indices that we want to set to true.
:type idx: 1D int array
"""
self.cut_flag = np.logical_or(self.cut_flag, np.in1d(self.all_idx, idx))
[docs]
def force_false(self, idx):
"""
Force certain events to false.
:param idx: All indices that we want to set to false.
:type idx: 1D int array
"""
self.cut_flag = np.logical_and(self.cut_flag, np.logical_not(np.in1d(self.all_idx, idx)))
[docs]
def get_flag(self):
"""
Return a bool array of all events, indicating which do survive all cuts.
:return: Indicates which events do survive all cuts.
:rtype: 1D int array
"""
return self.cut_flag
[docs]
def get_antiflag(self):
"""
Return a bool array of all events, indicating which do not survive all cuts.
:return: Indicates which events do not survive all cuts.
:rtype: 1D int array
"""
return np.logical_not(self.cut_flag)
[docs]
def get_idx(self):
"""
Return an int array of all event indices that survive all cuts.
:return: The indices that survive all cuts.
:rtype: 1D int array
"""
return self.all_idx[self.cut_flag]
[docs]
def get_antiidx(self):
"""
Return an int array of all event indices that do not survive all cuts.
:return: The indices that do not survive all cuts.
:rtype: 1D int array
"""
return self.all_idx[np.logical_not(self.cut_flag)]
[docs]
def total(self):
"""
Return the total number of events.
:return: The total number of events.
:rtype: int
"""
return self.get_flag().shape[0]
[docs]
def counts(self):
"""
Return the number of events that survive all cuts.
:return: The number of events that survive all cuts.
:rtype: int
"""
return self.get_idx().shape[0]