Beam Class Regression Tests¶
%load_ext autoreload
%autoreload 2
import numpy as np
from distgen import Generator
from distgen.physical_constants import PHYSICAL_CONSTANTS
from distgen.tools import check_abs_and_rel_tols
from glob import glob
yaml_files = list( glob('../examples/data/*.yaml') )
coordinates = {'x', 'y', 'z', 'px', 'py', 'pz',
'r', 'theta', 'pr', 'ptheta',
'xp', 'yp', 'thetax', 'thetay',
'gamma', 'energy', 'kinetic_energy',
'beta_x', 'beta_y', 'beta_z'}
def run_test_on_input_file(input_file, test):
G = Generator(input_file, verbose=0)
G['n_particle']=10_000
test(G.beam())
def run_test_on_input_files(input_files, test, verbose=False):
for input_file in input_files:
if(verbose):
print('Testing:', input_file)
run_test_on_input_file(input_file, test)
Statistical Tests¶
$\sum_i w_i = 1$
def test_weight_normalization(beam):
check_abs_and_rel_tols('macroparticle weights', np.sum(beam['w']), 1.0, abs_tol=1e-12, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_weight_normalization)
$\langle \mathcal{O}\rangle = \sum_i w_i \mathcal{O}_i$
def test_avg(beam):
for var in coordinates:
avg_beam, avg_numpy = beam.avg(var), np.sum(beam['w']*beam[var])
check_abs_and_rel_tols('beam.avg', avg_beam, avg_numpy, abs_tol=1e-12, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_avg)
$\sigma_{\mathcal{O}}^2 = \sum_i w_i (\mathcal{O}_i-\langle \mathcal{O}\rangle)^2 $
def test_std(beam):
for var in coordinates:
sigma2_beam = beam.std(var)**2
sigma2_numpy = np.sum(beam['w']*(beam[var]-beam.avg(var))**2)
check_abs_and_rel_tols('beam.std', sigma2_beam, sigma2_numpy, abs_tol=1e-8, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_std)
$r=\sqrt{ x^2 + y^2 }$
def test_r(beam):
check_abs_and_rel_tols('beam.r', beam['r'], np.sqrt(beam['x']**2 + beam['y']**2), abs_tol=1e-15, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_r)
$x=r\cos\theta$
def test_x(beam):
check_abs_and_rel_tols('beam.x = r cos(theta)', beam['x'], beam['r']*np.cos(beam['theta']), abs_tol=1e-12, rel_tol=1e-11)
run_test_on_input_files(yaml_files, test_x)
$y = r\sin\theta$
def test_y(beam):
check_abs_and_rel_tols('beam.y = r sin(theta)', beam['y'], beam['r']*np.sin(beam['theta']), abs_tol=1e-12, rel_tol=1e-11)
run_test_on_input_files(yaml_files, test_y)
$p_r = p_x\cos\theta + p_y\sin\theta$
def test_pr(beam):
check_abs_and_rel_tols('beam.pr = px cos(theta) + py sin(theta)', beam['pr'],
beam['px']*np.cos(beam['theta']) + beam['py']*np.sin(beam['theta']),
abs_tol=1e-12, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_pr)
$p_{\theta} = -p_x\sin\theta + p_y\cos\theta$
def test_ptheta(beam):
check_abs_and_rel_tols('beam.ptheta = -px sin(theta) + py cos(theta)',
beam['ptheta'],
-beam['px']*np.sin(beam['theta']) + beam['py']*np.cos(beam['theta']),
abs_tol=1e-12, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_ptheta)
Transverse Derivatives and Angles¶
Getting¶
$x^{\prime} = p_x/p_z$
def test_xp(beam):
check_abs_and_rel_tols('beam.xp = px/pz', beam['xp'], beam['px'].to(beam['pz'].units)/beam['pz'], abs_tol=1e-14, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_xp)
$y^{\prime} = p_y/p_z$
def test_yp(beam):
check_abs_and_rel_tols('beam.yp = py/pz', beam['yp'], beam['py'].to(beam['pz'].units)/beam['pz'], abs_tol=1e-14, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_yp)
$\theta_x = \arctan(p_x/p_z)$
def test_thetax(beam):
check_abs_and_rel_tols('beam.thetax = arctan(px/pz)',
beam['thetax'],
np.arctan2(beam['px'].to(beam['pz'].units), beam['pz']),
abs_tol=1e-15, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_thetax)
$\theta_y = \arctan(p_y/p_z)$
def test_thetay(beam):
check_abs_and_rel_tols('beam.thetay = arctan(py/pz)',
beam['thetay'],
np.arctan2(beam['py'].to(beam['pz'].units), beam['pz']),
abs_tol=1e-15, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_thetay)
def test_p(beam):
deviation = np.abs(beam['p'] - np.sqrt(beam['px']**2 +beam['py']**2 + beam['pz']**2))
check_abs_and_rel_tols('beam.p = sqrt(px^2 + py^2 + pz^2)',
beam['p'],
np.sqrt(beam['px']**2 +beam['py']**2 + beam['pz']**2),
abs_tol=1e-15, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_p)
$E = \sqrt{c^2|\vec{p}|^2 + (mc^2)^2}$
def test_energy(beam):
c = PHYSICAL_CONSTANTS['speed of light in vacuum']
check_abs_and_rel_tols('beam.energy = sqrt(c^2p^2 + (mc^2)^2)',
beam['energy'],
np.sqrt(c**2*beam['p']**2 + beam.mc2**2),
abs_tol=1e-9, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_energy)
$\gamma = \sqrt{1+\left(\frac{p}{mc}\right)^2}$, $E/mc^2$
def test_gamma(beam):
mc = beam.species_mass * PHYSICAL_CONSTANTS['speed of light in vacuum']
check_abs_and_rel_tols('beam.gamma = sqrt( 1 + (p/mc)^2 )',
beam['gamma'],
np.sqrt(1+ (beam['p']/mc).to_reduced_units()**2),
abs_tol=1e-10, rel_tol=1e-10)
check_abs_and_rel_tols('beam.gamma = E/mc^2',
beam['gamma'],
beam['energy']/beam.mc2,
abs_tol=1e-15, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_gamma)
$\beta = \frac{c|\vec{p}|}{E}$
def test_beta(beam):
check_abs_and_rel_tols('beam.beta = c|p|/E',
beam['beta'],
PHYSICAL_CONSTANTS['speed of light in vacuum']*beam.p/beam.energy,
abs_tol=1e-11, rel_tol=1e-14)
assert max(beam['beta']) < 1, 'max(beta) > 1, faster than light particle!'
run_test_on_input_files(yaml_files, test_beta)
$\beta_{x_i} = \frac{cp_{x_i}}{E}$, $\beta_x = x^{\prime}\beta_z$, $\beta_y = y^{\prime}\beta_z$, $\beta_z = \frac{\beta}{\sqrt{1+(x^{\prime})^2 +(y^{\prime})^2}}$
def test_beta_xi(beam):
for var in ['x', 'y', 'z']:
check_abs_and_rel_tols('beam.beta_xi = c pxi/E )',
beam[f'beta_{var}'],
(PHYSICAL_CONSTANTS['speed of light in vacuum']*beam[f'p{var}']/beam['energy']).to_reduced_units(),
abs_tol=1e-15, rel_tol=1e-15)
check_abs_and_rel_tols("beam.beta_z = sign(pz)*beta/sqrt( 1 + x'^2 + y'^2 )",
beam['beta_z'],
np.sign(beam['pz'])*beam['beta']/np.sqrt( 1 + beam['xp']**2 + beam['yp']**2 ),
abs_tol=1e-11, rel_tol=5e-9)
run_test_on_input_files(yaml_files, test_beta_xi)
KE = $mc^2(\gamma-1)$, $E-mc^2$
def test_kinetic_energy(beam):
if(PHYSICAL_CONSTANTS.species(beam['species'])['mass']>0):
check_abs_and_rel_tols('beam.kinetic_energy = mc2*(gamma-1)',
beam['kinetic_energy'],
beam.mc2*(beam['gamma']-1),
abs_tol=1e-9, rel_tol=1e-05)
check_abs_and_rel_tols('beam.kinetic_energy = E - mc2',
beam['kinetic_energy'],
beam['energy'] - beam.mc2,
abs_tol=1e-15, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_kinetic_energy)
def test_emitt_normalized(beam):
for var in ['x', 'y']:
mc = beam.species_mass * PHYSICAL_CONSTANTS['speed of light in vacuum']
stdx = beam.std(var)
stdp = (beam.std(f'p{var}')/mc).to_reduced_units()
dx = beam[var]-beam.avg(var)
dp = ( (beam[f'p{var}']-beam.avg(f'p{var}'))/mc ).to_reduced_units()
check_abs_and_rel_tols('beam.emitt (normalized)',
beam.emitt(var),
np.sqrt(stdx**2 * stdp**2 - (np.sum(beam['w']*dx*dp))**2),
abs_tol=1e-11, rel_tol=1e-11)
run_test_on_input_files(yaml_files, test_emitt_normalized)
$\epsilon_{x} = \sqrt{\sigma_x^2\sigma_{x^{\prime}}^2 - \langle \left(x-\langle x\rangle\right)\left(x^{\prime}-\langle x^{\prime}\rangle\right)\rangle^2 }$
def test_emitt_geometric(beam):
for var in ['x', 'y']:
stdx = beam.std(var)
stdp = (beam.std(f'{var}p'))
dx = beam[var]-beam.avg(var)
dp = (beam[f'{var}p']-beam.avg(f'{var}p'))
check_abs_and_rel_tols('beam.emitt (geometric)',
beam.emitt(var, 'geometric'),
np.sqrt(stdx**2 * stdp**2 - (np.sum(beam['w']*dx*dp))**2),
abs_tol=1e-14, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_emitt_geometric)
Twiss $\beta_{x_i} = \frac{\sigma_x^2}{\epsilon_x}$
def twiss_beta_xi(beam):
for var in ['x', 'y']:
stdx = beam.std(var)
epsx = beam.emitt(var,'geometric')
if(epsx>0):
check_abs_and_rel_tols('beam.Beta_xi (Twiss)',
beam.Beta(var),
stdx**2/epsx ,
abs_tol=1e-14, rel_tol=1e-15)
run_test_on_input_files(yaml_files, test_beta_xi)
Twiss $\alpha_{x_i} = -\frac{\langle(x-\langle x\rangle)(x^{\prime}-\langle x^{\prime}\rangle)\rangle}{\epsilon_x}$
def test_alpha_xi(beam):
for var in ['x', 'y']:
dx = beam[var]-beam.avg(var)
dp = (beam[f'{var}p']-beam.avg(f'{var}p'))
epsx = beam.emitt(var,'geometric')
if(epsx>0):
check_abs_and_rel_tols('beam.Alpha_xi (Twiss)',
beam.Alpha(var),
-sum(beam['w']*dx*dp)/epsx,
abs_tol=1e-14, rel_tol=1e-14)
run_test_on_input_files(yaml_files, test_alpha_xi)
Twiss $\gamma_{x_i} = \frac{\sigma_{x^{\prime}}^2}{\epsilon_x}$
def test_gamma_xi(beam):
for var in ['x', 'y']:
stdp = beam.std(f'{var}p')
epsx = beam.emitt(var,'geometric')
if(epsx>0):
check_abs_and_rel_tols('beam.Gamma_xi (Twiss)',
beam.Gamma(var),
stdp**2/epsx,
abs_tol=1e-14, rel_tol=1e-14)
run_test_on_input_files(yaml_files, test_gamma_xi)