Inheritance diagram for TrainDataSaver:

Public Member Functions
	__init__ (self, output_file, flag_file)

	initialize (self)

	event (self)

	terminate (self)

Public Attributes
	output_file = output_file
	Filename to save training data to.

	flag_list = ak.from_parquet(flag_file)
	Filename of the flag file indicating passing events.

	fast_mode = output_file.endswith(".parquet")
	Whether use fast mode or advanced mode.

	eventInfo = Belle2.PyStoreObj('EventMetaData')
	Initialise event metadata from data store.

	eventExtraInfo = Belle2.PyStoreObj('EventExtraInfo')
	Initialise event extra info from data store.

	df_dict = pd.DataFrame()
	Pandas dataframe to save particle features.

Detailed Description

Save MCParticles to Pandas Dataframe.

Arguments:
    output_file (str): Filename to save training data.
        Ending with ``parquet`` indicating fast mode, which will generate the final parquet file for training.
        Ending with ``h5`` indicating advanced mode, which will produce a temporary h5 file for further preprocessing.
    flag_file (str): Filename of the flag file indicating passing events.

Returns:
    None

Definition at line 68 of file NN_trainer_module.py.

Constructor & Destructor Documentation

◆ init()

__init__	(	self,
		output_file,
		flag_file )

Initialize the TrainDataSaver module.

:param output_file: Filename to save training data to.
:param flag_file: Filename of the flag file indicating passing events.

Definition at line 82 of file NN_trainer_module.py.

    ):
        """
        Initialize the TrainDataSaver module.
 
        :param output_file: Filename to save training data to.
        :param flag_file: Filename of the flag file indicating passing events.
        """
        super().__init__()
        
        self.output_file = output_file
        
        self.flag_list = ak.from_parquet(flag_file)
        
        self.fast_mode = output_file.endswith(".parquet")
 
        # delete output file if it already exists, since we will append later
        if os.path.exists(output_file):
            os.remove(output_file)
 

Member Function Documentation

◆ event()

event ( self )

Process each event and append event information to the dictionary.

Definition at line 116 of file NN_trainer_module.py.

    def event(self):
        """
        Process each event and append event information to the dictionary.
        """
        evtNum = self.eventInfo.getEvent()
        self.df_dict = pd.concat([
            self.df_dict,
            load_particle_list(mcplist=Belle2.PyStoreArray("MCParticles"), evtNum=evtNum, label=(evtNum in self.flag_list))
            ])
 

◆ initialize()

initialize ( self )

Initialize the data store and the dictionary to save particle features before processing events.

Definition at line 105 of file NN_trainer_module.py.

    def initialize(self):
        """
        Initialize the data store and the dictionary to save particle features before processing events.
        """
        
        self.eventInfo = Belle2.PyStoreObj('EventMetaData')
        
        self.eventExtraInfo = Belle2.PyStoreObj('EventExtraInfo')
        
        self.df_dict = pd.DataFrame()
 

◆ terminate()

terminate ( self )

Append events on disk in either of the two different ways and free memory.

In fast mode, the dataframe containing particle-level information and skim labels is preprocessed
and saved as a parquet file which is ready for NN training.

In advanced mode, the dataframe is saved as a h5 file and waits for combination with event-level information
before preprocessing.

Definition at line 126 of file NN_trainer_module.py.

    def terminate(self):
        """
        Append events on disk in either of the two different ways and free memory.
 
        In fast mode, the dataframe containing particle-level information and skim labels is preprocessed
        and saved as a parquet file which is ready for NN training.
 
        In advanced mode, the dataframe is saved as a h5 file and waits for combination with event-level information
        before preprocessing.
        """
        if self.fast_mode:
            ak.to_parquet(preprocessed(self.df_dict), self.output_file)
        else:
            self.df_dict.to_hdf(self.output_file, key='mc_information', mode='a', format='table', append=True)
        self.df_dict = pd.DataFrame()
 
 

Member Data Documentation

◆ df_dict

df_dict = pd.DataFrame()

Pandas dataframe to save particle features.

Definition at line 114 of file NN_trainer_module.py.

◆ eventExtraInfo

eventExtraInfo = Belle2.PyStoreObj('EventExtraInfo')

Initialise event extra info from data store.

Definition at line 112 of file NN_trainer_module.py.

◆ eventInfo

eventInfo = Belle2.PyStoreObj('EventMetaData')

Initialise event metadata from data store.

Definition at line 110 of file NN_trainer_module.py.

◆ fast_mode

fast_mode = output_file.endswith(".parquet")

Whether use fast mode or advanced mode.

Definition at line 99 of file NN_trainer_module.py.

◆ flag_list

flag_list = ak.from_parquet(flag_file)

Filename of the flag file indicating passing events.

Definition at line 97 of file NN_trainer_module.py.

◆ output_file

output_file = output_file

Filename to save training data to.

Definition at line 95 of file NN_trainer_module.py.

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

generators/scripts/smartBKG/b2modules/NN_trainer_module.py

Public Member Functions

Public Attributes