From e11d00a64d7ef25cdc0a345c17424e3b9668793c Mon Sep 17 00:00:00 2001 From: Alex Bojanich Date: Wed, 4 Sep 2024 22:25:35 -0700 Subject: [PATCH 1/5] Add Python CSR wake plotting test. --- docs/source/index.rst | 21 ++++-- tests/python/test_wake.py | 145 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 161 insertions(+), 5 deletions(-) create mode 100644 tests/python/test_wake.py diff --git a/docs/source/index.rst b/docs/source/index.rst index 6ab1b75e0..967641494 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -3,23 +3,34 @@ ImpactX ------- -ImpactX is an s-based beam dynamics code including space charge effects. +ImpactX enables high-performance modeling of beam dynamics in particle accelerators with collective effects. + This is the next generation of the `IMPACT-Z `__ code. +ImpactX runs on modern GPUs or CPUs alike, provides user-friendly interfaces suitable for AI/ML workflows, has many :ref:`benchmarks ` to ensure its correctness, and an extensive documentation. + +As a beam dynamics code, ImpactX uses the reference trajectory path length :math:`s` as the independent variable of motion to achieve large speedups. +The code includes the effects of externally applied fields from magnets and accelerating cavities as well as the effect of self-fields (space charge fields, CSR, wakefields, ...). +All particle tracking models are symplectic, and space charge is included by solving the Poisson equation in the beam rest frame. +The code may be used to model the dynamics of beams in both linear and ring accelerators. +See our :ref:`theory chapter ` for details on our models, assumptions and concepts. + +ImpactX is part of the `Beam, Plasma & Accelerator Simulation Toolkit (BLAST) `__. +Please see this page for more information to select the code most appropriate for your application. -Please contact us with any questions on it or if you like to contribute to its development. .. _contact: Contact us ^^^^^^^^^^ -If you are starting using ImpactX, or if you have a user question, please pop in our `discussions page `__ and get in touch with the community. +We organize support as an open community, helping each other. -The `ImpactX GitHub repo `__ is the main communication platform. +The `ImpactX GitHub repository `__ is our open communication platform. Have a look at the action icons on the top right of the web page: feel free to watch the repo if you want to receive updates, or to star the repo to support the project. For bug reports or to request new features, you can also open a new `issue `__. -We also have a `discussion page `__ on which you can find already answered questions, add new questions, get help with installation procedures, discuss ideas or share comments. +If you are starting to use ImpactX and have questions not answered by this documentation, have a look on our `discussion page `__. +There, you can find already answered questions, add your own, get help with installation procedures, help others and share ideas. .. raw:: html diff --git a/tests/python/test_wake.py b/tests/python/test_wake.py new file mode 100644 index 000000000..b0fe86be3 --- /dev/null +++ b/tests/python/test_wake.py @@ -0,0 +1,145 @@ +import sys +import matplotlib.pyplot as plt +import numpy as np + +np.set_printoptions(threshold=sys.maxsize) +from conftest import basepath +from impactx import ImpactX, amr, wakeconvolution +from amrex.space3d import PODVector_real_std # Import amrex array + +def test_wake(save_png=True): + amr.initialize([]) + try: + sim = ImpactX() + sim.n_cell = [16, 24, 32] + sim.load_inputs_file(basepath + "/../../examples/chicane/input_chicane_csr.in") + sim.slice_step_diagnostics = False + + sim.init_grids() + sim.init_beam_distribution_from_inputs() + sim.init_lattice_elements_from_inputs() + sim.evolve() + + sim.deposit_charge() + + element_has_csr = True + R = 10.35 # Units [m] + sigma_t = 1.9975134930563207e-05 + + if element_has_csr: + pc = sim.particle_container() + x_min, y_min, t_min, x_max, y_max, t_max = pc.min_and_max_positions() + + is_unity_particle_weight = False + GetNumberDensity = True + + num_bins = 150 + bin_min = t_min + bin_max = t_max + bin_size = (bin_max - bin_min) / (num_bins - 1) + print("t_min, t_max, sigma_t = ") + print(t_min,t_max,sigma_t) + + # Create charge_distribution and slopes as PODVector_real_std + charge_distribution = PODVector_real_std(num_bins + 1) + slopes = PODVector_real_std(num_bins) + + # Call deposit_charge with the correct type + wakeconvolution.deposit_charge( + pc, + charge_distribution, + bin_min, + bin_size, + is_unity_particle_weight, + ) + + # Call derivative_charge with the correct types + wakeconvolution.derivative_charge( + charge_distribution, + slopes, + bin_size, + GetNumberDensity, + ) + + # Convert charge distribution to numpy array and plot it + charge_distribution_np = charge_distribution.to_numpy() + charge_distribution_s_values = np.linspace(bin_min, bin_max, len(charge_distribution_np)) + plt.figure() + print("Length density s values = ") + print(len(charge_distribution_np)) + print("Length density values = ") + print(len(charge_distribution_np)) + plt.plot(charge_distribution_s_values, charge_distribution_np) + plt.xlabel("Longitudinal Position s (m)") + plt.ylabel("Charge density") + plt.title("Charge density vs Longitudinal Position") + if save_png: + plt.savefig("density.png") + else: + plt.show() + + # Convert slopes to numpy array and plot it + slopes_np = slopes.to_numpy() + slopes_s_values = np.linspace(bin_min, bin_max, len(slopes_np)) + plt.figure() + plt.plot(slopes_s_values, slopes_np) + plt.xlabel("Longitudinal Position s (m)") + plt.ylabel("Slopes") + plt.title("Slopes vs Longitudinal Position") + if save_png: + plt.savefig("slopes.png") + else: + plt.show() + + # Create wake_function as PODVector_real_std with size 2 * len(slopes) + wake_function = PODVector_real_std(2 * len(slopes)) + wake_s_values = np.linspace(bin_min, bin_max, 2 * len(slopes)) + for i in range(len(wake_function)): + if (i < num_bins): + s = (i * bin_size) + else: + s = (i-2*num_bins)*bin_size + wake_s_values[i] = s + wake_function[i] = wakeconvolution.w_l_csr(s, R, bin_size) + + + # Convert wake_function to numpy array and plot it + wake_function_np = wake_function.to_numpy() + plt.figure() + plt.plot(wake_s_values, wake_function_np) + plt.xlabel("Longitudinal Position s (m)") + plt.ylabel("Wake Function") + plt.title("Wake Function vs Longitudinal Position") + if save_png: + plt.savefig("wake_function.png") + else: + plt.show() + + # Call convolve_fft with the correct types and capture the output + convolved_wakefield = wakeconvolution.convolve_fft( + slopes, + wake_function, + bin_size + ) + + # Convert the result to numpy array + convolved_wakefield_np = convolved_wakefield.to_numpy() + + # Adjust the s_values to match the length of convolved_wakefield_np + s_values = np.linspace(bin_min, bin_max, len(convolved_wakefield_np)) + normalized_s_values = s_values / sigma_t + + plt.plot(normalized_s_values, convolved_wakefield_np) + plt.xlabel("Longitudinal Position s/sigma_s") + plt.ylabel("Wakefield (V C/m)") + plt.title("Convolved CSR Wakefield vs Longitudinal Position") + if save_png: + plt.savefig("convolved_wakefield.png") + else: + plt.show() + plt.close("all") + finally: + sim.finalize() + +if __name__ == "__main__": + test_wake() From ce038bda5840a0274f65069744b11ce8fb070c5c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 5 Sep 2024 22:20:57 +0000 Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/python/test_wake.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/tests/python/test_wake.py b/tests/python/test_wake.py index b0fe86be3..8474cd5ae 100644 --- a/tests/python/test_wake.py +++ b/tests/python/test_wake.py @@ -1,11 +1,14 @@ import sys + import matplotlib.pyplot as plt import numpy as np np.set_printoptions(threshold=sys.maxsize) from conftest import basepath -from impactx import ImpactX, amr, wakeconvolution + from amrex.space3d import PODVector_real_std # Import amrex array +from impactx import ImpactX, amr, wakeconvolution + def test_wake(save_png=True): amr.initialize([]) @@ -38,7 +41,7 @@ def test_wake(save_png=True): bin_max = t_max bin_size = (bin_max - bin_min) / (num_bins - 1) print("t_min, t_max, sigma_t = ") - print(t_min,t_max,sigma_t) + print(t_min, t_max, sigma_t) # Create charge_distribution and slopes as PODVector_real_std charge_distribution = PODVector_real_std(num_bins + 1) @@ -63,7 +66,9 @@ def test_wake(save_png=True): # Convert charge distribution to numpy array and plot it charge_distribution_np = charge_distribution.to_numpy() - charge_distribution_s_values = np.linspace(bin_min, bin_max, len(charge_distribution_np)) + charge_distribution_s_values = np.linspace( + bin_min, bin_max, len(charge_distribution_np) + ) plt.figure() print("Length density s values = ") print(len(charge_distribution_np)) @@ -82,7 +87,7 @@ def test_wake(save_png=True): slopes_np = slopes.to_numpy() slopes_s_values = np.linspace(bin_min, bin_max, len(slopes_np)) plt.figure() - plt.plot(slopes_s_values, slopes_np) + plt.plot(slopes_s_values, slopes_np) plt.xlabel("Longitudinal Position s (m)") plt.ylabel("Slopes") plt.title("Slopes vs Longitudinal Position") @@ -95,14 +100,13 @@ def test_wake(save_png=True): wake_function = PODVector_real_std(2 * len(slopes)) wake_s_values = np.linspace(bin_min, bin_max, 2 * len(slopes)) for i in range(len(wake_function)): - if (i < num_bins): - s = (i * bin_size) + if i < num_bins: + s = i * bin_size else: - s = (i-2*num_bins)*bin_size + s = (i - 2 * num_bins) * bin_size wake_s_values[i] = s wake_function[i] = wakeconvolution.w_l_csr(s, R, bin_size) - # Convert wake_function to numpy array and plot it wake_function_np = wake_function.to_numpy() plt.figure() @@ -117,9 +121,7 @@ def test_wake(save_png=True): # Call convolve_fft with the correct types and capture the output convolved_wakefield = wakeconvolution.convolve_fft( - slopes, - wake_function, - bin_size + slopes, wake_function, bin_size ) # Convert the result to numpy array @@ -141,5 +143,6 @@ def test_wake(save_png=True): finally: sim.finalize() + if __name__ == "__main__": test_wake() From 32dc88604ddb13fa7f3bfb8e2be048aa2181f8fa Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Thu, 5 Sep 2024 15:43:42 -0700 Subject: [PATCH 3/5] Update test_wake.py Remove unnecessary print statements. --- tests/python/test_wake.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/python/test_wake.py b/tests/python/test_wake.py index 8474cd5ae..8b80090f9 100644 --- a/tests/python/test_wake.py +++ b/tests/python/test_wake.py @@ -40,8 +40,6 @@ def test_wake(save_png=True): bin_min = t_min bin_max = t_max bin_size = (bin_max - bin_min) / (num_bins - 1) - print("t_min, t_max, sigma_t = ") - print(t_min, t_max, sigma_t) # Create charge_distribution and slopes as PODVector_real_std charge_distribution = PODVector_real_std(num_bins + 1) @@ -70,10 +68,6 @@ def test_wake(save_png=True): bin_min, bin_max, len(charge_distribution_np) ) plt.figure() - print("Length density s values = ") - print(len(charge_distribution_np)) - print("Length density values = ") - print(len(charge_distribution_np)) plt.plot(charge_distribution_s_values, charge_distribution_np) plt.xlabel("Longitudinal Position s (m)") plt.ylabel("Charge density") From be5c7051d6870bc441a05ec93b949fd3da41504c Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Sun, 22 Sep 2024 22:44:39 -0700 Subject: [PATCH 4/5] Fix AMReX Init for `test_wake` --- tests/python/test_wake.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/python/test_wake.py b/tests/python/test_wake.py index 8b80090f9..5bcef8376 100644 --- a/tests/python/test_wake.py +++ b/tests/python/test_wake.py @@ -6,16 +6,16 @@ np.set_printoptions(threshold=sys.maxsize) from conftest import basepath +import impactx from amrex.space3d import PODVector_real_std # Import amrex array -from impactx import ImpactX, amr, wakeconvolution +from impactx import ImpactX, wakeconvolution def test_wake(save_png=True): - amr.initialize([]) try: sim = ImpactX() sim.n_cell = [16, 24, 32] - sim.load_inputs_file(basepath + "/../../examples/chicane/input_chicane_csr.in") + sim.load_inputs_file(basepath + "/examples/chicane/input_chicane_csr.in") sim.slice_step_diagnostics = False sim.init_grids() @@ -139,4 +139,8 @@ def test_wake(save_png=True): if __name__ == "__main__": + # Call MPI_Init and MPI_Finalize only once: + if impactx.Config.have_mpi: + from mpi4py import MPI # noqa + test_wake() From 1226130dda0de856c0753276dcaf69f8f5ca2ae9 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 9 Oct 2024 23:09:50 -0700 Subject: [PATCH 5/5] Finalize Test Standalone and PyTest --- tests/python/test_wake.py | 263 +++++++++++++++++++------------------- 1 file changed, 135 insertions(+), 128 deletions(-) mode change 100644 => 100755 tests/python/test_wake.py diff --git a/tests/python/test_wake.py b/tests/python/test_wake.py old mode 100644 new mode 100755 index 5bcef8376..658023f4f --- a/tests/python/test_wake.py +++ b/tests/python/test_wake.py @@ -1,3 +1,12 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2024 The ImpactX Community +# +# Authors: Alex Bojanich, Chad Mitchell, Axel Huebl +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + import sys import matplotlib.pyplot as plt @@ -6,141 +15,139 @@ np.set_printoptions(threshold=sys.maxsize) from conftest import basepath -import impactx -from amrex.space3d import PODVector_real_std # Import amrex array -from impactx import ImpactX, wakeconvolution +from impactx import Config, ImpactX, amr, wakeconvolution def test_wake(save_png=True): - try: - sim = ImpactX() - sim.n_cell = [16, 24, 32] - sim.load_inputs_file(basepath + "/examples/chicane/input_chicane_csr.in") - sim.slice_step_diagnostics = False - - sim.init_grids() - sim.init_beam_distribution_from_inputs() - sim.init_lattice_elements_from_inputs() - sim.evolve() - - sim.deposit_charge() - - element_has_csr = True - R = 10.35 # Units [m] - sigma_t = 1.9975134930563207e-05 - - if element_has_csr: - pc = sim.particle_container() - x_min, y_min, t_min, x_max, y_max, t_max = pc.min_and_max_positions() - - is_unity_particle_weight = False - GetNumberDensity = True - - num_bins = 150 - bin_min = t_min - bin_max = t_max - bin_size = (bin_max - bin_min) / (num_bins - 1) - - # Create charge_distribution and slopes as PODVector_real_std - charge_distribution = PODVector_real_std(num_bins + 1) - slopes = PODVector_real_std(num_bins) - - # Call deposit_charge with the correct type - wakeconvolution.deposit_charge( - pc, - charge_distribution, - bin_min, - bin_size, - is_unity_particle_weight, - ) - - # Call derivative_charge with the correct types - wakeconvolution.derivative_charge( - charge_distribution, - slopes, - bin_size, - GetNumberDensity, - ) - - # Convert charge distribution to numpy array and plot it - charge_distribution_np = charge_distribution.to_numpy() - charge_distribution_s_values = np.linspace( - bin_min, bin_max, len(charge_distribution_np) - ) - plt.figure() - plt.plot(charge_distribution_s_values, charge_distribution_np) - plt.xlabel("Longitudinal Position s (m)") - plt.ylabel("Charge density") - plt.title("Charge density vs Longitudinal Position") - if save_png: - plt.savefig("density.png") - else: - plt.show() - - # Convert slopes to numpy array and plot it - slopes_np = slopes.to_numpy() - slopes_s_values = np.linspace(bin_min, bin_max, len(slopes_np)) - plt.figure() - plt.plot(slopes_s_values, slopes_np) - plt.xlabel("Longitudinal Position s (m)") - plt.ylabel("Slopes") - plt.title("Slopes vs Longitudinal Position") - if save_png: - plt.savefig("slopes.png") - else: - plt.show() - - # Create wake_function as PODVector_real_std with size 2 * len(slopes) - wake_function = PODVector_real_std(2 * len(slopes)) - wake_s_values = np.linspace(bin_min, bin_max, 2 * len(slopes)) - for i in range(len(wake_function)): - if i < num_bins: - s = i * bin_size - else: - s = (i - 2 * num_bins) * bin_size - wake_s_values[i] = s - wake_function[i] = wakeconvolution.w_l_csr(s, R, bin_size) - - # Convert wake_function to numpy array and plot it - wake_function_np = wake_function.to_numpy() - plt.figure() - plt.plot(wake_s_values, wake_function_np) - plt.xlabel("Longitudinal Position s (m)") - plt.ylabel("Wake Function") - plt.title("Wake Function vs Longitudinal Position") - if save_png: - plt.savefig("wake_function.png") - else: - plt.show() - - # Call convolve_fft with the correct types and capture the output - convolved_wakefield = wakeconvolution.convolve_fft( - slopes, wake_function, bin_size - ) - - # Convert the result to numpy array - convolved_wakefield_np = convolved_wakefield.to_numpy() - - # Adjust the s_values to match the length of convolved_wakefield_np - s_values = np.linspace(bin_min, bin_max, len(convolved_wakefield_np)) - normalized_s_values = s_values / sigma_t - - plt.plot(normalized_s_values, convolved_wakefield_np) - plt.xlabel("Longitudinal Position s/sigma_s") - plt.ylabel("Wakefield (V C/m)") - plt.title("Convolved CSR Wakefield vs Longitudinal Position") - if save_png: - plt.savefig("convolved_wakefield.png") + sim = ImpactX() + sim.load_inputs_file(basepath + "/examples/chicane/input_chicane_csr.in") + sim.n_cell = [16, 24, 32] + sim.slice_step_diagnostics = False + + sim.init_grids() + sim.init_beam_distribution_from_inputs() + sim.init_lattice_elements_from_inputs() + sim.evolve() + + sim.deposit_charge() + + element_has_csr = True + R = 10.35 # Units [m] + sigma_t = 1.9975134930563207e-05 + + if element_has_csr: + pc = sim.particle_container() + x_min, y_min, t_min, x_max, y_max, t_max = pc.min_and_max_positions() + + is_unity_particle_weight = False + GetNumberDensity = True + + num_bins = 150 + bin_min = t_min + bin_max = t_max + bin_size = (bin_max - bin_min) / (num_bins - 1) + + # Create charge_distribution and slopes as PODVector_real_std + charge_distribution = amr.PODVector_real_std(num_bins + 1) + slopes = amr.PODVector_real_std(num_bins) + + # Call deposit_charge with the correct type + wakeconvolution.deposit_charge( + pc, + charge_distribution, + bin_min, + bin_size, + is_unity_particle_weight, + ) + + # Call derivative_charge with the correct types + wakeconvolution.derivative_charge( + charge_distribution, + slopes, + bin_size, + GetNumberDensity, + ) + + # Convert charge distribution to numpy array and plot it + charge_distribution_np = charge_distribution.to_numpy() + charge_distribution_s_values = np.linspace( + bin_min, bin_max, len(charge_distribution_np) + ) + plt.figure() + plt.plot(charge_distribution_s_values, charge_distribution_np) + plt.xlabel("Longitudinal Position s (m)") + plt.ylabel("Charge density") + plt.title("Charge density vs Longitudinal Position") + if save_png: + plt.savefig("density.png") + else: + plt.show() + + # Convert slopes to numpy array and plot it + slopes_np = slopes.to_numpy() + slopes_s_values = np.linspace(bin_min, bin_max, len(slopes_np)) + plt.figure() + plt.plot(slopes_s_values, slopes_np) + plt.xlabel("Longitudinal Position s (m)") + plt.ylabel("Slopes") + plt.title("Slopes vs Longitudinal Position") + if save_png: + plt.savefig("slopes.png") + else: + plt.show() + + # Create wake_function as PODVector_real_std with size 2 * len(slopes) + wake_function = amr.PODVector_real_std(2 * len(slopes)) + wake_s_values = np.linspace(bin_min, bin_max, 2 * len(slopes)) + for i in range(len(wake_function)): + if i < num_bins: + s = i * bin_size else: - plt.show() - plt.close("all") - finally: - sim.finalize() + s = (i - 2 * num_bins) * bin_size + wake_s_values[i] = s + wake_function[i] = wakeconvolution.w_l_csr(s, R, bin_size) + + # Convert wake_function to numpy array and plot it + wake_function_np = wake_function.to_numpy() + plt.figure() + plt.plot(wake_s_values, wake_function_np) + plt.xlabel("Longitudinal Position s (m)") + plt.ylabel("Wake Function") + plt.title("Wake Function vs Longitudinal Position") + if save_png: + plt.savefig("wake_function.png") + else: + plt.show() + + # Call convolve_fft with the correct types and capture the output + convolved_wakefield = wakeconvolution.convolve_fft( + slopes, wake_function, bin_size + ) + + # Convert the result to numpy array + convolved_wakefield_np = convolved_wakefield.to_numpy() + + # Adjust the s_values to match the length of convolved_wakefield_np + s_values = np.linspace(bin_min, bin_max, len(convolved_wakefield_np)) + normalized_s_values = s_values / sigma_t + + plt.plot(normalized_s_values, convolved_wakefield_np) + plt.xlabel("Longitudinal Position s/sigma_s") + plt.ylabel("Wakefield (V C/m)") + plt.title("Convolved CSR Wakefield vs Longitudinal Position") + if save_png: + plt.savefig("convolved_wakefield.png") + else: + plt.show() + plt.close("all") + + sim.finalize() if __name__ == "__main__": # Call MPI_Init and MPI_Finalize only once: - if impactx.Config.have_mpi: + if Config.have_mpi: from mpi4py import MPI # noqa - test_wake() + amr.initialize([]) + test_wake(save_png=False)