TrainDataSaver Class Reference

Inheritance diagram for TrainDataSaver:

Collaboration diagram for TrainDataSaver:

Public Member Functions
def	__init__ (self, output_file, flag_file)
def	initialize (self)
def	event (self)
def	terminate (self)

Public Attributes
	output_file
	Filename to save training data to.
	flag_list
	Filename of the flag file indicating passing events.
	fast_mode
	Whether use fast mode or advanced mode.
	eventInfo
	Initialise event metadata from data store.
	eventExtraInfo
	Initialise event extra info from data store.
	df_dict
	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 temperary 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__()

def __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 apend later
        if os.path.exists(output_file):
            os.remove(output_file)
 

Member Function Documentation

event()

def 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()

def 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()

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.

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

Pandas dataframe to save particle features.

Definition at line 114 of file NN_trainer_module.py.

eventExtraInfo

eventExtraInfo

Initialise event extra info from data store.

Definition at line 112 of file NN_trainer_module.py.

eventInfo

eventInfo

Initialise event metadata from data store.

Definition at line 110 of file NN_trainer_module.py.

fast_mode

fast_mode

Whether use fast mode or advanced mode.

Definition at line 99 of file NN_trainer_module.py.

flag_list

flag_list

Filename of the flag file indicating passing events.

Definition at line 97 of file NN_trainer_module.py.

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