-
Notifications
You must be signed in to change notification settings - Fork 14
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #69 from AngelFP/new_pipeline
Implement new tracking pipeline
- Loading branch information
Showing
45 changed files
with
3,476 additions
and
1,522 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,84 @@ | ||
| from pytest import approx | ||
| import numpy as np | ||
| import matplotlib.pyplot as plt | ||
| from aptools.plotting.quick_diagnostics import slice_analysis | ||
|
|
||
| from wake_t import PlasmaStage, GaussianPulse | ||
| from wake_t.utilities.bunch_generation import get_matched_bunch | ||
| from wake_t.diagnostics import analyze_bunch, analyze_bunch_list | ||
|
|
||
|
|
||
| def test_custom_blowout_wakefield(make_plots=False): | ||
| """Check that the `'custom_blowout'` wakefield model works as expected. | ||
| A matched bunch with high emittance is propagated through a 3 cm plasma | ||
| stage. Due to the high emittance and Ez slope, the beam develops not only | ||
| a large chirp but also an uncorrelated energy spread. The test checks that | ||
| the final projected energy spread is always the same. | ||
| """ | ||
| # Set numpy random seed to get reproducible results | ||
| np.random.seed(1) | ||
|
|
||
| # Create laser driver. | ||
| laser = GaussianPulse(100e-6, l_0=800e-9, w_0=50e-6, a_0=3, | ||
| tau=30e-15, z_foc=0.) | ||
|
|
||
| # Create bunch (matched to a blowout at a density of 10^{23} m^{-3}). | ||
| en = 10e-6 # m | ||
| ene = 200 # units of beta*gamma | ||
| ene_sp = 0.3 # % | ||
| xi_c = laser.xi_c - 55e-6 # m | ||
| s_t = 10 # fs | ||
| q_tot = 100 # pC | ||
| n_part = 3e4 | ||
| k_x = 1e6 | ||
| bunch = get_matched_bunch(en, en, ene, ene_sp, s_t, xi_c, q_tot, n_part, | ||
| k_x=k_x) | ||
|
|
||
| # Create plasma stage. | ||
| plasma = PlasmaStage( | ||
| 3e-2, 1e23, laser=laser, wakefield_model='custom_blowout', | ||
| lon_field=-10e9, lon_field_slope=1e15, foc_strength=k_x, | ||
| xi_fields=xi_c, n_out=50, bunch_pusher='boris') | ||
|
|
||
| # Do tracking. | ||
| bunch_list = plasma.track(bunch) | ||
|
|
||
| bunch_params = analyze_bunch(bunch) | ||
| rel_ene_sp = bunch_params['rel_ene_spread'] | ||
| assert approx(rel_ene_sp, rel=1e-10) == 0.21192531095086517 | ||
|
|
||
| if make_plots: | ||
| # Analyze bunch evolution. | ||
| params_evolution = analyze_bunch_list(bunch_list) | ||
|
|
||
|
|
||
| # Quick plot of results. | ||
| z = params_evolution['prop_dist'] * 1e2 | ||
| fig_1 = plt.figure() | ||
| plt.subplot(411) | ||
| plt.plot(z, params_evolution['beta_x']*1e3) | ||
| plt.tick_params(axis='x', which='both', labelbottom=False) | ||
| plt.ylabel("$\\beta_x$ [mm]") | ||
| plt.subplot(412) | ||
| plt.plot(z, params_evolution['emitt_x']*1e6) | ||
| plt.tick_params(axis='x', which='both', labelbottom=False) | ||
| plt.ylabel("$\\epsilon_{nx}$ [$\\mu$m]") | ||
| plt.subplot(413) | ||
| plt.plot(z, params_evolution['rel_ene_spread']*100) | ||
| plt.tick_params(axis='x', which='both', labelbottom=False) | ||
| plt.ylabel("$\\frac{\\Delta \\gamma}{\\gamma}$ [%]") | ||
| plt.subplot(414) | ||
| plt.plot(z, params_evolution['avg_ene']) | ||
| plt.xlabel("z [mm]") | ||
| plt.ylabel("$\\gamma$") | ||
| plt.tight_layout() | ||
| fig_2 = plt.figure() | ||
| slice_analysis( | ||
| bunch.x, bunch.y, bunch.xi, bunch.px, bunch.py, bunch.pz, | ||
| bunch.q, fig=fig_2) | ||
| plt.show() | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| test_custom_blowout_wakefield(make_plots=True) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| from copy import deepcopy | ||
|
|
||
| import numpy as np | ||
| import scipy.constants as ct | ||
|
|
||
| from wake_t import GaussianPulse, PlasmaStage, Beamline | ||
| from wake_t.utilities.bunch_generation import get_matched_bunch | ||
|
|
||
|
|
||
| def test_variable_parabolic_coefficient(): | ||
| """ | ||
| Checks that a z-dependent parabolic coefficient works as expected. | ||
| A piecewise function for the coefficient is defined, which should be | ||
| equivalent to tracking multiple plasma stages each having a different | ||
| parabolic coefficient. | ||
| """ | ||
| # Plasma density. | ||
| n_p = 1e23 | ||
|
|
||
| # Grid and field parameters. | ||
| xi_max = 0 | ||
| xi_min = -100e-6 | ||
| r_max = 100e-6 | ||
| Nxi = 100 | ||
| Nr = 100 | ||
| dz_fields = 1e-3 | ||
|
|
||
| # Laser parameters. | ||
| a0 = 1 | ||
| w0 = 30e-6 | ||
| tau = 25e-15 | ||
| l0 = 0.8e-6 | ||
|
|
||
| # Guiding plasma channel | ||
| r_e = ct.e**2 / (4. * np.pi * ct.epsilon_0 * ct.m_e * ct.c**2) | ||
| rel_delta_n_over_w2 = 1. / (np.pi * r_e * w0**4 * n_p) | ||
|
|
||
| # Length and parabolic coefficient of each section (stretch). | ||
| L_stretches = [1e-2, 1e-2, 1e-2, 1e-2, 1e-2] | ||
| pc_stretches = rel_delta_n_over_w2 * np.array([1, 2, 1, 0.5, 2]) | ||
|
|
||
| # Define z-dependent parabolic coefficient. | ||
| def parabolic_coefficient(z): | ||
| z_stretch_end = np.cumsum(L_stretches) | ||
| i = np.sum(np.float32(z_stretch_end) <= np.float32(z)) | ||
| if i == len(L_stretches): | ||
| i -= 1 | ||
| return pc_stretches[i] | ||
|
|
||
| # Create identical laser pulses for each case. | ||
| laser = GaussianPulse(-50e-6, a0, w0, tau, z_foc=0., l_0=l0) | ||
| laser_1 = deepcopy(laser) | ||
| laser_2 = deepcopy(laser) | ||
|
|
||
| # Create identical bunches for each case. | ||
| bunch = get_matched_bunch( | ||
| 1e-6, 1e-6, 200, 1, 3, laser.xi_c - 30e-6, 1e-6, 1e4, n_p=n_p) | ||
| bunch_1 = deepcopy(bunch) | ||
| bunch_2 = deepcopy(bunch) | ||
|
|
||
| # Create single plasma stage (containing all sections). | ||
| plasma_single = PlasmaStage( | ||
| np.sum(L_stretches), n_p, wakefield_model='quasistatic_2d', n_out=10, | ||
| xi_min=xi_min, xi_max=xi_max, r_max=r_max, r_max_plasma=r_max, | ||
| laser=laser_1, laser_evolution=True, n_r=Nr, n_xi=Nxi, ppc=2, | ||
| dz_fields=dz_fields, parabolic_coefficient=parabolic_coefficient) | ||
|
|
||
| # Track single plasma. | ||
| plasma_single.track(bunch_1) | ||
|
|
||
| # Create set of plasma stages, one per section. | ||
| sub_stages = [] | ||
| for i, (l, pc) in enumerate(zip(L_stretches, pc_stretches)): | ||
| stage = PlasmaStage( | ||
| l, n_p, wakefield_model='quasistatic_2d', n_out=3, | ||
| xi_min=xi_min, xi_max=xi_max, r_max=r_max, r_max_plasma=r_max, | ||
| laser=laser_2, laser_evolution=True, n_r=Nr, n_xi=Nxi, ppc=2, | ||
| dz_fields=dz_fields, parabolic_coefficient=pc) | ||
| sub_stages.append(stage) | ||
| plasma_multi = Beamline(sub_stages) | ||
|
|
||
| # Track through set of stages. | ||
| plasma_multi.track(bunch_2) | ||
|
|
||
| # Get final envelope of both lasers. | ||
| a_env_1 = laser_1.get_envelope() | ||
| a_mod_1 = np.abs(a_env_1) | ||
| a_phase_1 = np.angle(a_env_1) | ||
| a_env_2 = laser_2.get_envelope() | ||
| a_mod_2 = np.abs(a_env_2) | ||
| a_phase_2 = np.angle(a_env_2) | ||
|
|
||
| # Check that both envelopes are equal. | ||
| np.testing.assert_almost_equal(a_mod_2, a_mod_1) | ||
| np.testing.assert_almost_equal(a_phase_2, a_phase_1) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| test_variable_parabolic_coefficient() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,83 @@ | ||
| from pytest import approx | ||
| import numpy as np | ||
| import matplotlib.pyplot as plt | ||
| from aptools.plotting.quick_diagnostics import slice_analysis | ||
|
|
||
| from wake_t import PlasmaStage, GaussianPulse | ||
| from wake_t.utilities.bunch_generation import get_matched_bunch | ||
| from wake_t.diagnostics import analyze_bunch, analyze_bunch_list | ||
|
|
||
|
|
||
| def test_simple_blowout_wakefield(make_plots=False): | ||
| """Check that the `'simple_blowout'` wakefield model works as expected. | ||
| A matched bunch with high emittance is propagated through a 3 cm plasma | ||
| stage. Due to the high emittance and Ez slope, the beam develops not only | ||
| a large chirp but also an uncorrelated energy spread. The test checks that | ||
| the final projected energy spread is always the same. | ||
| """ | ||
| # Set numpy random seed to get reproducible results | ||
| np.random.seed(1) | ||
|
|
||
| # Create laser driver. | ||
| laser = GaussianPulse(100e-6, l_0=800e-9, w_0=50e-6, a_0=3, | ||
| tau=30e-15, z_foc=0.) | ||
|
|
||
|
|
||
| # Create bunch (matched to a blowout at a density of 10^{23} m^{-3}). | ||
| en = 10e-6 # m | ||
| ene = 200 # units of beta*gamma | ||
| ene_sp = 0.3 # % | ||
| xi_c = laser.xi_c - 55e-6 # m | ||
| s_t = 10 # fs | ||
| q_tot = 100 # pC | ||
| n_part = 3e4 | ||
| bunch = get_matched_bunch(en, en, ene, ene_sp, s_t, xi_c, q_tot, n_part, | ||
| n_p=1e23) | ||
|
|
||
| # Create plasma stage. | ||
| plasma = PlasmaStage( | ||
| 3e-2, 1e23, laser=laser, wakefield_model='simple_blowout', n_out=50, | ||
| bunch_pusher='boris') | ||
|
|
||
| # Do tracking. | ||
| bunch_list = plasma.track(bunch) | ||
|
|
||
| bunch_params = analyze_bunch(bunch) | ||
| rel_ene_sp = bunch_params['rel_ene_spread'] | ||
| assert approx(rel_ene_sp, rel=1e-10) == 0.36376504595288617 | ||
|
|
||
| if make_plots: | ||
| # Analyze bunch evolution. | ||
| params_evolution = analyze_bunch_list(bunch_list) | ||
|
|
||
|
|
||
| # Quick plot of results. | ||
| z = params_evolution['prop_dist'] * 1e2 | ||
| fig_1 = plt.figure() | ||
| plt.subplot(411) | ||
| plt.plot(z, params_evolution['beta_x']*1e3) | ||
| plt.tick_params(axis='x', which='both', labelbottom=False) | ||
| plt.ylabel("$\\beta_x$ [mm]") | ||
| plt.subplot(412) | ||
| plt.plot(z, params_evolution['emitt_x']*1e6) | ||
| plt.tick_params(axis='x', which='both', labelbottom=False) | ||
| plt.ylabel("$\\epsilon_{nx}$ [$\\mu$m]") | ||
| plt.subplot(413) | ||
| plt.plot(z, params_evolution['rel_ene_spread']*100) | ||
| plt.tick_params(axis='x', which='both', labelbottom=False) | ||
| plt.ylabel("$\\frac{\\Delta \\gamma}{\\gamma}$ [%]") | ||
| plt.subplot(414) | ||
| plt.plot(z, params_evolution['avg_ene']) | ||
| plt.xlabel("z [mm]") | ||
| plt.ylabel("$\\gamma$") | ||
| plt.tight_layout() | ||
| fig_2 = plt.figure() | ||
| slice_analysis( | ||
| bunch.x, bunch.y, bunch.xi, bunch.px, bunch.py, bunch.pz, | ||
| bunch.q, fig=fig_2) | ||
| plt.show() | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| test_simple_blowout_wakefield(make_plots=True) |
Oops, something went wrong.