13.3. Selective background simulation using graph neural networks#

Using graph neural network with attention mechanism to predict whether a generated event will pass the skim after detector simulation and reconstruction. Selection and weighting methods are then invested to choose proper events according to their scores while trying to avoid bias. The computational resource used for steps between generation and skim will be saved in this way.

Currently only the practical part is available e.g. the well trained neural network (parameters saved in global tag SmartBKG_GATGAP with payload GATGAPgen.pth) to filter out events that can pass FEI hadronic B0 skim. The corresponding neural network built with PyTorch is stored in generators/smartBKG/model/gatgap.py while generators/smartBKG/NN_filter_module.py is a wrapper (basf2.module) suited in basf2 framework.

class smartBKG.NN_filter_module.NNFilterModule(model_file=None, model_config={'attention_heads': 4, 'n_layers': 6, 'units': 32, 'use_gap': True}, preproc_config={'cuts': ['abs(x) < 10', 'abs(y) < 10', 'abs(z) < 10'], 'features': ['PDG', 'prodTime', 'energy', 'x', 'y', 'z', 'px', 'py', 'pz']}, threshold=None, extra_info_var='NN_prediction', global_tag='SmartBKG_GATGAP', payload='GATGAPgen.pth')[source]#
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

Parameters

  • model_file (str) – Path to saved model

  • model_config (dict) – Parameters to build the model

  • preproc_config (dict) – Parameters to provide information 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.