Newer
Older
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
import numpy as np
import matplotlib.pyplot as plt
import os
import argparse
def extract_x0_values(filename):
x0_values = []
counter = 0
with open(filename, 'r') as file:
for line in file:
if '| 0 Average Material' in line:
counter += 1
items = [item.strip() for item in line.split('|')]
x0_value = float(items[1].split()[10])
x0_values.append(x0_value)
print("processing: ", counter, " x0: ", x0_value)
return x0_values
def plot_all_in_one(x0_values, output_file, bins, bias, detector):
"""Plot all charts in one figure"""
theta_bins = bins + 1
phi_bins = bins*2 + 1
range_theta = [bias, theta_bins-bias]
# Adjust figure size and ratio, increase height to accommodate title
fig = plt.figure(figsize=(28, 8))
# Adjust spacing between subplots and top margin
gs = fig.add_gridspec(1, 3, width_ratios=[1.2, 1, 1], wspace=0.25, top=0.85)
# Add three subplots
ax1 = fig.add_subplot(gs[0]) # Heat map
ax2 = fig.add_subplot(gs[1]) # Phi projection
ax3 = fig.add_subplot(gs[2]) # Theta projection
# Draw heat map
X0_matrix = np.array(x0_values).reshape(theta_bins, phi_bins)[range_theta[0]:range_theta[1], :]
phi = np.linspace(-180, 180, phi_bins)
theta = np.linspace(range_theta[0]*180/theta_bins, range_theta[1]*180/theta_bins, range_theta[1]-range_theta[0])
THETA, PHI = np.meshgrid(theta, phi)
im = ax1.pcolormesh(THETA, PHI, X0_matrix.T, shading='auto', cmap='viridis')
cbar = plt.colorbar(im, ax=ax1)
cbar.set_label('X/X0', fontsize=12)
ax1.set_title('Material Budget Distribution', fontsize=12, pad=10)
ax1.set_xlabel('Theta (angle)', fontsize=10)
ax1.set_ylabel('Phi (angle)', fontsize=10)
# Draw phi projection
phi_projection = np.sum(X0_matrix, axis=0) / (range_theta[1]-range_theta[0])
ax2.plot(phi, phi_projection, 'b-', linewidth=2)
ax2.fill_between(phi, phi_projection, alpha=0.3)
ax2.set_title('Projection along Phi', fontsize=12, pad=10)
ax2.set_xlabel('Phi (degree)', fontsize=10)
ax2.set_ylabel('Average X/X0', fontsize=10)
ax2.grid(True, linestyle='--', alpha=0.7)
ax2.spines['top'].set_visible(False)
ax2.spines['right'].set_visible(False)
# Draw theta projection
theta_projection = np.sum(X0_matrix, axis=1) / phi_bins
ax3.plot(theta, theta_projection, 'r-', linewidth=2)
ax3.fill_between(theta, theta_projection, alpha=0.3)
ax3.set_title('Projection along Theta', fontsize=12, pad=10)
ax3.set_xlabel('Theta (degree)', fontsize=10)
ax3.set_ylabel('Average X/X0', fontsize=10)
ax3.grid(True, linestyle='--', alpha=0.7)
ax3.spines['top'].set_visible(False)
ax3.spines['right'].set_visible(False)
# Adjust main title position
plt.suptitle('Material Budget Analysis for {}'.format(detector),
fontsize=16, # Increase font size
y=0.95) # Adjust title vertical position
plt.tight_layout()
plt.savefig(output_file, dpi=600, bbox_inches='tight')
plt.show()
def main():
parser = argparse.ArgumentParser(description='Process material scan data and generate plots.')
parser.add_argument('--input_file', help='Path to the input text file')
parser.add_argument('--output_file', default='material_budget.png', help='Path to the output png file')
parser.add_argument('--detector', default='all_world', help='Detector name')
parser.add_argument('--bins', type=int, default=100, help='theta bins is "bins"+1, phi bins is "bins"*2+1, (default: 100)')
parser.add_argument('--bias', type=int, default=0, help='Bias value for theta range (default: 0)')
args = parser.parse_args()
# all_world Coil Ecal FTD Hcal Lumical Muon OTK ParaffinEndcap SIT TPC VXD
# bias detectors: BeamPipe onlyTracker
# Extract detector name from filename
x0_values = extract_x0_values(args.input_file)
plot_all_in_one(x0_values, args.output_file, args.bins, args.bias, args.detector)
if __name__ == '__main__':
main()