-
Notifications
You must be signed in to change notification settings - Fork 20
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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 <abojanich@lbl.gov> Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
- Loading branch information
1 parent
08172bf
commit 25b91f6
Showing
1 changed file
with
153 additions
and
0 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |