From 4df6a60e352252ec44bbbba8956bd6b5d6feb495 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Szymon=20=C5=81opaciuk?= Date: Tue, 18 Apr 2023 11:01:11 +0200 Subject: [PATCH] use for_all_test_contexts in tests --- .gitignore | 1 + tests/test_beambeam2d.py | 287 ++++++----- tests/test_beambeam3d.py | 228 +++++---- tests/test_beamstrahlung.py | 608 ++++++++++++------------ tests/test_cerrf.py | 23 +- tests/test_constant_charge_slicing.py | 2 +- tests/test_electroncloud.py | 164 +++---- tests/test_electronlens_interpolated.py | 105 ++-- tests/test_mean_std.py | 29 +- tests/test_profiles.py | 59 ++- tests/test_spacecharge.py | 437 +++++++++-------- 11 files changed, 953 insertions(+), 990 deletions(-) diff --git a/.gitignore b/.gitignore index 9f5a36af..caa23f47 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ norepository_* build CyFPPS.cpp *.egg-info +.idea/** diff --git a/tests/test_beambeam2d.py b/tests/test_beambeam2d.py index b9965902..5bc9f726 100644 --- a/tests/test_beambeam2d.py +++ b/tests/test_beambeam2d.py @@ -5,156 +5,155 @@ import numpy as np -import xobjects as xo -import xtrack as xt import xpart as xp import ducktrack as dtk +from xobjects.test_helpers import for_all_test_contexts + from scipy.constants import m_p as pmass_kg from scipy.constants import e as qe from scipy.constants import c as clight pmass = pmass_kg*clight**2/qe -def test_beambeam(): - for context in xo.context.get_test_contexts(): - print(repr(context)) - - ################################# - # Generate particles and probes # - ################################# - - n_macroparticles_b1 = int(1e6) - bunch_intensity_b1 = 2.5e11 - sigma_x_b1 = 3e-3 - sigma_y_b1 = 2e-3 - mean_x_b1 = 1.3e-3 - mean_y_b1 = -1.2e-3 - - n_macroparticles_b2 = int(1e6) - bunch_intensity_b2 = 3e11 - sigma_x_b2 = 1.7e-3 - sigma_y_b2 = 2.1e-3 - mean_x_b2 = -1e-3 - mean_y_b2 = 1.4e-3 - - sigma_z = 30e-2 - p0c = 25.92e9 - mass = 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_b1, r_probes, _, _, _ - ) = generate_particles_object( - n_macroparticles_b1, - bunch_intensity_b1, - sigma_x_b1, - sigma_y_b1, - sigma_z, - p0c, - mass, - n_probes, - r_max_probes, - z_probes, - theta_probes) - # Move to right context - particles_b1 = xp.Particles(_context=context, - **particles_b1.to_dict()) - particles_b1.x += mean_x_b1 - particles_b1.y += mean_y_b1 - - (particles_b2, r_probes, _, _, _ - ) = generate_particles_object( - n_macroparticles_b2, - bunch_intensity_b2, - sigma_x_b2, - sigma_y_b2, - sigma_z, - p0c, - mass, - n_probes, - r_max_probes, - z_probes, - theta_probes) - particles_b2 = xp.Particles(_context=context, - **particles_b2.to_dict()) - particles_b2.x += mean_x_b2 - particles_b2.y += mean_y_b2 - - ############# - # Beam-beam # - ############# - - from xfields import BeamBeamBiGaussian2D, mean_and_std - - # if beta0 is array I just take the first - beta0_b2 = context.nparray_from_context_array(particles_b2.beta0)[0] - beta0_b1 = context.nparray_from_context_array(particles_b1.beta0)[0] - - bbeam_b1 = BeamBeamBiGaussian2D( - _context=context, - n_particles=bunch_intensity_b2, - q0 = particles_b2.q0, - beta0=beta0_b2, - sigma_x=1., # dummy - sigma_y=1., # dummy - mean_x=1., # dummy - mean_y=1., # dummy - min_sigma_diff=1e-10) - - # Measure beam properties - mean_x_meas, sigma_x_meas = mean_and_std(particles_b2.x) - mean_y_meas, sigma_y_meas = mean_and_std(particles_b2.y) - # Update bb lens - bbeam_b1.sigma_x = sigma_x_meas - bbeam_b1.sigma_y = sigma_y_meas - bbeam_b1.mean_x = mean_x_meas - bbeam_b1.mean_y = mean_y_meas - #Track - bbeam_b1.track(particles_b1) - - ############################# - # Compare against ducktrack # - ############################# - - p2np = context.nparray_from_context_array - x_probes = p2np(particles_b1.x[:n_probes]) - y_probes = p2np(particles_b1.y[:n_probes]) - z_probes = p2np(particles_b1.zeta[:n_probes]) - - from ducktrack.elements import BeamBeam4D - bb_b1_dtk= BeamBeam4D( - charge = bunch_intensity_b2, - sigma_x=sigma_x_b2, - sigma_y=sigma_y_b2, - x_bb=mean_x_b2, - y_bb=mean_y_b2, - beta_r=np.float64(beta0_b2)) - - p_dtk = dtk.TestParticles(p0c=p0c, - mass=mass, - x=x_probes.copy(), - y=y_probes.copy(), - zeta=z_probes.copy()) - - bb_b1_dtk.track(p_dtk) - - assert np.allclose(p_dtk.px, - p2np(particles_b1.px[:n_probes]), - atol=2e-2*np.max(np.abs(p_dtk.px))) - assert np.allclose(p_dtk.py, - p2np(particles_b1.py[:n_probes]), - atol=2e-2*np.max(np.abs(p_dtk.py))) - - - # Some test on scale_strength - bbeam_b1.scale_strength = 0 - p_before = particles_b1.copy() - - bbeam_b1.track(particles_b1) - - assert np.allclose(p2np(p_before.px), p2np(particles_b1.px), atol=1e-14) - assert np.allclose(p2np(p_before.py), p2np(particles_b1.py), atol=1e-14) + +@for_all_test_contexts +def test_beambeam(test_context): + ################################# + # Generate particles and probes # + ################################# + + n_macroparticles_b1 = int(1e6) + bunch_intensity_b1 = 2.5e11 + sigma_x_b1 = 3e-3 + sigma_y_b1 = 2e-3 + mean_x_b1 = 1.3e-3 + mean_y_b1 = -1.2e-3 + + n_macroparticles_b2 = int(1e6) + bunch_intensity_b2 = 3e11 + sigma_x_b2 = 1.7e-3 + sigma_y_b2 = 2.1e-3 + mean_x_b2 = -1e-3 + mean_y_b2 = 1.4e-3 + + sigma_z = 30e-2 + p0c = 25.92e9 + mass = 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_b1, r_probes, _, _, _ + ) = generate_particles_object( + n_macroparticles_b1, + bunch_intensity_b1, + sigma_x_b1, + sigma_y_b1, + sigma_z, + p0c, + mass, + n_probes, + r_max_probes, + z_probes, + theta_probes) + # Move to right context + particles_b1 = xp.Particles(_context=test_context, + **particles_b1.to_dict()) + particles_b1.x += mean_x_b1 + particles_b1.y += mean_y_b1 + + (particles_b2, r_probes, _, _, _ + ) = generate_particles_object( + n_macroparticles_b2, + bunch_intensity_b2, + sigma_x_b2, + sigma_y_b2, + sigma_z, + p0c, + mass, + n_probes, + r_max_probes, + z_probes, + theta_probes) + particles_b2 = xp.Particles(_context=test_context, + **particles_b2.to_dict()) + particles_b2.x += mean_x_b2 + particles_b2.y += mean_y_b2 + + ############# + # Beam-beam # + ############# + + from xfields import BeamBeamBiGaussian2D, mean_and_std + + # if beta0 is array I just take the first + beta0_b2 = test_context.nparray_from_context_array(particles_b2.beta0)[0] + beta0_b1 = test_context.nparray_from_context_array(particles_b1.beta0)[0] + + bbeam_b1 = BeamBeamBiGaussian2D( + _context=test_context, + n_particles=bunch_intensity_b2, + q0 = particles_b2.q0, + beta0=beta0_b2, + sigma_x=1., # dummy + sigma_y=1., # dummy + mean_x=1., # dummy + mean_y=1., # dummy + min_sigma_diff=1e-10) + + # Measure beam properties + mean_x_meas, sigma_x_meas = mean_and_std(particles_b2.x) + mean_y_meas, sigma_y_meas = mean_and_std(particles_b2.y) + # Update bb lens + bbeam_b1.sigma_x = sigma_x_meas + bbeam_b1.sigma_y = sigma_y_meas + bbeam_b1.mean_x = mean_x_meas + bbeam_b1.mean_y = mean_y_meas + #Track + bbeam_b1.track(particles_b1) + + ############################# + # Compare against ducktrack # + ############################# + + p2np = test_context.nparray_from_context_array + x_probes = p2np(particles_b1.x[:n_probes]) + y_probes = p2np(particles_b1.y[:n_probes]) + z_probes = p2np(particles_b1.zeta[:n_probes]) + + from ducktrack.elements import BeamBeam4D + bb_b1_dtk= BeamBeam4D( + charge = bunch_intensity_b2, + sigma_x=sigma_x_b2, + sigma_y=sigma_y_b2, + x_bb=mean_x_b2, + y_bb=mean_y_b2, + beta_r=np.float64(beta0_b2)) + + p_dtk = dtk.TestParticles(p0c=p0c, + mass=mass, + x=x_probes.copy(), + y=y_probes.copy(), + zeta=z_probes.copy()) + + bb_b1_dtk.track(p_dtk) + + assert np.allclose(p_dtk.px, + p2np(particles_b1.px[:n_probes]), + atol=2e-2*np.max(np.abs(p_dtk.px))) + assert np.allclose(p_dtk.py, + p2np(particles_b1.py[:n_probes]), + atol=2e-2*np.max(np.abs(p_dtk.py))) + + + # Some test on scale_strength + bbeam_b1.scale_strength = 0 + p_before = particles_b1.copy() + + bbeam_b1.track(particles_b1) + + assert np.allclose(p2np(p_before.px), p2np(particles_b1.px), atol=1e-14) + assert np.allclose(p2np(p_before.py), p2np(particles_b1.py), atol=1e-14) diff --git a/tests/test_beambeam3d.py b/tests/test_beambeam3d.py index 18fc3e7e..0704983f 100644 --- a/tests/test_beambeam3d.py +++ b/tests/test_beambeam3d.py @@ -1,8 +1,12 @@ +# copyright ################################# # +# This file is part of the Xfields Package. # +# Copyright (c) CERN, 2023. # +# ########################################### # + import numpy as np import pytest import xobjects as xo -import xtrack as xt import xfields as xf import xpart as xp @@ -10,121 +14,115 @@ import ducktrack as dtk -def test_compute_moments(): - for context in xo.context.get_test_contexts(): - - print(repr(context)) - - if isinstance(context, xo.ContextPyopencl): - print('Incompatible with OpenCL') - continue - - ########### - # ttbar 2 # - ########### - p0c = 182.5e9 # [eV] - mass0 = .511e6 # [eV] - - physemit_x = 1.46e-09 # [m] - physemit_y = 2.9e-12 # [m] - beta_x = 1 # [m] - beta_y = .0016 # [m] - sigma_x = np.sqrt(physemit_x*beta_x) # [m] - sigma_px = np.sqrt(physemit_x/beta_x) # [m] - sigma_y = np.sqrt(physemit_y*beta_y) # [m] - sigma_py = np.sqrt(physemit_y/beta_y) # [m] - sigma_z_tot = .00254 # [m] sr+bs - sigma_delta_tot = .00192 # [m] - n_macroparticles_b1 = int(1e6) - - n_slices = 2 - threshold_num_macroparticles=20 - - ############# - # particles # - ############# - - #e- - part_range = np.linspace(-5*sigma_z_tot,5*sigma_z_tot,n_macroparticles_b1) - particles_b0 = xp.Particles( - _context = context, - q0 = -1, - p0c = p0c, - mass0 = mass0, - x = part_range, - zeta = part_range, - ) - - slicer = xf.TempSlicer(n_slices=n_slices, sigma_z=sigma_z_tot, mode="unibin") - - particles_b1 = particles_b0.copy() - particles_b2 = particles_b0.copy() - particles_b2.state[:int(n_macroparticles_b1/4)] = 0 # set 1/4 of the particles to lost - - # compute slice moments: lost particles are labeled with state=0 and their slice idx will be set to -1 - slice_moments_b1 = slicer.compute_moments(particles_b1, threshold_num_macroparticles=threshold_num_macroparticles) - slice_moments_b2 = slicer.compute_moments(particles_b2, threshold_num_macroparticles=threshold_num_macroparticles) - - other_beam_num_particles_b1 = slice_moments_b1[:n_slices] - x_center_b1 = slice_moments_b1[ n_slices:2*n_slices] - Sigma_11_b1 = slice_moments_b1[ 7*n_slices: 8*n_slices] - - other_beam_num_particles_b2 = slice_moments_b2[:n_slices] - x_center_b2 = slice_moments_b2[ n_slices:2*n_slices] - Sigma_11_b2 = slice_moments_b2[ 7*n_slices: 8*n_slices] - - # check if all lost particles have slice idx = -1 - assert np.all(particles_b2.slice[particles_b2.state == 0] == -1) - - # check if the mean and std of the alive particles in each slice agrees with compute_moments - for s in range(n_slices): - slice_b1 = particles_b1.x[particles_b1.slice==s] - num_parts_slice_b1 = len(slice_b1) - mean_b1 = np.mean(slice_b1) - diff_b1 = slice_b1 - mean_b1 - sigma_b1 = float((diff_b1**2).sum()) / len(slice_b1) - - if num_parts_slice_b1 > threshold_num_macroparticles: - assert num_parts_slice_b1 == other_beam_num_particles_b1[s] - assert mean_b1 == x_center_b1[s] - assert sigma_b1 == Sigma_11_b1[s] - else: - print(f"Slice {s} has insufficient ({num_parts_slice_b1}) particles! Need at least {threshold_num_macroparticles}.") - - slice_b2 = particles_b2.x[particles_b2.slice==s] - num_parts_slice_b2 = len(slice_b2) - mean_b2 = np.mean(slice_b2) - diff_b2 = slice_b2 - mean_b2 - sigma_b2 = float((diff_b2**2).sum()) / len(slice_b2) - if num_parts_slice_b2 > threshold_num_macroparticles: - assert num_parts_slice_b2 == other_beam_num_particles_b2[s] - assert mean_b2 == x_center_b2[s] - assert sigma_b2 == Sigma_11_b2[s] - else: - print(f"Slice {s} has insufficient ({num_parts_slice_b2}) particles! Need at least {threshold_num_macroparticles}.") - - # compute number of particles in bunch head slice - slice_b1 = particles_b1.x[particles_b1.slice==0] - num_parts_slice_b1 = len(slice_b1) - x_center_b1_before = x_center_b1 - print(x_center_b1) - - # lose a particle in the head of the bunch (slice 0) during tracking, slice moment is messed up - lost_idx = -1 - particles_b1.x[lost_idx] = 1e34 - slice_moments_b1 = slicer.compute_moments(particles_b1, threshold_num_macroparticles=threshold_num_macroparticles) - x_center_b1 = slice_moments_b1[ n_slices:2*n_slices] - print(x_center_b1) - - assert x_center_b1[0] > 1e20 - # update particle status, recompute slices - particles_b1.state[lost_idx] = 0 - slice_moments_b1 = slicer.compute_moments(particles_b1, threshold_num_macroparticles=threshold_num_macroparticles) - x_center_b1 = slice_moments_b1[ n_slices:2*n_slices] - print(x_center_b1) - - assert np.abs((x_center_b1[0]-x_center_b1_before[0])/x_center_b1_before[0]) < 1e-5 +@for_all_test_contexts(excluding='ContextPyopencl') +def test_compute_moments(test_context): + ########### + # ttbar 2 # + ########### + p0c = 182.5e9 # [eV] + mass0 = .511e6 # [eV] + + physemit_x = 1.46e-09 # [m] + physemit_y = 2.9e-12 # [m] + beta_x = 1 # [m] + beta_y = .0016 # [m] + sigma_x = np.sqrt(physemit_x*beta_x) # [m] + sigma_px = np.sqrt(physemit_x/beta_x) # [m] + sigma_y = np.sqrt(physemit_y*beta_y) # [m] + sigma_py = np.sqrt(physemit_y/beta_y) # [m] + sigma_z_tot = .00254 # [m] sr+bs + sigma_delta_tot = .00192 # [m] + n_macroparticles_b1 = int(1e6) + + n_slices = 2 + threshold_num_macroparticles=20 + + ############# + # particles # + ############# + + #e- + part_range = np.linspace(-5*sigma_z_tot,5*sigma_z_tot,n_macroparticles_b1) + particles_b0 = xp.Particles( + _context = test_context, + q0 = -1, + p0c = p0c, + mass0 = mass0, + x = part_range, + zeta = part_range, + ) + + slicer = xf.TempSlicer(n_slices=n_slices, sigma_z=sigma_z_tot, mode="unibin") + + particles_b1 = particles_b0.copy() + particles_b2 = particles_b0.copy() + particles_b2.state[:int(n_macroparticles_b1/4)] = 0 # set 1/4 of the particles to lost + + # compute slice moments: lost particles are labeled with state=0 and their slice idx will be set to -1 + slice_moments_b1 = slicer.compute_moments(particles_b1, threshold_num_macroparticles=threshold_num_macroparticles) + slice_moments_b2 = slicer.compute_moments(particles_b2, threshold_num_macroparticles=threshold_num_macroparticles) + + other_beam_num_particles_b1 = slice_moments_b1[:n_slices] + x_center_b1 = slice_moments_b1[ n_slices:2*n_slices] + Sigma_11_b1 = slice_moments_b1[ 7*n_slices: 8*n_slices] + + other_beam_num_particles_b2 = slice_moments_b2[:n_slices] + x_center_b2 = slice_moments_b2[ n_slices:2*n_slices] + Sigma_11_b2 = slice_moments_b2[ 7*n_slices: 8*n_slices] + + # check if all lost particles have slice idx = -1 + assert np.all(particles_b2.slice[particles_b2.state == 0] == -1) + + # check if the mean and std of the alive particles in each slice agrees with compute_moments + for s in range(n_slices): + slice_b1 = particles_b1.x[particles_b1.slice==s] + num_parts_slice_b1 = len(slice_b1) + mean_b1 = np.mean(slice_b1) + diff_b1 = slice_b1 - mean_b1 + sigma_b1 = float((diff_b1**2).sum()) / len(slice_b1) + + if num_parts_slice_b1 > threshold_num_macroparticles: + assert num_parts_slice_b1 == other_beam_num_particles_b1[s] + assert mean_b1 == x_center_b1[s] + assert sigma_b1 == Sigma_11_b1[s] + else: + print(f"Slice {s} has insufficient ({num_parts_slice_b1}) particles! Need at least {threshold_num_macroparticles}.") + + slice_b2 = particles_b2.x[particles_b2.slice==s] + num_parts_slice_b2 = len(slice_b2) + mean_b2 = np.mean(slice_b2) + diff_b2 = slice_b2 - mean_b2 + sigma_b2 = float((diff_b2**2).sum()) / len(slice_b2) + if num_parts_slice_b2 > threshold_num_macroparticles: + assert num_parts_slice_b2 == other_beam_num_particles_b2[s] + assert mean_b2 == x_center_b2[s] + assert sigma_b2 == Sigma_11_b2[s] + else: + print(f"Slice {s} has insufficient ({num_parts_slice_b2}) particles! Need at least {threshold_num_macroparticles}.") + + # compute number of particles in bunch head slice + slice_b1 = particles_b1.x[particles_b1.slice==0] + num_parts_slice_b1 = len(slice_b1) + x_center_b1_before = x_center_b1 + print(x_center_b1) + + # lose a particle in the head of the bunch (slice 0) during tracking, slice moment is messed up + lost_idx = -1 + particles_b1.x[lost_idx] = 1e34 + slice_moments_b1 = slicer.compute_moments(particles_b1, threshold_num_macroparticles=threshold_num_macroparticles) + x_center_b1 = slice_moments_b1[ n_slices:2*n_slices] + print(x_center_b1) + + assert x_center_b1[0] > 1e20 + + # update particle status, recompute slices + particles_b1.state[lost_idx] = 0 + slice_moments_b1 = slicer.compute_moments(particles_b1, threshold_num_macroparticles=threshold_num_macroparticles) + x_center_b1 = slice_moments_b1[ n_slices:2*n_slices] + print(x_center_b1) + + assert np.abs((x_center_b1[0]-x_center_b1_before[0])/x_center_b1_before[0]) < 1e-5 def sigma_configurations(): print('decoupled round beam') diff --git a/tests/test_beamstrahlung.py b/tests/test_beamstrahlung.py index 3398b930..e22b8599 100644 --- a/tests/test_beamstrahlung.py +++ b/tests/test_beamstrahlung.py @@ -8,320 +8,306 @@ import xfields as xf import xpart as xp +from xobjects.test_helpers import for_all_test_contexts + test_data_folder = pathlib.Path( __file__).parent.joinpath('../test_data').absolute() -def test_beambeam3d_beamstrahlung_single_collision(): - for context in xo.context.get_test_contexts(): - - if isinstance(context, xo.ContextPyopencl): - print('Incompatible with OpenCL') - continue - - print(repr(context)) - - ########### - # ttbar 2 # - ########### - bunch_intensity = 2.3e11 # [1] - energy = 182.5 # [GeV] - p0c = 182.5e9 # [eV] - mass0 = .511e6 # [eV] - phi = 15e-3 # [rad] half xing - physemit_x = 1.46e-09 # [m] - physemit_y = 2.9e-12 # [m] - beta_x = 1 # [m] - beta_y = .0016 # [m] - sigma_x = np.sqrt(physemit_x*beta_x) # [m] - sigma_px = np.sqrt(physemit_x/beta_x) # [m] - sigma_y = np.sqrt(physemit_y*beta_y) # [m] - sigma_py = np.sqrt(physemit_y/beta_y) # [m] - sigma_z_tot = .00254 # [m] sr+bs - sigma_delta_tot = .00192 # [m] - n_macroparticles_b1 = int(1e6) - - n_slices = 100 - - ############# - # particles # - ############# - - #e- - particles_b1 = xp.Particles( - _context = context, - q0 = -1, - p0c = p0c, - mass0 = mass0, - x = sigma_x *np.random.randn(n_macroparticles_b1), - y = sigma_y *np.random.randn(n_macroparticles_b1), - zeta = sigma_z_tot *np.random.randn(n_macroparticles_b1), - px = sigma_px *np.random.randn(n_macroparticles_b1), - py = sigma_py *np.random.randn(n_macroparticles_b1), - delta = sigma_delta_tot*np.random.randn(n_macroparticles_b1), - ) - particles_b1.name = "b1" - - slicer = xf.TempSlicer(n_slices=n_slices, sigma_z=sigma_z_tot, mode="unicharge") - - el_beambeam_b1 = xf.BeamBeamBiGaussian3D( - _context=context, - config_for_update = None, - other_beam_q0=1, - phi=phi, - alpha=0, - # decide between round or elliptical kick formula - min_sigma_diff = 1e-28, - # slice intensity [num. real particles] n_slices inferred from length of this - slices_other_beam_num_particles = slicer.bin_weights * bunch_intensity, - # unboosted strong beam moments - slices_other_beam_zeta_center = slicer.bin_centers, - slices_other_beam_Sigma_11 = n_slices*[sigma_x**2], - slices_other_beam_Sigma_22 = n_slices*[sigma_px**2], - slices_other_beam_Sigma_33 = n_slices*[sigma_y**2], - slices_other_beam_Sigma_44 = n_slices*[sigma_py**2], - # only if BS on - slices_other_beam_zeta_bin_width_star_beamstrahlung = slicer.bin_widths_beamstrahlung / np.cos(phi), # boosted dz - # has to be set - slices_other_beam_Sigma_12 = n_slices*[0], - slices_other_beam_Sigma_34 = n_slices*[0], - ) - - ######################### - # track for 1 collision # - ######################### - - line = xt.Line(elements = [el_beambeam_b1]) - line.build_tracker(_context=context) - - assert line._needs_rng == False - - line.configure_radiation(model_beamstrahlung='quantum') - assert line._needs_rng == True - - record = line.start_internal_logging_for_elements_of_type( - xf.BeamBeamBiGaussian3D, capacity={"beamstrahlungtable": int(3e5)}) - line.track(particles_b1, num_turns=1) - line.stop_internal_logging_for_elements_of_type(xf.BeamBeamBiGaussian3D) - - record.move(_context=xo.context_default) - - ########################################### - # test 1: compare spectrum with guineapig # - ########################################### - - fname = (test_data_folder - / "beamstrahlung/guineapig_ttbar2_beamstrahlung_photon_energies_gev.txt") - guinea_photons = np.loadtxt(fname) # contains about 250k photons emitted from 1e6 macroparticles in 1 collision - n_bins = 10 - bins = np.logspace( - np.log10(1e-14), np.log10(1e1), n_bins) - xsuite_hist = np.histogram(record.beamstrahlungtable.photon_energy/1e9, - bins=bins)[0] - guinea_hist = np.histogram(guinea_photons, bins=bins)[0] - - bin_rel_errors = np.abs(xsuite_hist - guinea_hist) / guinea_hist - print(f"bin relative errors [1]: {bin_rel_errors}") - - # test if relative error in the last 5 bins is smaller than 1e-1 - assert np.allclose(xsuite_hist[-5:], guinea_hist[-5:], rtol=1e-1, atol=0) - - ############################################ - # test 2: compare beamstrahlung parameters # - ############################################ - - # average and maximum BS parameter - # https://www.researchgate.net/publication/2278298_Beam-Beam_Phenomena_In_Linear_Colliders - # page 20 - r0 = cst.e**2/(4*np.pi*cst.epsilon_0*cst.m_e*cst.c**2) # - if pp - - upsilon_max = ( - 2 * r0**2 * energy/(mass0*1e-9) * bunch_intensity - / (1/137*sigma_z_tot*(sigma_x + 1.85*sigma_y))) - upsilon_avg = (5/6 * r0**2 * energy/(mass0*1e-9) * bunch_intensity - / (1/137*sigma_z_tot*(sigma_x + sigma_y))) - - # get rid of padded zeros in table - photon_critical_energy = np.array( - sorted(set(record.beamstrahlungtable.photon_critical_energy))[1:]) - primary_energy = np.array( - sorted(set( record.beamstrahlungtable.primary_energy))[1:]) - - upsilon_avg_sim = np.mean(0.67 * photon_critical_energy / primary_energy) - upsilon_max_sim = np.max(0.67 * photon_critical_energy / primary_energy) - - print("Y max. [1]:", upsilon_max) - print("Y avg. [1]:", upsilon_avg) - print("Y max. [1]:", upsilon_max_sim) - print("Y avg. [1]:", upsilon_avg_sim) - print("Y max. ratio [1]:", upsilon_max_sim / upsilon_max) - print("Y avg. ratio [1]:", upsilon_avg_sim / upsilon_avg) - - assert np.abs(upsilon_max_sim / upsilon_max - 1) < 2e-2 - assert np.abs(upsilon_avg_sim / upsilon_avg - 1) < 5e-2 - - - -def test_beambeam3d_collective_beamstrahlung_single_collision(): - for context in xo.context.get_test_contexts(): - - print(repr(context)) - - if isinstance(context, xo.ContextPyopencl): - print("skipping test for pyopencl context") - continue - - ########### - # ttbar 2 # - ########### - bunch_intensity = 2.3e11 # [1] - energy = 182.5 # [GeV] - p0c = 182.5e9 # [eV] - mass0 = .511e6 # [eV] - phi = 15e-3 # [rad] half xing - u_sr = 9.2 # [GeV] - u_bs = .0114 # [GeV] - k2_factor = .4 # [1] - qx = .554 # [1] half arc - qy = .588 # [1] - qs = .0436 # [1] - physemit_x = 1.46e-09 # [m] - physemit_y = 2.9e-12 # [m] - beta_x = 1 # [m] - beta_y = .0016 # [m] - sigma_x = np.sqrt(physemit_x*beta_x) # [m] - sigma_px = np.sqrt(physemit_x/beta_x) # [m] - sigma_y = np.sqrt(physemit_y*beta_y) # [m] - sigma_py = np.sqrt(physemit_y/beta_y) # [m] - sigma_z = .00194 # [m] sr - sigma_z_tot = .00254 # [m] sr+bs - sigma_delta = .0015 # [m] - sigma_delta_tot = .00192 # [m] - beta_s = sigma_z/sigma_delta # [m] - physemit_s = sigma_z*sigma_delta # [m] - physemit_s_tot = sigma_z_tot*sigma_delta_tot # [m] - n_macroparticles_b1 = int(1e6) - n_macroparticles_b2 = int(1e6) - - n_slices = 100 - - ############# - # particles # - ############# - - #e- - particles_b1 = xp.Particles( - _context = context, - q0 = -1, - p0c = p0c, - mass0 = mass0, - x = sigma_x *np.random.randn(n_macroparticles_b1), - y = sigma_y *np.random.randn(n_macroparticles_b1), - zeta = sigma_z_tot *np.random.randn(n_macroparticles_b1), - px = sigma_px *np.random.randn(n_macroparticles_b1), - py = sigma_py *np.random.randn(n_macroparticles_b1), - delta = sigma_delta_tot*np.random.randn(n_macroparticles_b1), +@for_all_test_contexts(excluding='ContextPyopencl') +def test_beambeam3d_beamstrahlung_single_collision(test_context): + ########### + # ttbar 2 # + ########### + bunch_intensity = 2.3e11 # [1] + energy = 182.5 # [GeV] + p0c = 182.5e9 # [eV] + mass0 = .511e6 # [eV] + phi = 15e-3 # [rad] half xing + physemit_x = 1.46e-09 # [m] + physemit_y = 2.9e-12 # [m] + beta_x = 1 # [m] + beta_y = .0016 # [m] + sigma_x = np.sqrt(physemit_x*beta_x) # [m] + sigma_px = np.sqrt(physemit_x/beta_x) # [m] + sigma_y = np.sqrt(physemit_y*beta_y) # [m] + sigma_py = np.sqrt(physemit_y/beta_y) # [m] + sigma_z_tot = .00254 # [m] sr+bs + sigma_delta_tot = .00192 # [m] + n_macroparticles_b1 = int(1e6) + + n_slices = 100 + + ############# + # particles # + ############# + + #e- + particles_b1 = xp.Particles( + _context = test_context, + q0 = -1, + p0c = p0c, + mass0 = mass0, + x = sigma_x *np.random.randn(n_macroparticles_b1), + y = sigma_y *np.random.randn(n_macroparticles_b1), + zeta = sigma_z_tot *np.random.randn(n_macroparticles_b1), + px = sigma_px *np.random.randn(n_macroparticles_b1), + py = sigma_py *np.random.randn(n_macroparticles_b1), + delta = sigma_delta_tot*np.random.randn(n_macroparticles_b1), + ) + + particles_b1.name = "b1" + + slicer = xf.TempSlicer(n_slices=n_slices, sigma_z=sigma_z_tot, mode="unicharge") + + el_beambeam_b1 = xf.BeamBeamBiGaussian3D( + _context=test_context, + config_for_update = None, + other_beam_q0=1, + phi=phi, + alpha=0, + # decide between round or elliptical kick formula + min_sigma_diff = 1e-28, + # slice intensity [num. real particles] n_slices inferred from length of this + slices_other_beam_num_particles = slicer.bin_weights * bunch_intensity, + # unboosted strong beam moments + slices_other_beam_zeta_center = slicer.bin_centers, + slices_other_beam_Sigma_11 = n_slices*[sigma_x**2], + slices_other_beam_Sigma_22 = n_slices*[sigma_px**2], + slices_other_beam_Sigma_33 = n_slices*[sigma_y**2], + slices_other_beam_Sigma_44 = n_slices*[sigma_py**2], + # only if BS on + slices_other_beam_zeta_bin_width_star_beamstrahlung = slicer.bin_widths_beamstrahlung / np.cos(phi), # boosted dz + # has to be set + slices_other_beam_Sigma_12 = n_slices*[0], + slices_other_beam_Sigma_34 = n_slices*[0], + ) + + ######################### + # track for 1 collision # + ######################### + + line = xt.Line(elements = [el_beambeam_b1]) + line.build_tracker(_context=test_context) + + assert line._needs_rng == False + + line.configure_radiation(model_beamstrahlung='quantum') + assert line._needs_rng == True + + record = line.start_internal_logging_for_elements_of_type( + xf.BeamBeamBiGaussian3D, capacity={"beamstrahlungtable": int(3e5)}) + line.track(particles_b1, num_turns=1) + line.stop_internal_logging_for_elements_of_type(xf.BeamBeamBiGaussian3D) + + record.move(_context=xo.context_default) + + ########################################### + # test 1: compare spectrum with guineapig # + ########################################### + + fname = (test_data_folder + / "beamstrahlung/guineapig_ttbar2_beamstrahlung_photon_energies_gev.txt") + guinea_photons = np.loadtxt(fname) # contains about 250k photons emitted from 1e6 macroparticles in 1 collision + n_bins = 10 + bins = np.logspace( + np.log10(1e-14), np.log10(1e1), n_bins) + xsuite_hist = np.histogram(record.beamstrahlungtable.photon_energy/1e9, + bins=bins)[0] + guinea_hist = np.histogram(guinea_photons, bins=bins)[0] + + bin_rel_errors = np.abs(xsuite_hist - guinea_hist) / guinea_hist + print(f"bin relative errors [1]: {bin_rel_errors}") + + # test if relative error in the last 5 bins is smaller than 1e-1 + assert np.allclose(xsuite_hist[-5:], guinea_hist[-5:], rtol=1e-1, atol=0) + + ############################################ + # test 2: compare beamstrahlung parameters # + ############################################ + + # average and maximum BS parameter + # https://www.researchgate.net/publication/2278298_Beam-Beam_Phenomena_In_Linear_Colliders + # page 20 + r0 = cst.e**2/(4*np.pi*cst.epsilon_0*cst.m_e*cst.c**2) # - if pp + + upsilon_max = ( + 2 * r0**2 * energy/(mass0*1e-9) * bunch_intensity + / (1/137*sigma_z_tot*(sigma_x + 1.85*sigma_y))) + upsilon_avg = (5/6 * r0**2 * energy/(mass0*1e-9) * bunch_intensity + / (1/137*sigma_z_tot*(sigma_x + sigma_y))) + + # get rid of padded zeros in table + photon_critical_energy = np.array( + sorted(set(record.beamstrahlungtable.photon_critical_energy))[1:]) + primary_energy = np.array( + sorted(set( record.beamstrahlungtable.primary_energy))[1:]) + + upsilon_avg_sim = np.mean(0.67 * photon_critical_energy / primary_energy) + upsilon_max_sim = np.max(0.67 * photon_critical_energy / primary_energy) + + print("Y max. [1]:", upsilon_max) + print("Y avg. [1]:", upsilon_avg) + print("Y max. [1]:", upsilon_max_sim) + print("Y avg. [1]:", upsilon_avg_sim) + print("Y max. ratio [1]:", upsilon_max_sim / upsilon_max) + print("Y avg. ratio [1]:", upsilon_avg_sim / upsilon_avg) + + assert np.abs(upsilon_max_sim / upsilon_max - 1) < 2e-2 + assert np.abs(upsilon_avg_sim / upsilon_avg - 1) < 5e-2 + + +@for_all_test_contexts(excluding='ContextPyopencl') +def test_beambeam3d_collective_beamstrahlung_single_collision(test_context): + ########### + # ttbar 2 # + ########### + bunch_intensity = 2.3e11 # [1] + energy = 182.5 # [GeV] + p0c = 182.5e9 # [eV] + mass0 = .511e6 # [eV] + phi = 15e-3 # [rad] half xing + u_sr = 9.2 # [GeV] + u_bs = .0114 # [GeV] + k2_factor = .4 # [1] + qx = .554 # [1] half arc + qy = .588 # [1] + qs = .0436 # [1] + physemit_x = 1.46e-09 # [m] + physemit_y = 2.9e-12 # [m] + beta_x = 1 # [m] + beta_y = .0016 # [m] + sigma_x = np.sqrt(physemit_x*beta_x) # [m] + sigma_px = np.sqrt(physemit_x/beta_x) # [m] + sigma_y = np.sqrt(physemit_y*beta_y) # [m] + sigma_py = np.sqrt(physemit_y/beta_y) # [m] + sigma_z = .00194 # [m] sr + sigma_z_tot = .00254 # [m] sr+bs + sigma_delta = .0015 # [m] + sigma_delta_tot = .00192 # [m] + beta_s = sigma_z/sigma_delta # [m] + physemit_s = sigma_z*sigma_delta # [m] + physemit_s_tot = sigma_z_tot*sigma_delta_tot # [m] + n_macroparticles_b1 = int(1e6) + n_macroparticles_b2 = int(1e6) + + n_slices = 100 + + ############# + # particles # + ############# + + #e- + particles_b1 = xp.Particles( + _context = test_context, + q0 = -1, + p0c = p0c, + mass0 = mass0, + x = sigma_x *np.random.randn(n_macroparticles_b1), + y = sigma_y *np.random.randn(n_macroparticles_b1), + zeta = sigma_z_tot *np.random.randn(n_macroparticles_b1), + px = sigma_px *np.random.randn(n_macroparticles_b1), + py = sigma_py *np.random.randn(n_macroparticles_b1), + delta = sigma_delta_tot*np.random.randn(n_macroparticles_b1), + ) + + particles_b1.name = "b1" + + slicer = xf.TempSlicer(n_slices=n_slices, sigma_z=sigma_z_tot, mode="unicharge") + + # this is different w.r.t WS test + config_for_update=xf.ConfigForUpdateBeamBeamBiGaussian3D( + pipeline_manager=None, + element_name="beambeam", + slicer=slicer, + update_every=None, # Never updates (test in weakstrong mode) ) - particles_b1.name = "b1" - - slicer = xf.TempSlicer(n_slices=n_slices, sigma_z=sigma_z_tot, mode="unicharge") - - # this is different w.r.t WS test - config_for_update=xf.ConfigForUpdateBeamBeamBiGaussian3D( - pipeline_manager=None, - element_name="beambeam", - slicer=slicer, - update_every=None, # Never updates (test in weakstrong mode) - ) - - el_beambeam_b1 = xf.BeamBeamBiGaussian3D( - _context=context, - config_for_update = config_for_update, - other_beam_q0=1, - phi=phi, - alpha=0, - # decide between round or elliptical kick formula - min_sigma_diff = 1e-28, - # slice intensity [num. real particles] n_slices inferred from length of this - slices_other_beam_num_particles = slicer.bin_weights * bunch_intensity, - # unboosted strong beam moments - slices_other_beam_zeta_center = slicer.bin_centers, - slices_other_beam_Sigma_11 = n_slices*[sigma_x**2], - slices_other_beam_Sigma_22 = n_slices*[sigma_px**2], - slices_other_beam_Sigma_33 = n_slices*[sigma_y**2], - slices_other_beam_Sigma_44 = n_slices*[sigma_py**2], - # only if BS on - slices_other_beam_zeta_bin_width_star_beamstrahlung = slicer.bin_widths_beamstrahlung / np.cos(phi), # boosted dz - # has to be set - slices_other_beam_Sigma_12 = n_slices*[0], - slices_other_beam_Sigma_34 = n_slices*[0], - ) - - el_beambeam_b1.name = "beambeam" - - ######################### - # track for 1 collision # - ######################### - - line = xt.Line(elements = [el_beambeam_b1]) - line.build_tracker(_context=context) - - assert line._needs_rng == False - - line.configure_radiation(model_beamstrahlung='quantum') - assert line._needs_rng == True - - record = line.start_internal_logging_for_elements_of_type(xf.BeamBeamBiGaussian3D, capacity={"beamstrahlungtable": int(3e5)}) - line.track(particles_b1, num_turns=1) - line.stop_internal_logging_for_elements_of_type(xf.BeamBeamBiGaussian3D) - - record.move(_context=xo.context_default) - - ########################################### - # test 1: compare spectrum with guineapig # - ########################################### - - fname = test_data_folder / "beamstrahlung/guineapig_ttbar2_beamstrahlung_photon_energies_gev.txt" - guinea_photons = np.loadtxt(fname) # contains about 250k photons emitted from 1e6 macroparticles in 1 collision - n_bins = 10 - bins = np.logspace(np.log10(1e-14), np.log10(1e1), n_bins) - xsuite_hist = np.histogram(record.beamstrahlungtable.photon_energy/1e9, bins=bins)[0] - guinea_hist = np.histogram(guinea_photons, bins=bins)[0] - - bin_rel_errors = np.abs(xsuite_hist - guinea_hist) / guinea_hist - print(f"bin relative errors [1]: {bin_rel_errors}") - - # test if relative error in the last 5 bins is smaller than 1e-1 - assert np.allclose(xsuite_hist[-5:], guinea_hist[-5:], rtol=1e-1, atol=0) - - ############################################ - # test 2: compare beamstrahlung parameters # - ############################################ - - # average and maximum BS parameter - # https://www.researchgate.net/publication/2278298_Beam-Beam_Phenomena_In_Linear_Colliders - # page 20 - r0 = cst.e**2/(4*np.pi*cst.epsilon_0*cst.m_e*cst.c**2) # - if pp - - upsilon_max = 2 * r0**2 * energy/(mass0*1e-9) * bunch_intensity / (1/137*sigma_z_tot*(sigma_x + 1.85*sigma_y)) - upsilon_avg = 5/6 * r0**2 * energy/(mass0*1e-9) * bunch_intensity / (1/137*sigma_z_tot*(sigma_x + sigma_y)) - - # get rid of padded zeros in table - photon_critical_energy = np.array(sorted(set(record.beamstrahlungtable.photon_critical_energy))[1:]) - primary_energy = np.array(sorted(set( record.beamstrahlungtable.primary_energy))[1:]) - - upsilon_avg_sim = np.mean(0.67 * photon_critical_energy / primary_energy) - upsilon_max_sim = np.max(0.67 * photon_critical_energy / primary_energy) - - print("Y max. [1]:", upsilon_max) - print("Y avg. [1]:", upsilon_avg) - print("Y max. [1]:", upsilon_max_sim) - print("Y avg. [1]:", upsilon_avg_sim) - print("Y max. ratio [1]:", upsilon_max_sim / upsilon_max) - print("Y avg. ratio [1]:", upsilon_avg_sim / upsilon_avg) - - assert np.abs(upsilon_max_sim / upsilon_max - 1) < 2e-2 - assert np.abs(upsilon_avg_sim / upsilon_avg - 1) < 5e-2 - - + el_beambeam_b1 = xf.BeamBeamBiGaussian3D( + _context=test_context, + config_for_update = config_for_update, + other_beam_q0=1, + phi=phi, + alpha=0, + # decide between round or elliptical kick formula + min_sigma_diff = 1e-28, + # slice intensity [num. real particles] n_slices inferred from length of this + slices_other_beam_num_particles = slicer.bin_weights * bunch_intensity, + # unboosted strong beam moments + slices_other_beam_zeta_center = slicer.bin_centers, + slices_other_beam_Sigma_11 = n_slices*[sigma_x**2], + slices_other_beam_Sigma_22 = n_slices*[sigma_px**2], + slices_other_beam_Sigma_33 = n_slices*[sigma_y**2], + slices_other_beam_Sigma_44 = n_slices*[sigma_py**2], + # only if BS on + slices_other_beam_zeta_bin_width_star_beamstrahlung = slicer.bin_widths_beamstrahlung / np.cos(phi), # boosted dz + # has to be set + slices_other_beam_Sigma_12 = n_slices*[0], + slices_other_beam_Sigma_34 = n_slices*[0], + ) + + el_beambeam_b1.name = "beambeam" + + ######################### + # track for 1 collision # + ######################### + + line = xt.Line(elements = [el_beambeam_b1]) + line.build_tracker(_context=test_context) + + assert line._needs_rng == False + + line.configure_radiation(model_beamstrahlung='quantum') + assert line._needs_rng == True + + record = line.start_internal_logging_for_elements_of_type(xf.BeamBeamBiGaussian3D, capacity={"beamstrahlungtable": int(3e5)}) + line.track(particles_b1, num_turns=1) + line.stop_internal_logging_for_elements_of_type(xf.BeamBeamBiGaussian3D) + + record.move(_context=xo.context_default) + + ########################################### + # test 1: compare spectrum with guineapig # + ########################################### + + fname = test_data_folder / "beamstrahlung/guineapig_ttbar2_beamstrahlung_photon_energies_gev.txt" + guinea_photons = np.loadtxt(fname) # contains about 250k photons emitted from 1e6 macroparticles in 1 collision + n_bins = 10 + bins = np.logspace(np.log10(1e-14), np.log10(1e1), n_bins) + xsuite_hist = np.histogram(record.beamstrahlungtable.photon_energy/1e9, bins=bins)[0] + guinea_hist = np.histogram(guinea_photons, bins=bins)[0] + + bin_rel_errors = np.abs(xsuite_hist - guinea_hist) / guinea_hist + print(f"bin relative errors [1]: {bin_rel_errors}") + + # test if relative error in the last 5 bins is smaller than 1e-1 + assert np.allclose(xsuite_hist[-5:], guinea_hist[-5:], rtol=1e-1, atol=0) + + ############################################ + # test 2: compare beamstrahlung parameters # + ############################################ + + # average and maximum BS parameter + # https://www.researchgate.net/publication/2278298_Beam-Beam_Phenomena_In_Linear_Colliders + # page 20 + r0 = cst.e**2/(4*np.pi*cst.epsilon_0*cst.m_e*cst.c**2) # - if pp + + upsilon_max = 2 * r0**2 * energy/(mass0*1e-9) * bunch_intensity / (1/137*sigma_z_tot*(sigma_x + 1.85*sigma_y)) + upsilon_avg = 5/6 * r0**2 * energy/(mass0*1e-9) * bunch_intensity / (1/137*sigma_z_tot*(sigma_x + sigma_y)) + + # get rid of padded zeros in table + photon_critical_energy = np.array(sorted(set(record.beamstrahlungtable.photon_critical_energy))[1:]) + primary_energy = np.array(sorted(set( record.beamstrahlungtable.primary_energy))[1:]) + + upsilon_avg_sim = np.mean(0.67 * photon_critical_energy / primary_energy) + upsilon_max_sim = np.max(0.67 * photon_critical_energy / primary_energy) + + print("Y max. [1]:", upsilon_max) + print("Y avg. [1]:", upsilon_avg) + print("Y max. [1]:", upsilon_max_sim) + print("Y avg. [1]:", upsilon_avg_sim) + print("Y max. ratio [1]:", upsilon_max_sim / upsilon_max) + print("Y avg. ratio [1]:", upsilon_avg_sim / upsilon_avg) + + assert np.abs(upsilon_max_sim / upsilon_max - 1) < 2e-2 + assert np.abs(upsilon_avg_sim / upsilon_avg - 1) < 5e-2 diff --git a/tests/test_cerrf.py b/tests/test_cerrf.py index 135c306a..93d8b2ab 100644 --- a/tests/test_cerrf.py +++ b/tests/test_cerrf.py @@ -9,6 +9,7 @@ import xobjects as xo from xobjects.context import available from xfields.general import _pkg_root +from xobjects.test_helpers import for_all_test_contexts @pytest.fixture @@ -80,12 +81,8 @@ def compute(self): return FaddeevaCalculator -@pytest.mark.parametrize( - 'context', - xo.context.get_test_contexts(), - ids=[str(ctx) for ctx in xo.context.get_test_contexts()], -) -def test_faddeeva_w_q1(faddeeva_calculator, context): +@for_all_test_contexts +def test_faddeeva_w_q1(faddeeva_calculator, test_context): FaddeevaCalculator = faddeeva_calculator # Generate the test grid @@ -100,7 +97,7 @@ def test_faddeeva_w_q1(faddeeva_calculator, context): # Calculate the values based on the grid z = (re_absc + 1j * im_absc).reshape(n_re * n_im) - calculator = FaddeevaCalculator(z=z, _context=context) + calculator = FaddeevaCalculator(z=z, _context=test_context) calculator.compute() # Using scipy's wofz implemenation of the Faddeeva method. This is @@ -127,20 +124,14 @@ def test_faddeeva_w_q1(faddeeva_calculator, context): assert d_abs_im.max() < 0.5e-9 -@pytest.mark.parametrize( - 'context', - xo.context.get_test_contexts(), - ids=[str(ctx) for ctx in xo.context.get_test_contexts()], -) -def test_faddeeva_w_all_quadrants(faddeeva_calculator, context): +@for_all_test_contexts +def test_faddeeva_w_all_quadrants(faddeeva_calculator, test_context): FaddeevaCalculator = faddeeva_calculator x0 = 5.33 y0 = 4.29 num_args = 10000 - ctx = context - re_max = np.float64(np.sqrt(2.0) * x0) im_max = np.float64(np.sqrt(2.0) * y0) @@ -163,7 +154,7 @@ def test_faddeeva_w_all_quadrants(faddeeva_calculator, context): z = re_absc + 1j * im_absc # Calculate the values based on the grid - calculator = FaddeevaCalculator(z=z, _context=context) + calculator = FaddeevaCalculator(z=z, _context=test_context) calculator.compute() # Create comparison data for veryfing the correctness of faddeeva_w(). diff --git a/tests/test_constant_charge_slicing.py b/tests/test_constant_charge_slicing.py index 79df1b8c..56e36f34 100644 --- a/tests/test_constant_charge_slicing.py +++ b/tests/test_constant_charge_slicing.py @@ -3,10 +3,10 @@ import xfields as xf - constant_charge_slicing_gaussian = ( xf.config_tools.beambeam_config_tools.config_tools.constant_charge_slicing_gaussian) + def test_constant_charge_slicing_gaussian(): sigma_z = 0.4 diff --git a/tests/test_electroncloud.py b/tests/test_electroncloud.py index e4e17705..0dbd79c7 100644 --- a/tests/test_electroncloud.py +++ b/tests/test_electroncloud.py @@ -9,96 +9,96 @@ import xpart as xp import xfields as xf +from xobjects.test_helpers import for_all_test_contexts -def test_tricubic_interpolation(): - for context in xo.context.get_test_contexts(): - print(f"Test {context.__class__}") - scale=0.05 - ff = lambda x, y, z: sum([ scale * x**i * y**j * z**k - for i in range(4) for j in range(4) for k in range(4)]) - dfdx = lambda x, y, z: sum([ i * scale * x**(i-1) * y**j * z**k - for i in range(1,4) for j in range(4) for k in range(4)]) - dfdy = lambda x, y, z: sum([ j * scale * x**i * y**(j-1) * z**k - for i in range(4) for j in range(1,4) for k in range(4)]) - dfdz = lambda x, y, z: sum([ k * scale * x**i * y**j * z**(k-1) - for i in range(4) for j in range(4) for k in range(1,4)]) - dfdxy = lambda x, y, z: sum([ i * j * scale * x**(i-1) * y**(j-1) * z**k - for i in range(1,4) for j in range(1,4) for k in range(4)]) - dfdxz = lambda x, y, z: sum([ i * k * scale * x**(i-1) * y**j * z**(k-1) - for i in range(1,4) for j in range(4) for k in range(1,4)]) - dfdyz = lambda x, y, z: sum([ j * k * scale * x**i * y**(j-1) * z**(k-1) - for i in range(4) for j in range(1,4) for k in range(1,4)]) - dfdxyz = lambda x, y, z: sum([ i * j * k * scale * x**(i-1) * y**(j-1) * z**(k-1) - for i in range(1,4) for j in range(1,4) for k in range(1,4)]) +@for_all_test_contexts +def test_tricubic_interpolation(test_context): + scale = 0.05 + ff = lambda x, y, z: sum([ scale * x**i * y**j * z**k + for i in range(4) for j in range(4) for k in range(4)]) + dfdx = lambda x, y, z: sum([ i * scale * x**(i-1) * y**j * z**k + for i in range(1,4) for j in range(4) for k in range(4)]) + dfdy = lambda x, y, z: sum([ j * scale * x**i * y**(j-1) * z**k + for i in range(4) for j in range(1,4) for k in range(4)]) + dfdz = lambda x, y, z: sum([ k * scale * x**i * y**j * z**(k-1) + for i in range(4) for j in range(4) for k in range(1,4)]) + dfdxy = lambda x, y, z: sum([ i * j * scale * x**(i-1) * y**(j-1) * z**k + for i in range(1,4) for j in range(1,4) for k in range(4)]) + dfdxz = lambda x, y, z: sum([ i * k * scale * x**(i-1) * y**j * z**(k-1) + for i in range(1,4) for j in range(4) for k in range(1,4)]) + dfdyz = lambda x, y, z: sum([ j * k * scale * x**i * y**(j-1) * z**(k-1) + for i in range(4) for j in range(1,4) for k in range(1,4)]) + dfdxyz = lambda x, y, z: sum([ i * j * k * scale * x**(i-1) * y**(j-1) * z**(k-1) + for i in range(1,4) for j in range(1,4) for k in range(1,4)]) - NN=21 - x_grid = np.linspace(-0.5, 0.5, NN) - y_grid = np.linspace(-0.5, 0.5, NN) - z_grid = np.linspace(-0.5, 0.5, NN) + NN=21 + x_grid = np.linspace(-0.5, 0.5, NN) + y_grid = np.linspace(-0.5, 0.5, NN) + z_grid = np.linspace(-0.5, 0.5, NN) - fieldmap = xf.TriCubicInterpolatedFieldMap(_context=context, - x_grid=x_grid, y_grid=y_grid, z_grid=z_grid) - ecloud = xf.ElectronCloud(length=1, fieldmap=fieldmap, - _buffer=fieldmap._buffer) + fieldmap = xf.TriCubicInterpolatedFieldMap(_context=test_context, + x_grid=x_grid, y_grid=y_grid, z_grid=z_grid) + ecloud = xf.ElectronCloud(length=1, fieldmap=fieldmap, + _buffer=fieldmap._buffer) - x0 = fieldmap.x_grid[0] - y0 = fieldmap.y_grid[0] - z0 = fieldmap.z_grid[0] - dx = fieldmap.dx - dy = fieldmap.dy - dz = fieldmap.dz - nx = fieldmap.nx - ny = fieldmap.ny - for ix in range(NN): - for iy in range(NN): - for iz in range(NN): - index = 0 + 8 * ix + 8 * nx * iy + 8 * nx * ny * iz - fieldmap._phi_taylor[index + 0] = ff(x_grid[ix], y_grid[iy], z_grid[iz]) - fieldmap._phi_taylor[index + 1] = dfdx(x_grid[ix], y_grid[iy], z_grid[iz]) * dx - fieldmap._phi_taylor[index + 2] = dfdy(x_grid[ix], y_grid[iy], z_grid[iz]) * dy - fieldmap._phi_taylor[index + 3] = dfdz(x_grid[ix], y_grid[iy], z_grid[iz]) * dz - fieldmap._phi_taylor[index + 4] = dfdxy(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dy - fieldmap._phi_taylor[index + 5] = dfdxz(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dz - fieldmap._phi_taylor[index + 6] = dfdyz(x_grid[ix], y_grid[iy], z_grid[iz]) * dy * dz - fieldmap._phi_taylor[index + 7] = dfdxyz(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dy * dz + x0 = fieldmap.x_grid[0] + y0 = fieldmap.y_grid[0] + z0 = fieldmap.z_grid[0] + dx = fieldmap.dx + dy = fieldmap.dy + dz = fieldmap.dz + nx = fieldmap.nx + ny = fieldmap.ny + for ix in range(NN): + for iy in range(NN): + for iz in range(NN): + index = 0 + 8 * ix + 8 * nx * iy + 8 * nx * ny * iz + fieldmap._phi_taylor[index + 0] = ff(x_grid[ix], y_grid[iy], z_grid[iz]) + fieldmap._phi_taylor[index + 1] = dfdx(x_grid[ix], y_grid[iy], z_grid[iz]) * dx + fieldmap._phi_taylor[index + 2] = dfdy(x_grid[ix], y_grid[iy], z_grid[iz]) * dy + fieldmap._phi_taylor[index + 3] = dfdz(x_grid[ix], y_grid[iy], z_grid[iz]) * dz + fieldmap._phi_taylor[index + 4] = dfdxy(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dy + fieldmap._phi_taylor[index + 5] = dfdxz(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dz + fieldmap._phi_taylor[index + 6] = dfdyz(x_grid[ix], y_grid[iy], z_grid[iz]) * dy * dz + fieldmap._phi_taylor[index + 7] = dfdxyz(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dy * dz - n_parts = 1000 - rng = default_rng(12345) - x_test = rng.random(n_parts) * 1.2 - 0.6 - y_test = rng.random(n_parts) * 1.2 - 0.6 - tau_test = rng.random(n_parts) * 1.2 - 0.6 + n_parts = 1000 + rng = default_rng(12345) + x_test = rng.random(n_parts) * 1.2 - 0.6 + y_test = rng.random(n_parts) * 1.2 - 0.6 + tau_test = rng.random(n_parts) * 1.2 - 0.6 - p0c = 450e9 - testp0 = xp.Particles(p0c=p0c) - beta0 = testp0.beta0 - part = xp.Particles(_context=context, x=x_test, y=y_test, - zeta=beta0*tau_test, p0c=p0c) - ecloud.track(part) + p0c = 450e9 + testp0 = xp.Particles(p0c=p0c) + beta0 = testp0.beta0 + part = xp.Particles(_context=test_context, x=x_test, y=y_test, + zeta=beta0*tau_test, p0c=p0c) + ecloud.track(part) - part.move(_context=xo.ContextCpu()) - mask_p = part.state != -11 - true_px = np.array([-dfdx(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p], - part.zeta[mask_p] / part.beta0[mask_p])]) - true_py = np.array([-dfdy(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p], - part.zeta[mask_p] / part.beta0[mask_p])]) - true_ptau = np.array([-dfdz(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p], - part.zeta[mask_p] / part.beta0[mask_p])]) + part.move(_context=xo.ContextCpu()) + mask_p = part.state != -11 + true_px = np.array([-dfdx(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p], + part.zeta[mask_p] / part.beta0[mask_p])]) + true_py = np.array([-dfdy(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p], + part.zeta[mask_p] / part.beta0[mask_p])]) + true_ptau = np.array([-dfdz(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p], + part.zeta[mask_p] / part.beta0[mask_p])]) - # print(true_px[:5]) - # print(part.ptau[:5]) - # print(part.state[:5]) - # print(f"px kick diff.: {np.mean(part.px[mask_p]-true_px):.2e} +- {np.std(part.px[mask_p] - true_px):.2e}") - # print(f"py kick diff.: {np.mean(part.py[mask_p]-true_py):.2e} +- {np.std(part.py[mask_p] - true_py):.2e}") - # print(f"ptau kick diff.: {np.mean(part.ptau[mask_p]-true_ptau):.2e} +- {np.std(part.ptau[mask_p] - true_ptau):.2e}") - # print(np.allclose(part.px[mask_p], true_px, atol=1.e-13, rtol=1.e-13)) - # print(np.allclose(part.py[mask_p], true_py, atol=1.e-13, rtol=1.e-13)) - # print(np.allclose(part.ptau[mask_p], true_ptau, atol=1.e-13, rtol=1.e-13)) - # print(np.max(np.abs(part.px[mask_p]- true_px))) - # print(np.max(np.abs(part.py[mask_p]- true_py))) - # print(np.max(np.abs(part.ptau[mask_p]- true_ptau))) + # print(true_px[:5]) + # print(part.ptau[:5]) + # print(part.state[:5]) + # print(f"px kick diff.: {np.mean(part.px[mask_p]-true_px):.2e} +- {np.std(part.px[mask_p] - true_px):.2e}") + # print(f"py kick diff.: {np.mean(part.py[mask_p]-true_py):.2e} +- {np.std(part.py[mask_p] - true_py):.2e}") + # print(f"ptau kick diff.: {np.mean(part.ptau[mask_p]-true_ptau):.2e} +- {np.std(part.ptau[mask_p] - true_ptau):.2e}") + # print(np.allclose(part.px[mask_p], true_px, atol=1.e-13, rtol=1.e-13)) + # print(np.allclose(part.py[mask_p], true_py, atol=1.e-13, rtol=1.e-13)) + # print(np.allclose(part.ptau[mask_p], true_ptau, atol=1.e-13, rtol=1.e-13)) + # print(np.max(np.abs(part.px[mask_p]- true_px))) + # print(np.max(np.abs(part.py[mask_p]- true_py))) + # print(np.max(np.abs(part.ptau[mask_p]- true_ptau))) - assert np.allclose(part.px[mask_p], true_px, atol=1.e-13, rtol=1.e-13) - assert np.allclose(part.py[mask_p], true_py, atol=1.e-13, rtol=1.e-13) - assert np.allclose(part.ptau[mask_p], true_ptau, atol=1.e-13, rtol=1.e-13) + assert np.allclose(part.px[mask_p], true_px, atol=1.e-13, rtol=1.e-13) + assert np.allclose(part.py[mask_p], true_py, atol=1.e-13, rtol=1.e-13) + assert np.allclose(part.ptau[mask_p], true_ptau, atol=1.e-13, rtol=1.e-13) diff --git a/tests/test_electronlens_interpolated.py b/tests/test_electronlens_interpolated.py index c25589c2..4215d413 100644 --- a/tests/test_electronlens_interpolated.py +++ b/tests/test_electronlens_interpolated.py @@ -3,73 +3,72 @@ # Copyright (c) CERN, 2021. # # ########################################### # - import numpy as np -import xobjects as xo import xpart as xp import xtrack as xt import xfields as xf +from xobjects.test_helpers import for_all_test_contexts + -def test_electronlens_interpolated(): - for context in xo.context.get_test_contexts(): - print(f"Test {context.__class__}") - outer_radius = 3.e-3 - inner_radius = 1.5e-3 +@for_all_test_contexts +def test_electronlens_interpolated(test_context): + outer_radius = 3.e-3 + inner_radius = 1.5e-3 - x_center = 0. - y_center = 0. + x_center = 0. + y_center = 0. - x_range = (-0.5e-2, 0.5e-2) - y_range = (-0.55e-2, 0.55e-2) - nx = 101 - ny = 151 - x_grid = np.linspace(x_range[0], x_range[1], nx) - y_grid = np.linspace(y_range[0], y_range[1], ny) - dx=x_grid[1] - x_grid[0] - dy=y_grid[1] - y_grid[0] + x_range = (-0.5e-2, 0.5e-2) + y_range = (-0.55e-2, 0.55e-2) + nx = 101 + ny = 151 + x_grid = np.linspace(x_range[0], x_range[1], nx) + y_grid = np.linspace(y_range[0], y_range[1], ny) + dx=x_grid[1] - x_grid[0] + dy=y_grid[1] - y_grid[0] - X, Y = np.meshgrid(x_grid, y_grid, indexing="ij") - rho = np.zeros_like(X) - rho[:] = 0 - R = np.sqrt( (X - x_center)**2 + (Y - y_center)**2) - rho[ (R > inner_radius) & (R < outer_radius) ] = 1. - #rho[ X < 0 ] = 0 - norm_rho = np.sum(rho[:,:])*dx*dy - rho[:] /= norm_rho + X, Y = np.meshgrid(x_grid, y_grid, indexing="ij") + rho = np.zeros_like(X) + rho[:] = 0 + R = np.sqrt( (X - x_center)**2 + (Y - y_center)**2) + rho[ (R > inner_radius) & (R < outer_radius) ] = 1. + #rho[ X < 0 ] = 0 + norm_rho = np.sum(rho[:,:])*dx*dy + rho[:] /= norm_rho - elens = xf.ElectronLensInterpolated(current=1, length=1, voltage=15e3, - x_grid=x_grid, y_grid=y_grid, rho=rho) + elens = xf.ElectronLensInterpolated(current=1, length=1, voltage=15e3, + x_grid=x_grid, y_grid=y_grid, rho=rho) - elens_ideal = xt.Elens(current=1, elens_length=1, voltage=15e3, - inner_radius=inner_radius, outer_radius=outer_radius) + elens_ideal = xt.Elens(current=1, elens_length=1, voltage=15e3, + inner_radius=inner_radius, outer_radius=outer_radius) - npart = 1000 - x_init = np.linspace(-0.42e-2, 0.42e-2, npart) - y_init = np.linspace(-0.42e-2, 0.42e-2, npart) - X_init, Y_init = np.meshgrid(x_init, y_init) - X_init = X_init.flatten() - Y_init = Y_init.flatten() + npart = 1000 + x_init = np.linspace(-0.42e-2, 0.42e-2, npart) + y_init = np.linspace(-0.42e-2, 0.42e-2, npart) + X_init, Y_init = np.meshgrid(x_init, y_init) + X_init = X_init.flatten() + Y_init = Y_init.flatten() - part = xp.Particles(x=X_init[:], y=Y_init[:], - zeta=[0], p0c=450e9 - ) + part = xp.Particles(x=X_init[:], y=Y_init[:], + zeta=[0], p0c=450e9 + ) - part_ideal = xp.Particles(x=X_init[:], y=Y_init[:], - zeta=[0], p0c=450e9 - ) + part_ideal = xp.Particles(x=X_init[:], y=Y_init[:], + zeta=[0], p0c=450e9 + ) - elens.track(part) - elens_ideal.track(part_ideal) + elens.track(part) + elens_ideal.track(part_ideal) - sort_mask = np.argsort(part.particle_id) - sort_ideal_mask = np.argsort(part_ideal.particle_id) - print(np.max(np.abs(part.px[sort_mask] - part_ideal.px[sort_ideal_mask]))) - print(np.max(np.abs(part.py[sort_mask] - part_ideal.py[sort_ideal_mask]))) - assert np.allclose(part.px[sort_mask], part_ideal.px[sort_ideal_mask], - atol=1.e-8, rtol=1.e-15) - assert np.allclose(part.py[sort_mask], part_ideal.py[sort_ideal_mask], - atol=1.e-8, rtol=1.e-15) - assert np.all(part.delta == 0.) - assert np.all(part.ptau == 0.) + sort_mask = np.argsort(part.particle_id) + sort_ideal_mask = np.argsort(part_ideal.particle_id) + print(np.max(np.abs(part.px[sort_mask] - part_ideal.px[sort_ideal_mask]))) + print(np.max(np.abs(part.py[sort_mask] - part_ideal.py[sort_ideal_mask]))) + assert np.allclose(part.px[sort_mask], part_ideal.px[sort_ideal_mask], + atol=1.e-8, rtol=1.e-15) + assert np.allclose(part.py[sort_mask], part_ideal.py[sort_ideal_mask], + atol=1.e-8, rtol=1.e-15) + assert np.all(part.delta == 0.) + assert np.all(part.ptau == 0.) diff --git a/tests/test_mean_std.py b/tests/test_mean_std.py index e23211c7..64955ddb 100644 --- a/tests/test_mean_std.py +++ b/tests/test_mean_std.py @@ -5,24 +5,23 @@ import numpy as np -import xobjects as xo import xfields as xf -def test_mean_and_std(): +from xobjects.test_helpers import for_all_test_contexts - for ctx in xo.context.get_test_contexts(): - print(f"Test {ctx.__class__}") - n_x=100 - a_host = np.array(np.random.rand(n_x)) - a_dev = ctx.nparray_to_context_array(a_host) +@for_all_test_contexts +def test_mean_and_std(test_context): + n_x = 100 + a_host = np.array(np.random.rand(n_x)) + a_dev = test_context.nparray_to_context_array(a_host) - mm, ss = xf.mean_and_std(a_dev) - assert np.isclose(mm, np.mean(a_host)) - assert np.isclose(ss, np.std(a_host)) + mm, ss = xf.mean_and_std(a_dev) + assert np.isclose(mm, np.mean(a_host)) + assert np.isclose(ss, np.std(a_host)) - weights_host = np.zeros_like(a_host)+.2 - weights_dev = ctx.nparray_to_context_array(weights_host) - mm, ss = xf.mean_and_std(a_dev, weights=weights_dev) - assert np.isclose(mm, np.mean(a_host)) - assert np.isclose(ss, np.std(a_host)) + weights_host = np.zeros_like(a_host)+.2 + weights_dev = test_context.nparray_to_context_array(weights_host) + mm, ss = xf.mean_and_std(a_dev, weights=weights_dev) + assert np.isclose(mm, np.mean(a_host)) + assert np.isclose(ss, np.std(a_host)) diff --git a/tests/test_profiles.py b/tests/test_profiles.py index 2ec1844e..9f020ae7 100644 --- a/tests/test_profiles.py +++ b/tests/test_profiles.py @@ -5,35 +5,32 @@ import numpy as np -import xobjects as xo - from xfields import LongitudinalProfileQGaussian - -def test_qgauss(): - for ctx in xo.context.get_test_contexts(): - print(repr(ctx)) - - z0 = 0.1 - sigma_z = 0.2 - npart = 1e11 - - for qq in [0, 0.5, 0.95, 1.05, 1.3]: - lprofile = LongitudinalProfileQGaussian( - _context=ctx, - number_of_particles=npart, - sigma_z=sigma_z, - z0=z0, - q_parameter=qq) - - z = np.linspace(-10., 10., 10000) - z_dev = ctx.nparray_to_context_array(z) - lden_dev = lprofile.line_density(z_dev) - lden = ctx.nparray_from_context_array(lden_dev) - - area = np.trapz(lden, z) - z_mean = np.trapz(lden*z/area, z) - z_std = np.sqrt(np.trapz(lden*(z-z_mean)**2/area, z)) - assert np.isclose(area, npart) - assert np.isclose(z_mean, z0) - assert np.isclose(z_std, sigma_z) - +from xobjects.test_helpers import for_all_test_contexts + + +@for_all_test_contexts +def test_qgauss(test_context): + z0 = 0.1 + sigma_z = 0.2 + npart = 1e11 + + for qq in [0, 0.5, 0.95, 1.05, 1.3]: + lprofile = LongitudinalProfileQGaussian( + _context=test_context, + number_of_particles=npart, + sigma_z=sigma_z, + z0=z0, + q_parameter=qq) + + z = np.linspace(-10., 10., 10000) + z_dev = test_context.nparray_to_context_array(z) + lden_dev = lprofile.line_density(z_dev) + lden = test_context.nparray_from_context_array(lden_dev) + + area = np.trapz(lden, z) + z_mean = np.trapz(lden*z/area, z) + z_std = np.sqrt(np.trapz(lden*(z-z_mean)**2/area, z)) + assert np.isclose(area, npart) + assert np.isclose(z_mean, z0) + assert np.isclose(z_std, sigma_z) diff --git a/tests/test_spacecharge.py b/tests/test_spacecharge.py index 9f1eb88b..cb0ac239 100644 --- a/tests/test_spacecharge.py +++ b/tests/test_spacecharge.py @@ -4,235 +4,228 @@ # ########################################### # import numpy as np +import pytest + from scipy.constants import m_p as pmass_kg from scipy.constants import e as qe from scipy.constants import c as clight -import xobjects as xo -import xtrack as xt import xpart as xp import ducktrack as dtk +from xobjects.test_helpers import for_all_test_contexts + pmass = pmass_kg*clight**2/qe -def test_spacecharge_gauss_qgauss(): - for frozen in [True, False]: - for context in xo.context.get_test_contexts(): - print(f"Test {context.__class__}") - - ################################# - # 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 = 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_dtk, r_probes, _, _, _) = generate_particles_object( - n_macroparticles, - bunch_intensity, - sigma_x, - sigma_y, - sigma_z, - p0c, - mass, - n_probes, - r_max_probes, - z_probes, - theta_probes) - particles = xp.Particles( - _context=context, **particles_dtk.to_dict()) - - particles.x += x0 - particles.y += y0 - - ################ - # Space charge # - ################ - - from xfields import LongitudinalProfileQGaussian - lprofile = LongitudinalProfileQGaussian( - _context=context, - number_of_particles=bunch_intensity, - sigma_z=sigma_z, - z0=0., - q_parameter=1. # there is a bug in ducktrack, - # only q=1 can be tested - ) - - 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 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]) - - scdtk = dtk.SCQGaussProfile( - number_of_particles = bunch_intensity, - bunchlength_rms=sigma_z, - sigma_x=sigma_x, - sigma_y=sigma_y, - length=scgauss.length, - q_parameter=scgauss.longitudinal_profile.q_parameter, - x_co=x0, - y_co=y0) - - p_dtk = dtk.TestParticles(p0c=p0c, - mass=mass, - x=x_probes.copy(), - y=y_probes.copy(), - zeta=z_probes.copy()) - - scdtk.track(p_dtk) - - assert np.allclose( - p2np(particles.px[:n_probes]), - p_dtk.px, - atol={True:1e-7, False:1e2}[frozen] - * np.max(np.abs(p_dtk.px))) - assert np.allclose( - p2np(particles.py[:n_probes]), - p_dtk.py, - atol={True:1e-7, False:1e2}[frozen] - * np.max(np.abs(p_dtk.py))) - - -def test_spacecharge_pic(): - for solver in ['FFTSolver2p5D', 'FFTSolver3D']: - for context in xo.context.get_test_contexts(): - print(f"Test {context.__class__}") - - print(repr(context)) - - ################################# - # Generate particles and probes # - ################################# - - n_macroparticles = int(5e6) - bunch_intensity = 2.5e11 - sigma_x = 3e-3 - sigma_y = 2e-3 - sigma_z = 30e-2 - p0c = 25.92e9 - mass = 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_gen, r_probes, x_probes, - y_probes, z_probes) = generate_particles_object( - n_macroparticles, - bunch_intensity, - sigma_x, - sigma_y, - sigma_z, - p0c, - mass, - n_probes, - r_max_probes, - z_probes, - theta_probes) - # Transfer particles to context - particles = xp.Particles( - _context=context, **particles_gen.to_dict()) - - ###################### - # Space charge (PIC) # - ###################### - - x_lim = 4.*sigma_x - y_lim = 4.*sigma_y - z_lim = 4.*sigma_z - - from xfields import SpaceCharge3D - - spcharge = SpaceCharge3D( - _context=context, - length=1, update_on_track=True, apply_z_kick=False, - x_range=(-x_lim, x_lim), - y_range=(-y_lim, y_lim), - z_range=(-z_lim, z_lim), - nx=128, ny=128, nz=25, - solver=solver, - gamma0=particles_gen.gamma0[0], - ) - - spcharge.track(particles) - - ############################# - # Compare against ducktrack # - ############################# - - p2np = context.nparray_from_context_array - - scdtk = dtk.SCQGaussProfile( - number_of_particles = bunch_intensity, - bunchlength_rms=sigma_z, - sigma_x=sigma_x, - sigma_y=sigma_y, - length=spcharge.length, - x_co=0., - y_co=0.) - - p_dtk = dtk.TestParticles(p0c=p0c, - mass=mass, - x=x_probes.copy(), - y=y_probes.copy(), - zeta=z_probes.copy()) - - 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_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_dtk.py[mask_inside_grid], - atol=3e-2*np.max(np.abs(p_dtk.py[mask_inside_grid]))) + +@pytest.mark.parametrize('frozen', [True, False]) +@for_all_test_contexts +def test_spacecharge_gauss_qgauss(frozen, test_context): + 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 = 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_dtk, r_probes, _, _, _) = generate_particles_object( + n_macroparticles, + bunch_intensity, + sigma_x, + sigma_y, + sigma_z, + p0c, + mass, + n_probes, + r_max_probes, + z_probes, + theta_probes) + particles = xp.Particles( + _context=test_context, **particles_dtk.to_dict()) + + particles.x += x0 + particles.y += y0 + + ################ + # Space charge # + ################ + + from xfields import LongitudinalProfileQGaussian + lprofile = LongitudinalProfileQGaussian( + _context=test_context, + number_of_particles=bunch_intensity, + sigma_z=sigma_z, + z0=0., + q_parameter=1. # there is a bug in ducktrack, + # only q=1 can be tested + ) + + 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=test_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 ducktrack # + ############################# + + p2np = test_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]) + + scdtk = dtk.SCQGaussProfile( + number_of_particles = bunch_intensity, + bunchlength_rms=sigma_z, + sigma_x=sigma_x, + sigma_y=sigma_y, + length=scgauss.length, + q_parameter=scgauss.longitudinal_profile.q_parameter, + x_co=x0, + y_co=y0) + + p_dtk = dtk.TestParticles(p0c=p0c, + mass=mass, + x=x_probes.copy(), + y=y_probes.copy(), + zeta=z_probes.copy()) + + scdtk.track(p_dtk) + + assert np.allclose( + p2np(particles.px[:n_probes]), + p_dtk.px, + atol={True:1e-7, False:1e2}[frozen] + * np.max(np.abs(p_dtk.px))) + assert np.allclose( + p2np(particles.py[:n_probes]), + p_dtk.py, + atol={True:1e-7, False:1e2}[frozen] + * np.max(np.abs(p_dtk.py))) + + +@pytest.mark.parametrize('solver', ['FFTSolver2p5D', 'FFTSolver3D']) +@for_all_test_contexts +def test_spacecharge_pic(solver, test_context): + ################################# + # Generate particles and probes # + ################################# + + n_macroparticles = int(5e6) + bunch_intensity = 2.5e11 + sigma_x = 3e-3 + sigma_y = 2e-3 + sigma_z = 30e-2 + p0c = 25.92e9 + mass = 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_gen, r_probes, x_probes, + y_probes, z_probes) = generate_particles_object( + n_macroparticles, + bunch_intensity, + sigma_x, + sigma_y, + sigma_z, + p0c, + mass, + n_probes, + r_max_probes, + z_probes, + theta_probes) + # Transfer particles to context + particles = xp.Particles( + _context=test_context, **particles_gen.to_dict()) + + ###################### + # Space charge (PIC) # + ###################### + + x_lim = 4.*sigma_x + y_lim = 4.*sigma_y + z_lim = 4.*sigma_z + + from xfields import SpaceCharge3D + + spcharge = SpaceCharge3D( + _context=test_context, + length=1, update_on_track=True, apply_z_kick=False, + x_range=(-x_lim, x_lim), + y_range=(-y_lim, y_lim), + z_range=(-z_lim, z_lim), + nx=128, ny=128, nz=25, + solver=solver, + gamma0=particles_gen.gamma0[0], + ) + + spcharge.track(particles) + + ############################# + # Compare against ducktrack # + ############################# + + p2np = test_context.nparray_from_context_array + + scdtk = dtk.SCQGaussProfile( + number_of_particles = bunch_intensity, + bunchlength_rms=sigma_z, + sigma_x=sigma_x, + sigma_y=sigma_y, + length=spcharge.length, + x_co=0., + y_co=0.) + + p_dtk = dtk.TestParticles(p0c=p0c, + mass=mass, + x=x_probes.copy(), + y=y_probes.copy(), + zeta=z_probes.copy()) + + 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_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_dtk.py[mask_inside_grid], + atol=3e-2*np.max(np.abs(p_dtk.py[mask_inside_grid])))