Skip to content

Commit

Permalink
tests are clean
Browse files Browse the repository at this point in the history
  • Loading branch information
@giadarol
giadarol committed Nov 10, 2021
1 parent b53d706 commit f3ef6c3
Showing 1 changed file with 28 additions and 144 deletions.
172 changes: 28 additions & 144 deletions tests/test_spacecharge.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import xtrack as xt
import xpart as xp

import xslowtrack as xst
import ducktrack as dtk

pmass = pmass_kg*clight**2/qe

Expand Down Expand Up @@ -35,7 +35,7 @@ def test_spacecharge_gauss_qgauss():
n_probes = 1000

from xfields.test_support.temp_makepart import generate_particles_object
(particles_pyst, r_probes, _, _, _) = generate_particles_object(
(particles_dtk, r_probes, _, _, _) = generate_particles_object(
n_macroparticles,
bunch_intensity,
sigma_x,
Expand All @@ -48,7 +48,7 @@ def test_spacecharge_gauss_qgauss():
z_probes,
theta_probes)
particles = xp.Particles(
_context=context, **particles_pyst.to_dict())
_context=context, **particles_dtk.to_dict())

particles.x += x0
particles.y += y0
Expand All @@ -63,7 +63,7 @@ def test_spacecharge_gauss_qgauss():
number_of_particles=bunch_intensity,
sigma_z=sigma_z,
z0=0.,
q_parameter=1. # there is a bug in xline,
q_parameter=1. # there is a bug in ducktrack,
# only q=1 can be tested
)

Expand Down Expand Up @@ -94,16 +94,16 @@ def test_spacecharge_gauss_qgauss():

scgauss.track(particles)

##############################
# Compare against xline #
##############################
#############################
# Compare against ducktrack #
#############################

p2np = context.nparray_from_context_array
x_probes = p2np(particles.x[:n_probes])
y_probes = p2np(particles.y[:n_probes])
z_probes = p2np(particles.zeta[:n_probes])

scpyst = xst.SCQGaussProfile(
scdtk = dtk.SCQGaussProfile(
number_of_particles = bunch_intensity,
bunchlength_rms=sigma_z,
sigma_x=sigma_x,
Expand All @@ -113,141 +113,24 @@ def test_spacecharge_gauss_qgauss():
x_co=x0,
y_co=y0)

p_pyst = xst.TestParticles(p0c=p0c,
p_dtk = dtk.TestParticles(p0c=p0c,
mass=mass,
x=x_probes.copy(),
y=y_probes.copy(),
zeta=z_probes.copy())

scpyst.track(p_pyst)
scdtk.track(p_dtk)

assert np.allclose(
p2np(particles.px[:n_probes]),
p_pyst.px,
p_dtk.px,
atol={True:1e-7, False:1e2}[frozen]
* np.max(np.abs(p_pyst.px)))
* np.max(np.abs(p_dtk.px)))
assert np.allclose(
p2np(particles.py[:n_probes]),
p_pyst.py,
p_dtk.py,
atol={True:1e-7, False:1e2}[frozen]
* np.max(np.abs(p_pyst.py)))

# TODO: re-enable when switch for longitudinal profiles will be introduced
#def test_spacecharge_gauss_coast():
# for frozen in [True, False]:
# for CTX in xo.ContextCpu, xo.ContextPyopencl, xo.ContextCupy:
# if CTX not in xo.context.available:
# continue
#
# context = CTX()
#
# #################################
# # Generate particles and probes #
# #################################
#
# n_macroparticles = int(1e6)
# bunch_intensity = 2.5e11
# sigma_x = 3e-3
# sigma_y = 2e-3
# sigma_z = 30e-2
# x0 = 1e-3
# y0 = -4e-3
# p0c = 25.92e9
# mass = Particles.pmass,
# theta_probes = 30 * np.pi/180
# r_max_probes = 2e-2
# z_probes = 1.2*sigma_z
# n_probes = 1000
#
# from xfields.test_support.temp_makepart import generate_particles_object
# (particles, r_probes, _, _, _) = generate_particles_object(
# context,
# n_macroparticles,
# bunch_intensity,
# sigma_x,
# sigma_y,
# sigma_z,
# p0c,
# mass,
# n_probes,
# r_max_probes,
# z_probes,
# theta_probes)
#
# particles.x += x0
# particles.y += y0
#
# ################
# # Space charge #
# ################
#
# beam_line_density = 1e13
# from xfields import LongitudinalProfileCoasting
# lprofile = LongitudinalProfileCoasting(
# context=context,
# beam_line_density=beam_line_density)
#
# from xfields import SpaceChargeBiGaussian
# # Just not to fool myself in the test
# if frozen:
# x0_init = x0
# y0_init = y0
# sx_init = sigma_x
# sy_init = sigma_y
# else:
# x0_init = None
# y0_init = None
# sx_init = None
# sy_init = None
#
# scgauss = SpaceChargeBiGaussian(
# context=context,
# update_on_track=not(frozen),
# length=1.,
# apply_z_kick=False,
# longitudinal_profile=lprofile,
# mean_x=x0_init,
# mean_y=y0_init,
# sigma_x=sx_init,
# sigma_y=sy_init,
# min_sigma_diff=1e-10)
#
# scgauss.track(particles)
#
# ##############################
# # Compare against xline #
# ##############################
#
# p2np = context.nparray_from_context_array
# x_probes = p2np(particles.x[:n_probes])
# y_probes = p2np(particles.y[:n_probes])
# z_probes = p2np(particles.zeta[:n_probes])
#
# from xline.elements import SCCoasting
# scpyst = SCCoasting(
# number_of_particles = beam_line_density,
# circumference = 1.,
# sigma_x=sigma_x,
# sigma_y=sigma_y,
# length=scgauss.length,
# x_co=x0,
# y_co=y0)
#
# p_pyst = Particles(p0c=p0c,
# mass=mass,
# x=x_probes.copy(),
# y=y_probes.copy(),
# zeta=z_probes.copy())
#
# scpyst.track(p_pyst)
#
#
# assert np.allclose(
# p2np(particles.px[:n_probes]),
# p_pyst.px, atol=1e-2*np.max(np.abs(p_pyst.px)))
# assert np.allclose(
# p2np(particles.py[:n_probes]),
# p_pyst.py, atol=1e-2*np.max(np.abs(p_pyst.py)))
* np.max(np.abs(p_dtk.py)))


def test_spacecharge_pic():
Expand All @@ -274,7 +157,7 @@ def test_spacecharge_pic():
n_probes = 1000

from xfields.test_support.temp_makepart import generate_particles_object
(particles_pyst, r_probes, x_probes,
(particles_gen, r_probes, x_probes,
y_probes, z_probes) = generate_particles_object(
n_macroparticles,
bunch_intensity,
Expand All @@ -287,8 +170,9 @@ def test_spacecharge_pic():
r_max_probes,
z_probes,
theta_probes)
# Transfer particles to context
particles = xp.Particles(
_context=context, **particles_pyst.to_dict())
_context=context, **particles_gen.to_dict())

######################
# Space charge (PIC) #
Expand All @@ -308,18 +192,18 @@ def test_spacecharge_pic():
z_range=(-z_lim, z_lim),
nx=128, ny=128, nz=25,
solver=solver,
gamma0=particles_pyst.gamma0[0],
gamma0=particles_gen.gamma0[0],
)

spcharge.track(particles)

##############################
# Compare against xline #
##############################
#############################
# Compare against ducktrack #
#############################

p2np = context.nparray_from_context_array

scpyst = xst.SCQGaussProfile(
scdtk = dtk.SCQGaussProfile(
number_of_particles = bunch_intensity,
bunchlength_rms=sigma_z,
sigma_x=sigma_x,
Expand All @@ -328,22 +212,22 @@ def test_spacecharge_pic():
x_co=0.,
y_co=0.)

p_pyst = xst.TestParticles(p0c=p0c,
p_dtk = dtk.TestParticles(p0c=p0c,
mass=mass,
x=x_probes.copy(),
y=y_probes.copy(),
zeta=z_probes.copy())

scpyst.track(p_pyst)
scdtk.track(p_dtk)

mask_inside_grid = ((np.abs(x_probes)<0.9*x_lim) &
(np.abs(y_probes)<0.9*y_lim))

assert np.allclose(
p2np(particles.px[:n_probes])[mask_inside_grid],
p_pyst.px[mask_inside_grid],
atol=3e-2*np.max(np.abs(p_pyst.px[mask_inside_grid])))
p_dtk.px[mask_inside_grid],
atol=3e-2*np.max(np.abs(p_dtk.px[mask_inside_grid])))
assert np.allclose(
p2np(particles.py[:n_probes])[mask_inside_grid],
p_pyst.py[mask_inside_grid],
atol=3e-2*np.max(np.abs(p_pyst.py[mask_inside_grid])))
p_dtk.py[mask_inside_grid],
atol=3e-2*np.max(np.abs(p_dtk.py[mask_inside_grid])))

0 comments on commit f3ef6c3

Please sign in to comment.