Finalize Test

Standalone and PyTest

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)