Par04/training_vae/root2h5.py

0001 """
0002 Converts an EDM4HEP ROOT file to an HDF5 file, saving the shower energy in a 3D array and shower data in a 2D array.
0003 Units: Energy values are stored in MeV and angles are stored in radians.
0004
0005 """
0006
0007 #!/bin/env python
0008 import sys
0009 import argparse
0010 import numpy as np
0011 import os
0012 import uproot
0013 import h5py
0014
0015
0016 def parse_args(argv):
0017     p = argparse.ArgumentParser()
0018     p.add_argument("--outputDir", '-o', type=str, default="./", help="Path to the output directory")
0019     p.add_argument("--inputFile", '-i', type=str, required=True, help="Name of the EDM4hep file to translate")
0020     p.add_argument("--numR", type=int, default=18, help="Number of cells in R")
0021     p.add_argument("--numPhi", type=int, default=50, help="Number of cells in phi")
0022     p.add_argument("--numZ", type=int, default=45, help="Number of cells in z")
0023     p.add_argument("--samplingFraction", type=float, default=1., help="Sampling fraction to use to rescale cell energy. Defined as f=active/(active+absorber)")
0024     args = p.parse_args()
0025     return args
0026
0027
0028 def main(argv):
0029     # Parse commandline arguments
0030     args = parse_args(argv)
0031     input_file = args.inputFile
0032     output_dir = args.outputDir
0033
0034     # Number of cells in the r, phi & z directions
0035     num_cells_R = args.numR
0036     num_cells_phi = args.numPhi
0037     num_cells_z = args.numZ
0038
0039     # Sampling fraction that rescales energy of each cell
0040     sampling_fraction = args.samplingFraction
0041
0042     if os.stat(input_file).st_size > 0:
0043         h5_file = h5py.File(
0044             f"{output_dir}/{os.path.splitext(os.path.basename(input_file))[0]}.h5", "w"
0045         )
0046         print(f"Creating output file {output_dir}/{os.path.splitext(os.path.basename(input_file))[0]}.h5")
0047         # Read Root file
0048         file = uproot.open(input_file)
0049         energy_particle = file["global"]["EnergyMC"].array()
0050         # For future once theta,phi are implemented in Par04 event/run action
0051         #phi_particle = file["global"]["PhiMC"].array()
0052         #theta_particle = file["global"]["ThetaMC"].array()
0053         cell_r = file["virtualReadout"]["rhoCell"].array()
0054         cell_phi = file["virtualReadout"]["phiCell"].array()
0055         cell_energy = file["virtualReadout"]["EnergyCell"].array()
0056         cell_z = file["virtualReadout"]["zCell"].array()
0057         all_events = []
0058         num_showers = len(energy_particle)
0059         # loop over events
0060         for event in range(num_showers):
0061             # Initialize a 3D array with shape nb_events, nb_cells in x,y,z (rho,phi,z)
0062             shower = np.zeros((num_cells_R, num_cells_phi, num_cells_z))
0063             for cell in range(len(cell_r[event])):
0064                 # This if statement is added to avoid having cells outside of desired cylinder size
0065                 if (
0066                     (cell_r[event][cell] < num_cells_R)
0067                     and (cell_phi[event][cell] < num_cells_phi)
0068                     and (cell_z[event][cell] < num_cells_z)
0069                 ):
0070                     shower[cell_r[event][cell]][cell_phi[event][cell]][
0071                         cell_z[event][cell]
0072                     ] = cell_energy[event][cell]
0073             all_events.append(shower)
0074         # Save dataset
0075         print(f"Creating datasets with shape {np.shape(energy_particle)} and {np.shape(all_events)} ")
0076         h5_file.create_dataset("incident_energy", data=energy_particle, compression="gzip", compression_opts=9,)
0077         # For future once theta,phi are implemented in Par04 event/run action
0078         #h5_file.create_dataset("incident_phi", data=phi_particle, compression="gzip", compression_opts=9,)
0079         #h5_file.create_dataset("incident_theta", data=theta_particle, compression="gzip", compression_opts=9,)
0080         h5_file.create_dataset("showers", data=all_events, compression="gzip", compression_opts=9,)
0081     h5_file.close()
0082
0083
0084 if __name__ == "__main__":
0085     exit(main(sys.argv[1:]))