-
Notifications
You must be signed in to change notification settings - Fork 194
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Updated visualization scripts for beam-beam collision example (#4797)
* add viz scripts to beam-beam collision example * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update plot_reduced.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Modified pull request * Modified README.rst * Modified README.rst * updated vizs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * restore plot fields * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rm EOL white spaces * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add missing text back. * Change figure link --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Arianna Formenti <[email protected]> Co-authored-by: Remi Lehe <[email protected]>
- Loading branch information
4 people
authored
Sep 17, 2024
1 parent
55e86de
commit 6463b1f
Showing
4 changed files
with
233 additions
and
5 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
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
139 changes: 139 additions & 0 deletions
139
Examples/Physics_applications/beam_beam_collision/plot_fields.py
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,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() |
48 changes: 48 additions & 0 deletions
48
Examples/Physics_applications/beam_beam_collision/plot_reduced.py
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,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") |