Skip to content

Commit 6f31a47

Browse files
MMannieCorbisierMMannieCorbisier
committed
adding of time and position tests
1 parent 0a3c31d commit 6f31a47

File tree

1 file changed

+107
-238
lines changed

1 file changed

+107
-238
lines changed

‎01_element/test_drift2.py‎

Lines changed: 107 additions & 238 deletions
Original file line numberDiff line numberDiff line change
@@ -4,264 +4,133 @@
44
import uproot
55
import numpy as np
66
import scipy.constants as const
7-
8-
class TestClass:
7+
from scipy import stats
8+
9+
def getSamplerData(f,samplerName, variable, mask) :
10+
d = f[samplerName+variable].array(library='np')
11+
elem = d[0]
12+
if np.ndim(elem) == 0:
13+
d_out=np.array(d)
14+
else:
15+
d_out = np.array([x[0] for x in d if x.size>0])
16+
d_out = d_out[mask]
917

10-
def test_LowEnergy_electron(self) :
11-
os.chdir(os.path.dirname(__file__))
12-
base_name = "drift2"
13-
template_name = base_name+".tpl"
14-
gmad_name = base_name+".gmad"
15-
root_name = base_name+".root"
16-
17-
l = 2.0 #m
18-
ek_in=1e-3 #GeV
19-
phi_in=0
20-
theta_in=0
21-
n_in=2500
22-
p_type="e-"
23-
c=const.c
24-
e=const.e
18+
return d_out
19+
20+
def get_p_value(mu0,x, rtol=1e-6,atol=1e-12):
21+
mean = np.mean(x)
22+
sigma = np.std(x, ddof=1)
23+
test_value=0
24+
if abs(mean) > atol:
25+
if np.abs(sigma/mean) <rtol:
26+
test_value=1
27+
else:
28+
if sigma< atol:
29+
test_value=1
30+
31+
if test_value:
32+
if np.isclose(mean, mu0, rtol):
33+
return 1.0 # H0 perfect
34+
else:
35+
return 0.0 # H0 clearly wrong
36+
37+
t_stat, p_value = stats.ttest_1samp(x, mu0)
38+
39+
return p_value
40+
41+
def get_theoretical_beam_parameters(ek_in,p_type,sim_env):
42+
c=const.c
43+
e=const.e
44+
if p_type=="e-":
2545
m=const.electron_mass
26-
e_m=m*c**2
27-
E_tot=ek_in*e*1e9+e_m
28-
pz_rel=np.sqrt(E_tot**2-e_m**2)/c
29-
pz_rel=pz_rel*c/(e*1e9) # GeV/c
30-
data = {
31-
'LENGTH': str(l),
32-
'BEAM_ENERGY' : str(ek_in), #GeV
33-
'P_TYPE': p_type
34-
}
35-
36-
pybdsim.Run.RenderGmadJinjaTemplate(template_name,gmad_name,data)
37-
pybdsim.Run.Bdsim(gmad_name,base_name,n_in,1)
38-
39-
with uproot.open(root_name) as file:
40-
name="d1."
41-
f=file['Event'][name]
42-
a=f[name+'parentID'].array(library='np')
43-
parent_ID= np.array([x[0] for x in a if len(x)>0])
44-
mask=parent_ID==0 #from initial beam only
45-
n_out=np.sum(mask)
46-
47-
a=f[name+'theta'].array(library='np')
48-
theta_out= np.array([x[0] for x in a if x.size>0])
49-
theta_out=theta_out[mask]
50-
theta_out=np.mean(theta_out)
51-
52-
a=f[name+'p'].array(library='np')
53-
p_out= np.array([x[0] for x in a if len(x)>0])
54-
a=f[name+'zp'].array(library='np')
55-
pz_frac= np.array([x[0] for x in a if len(x)>0])
56-
pz_out=np.multiply(pz_frac,p_out)
57-
pz_out=pz_out[mask]
58-
pz_out=np.mean(pz_out)
59-
60-
61-
a=f[name+'phi'].array(library='np')
62-
phi_out= np.array([x[0] for x in a if x.size>0])
63-
phi_out=phi_out[mask]
64-
phi_out=np.mean(phi_out)
65-
66-
a=f[name+'kineticEnergy'].array(library='np') # GeV
67-
ek_out= np.array([x[0] for x in a if len(x)>0])
68-
ek_out=ek_out[mask]
69-
ek_out=np.mean(ek_out)
70-
71-
assert (pytest.approx(ek_out)==ek_in) & (pytest.approx(phi_out)==phi_in) & (pytest.approx(theta_out)==theta_in) & (n_in==n_out) & (pytest.approx(pz_out)==pz_rel)
72-
73-
74-
75-
def test_HighEnergy_electron(self) :
46+
elif p_type=="proton":
47+
m=const.proton_mass
48+
gamma=ek_in*e*1e9/(m*c**2)+1
49+
v=c*np.sqrt(1-1/gamma**2)
50+
time_out_th=sim_env["l"]/v*1e9 #s
51+
e_m=m*c**2
52+
E_tot=ek_in*e*1e9+e_m
53+
pz_rel=np.sqrt(E_tot**2-e_m**2)/c
54+
pz_rel=pz_rel*c/(e*1e9) # GeV/c
55+
56+
return E_tot, pz_rel, time_out_th
57+
58+
@pytest.fixture
59+
def sim_env():
60+
os.chdir(os.path.dirname(__file__))
61+
base_name = "drift2"
62+
return {
63+
"base_name": base_name,
64+
"template_name": base_name + ".tpl",
65+
"gmad_name": base_name + ".gmad",
66+
"root_name": base_name + ".root",
67+
"n_in": 2500,
68+
"phi_in": 0,
69+
"theta_in": 0,
70+
"l": 2 #m
71+
}
72+
73+
class TestClass:
7674

77-
os.chdir(os.path.dirname(__file__))
78-
base_name = "drift2"
79-
template_name = base_name+".tpl"
80-
gmad_name = base_name+".gmad"
81-
root_name = base_name+".root"
75+
alpha_test=0.05
76+
77+
@pytest.mark.parametrize(
78+
"ek_in, p_type",
79+
[
80+
(1e-3, "e-"),
81+
(1e3, "e-"), #GeV
82+
(1e-4, "proton"),
83+
(1.0, "proton"),
84+
]
85+
)
86+
87+
def test_drift2(self, sim_env, ek_in, p_type) :
8288

83-
l = 2.0
84-
ek_in=1e3 #GeV
85-
phi_in=0
86-
theta_in=0
87-
n_in=2500
88-
p_type="e-"
89-
c=const.c
90-
e=const.e
91-
m=const.electron_mass
92-
e_m=m*c**2
93-
E_tot=ek_in*e*1e9+ e_m
94-
pz_rel=np.sqrt(E_tot**2-e_m**2)/c
95-
pz_rel=pz_rel*c/(e*1e9)
89+
E_tot, pz_rel, time_out_th= get_theoretical_beam_parameters(ek_in,p_type,sim_env)
9690
data = {
97-
'LENGTH': str(l),
91+
'LENGTH': str(sim_env["l"]),
9892
'BEAM_ENERGY' : str(ek_in), #GeV
9993
'P_TYPE': p_type
10094
}
101-
102-
pybdsim.Run.RenderGmadJinjaTemplate(template_name,gmad_name,data)
103-
pybdsim.Run.Bdsim(gmad_name,base_name,n_in,1)
104-
105-
with uproot.open(root_name) as file:
106-
name="d1."
107-
f=file['Event'][name]
108-
a=f[name+'parentID'].array(library='np')
109-
parent_ID= np.array([x[0] for x in a if len(x)>0])
110-
mask=parent_ID==0 #from initial beam only
111-
n_out=np.sum(mask)
112-
113-
a=f[name+'theta'].array(library='np')
114-
theta_out= np.array([x[0] for x in a if x.size>0])
115-
theta_out=theta_out[mask]
116-
theta_out=np.mean(theta_out)
117-
118-
a=f[name+'phi'].array(library='np')
119-
phi_out= np.array([x[0] for x in a if x.size>0])
120-
phi_out=phi_out[mask]
121-
phi_out=np.mean(phi_out)
122-
123-
a=f[name+'p'].array(library='np')
124-
p_out= np.array([x[0] for x in a if len(x)>0])
125-
a=file['Event'][name][name+'zp'].array(library='np')
126-
pz_frac= np.array([x[0] for x in a if len(x)>0])
127-
pz_out=np.multiply(pz_frac,p_out)
128-
pz_out=pz_out[mask]
129-
pz_out=np.mean(pz_out)
130-
131-
a=f[name+'kineticEnergy'].array(library='np') # in GeV
132-
ek_out= np.array([x[0] for x in a if len(x)>0])
133-
ek_out=ek_out[mask]
134-
ek_out=np.mean(ek_out)
135-
136-
assert (pytest.approx(ek_out)==ek_in) & (pytest.approx(phi_out)==phi_in) & (pytest.approx(theta_out)==theta_in) & (n_in==n_out) & (pytest.approx(pz_out)==pz_rel)
137-
13895

139-
def test_LowEnergy_proton(self) :
140-
os.chdir(os.path.dirname(__file__))
141-
base_name = "drift2"
142-
template_name = base_name+".tpl"
143-
gmad_name = base_name+".gmad"
144-
root_name = base_name+".root"
145-
146-
l = 2.0
147-
ek_in=1e-4 #GeV
148-
phi_in=0
149-
theta_in=0
150-
n_in=2500
151-
p_type="proton"
152-
c=const.c
153-
e=const.e
154-
m=const.proton_mass
155-
e_m=m*c**2
156-
E_tot=ek_in*e*1e9+ e_m
157-
pz_rel=np.sqrt(E_tot**2-e_m**2)/c
158-
pz_rel=pz_rel*c/(e*1e9) #GeV
159-
data = {
160-
'LENGTH': str(l),
161-
'BEAM_ENERGY' : str(ek_in), #GeV
162-
'P_TYPE': p_type
163-
}
96+
pybdsim.Run.RenderGmadJinjaTemplate(sim_env["template_name"],sim_env["gmad_name"],data)
97+
pybdsim.Run.Bdsim(sim_env["gmad_name"],sim_env["base_name"],sim_env["n_in"],1)
16498

165-
pybdsim.Run.RenderGmadJinjaTemplate(template_name,gmad_name,data)
166-
pybdsim.Run.Bdsim(gmad_name,base_name,n_in,1)
167-
168-
with uproot.open(root_name) as file:
169-
name="d1."
170-
f=file['Event'][name]
171-
a=f[name+'parentID'].array(library='np')
99+
with uproot.open(sim_env["root_name"]) as file:
100+
samplerName="d1."
101+
f=file['Event'][samplerName]
102+
a=f[samplerName+'parentID'].array(library='np')
172103
parent_ID= np.array([x[0] for x in a if len(x)>0])
173104
mask=parent_ID==0 #from initial beam only
174105
n_out=np.sum(mask)
175106

176-
a=f[name+'theta'].array(library='np')
177-
theta_out= np.array([x[0] for x in a if x.size>0])
178-
theta_out=theta_out[mask]
179-
theta_out=np.mean(theta_out)
180-
181-
a=f[name+'phi'].array(library='np')
182-
phi_out= np.array([x[0] for x in a if x.size>0])
183-
phi_out=phi_out[mask]
184-
phi_out=np.mean(phi_out)
107+
theta_out=getSamplerData(f,samplerName,"theta", mask)
108+
p_theta=get_p_value(sim_env["theta_in"],theta_out)
185109

186-
a=f[name+'p'].array(library='np')
110+
a=f[samplerName+'p'].array(library='np')
187111
p_out= np.array([x[0] for x in a if len(x)>0])
188-
a=f[name+'zp'].array(library='np')
112+
a=f[samplerName+'zp'].array(library='np')
189113
pz_frac= np.array([x[0] for x in a if len(x)>0])
190114
pz_out=np.multiply(pz_frac,p_out)
191115
pz_out=pz_out[mask]
192-
pz_out=np.mean(pz_out)
193-
194-
a=f[name+'kineticEnergy'].array(library='np') # in GeV
195-
ek_out= np.array([x[0] for x in a if len(x)>0])
196-
ek_out=ek_out[mask]
197-
ek_out=np.mean(ek_out)
198-
199-
assert (pytest.approx(ek_out)==ek_in) & (pytest.approx(phi_out)==phi_in) & (pytest.approx(theta_out)==theta_in) & (n_in==n_out) & (pytest.approx(pz_out)==pz_rel)
200-
201-
202-
203-
def test_HighEnergy_proton(self) :
204-
205-
os.chdir(os.path.dirname(__file__))
206-
base_name = "drift2"
207-
template_name = base_name+".tpl"
208-
gmad_name = base_name+".gmad"
209-
root_name = base_name+".root"
210-
211-
l = 2.0
212-
ek_in=1 #GeV
213-
phi_in=0
214-
theta_in=0
215-
n_in=2500
216-
p_type="proton"
217-
c=const.c
218-
e=const.e
219-
m=const.proton_mass
220-
e_m=m*c**2
221-
E_tot=ek_in*e*1e9+ e_m
222-
pz_rel=np.sqrt(E_tot**2-e_m**2)/c
223-
pz_rel=pz_rel*c/(e*1e9)
224-
data = {
225-
'LENGTH': str(l),
226-
'BEAM_ENERGY' : str(ek_in), #GeV
227-
'P_TYPE': p_type
228-
}
229-
230-
pybdsim.Run.RenderGmadJinjaTemplate(template_name,gmad_name,data)
231-
232-
pybdsim.Run.Bdsim(gmad_name,base_name,n_in,1)
233-
234-
with uproot.open(root_name) as file:
235-
name="d1."
236-
f=file['Event'][name]
237-
a=f[name+'parentID'].array(library='np')
238-
parent_ID= np.array([x[0] for x in a if len(x)>0])
239-
mask=parent_ID==0 #from initial beam only
240-
n_out=np.sum(mask)
241-
242-
a=f[name+'theta'].array(library='np')
243-
theta_out= np.array([x[0] for x in a if x.size>0])
244-
theta_out=theta_out[mask]
245-
theta_out=np.mean(theta_out)
116+
p_pz=get_p_value(pz_rel,pz_out)
117+
118+
phi_out=getSamplerData(f,samplerName,"phi", mask)
119+
p_phi=get_p_value(sim_env["phi_in"],phi_out)
246120

247-
a=f[name+'phi'].array(library='np')
248-
phi_out= np.array([x[0] for x in a if x.size>0])
249-
phi_out=phi_out[mask]
250-
phi_out=np.mean(phi_out)
121+
ek_out=getSamplerData(f,samplerName,"kineticEnergy", mask)
122+
p_ek=get_p_value(ek_in,ek_out)
251123

252-
a=f[name+'p'].array(library='np')
253-
p_out= np.array([x[0] for x in a if len(x)>0])
254-
a=f[name+'zp'].array(library='np')
255-
pz_frac= np.array([x[0] for x in a if len(x)>0])
256-
pz_out=np.multiply(pz_frac,p_out)
257-
pz_out=pz_out[mask]
258-
pz_out=np.mean(pz_out)
124+
t_out=getSamplerData(f,samplerName,"T", mask)
125+
p_t=get_p_value(time_out_th,t_out)
259126

260-
a=f[name+'kineticEnergy'].array(library='np') # in GeV
261-
ek_out= np.array([x[0] for x in a if len(x)>0])
262-
ek_out=ek_out[mask]
263-
ek_out=np.mean(ek_out)
127+
s_sampler=getSamplerData(f,samplerName,"S", mask)
128+
p_l=get_p_value(sim_env["l"],s_sampler)
264129

265-
assert (pytest.approx(ek_out)==ek_in) & (pytest.approx(phi_out)==phi_in) & (pytest.approx(theta_out)==theta_in) & (n_in==n_out) & (pytest.approx(pz_out)==pz_rel)
266-
267-
130+
assert ( p_ek> self.alpha_test)
131+
assert (p_phi> self.alpha_test)
132+
assert (p_theta> self.alpha_test)
133+
assert (sim_env["n_in"]==n_out)
134+
assert (p_pz> self.alpha_test)
135+
assert ( p_t> self.alpha_test)
136+
assert ( p_l>self.alpha_test)

0 commit comments

Comments
 (0)