Simplification of the code

Author: MMannieCorbisier

01_element/test_drift2.py

@@ -4,56 +4,37 @@
 import uproot
 import numpy as np
 import scipy.constants as const
-from scipy import stats
 
-def getSamplerData(f,samplerName, variable, mask) :
+def getSamplerData(f, samplerName, variable, mask) :
     d =  f[samplerName+variable].array(library='np') 
     elem = d[0]
     if np.ndim(elem) == 0:
         d_out=np.array(d)
     else:
         d_out =  np.array([x[0] for x in d if x.size>0])
     d_out = d_out[mask]
+    d_out_mean=np.mean(d_out)
 
-    return d_out
+    return d_out_mean
 
-def get_p_value(mu0,x, rtol=1e-6,atol=1e-12):
-    mean = np.mean(x)
-    sigma = np.std(x, ddof=1)
-    test_value=0
-    if abs(mean) > atol:
-        if np.abs(sigma/mean) <rtol:
-            test_value=1
-    else:
-        if sigma< atol:
-            test_value=1
-            
-    if test_value:
-        if np.isclose(mean, mu0, rtol):
-            return 0.0   # H0 perfect
-        else:
-            return 1.0   # H0 clearly wrong
-            
-    t_stat, p_value = stats.ttest_1samp(x, mu0)
-
-    return p_value
 
-def get_theoretical_beam_parameters(ek_in,p_type,sim_env):
+def get_theoretical_beam_parameters(Ek_in, p_type, sim_env):
     c=const.c
     e=const.e
     if p_type=="e-":
         m=const.electron_mass
     elif p_type=="proton":
         m=const.proton_mass
-    gamma=ek_in*e*1e9/(m*c**2)+1
+    gamma=Ek_in*e*1e9/(m*c**2)+1
     v=c*np.sqrt(1-1/gamma**2)
     time_out_th=sim_env["l"]/v*1e9 #s
-    e_m=m*c**2
-    E_tot=ek_in*e*1e9+e_m 
-    pz_rel=np.sqrt(E_tot**2-e_m**2)/c
-    pz_rel=pz_rel*c/(e*1e9) # GeV/c
+    E_m=m*c**2
+    E_tot=Ek_in*e*1e9+E_m 
+    pz_in=np.sqrt(E_tot**2-E_m**2)/c
+    pz_in=pz_in*c/(e*1e9) # GeV/c
+    
+    return E_tot, pz_in, time_out_th
 
-    return E_tot, pz_rel, time_out_th
 
 @pytest.fixture
 def sim_env():
@@ -70,12 +51,11 @@ def sim_env():
         "l": 2 #m      
     }
 
+    
 class TestClass:
 
-    alpha_test=0.05
-    
     @pytest.mark.parametrize(
-    "ek_in, p_type",
+    "Ek_in, p_type",
     [
         (1e-3, "e-"),
         (1e3, "e-"), #GeV
@@ -84,12 +64,12 @@ class TestClass:
     ]
     )
 
-    def test_drift2(self, sim_env, ek_in, p_type) :
-
-        E_tot, pz_rel, time_out_th= get_theoretical_beam_parameters(ek_in,p_type,sim_env)
+    def test_drift2(self, sim_env, Ek_in, p_type) :
+    
+        E_tot, pz_in, time_out_th= get_theoretical_beam_parameters(Ek_in,p_type,sim_env)
         data = {
             'LENGTH': str(sim_env["l"]),
-            'BEAM_ENERGY' : str(ek_in), #GeV
+            'BEAM_ENERGY' : str(Ek_in), #GeV
             'P_TYPE': p_type
         }
 
@@ -104,34 +84,28 @@ def test_drift2(self, sim_env, ek_in, p_type) :
             mask=parent_ID==0 #from initial beam only
             n_out=np.sum(mask)
 
-            theta_out=getSamplerData(f,samplerName,"theta", mask)
-            p_theta=get_p_value(sim_env["theta_in"],theta_out)
+            theta_out_mean=getSamplerData(f,samplerName,"theta", mask)
 
             a=f[samplerName+'p'].array(library='np')
             p_out= np.array([x[0] for x in a if len(x)>0])
             a=f[samplerName+'zp'].array(library='np')
             pz_frac= np.array([x[0] for x in a if len(x)>0])
             pz_out=np.multiply(pz_frac,p_out)
             pz_out=pz_out[mask]
-            p_pz=get_p_value(pz_rel,pz_out)
+            pz_out_mean=np.mean(pz_out)
 
-            phi_out=getSamplerData(f,samplerName,"phi", mask)
-            p_phi=get_p_value(sim_env["phi_in"],phi_out)
+            phi_out_mean=getSamplerData(f,samplerName,"phi", mask)
 
-            ek_out=getSamplerData(f,samplerName,"kineticEnergy", mask)
-            p_ek=get_p_value(ek_in,ek_out)
+            Ek_out_mean=getSamplerData(f,samplerName,"kineticEnergy", mask)
 
-            t_out=getSamplerData(f,samplerName,"T", mask)
-            p_t=get_p_value(time_out_th,t_out)
+            t_out_mean=getSamplerData(f,samplerName,"T", mask)
 
-            s_sampler=getSamplerData(f,samplerName,"S", mask)
-            p_l=get_p_value(sim_env["l"],s_sampler)
-
-        assert ( p_ek< self.alpha_test)
-        assert (p_phi< self.alpha_test)
-        assert (p_theta< self.alpha_test)
-        assert (sim_env["n_in"]==n_out)
-        assert (p_pz< self.alpha_test)
-        assert ( p_t< self.alpha_test)
-        assert ( p_l< self.alpha_test)
+            s_sampler_mean=getSamplerData(f,samplerName,"S", mask)
 
+        assert Ek_out_mean==pytest.approx(Ek_in,rel=1e-9)
+        assert phi_out_mean==pytest.approx(sim_env["phi_in"], rel=1e-9)
+        assert theta_out_mean==pytest.approx(sim_env["theta_in"],rel=1e-9)
+        assert sim_env["n_in"]==n_out
+        assert pz_out_mean==pytest.approx(pz_in,1e-5)
+        assert t_out_mean==pytest.approx(time_out_th,1e-6)
+        assert s_sampler_mean==pytest.approx(sim_env["l"],1-9)