NNFilterModule Class Reference

Inheritance diagram for NNFilterModule:

Collaboration diagram for NNFilterModule:

Public Member Functions
def	__init__ (self, model_file=None, model_config=MODEL_CONFIG, preproc_config=PREPROC_CONFIG, threshold=None, extra_info_var="NN_prediction", global_tag="SmartBKG_GATGAP", payload="GATGAPgen.pth")
def	initialize (self)
def	event (self)

Public Attributes
	model_file
	Path to the saved model file.
	model_config
	Parameters for building the model.
	preproc_config
	Parameters for preprocessing.
	threshold
	Threshold for event selection using reweighting method, value None indicating sampling method.
	extra_info_var
	Name of eventExtraInfo to save model prediction to.
	payload
	Payload for the well-trained model in global tag.
	model
	model with trained parameters
	EventExtraInfo
	StoreArray to save weights to.
	EventInfo
	Initialise event metadata from data store.
	out_features
	node features

Detailed Description

Goals:
   1. Build a graph from an event composed of MCParticles
   2. Apply the well-trained model for reweighting or sampling method to get a score
   3. Execute reweighting or sampling process to get a weight

Arguments:
   model_file(str): Path to the saved model
   model_config(dict): Parameters to build the model
   preproc_config(dict): Parameters for preprocessing
   threshold(float): Threshold for event selection using reweighting method, value *None* indicating sampling mehtod
   extra_info_var(str): Name of eventExtraInfo to save model prediction to
   global_tag(str): Tag in ConditionDB where the well trained model was stored
   payload(str): Payload for the well trained model in global tag

Returns:
   Pass or rejected according to random sampling or selection with the given threshold

Note:
    Score after the NN filter indicating the probability of the event to pass is saved
    under ``EventExtraInfo.extra_info_var``.

    Use ``eventExtraInfo(extra_info_var)`` in ``modularAnalysis.variablesToNtuple`` or
    ``additionalBranches=["EventExtraInfo"]`` in ``mdst.add_mdst_output`` to have access to the scores.

Definition at line 23 of file NN_filter_module.py.

Constructor & Destructor Documentation

__init__()

def __init__	(		self,
			model_file = None,
			model_config = MODEL_CONFIG,
			preproc_config = PREPROC_CONFIG,
			threshold = None,
			extra_info_var = "NN_prediction",
			global_tag = "SmartBKG_GATGAP",
			payload = "GATGAPgen.pth"
	)

Initialise the class.
:param model_file: Path to the saved model file.
:param model_config: Parameters for building the model.
:param preproc_config: Parameters for preprocessing.
:param threshold: Threshold for event selection using reweighting method, value *None* indicating sampling mehtod.
:param extra_info_var: Name of eventExtraInfo to save model prediction to.
:param global_tag: Tag in ConditionDB where the well-trained model was stored.
:param payload: Payload for the well-trained model in global tag.

Definition at line 50 of file NN_filter_module.py.

    ):
        """
        Initialise the class.
        :param model_file: Path to the saved model file.
        :param model_config: Parameters for building the model.
        :param preproc_config: Parameters for preprocessing.
        :param threshold: Threshold for event selection using reweighting method, value *None* indicating sampling mehtod.
        :param extra_info_var: Name of eventExtraInfo to save model prediction to.
        :param global_tag: Tag in ConditionDB where the well-trained model was stored.
        :param payload: Payload for the well-trained model in global tag.
        """
        super().__init__()
        
        self.model_file = model_file
        
        self.model_config = model_config
        
        self.preproc_config = preproc_config
        
        self.threshold = threshold
        
        self.extra_info_var = extra_info_var
        
        self.payload = payload
 
        # set additional database conditions for trained neural network
        b2.conditions.prepend_globaltag(global_tag)
 

Member Function Documentation

event()

def event

(

self

)

Collect information from database, build graphs, make predictions and select through sampling or threshold

Definition at line 112 of file NN_filter_module.py.

    def event(self):
        """
        Collect information from database, build graphs, make predictions and select through sampling or threshold
        """
        # Need to create the eventExtraInfo entry for each event
        self.EventExtraInfo.create()
        df_dict = load_particle_list(mcplist=Belle2.PyStoreArray("MCParticles"), evtNum=self.EventInfo.getEvent(), label=True)
        single_input = preprocessed(df_dict, particle_selection=self.preproc_config['cuts'])
        graph = ArrayDataset(single_input, batch_size=1)[0][0]
        # Output pass probability
        pred = torch.sigmoid(self.model(graph)).detach().numpy().squeeze()
 
        # Save the pass probability to EventExtraInfo
        self.EventExtraInfo.addExtraInfo(self.extra_info_var, pred)
 
        # Module returns bool of whether prediciton passes threshold for use in basf2 path flow control
        if self.threshold:
            self.return_value(int(pred >= self.threshold))
        else:
            self.return_value(int(pred >= np.random.rand()))

initialize()

def initialize

(

self

)

Initialise module before any events are processed

Definition at line 87 of file NN_filter_module.py.

    def initialize(self):
        """
        Initialise module before any events are processed
        """
        # read trained model parameters from database
        if not self.model_file:
            accessor = DBAccessorBase(DBStoreEntry.c_RawFile, self.payload, True)
            self.model_file = accessor.getFilename()
        trained_parameters = torch.load(self.model_file, map_location=DEVICE)
 
        
        self.model = GATGAPModel(**self.model_config)
        self.model.load_state_dict(trained_parameters['model_state_dict'])
 
        
        self.EventExtraInfo = Belle2.PyStoreObj('EventExtraInfo')
        if not self.EventExtraInfo.isValid():
            self.EventExtraInfo.registerInDataStore()
        
        self.EventInfo = Belle2.PyStoreObj('EventMetaData')
        
        self.out_features = self.preproc_config['features']
        if 'PDG' in self.preproc_config['features']:
            self.out_features.remove('PDG')
 

Member Data Documentation

EventExtraInfo

EventExtraInfo

StoreArray to save weights to.

Definition at line 102 of file NN_filter_module.py.

EventInfo

EventInfo

Initialise event metadata from data store.

Definition at line 106 of file NN_filter_module.py.

extra_info_var

extra_info_var

Name of eventExtraInfo to save model prediction to.

Definition at line 80 of file NN_filter_module.py.

model

model

model with trained parameters

Definition at line 98 of file NN_filter_module.py.

model_config

model_config

Parameters for building the model.

Definition at line 74 of file NN_filter_module.py.

model_file

model_file

Path to the saved model file.

Definition at line 72 of file NN_filter_module.py.

out_features

out_features

node features

Definition at line 108 of file NN_filter_module.py.

payload

payload

Payload for the well-trained model in global tag.

Definition at line 82 of file NN_filter_module.py.

preproc_config

preproc_config

Parameters for preprocessing.

Definition at line 76 of file NN_filter_module.py.

threshold

threshold

Threshold for event selection using reweighting method, value None indicating sampling method.

Definition at line 78 of file NN_filter_module.py.

---

The documentation for this class was generated from the following file:

• generators/scripts/smartBKG/b2modules/NN_filter_module.py