-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathefield2adc_alltraces_PM.py
executable file
·128 lines (99 loc) · 6.06 KB
/
efield2adc_alltraces_PM.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#!/usr/bin/python
import argparse
import grand.io.root_trees
from grand.io.root_trees import *
from electronic_chain import *
from misc_functions import *
from PM_functions.noisemodel import generate_galacticnoise
from PM_functions.LNA import LNA_get
from PM_functions.filters import filter_get
from PM_functions import manual_pipeline
import PM_functions.config as PM_config
def main():
# Tells other functions were to look for antenna length, coefficients, etc. files
PM_config.PM_files_path = "PM_functions/PM_files"
# Parse the command line arguments
clparser = argparse.ArgumentParser()
clparser.add_argument("filename", nargs="+")
clparser.add_argument("-os", "--output_file_suffix", default="_voltage_adc")
clparser.add_argument("-od", "--output_dir", default="")
clargs = clparser.parse_args()
# Loop through files
for in_root_file in clargs.filename:
# Removing the trees from the previous file from memory
# ToDo: This should not be up to a user, at least not in this ugly way
grand.io.root_trees.grand_tree_list = []
filename_only = os.path.split(in_root_file)[-1]
# If given, create the output directory and use it
if clargs.output_dir != "":
os.makedirs(clargs.output_dir, exist_ok=True)
out_root_file = clargs.output_dir+"/"+filename_only.split(".root")[0]+clargs.output_file_suffix+".root"
else:
out_root_file = filename_only.split(".root")[0] + clargs.output_file_suffix + ".root"
# Open the ROOT tree with simulation shower data
tshower = ShowerEventSimdataTree(in_root_file)
# Open the ROOT tree with traces
tefield = EfieldEventTree(in_root_file)
# Open the ROOT tree with simulation run info
trunefieldsimdata = EfieldRunSimdataTree(in_root_file)
trunefieldsimdata.get_entry(0)
# In this file antennas are misordered. This is the dictionary that gives ROOT file trace number for given trace file number
# trace_file_to_root = {95: 62, 94: 95, 93: 42, 92: 19, 91: 21, 90: 63, 89: 14, 88: 41, 87: 0, 86: 45, 85: 54, 84: 22, 83: 20, 82: 7, 81: 69, 80: 77, 79: 68, 78: 81, 77: 30, 76: 76, 75: 17, 74: 3, 73: 18, 72: 32, 71: 53, 70: 55, 69: 4, 68: 90, 67: 82, 66: 40, 65: 12, 64: 93, 63: 13, 62: 38, 61: 44, 60: 84, 59: 15, 58: 70, 57: 58, 56: 75, 55: 78, 54: 23, 53: 8, 52: 34, 51: 29, 50: 9, 49: 37, 48: 59, 47: 79, 46: 16, 45: 83, 44: 86, 43: 47, 42: 36, 41: 92, 40: 51, 39: 66, 38: 49, 37: 31, 36: 57, 35: 6, 34: 50, 33: 61, 32: 56, 31: 89, 30: 48, 29: 65, 28: 67, 27: 64, 26: 71, 25: 28, 24: 52, 23: 27, 22: 25, 21: 91, 20: 35, 19: 85, 18: 33, 17: 87, 16: 24, 15: 88, 14: 46, 13: 39, 12: 73, 11: 72, 10: 5, 9: 94, 8: 26, 7: 43, 6: 11, 5: 60, 4: 80, 3: 74, 2: 10, 1: 1, 0: 2}
# File for storing the voltage traces
tvoltage = VoltageEventTree(out_root_file)
# File for storing the ADC traces
tadc = ADCEventTree(out_root_file)
# Stuff necessary for creating the coefficients
sampling_time = trunefieldsimdata.t_bin_size
fs = 1 / sampling_time * 1000 # sampling frequency, MHZ
N = math.ceil(fs)
freqs_tot = np.fft.rfftfreq(N, d=sampling_time * 1e-9) / 1e6 # MHz
# Read LNA coefficient
LNA_coefficient = LNA_get(freqs_tot, 1)
# Read cable and filter coefficient
cable_coefficient, filter_coefficient = filter_get(freqs_tot, 1)
# Loop through all the events in the shower tree
for i in range(tshower.get_entries()):
tefield.get_entry(i)
# A bug in root_trees? Get entry should not be necessary
tshower.get_entry(i)
# Copy the tefield event contents that exist also in tvoltage and tadc
tvoltage.copy_contents(tefield)
tadc.copy_contents(tefield)
# Convert the shower angles into ZHAireS convention
e_theta = 180-tshower.shower_zenith
e_phi = tshower.shower_azimuth-180
if e_theta<0: e_theta=360-e_theta
if e_phi<0: e_phi=360+e_phi
print("theta is:", e_theta, "degree", tshower.shower_azimuth)
print("phi is:", e_phi, "degree", tshower.shower_zenith)
# Read galactic noise coefficients - recreated for every event
galactic_noise = generate_galacticnoise(tefield.du_count, sampling_time, 18)
# Time benchmarking init
time_passed(True)
# Stack the traces together for other functions
traces_t = np.stack([tefield.trace_x, tefield.trace_y, tefield.trace_z], axis=-2)
original_trace_length = traces_t.shape[-1]
# Execute the manual pipeline on the traces. With return_all_traces=True it stores and returns all the traces throughout the pipeline steps
all_traces_t, all_traces_f = manual_pipeline(traces_t, e_phi, e_theta, du_count=tefield.du_count, original_trace_length=original_trace_length, sampling_time=sampling_time, return_all_traces=True, galactic_noise_coefficient=galactic_noise, LNA_coefficient=LNA_coefficient, cable_coefficient=cable_coefficient, filter_coefficient=filter_coefficient)
# The last before last trace is the final voltage
# Store it in the tree
tvoltage.trace_x = all_traces_t[-2][:,0,:].astype(np.float32)
tvoltage.trace_y = all_traces_t[-2][:,1,:].astype(np.float32)
tvoltage.trace_z = all_traces_t[-2][:,2,:].astype(np.float32)
# The last trace is the final ADC
# Store it in the tree
tadc.trace_0 = all_traces_t[-1][:,0,:].astype(np.int16)
tadc.trace_1 = all_traces_t[-1][:,1,:].astype(np.int16)
tadc.trace_2 = all_traces_t[-1][:,2,:].astype(np.int16)
# Print how much time has passed since the time benchmarking init
print(time_passed())
# Fill the trees with the event
tvoltage.fill()
tadc.fill()
# Write the trees to their files
tvoltage.write()
tadc.write()
print("Wrote trees")
if __name__ == '__main__':
main()