From 25b91f60edc205559e472589d092982e845967a0 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Thu, 10 Oct 2024 00:00:09 -0700 Subject: [PATCH] Add Python CSR wake plotting test. (#696) * Add Python CSR wake plotting test. * Update test_wake.py: Remove unnecessary print statements. * Fix AMReX Init for `test_wake` * Finalize Test: Standalone and PyTest Co-authored-by: Alex Bojanich Co-authored-by: Axel Huebl --- tests/python/test_wake.py | 153 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100755 tests/python/test_wake.py diff --git a/tests/python/test_wake.py b/tests/python/test_wake.py new file mode 100755 index 000000000..658023f4f --- /dev/null +++ b/tests/python/test_wake.py @@ -0,0 +1,153 @@ +#!/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 +import numpy as np + +np.set_printoptions(threshold=sys.maxsize) +from conftest import basepath + +from impactx import Config, ImpactX, amr, wakeconvolution + + +def test_wake(save_png=True): + 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: + 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 Config.have_mpi: + from mpi4py import MPI # noqa + + amr.initialize([]) + test_wake(save_png=False)