Add Python CSR wake plotting test.

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 <https://github.com/impact-lbl/IMPACT-Z>`__ code.
+ImpactX runs on modern GPUs or CPUs alike, provides user-friendly interfaces suitable for AI/ML workflows, has many :ref:`benchmarks <usage-examples>` 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 <theory-concepts>` for details on our models, assumptions and concepts.
+
+ImpactX is part of the `Beam, Plasma & Accelerator Simulation Toolkit (BLAST) <https://blast.lbl.gov>`__.
+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 <https://github.com/ECP-WarpX/impactx/discussions>`__ and get in touch with the community.
+We organize support as an open community, helping each other.
 
-The `ImpactX GitHub repo <https://github.com/ECP-WarpX/impactx>`__ is the main communication platform.
+The `ImpactX GitHub repository <https://github.com/ECP-WarpX/impactx>`__ 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 <https://github.com/ECP-WarpX/impactx/issues>`__.
 
-We also have a `discussion page <https://github.com/ECP-WarpX/impactx/discussions>`__ 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 <https://github.com/ECP-WarpX/impactx/discussions>`__.
+There, you can find already answered questions, add your own, get help with installation procedures, help others and share ideas.
 
 .. raw:: html
 

tests/python/test_wake.py

@@ -0,0 +1,146 @@
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+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
+
+
+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")
+            else:
+                plt.show()
+            plt.close("all")
+    finally:
+        sim.finalize()
+
+
+if __name__ == "__main__":
+    # Call MPI_Init and MPI_Finalize only once:
+    if impactx.Config.have_mpi:
+        from mpi4py import MPI  # noqa
+
+    test_wake()