Belle II Software light-2406-ragdoll
test_VariablesToHDF5.py
1#!/usr/bin/env python3
2
3
10
11import os
12import basf2
13import pandas
14import b2test_utils
15from b2pandas_utils import VariablesToHDF5
16
17inputFile = b2test_utils.require_file('mdst14.root', 'validation')
18path = basf2.create_path()
19path.add_module('RootInput', inputFileName=inputFile)
20path.add_module('ParticleLoader', decayStrings=['e+'])
21
22# Write out electron id and momentum of all true electron candidates
23v2hdf5_e = VariablesToHDF5(
24 "e+:all", ['electronID', 'p', 'isSignal'], "particleDF.hdf5")
25path.add_module(v2hdf5_e)
26
27# event-wise mode is not supported at the moment. when it is add something like
28# the following comment and test file creation
29# Write out number of tracks and ecl-clusters in every event
30# v2hdf5_evt = VariablesToHDF5(
31# "", ['nTracks', 'nKLMClusters'], "eventDF.hdf5")
32# path.add_module(v2hdf5_evt)
33
34
36 basf2.process(path, 10) # v2hdf5 is a python module, so don't run over everything... remove this if it gets implemented in C++
37
38 # Testing
39 assert os.path.isfile('particleDF.hdf5'), "particleDF.hdf5 wasn't created"
40 df1 = pandas.read_hdf('particleDF.hdf5', 'e+:all')
41 assert len(df1) > 0, "electron dataframe contains zero entries"
42 assert 'electronID' in df1.columns, "electronID column is missing from electron dataframe"
43 assert 'p' in df1.columns, "p column is missing from electron dataframe"
44 assert 'evt' in df1.columns, "event number is missing from electron dataframe"
45 assert 'run' in df1.columns, "run number is missing from electron dataframe"
46 assert 'exp' in df1.columns, "experiment number is missing from electron dataframe"
47 assert 'icand' in df1.columns, "candidate number is missing from electron dataframe"
48 assert 'ncand' in df1.columns, "candidate count is missing from electron dataframe"
49
50 assert df1.run[0] == 0, "run number not as expected"
51 assert df1.exp[0] == 1003, "experiment number not as expected"
52 assert df1.evt[0] == 1, "event number not as expected"
def require_file(filename, data_type="", py_case=None)
Definition: __init__.py:54
def clean_working_directory()
Definition: __init__.py:189