Updated visualization scripts for beam-beam collision example

Examples/Physics_applications/beam_beam_collision/README.rst

@@ -11,7 +11,8 @@ We turn on the Quantum Synchrotron QED module for photon emission (also known as
 the Breit-Wheeler QED module for the generation of electron-positron pairs (also known as coherent pair generation in the collider community).
 
 To solve for the electromagnetic field we use the nodal version of the electrostatic relativistic solver.
-This solver computes the average velocity of each species, and solves the corresponding relativistic Poisson equation (see the WarpX documentation for `warpx.do_electrostatic = relativistic` for more detail). This solver accurately reproduced the subtle cancellation that occur for some component of the ``E + v x B`` terms which are crucial in simulations of relativistic particles.
+This solver computes the average velocity of each species, and solves the corresponding relativistic Poisson equation (see the WarpX documentation for `warpx.do_electrostatic = relativistic` for more detail).
+This solver accurately reproduces the subtle cancellation that occur for some component of ``E + v x B``, which are crucial in simulations of relativistic particles.
 
 
 This example is based on the following paper :cite:t:`ex-Yakimenko2019`.
@@ -26,20 +27,23 @@ For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with `
 
 .. literalinclude:: inputs_test_3d_beam_beam_collision
    :language: ini
-   :caption: You can copy this file from ``Examples/Physics_applications/beam-beam_collision/inputs_test_3d_beam_beam_collision``.
+   :caption: You can copy this file from ``Examples/Physics_applications/beam_beam_collision/inputs_test_3d_beam_beam_collision``.
 
 
 Visualize
 ---------
 
 The figure below shows the number of photons emitted per beam particle (left) and the number of secondary pairs generated per beam particle (right).
 
-We compare different results:
+We compare different results for the reduced diagnostics with the literature:
 * (red) simplified WarpX simulation as the example stored in the directory ``/Examples/Physics_applications/beam-beam_collision``;
 * (blue) large-scale WarpX simulation (high resolution and ad hoc generated tables ;
 * (black) literature results from :cite:t:`ex-Yakimenko2019`.
 
-The small-scale simulation has been performed with a resolution of ``nx = 64, ny = 64, nz = 64`` grid cells, while the large-scale one has a much higher resolution of ``nx = 512, ny = 512, nz = 1024``. Moreover, the large-scale simulation uses dedicated QED lookup tables instead of the builtin tables. To generate the tables within WarpX, the code must be compiled with the flag ``-DWarpX_QED_TABLE_GEN=ON``. For the large-scale simulation we have used the following options:
+The small-scale simulation has been performed with a resolution of ``nx = 64, ny = 64, nz = 64`` grid cells, while the large-scale one has a much higher resolution of ``nx = 512, ny = 512, nz = 1024``.
+Moreover, the large-scale simulation uses dedicated QED lookup tables instead of the builtin tables.
+To generate the tables within WarpX, the code must be compiled with the flag ``-DWarpX_QED_TABLE_GEN=ON``.
+For the large-scale simulation we have used the following options:
 
 .. code-block:: ini
 
@@ -63,8 +67,45 @@ The small-scale simulation has been performed with a resolution of ``nx = 64, ny
    qed_bw.tab_pair_frac_how_many=512
    qed_bw.save_table_in=my_bw_table.txt
 
+
 .. figure:: https://gist.github.com/user-attachments/assets/2dd43782-d039-4faa-9d27-e3cf8fb17352
    :alt: Beam-beam collision benchmark against :cite:t:`ex-Yakimenko2019`.
    :width: 100%
 
    Beam-beam collision benchmark against :cite:t:`ex-Yakimenko2019`.
+
+
+Below are two visualizations scripts that provide examples to graph the field and reduced diagnostics.
+They are available in the ``Examples/Physics_applications/beam-beam_collision/`` folder and can be run as simply as ``python3 plot_fields.py`` and ``python3 plot_reduced.py``.
+
+.. tab-set::
+
+   .. tab-item:: Field Diagnostics
+
+      This script visualizes the evolution of the fields (:math:`|E|, |B|, \rho`) during the collision between the two ultra-relativistic lepton beams.
+      The magnitude of E and B and the charge densities of the primary beams and of the secondary pairs are sliced along either one of the two transverse coordinates (:math:`x` and :math:`y`).
+
+      .. literalinclude:: plot_fields.py
+         :language: python3
+         :caption: You can copy this file from ``Examples/Physics_applications/beam-beam_collision/plot_fields.py``.
+
+      .. figure:: https://gist.github.com/user-attachments/assets/04c9c0ec-b580-446f-a11a-437c1b244a41
+         :alt: Slice across :math:`x` of different fields (:math:`|E|, |B|, \rho`) at timestep 45, in the middle of the collision.
+         :width: 100%
+
+         Slice across :math:`x` of different fields (:math:`|E|, |B|, \rho`) at timestep 45, in the middle of the collision.
+
+
+   .. tab-item:: Reduced Diagnostics
+
+      A similar script to the one below was used to produce the image showing the benchmark against :cite:t:`ex-Yakimenko2019`.
+
+      .. literalinclude:: plot_reduced.py
+         :language: python3
+         :caption: You can copy this file from ``Examples/Physics_applications/beam-beam_collision/plot_reduced.py``.
+
+      .. figure:: https://gist.github.com/user-attachments/assets/c280490a-f1f2-4329-ad3c-46817d245dc1
+         :alt: Photon and pair production rates in time throughout the collision.
+         :width: 100%
+
+         Photon and pair production rates in time throughout the collision.

Examples/Physics_applications/beam_beam_collision/inputs_test_3d_beam_beam_collision

@@ -211,7 +211,7 @@ warpx.do_qed_schwinger = 0.
 # FULL
 diagnostics.diags_names = diag1
 
-diag1.intervals = 0
+diag1.intervals = 15
 diag1.diag_type = Full
 diag1.write_species = 1
 diag1.fields_to_plot = Ex Ey Ez Bx By Bz rho_beam1 rho_beam2 rho_ele1 rho_pos1 rho_ele2 rho_pos2

Examples/Physics_applications/beam_beam_collision/plot_fields.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+
+import matplotlib.pyplot as plt
+import numpy as np
+from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable
+from openpmd_viewer import OpenPMDTimeSeries
+
+plt.rcParams.update({"font.size": 16})
+
+series = OpenPMDTimeSeries("./diags/diag1")
+steps = series.iterations
+
+
+for slice_axis in ["x", "y"]:  # slice the fields along x and y
+    for n in steps:  # loop through the available timesteps
+        fig, ax = plt.subplots(
+            ncols=2, nrows=2, figsize=(10, 6), dpi=300, sharex=True, sharey=True
+        )
+
+        # get E field
+        Ex, info = series.get_field(
+            field="E", coord="x", iteration=n, plot=False, slice_across=slice_axis
+        )
+        Ey, info = series.get_field(
+            field="E", coord="y", iteration=n, plot=False, slice_across=slice_axis
+        )
+        Ez, info = series.get_field(
+            field="E", coord="z", iteration=n, plot=False, slice_across=slice_axis
+        )
+        # get B field
+        Bx, info = series.get_field(
+            field="B", coord="x", iteration=n, plot=False, slice_across=slice_axis
+        )
+        By, info = series.get_field(
+            field="B", coord="y", iteration=n, plot=False, slice_across=slice_axis
+        )
+        Bz, info = series.get_field(
+            field="B", coord="z", iteration=n, plot=False, slice_across=slice_axis
+        )
+        # get charge densities
+        rho_beam1, info = series.get_field(
+            field="rho_beam1", iteration=n, plot=False, slice_across=slice_axis
+        )
+        rho_beam2, info = series.get_field(
+            field="rho_beam2", iteration=n, plot=False, slice_across=slice_axis
+        )
+        rho_ele1, info = series.get_field(
+            field="rho_ele1", iteration=n, plot=False, slice_across=slice_axis
+        )
+        rho_pos1, info = series.get_field(
+            field="rho_pos1", iteration=n, plot=False, slice_across=slice_axis
+        )
+        rho_ele2, info = series.get_field(
+            field="rho_ele2", iteration=n, plot=False, slice_across=slice_axis
+        )
+        rho_pos2, info = series.get_field(
+            field="rho_pos2", iteration=n, plot=False, slice_across=slice_axis
+        )
+
+        xmin = info.z.min()
+        xmax = info.z.max()
+        xlabel = "z [m]"
+
+        if slice_axis == "x":
+            ymin = info.y.min()
+            ymax = info.y.max()
+            ylabel = "y [m]"
+        elif slice_axis == "y":
+            ymin = info.x.min()
+            ymax = info.x.max()
+            ylabel = "x [m]"
+
+        # plot E magnitude
+        Emag = np.sqrt(Ex**2 + Ey**2 + Ez**2)
+        im = ax[0, 0].imshow(
+            np.transpose(Emag),
+            cmap="seismic",
+            extent=[xmin, xmax, ymin, ymax],
+            vmin=0,
+            vmax=np.max(np.abs(Emag)),
+        )
+        ax[0, 0].set_title("E [V/m]")
+        divider = make_axes_locatable(ax[0, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        fig.colorbar(im, cax=cax, orientation="vertical")
+
+        # plot B magnitude
+        Bmag = np.sqrt(Bx**2 + By**2 + Bz**2)
+        im = ax[1, 0].imshow(
+            np.transpose(Bmag),
+            cmap="seismic",
+            extent=[xmin, xmax, ymin, ymax],
+            vmin=0,
+            vmax=np.max(np.abs(Bmag)),
+        )
+        ax[1, 0].set_title("B [T]")
+        divider = make_axes_locatable(ax[1, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        fig.colorbar(im, cax=cax, orientation="vertical")
+
+        # plot beam densities
+        rho_beams = rho_beam1 + rho_beam2
+        im = ax[0, 1].imshow(
+            np.transpose(rho_beams),
+            cmap="seismic",
+            extent=[xmin, xmax, ymin, ymax],
+            vmin=-np.max(np.abs(rho_beams)),
+            vmax=np.max(np.abs(rho_beams)),
+        )
+        ax[0, 1].set_title(r"$\rho$ beams [C/m$^3$]")
+        divider = make_axes_locatable(ax[0, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        fig.colorbar(im, cax=cax, orientation="vertical")
+
+        # plot secondary densities
+        rho2 = rho_ele1 + rho_pos1 + rho_ele2 + rho_pos2
+        im = ax[1, 1].imshow(
+            np.transpose(rho2),
+            cmap="seismic",
+            extent=[xmin, xmax, ymin, ymax],
+            vmin=-np.max(np.abs(rho2)),
+            vmax=np.max(np.abs(rho2)),
+        )
+        ax[1, 1].set_title(r"$\rho$ secondaries [C/m$^3$]")
+        divider = make_axes_locatable(ax[1, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        fig.colorbar(im, cax=cax, orientation="vertical")
+
+        for a in ax[-1, :].reshape(-1):
+            a.set_xlabel(xlabel)
+        for a in ax[:, 0].reshape(-1):
+            a.set_ylabel(ylabel)
+
+        fig.suptitle(f"Iteration = {n}, time [s] = {series.current_t}", fontsize=20)
+        plt.tight_layout()
+
+        image_file_name = "FIELDS_" + slice_axis + f"_{n:03d}.png"
+        plt.savefig(image_file_name, dpi=100, bbox_inches="tight")
+        plt.close()

Examples/Physics_applications/beam_beam_collision/plot_reduced.py

@@ -0,0 +1,48 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from scipy.constants import c, nano, physical_constants
+
+r_e = physical_constants["classical electron radius"][0]
+my_dpi = 300
+sigmaz = 10 * nano
+
+fig, ax = plt.subplots(
+    ncols=2, nrows=1, figsize=(2000.0 / my_dpi, 1000.0 / my_dpi), dpi=my_dpi
+)
+
+rdir = "./diags/reducedfiles/"
+
+df_cr = pd.read_csv(f"{rdir}" + "ColliderRelevant_beam1_beam2.txt", sep=" ", header=0)
+df_pn = pd.read_csv(f"{rdir}" + "ParticleNumber.txt", sep=" ", header=0)
+
+
+times = df_cr[[col for col in df_cr.columns if "]time" in col]].to_numpy()
+steps = df_cr[[col for col in df_cr.columns if "]step" in col]].to_numpy()
+
+x = df_cr[[col for col in df_cr.columns if "]dL_dt" in col]].to_numpy()
+coll_index = np.argmax(x)
+coll_time = times[coll_index]
+
+# number of photons per beam particle
+np1 = df_pn[[col for col in df_pn.columns if "]pho1_weight" in col]].to_numpy()
+np2 = df_pn[[col for col in df_pn.columns if "]pho2_weight" in col]].to_numpy()
+Ne = df_pn[[col for col in df_pn.columns if "]beam1_weight" in col]].to_numpy()[0]
+Np = df_pn[[col for col in df_pn.columns if "]beam2_weight" in col]].to_numpy()[0]
+
+ax[0].plot((times - coll_time) / (sigmaz / c), (np1 + np2) / (Ne + Np), lw=2)
+ax[0].set_title(r"photon number/beam particle")
+
+# number of NLBW particles per beam particle
+e1 = df_pn[[col for col in df_pn.columns if "]ele1_weight" in col]].to_numpy()
+e2 = df_pn[[col for col in df_pn.columns if "]ele2_weight" in col]].to_numpy()
+
+ax[1].plot((times - coll_time) / (sigmaz / c), (e1 + e2) / (Ne + Np), lw=2)
+ax[1].set_title(r"NLBW particles/beam particle")
+
+for a in ax.reshape(-1):
+    a.set_xlabel(r"time [$\sigma_z/c$]")
+image_file_name = "reduced.png"
+plt.tight_layout()
+plt.savefig(image_file_name, dpi=300, bbox_inches="tight")
+plt.close("all")