From be7e1f7260b06a2b7dc66262b990c601d708088a Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 9 Oct 2024 23:09:50 -0700 Subject: [PATCH] Finalize Test Standalone and PyTest --- tests/python/test_wake.py | 252 +++++++++++++++++++------------------- 1 file changed, 125 insertions(+), 127 deletions(-) diff --git a/tests/python/test_wake.py b/tests/python/test_wake.py index 5bcef8376..63a517183 100644 --- a/tests/python/test_wake.py +++ b/tests/python/test_wake.py @@ -6,141 +6,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") + 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() - - # 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() + 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 + amr.initialize([]) test_wake()