"""Tools for visualization of simulation output."""
import matplotlib.pyplot as plt
import numpy as np
[docs]
def plotTrajectory(trajectory):
"""Plot the trajectory of a particle.
Args:
trajectory: dictionary with the positions and angles of the particle,
as returned by :py:meth:`~.computations.computeTrajectory`.
"""
fig, ((ax_Bx, ax_By), (ax_x, ax_y)) = plt.subplots(2, 2)
ax_Bx.plot(trajectory["z"], trajectory["Bx"], color="black")
ax_By.plot(trajectory["z"], trajectory["By"], color="black")
ax_x.plot(trajectory["z"], trajectory["x"])
ax_xp = ax_x.twinx()
ax_xp.plot(trajectory["z"], trajectory["xp"] * 1e3, color="C1")
ax_y.plot(trajectory["z"], trajectory["y"])
ax_yp = ax_y.twinx()
ax_yp.plot(trajectory["z"], trajectory["yp"] * 1e3, color="C1")
ax_Bx.set_ylabel("B$_x$ [T]")
ax_By.set_ylabel("B$_y$ [T]")
ax_x.set_ylabel("x [m]", color="C0")
ax_y.set_ylabel("y [m]", color="C0")
ax_x.set_xlabel("z [m]", color="C0")
ax_y.set_xlabel("z [m]", color="C0")
ax_xp.set_ylabel("x' [mrad]", color="C1")
ax_yp.set_ylabel("y' [mrad]", color="C1")
plt.tight_layout()
plt.show()
[docs]
def plotSpectrum(spectrum, axisScale=["lin", "lin"]):
fig, ax = plt.subplots(1)
ax.plot(spectrum["wlax"], spectrum["data"])
if axisScale[0] == "log":
ax.set_xscale("log")
if axisScale[1] == "log":
ax.set_yscale("log")
ax.set_xlabel("wavelength [nm]")
ax.set_ylabel("intensity [ph/(s mm^2 nm]")
ax.grid(which='both', axis='both')
plt.tight_layout()
plt.show()
[docs]
def plotI(intensity, colorScale="lin"):
"""Plot the two-dimensional distribution of the wavefront intensity.
Args:
intensity: the radiation intensity, as returned by
:py:meth:`~.wavefronts.Wavefront.getWfrI`.
colorScale (optional): 'lin' or 'log10' scale. Defaults to 'lin'.
"""
fig, ax = plt.subplots()
ax.set_xlabel("x position [mm]", fontsize=14)
ax.set_ylabel("y position [mm]", fontsize=14)
ax.tick_params(labelsize=14)
extent = np.array((intensity["xax"][0], intensity["xax"][-1], \
intensity["yax"][0], intensity["yax"][-1])) * 1e3
if colorScale == "lin":
im = ax.imshow(intensity["data"], extent=extent,
origin="lower", cmap="jet")
elif colorScale == "log10":
im = ax.imshow(np.log10(intensity["data"]), extent=extent,
origin="lower", cmap="jet")
else:
print("Color scale should be one of 'lin' or 'log10'!")
return
plt.tight_layout()
plt.show()
[docs]
def plotPwrDensity(pwrDensity, colorScale="lin"):
"""Plot the two-dimensional distribution of the power density.
Args:
pwrDensity: the radiation power density, as returned by
:py:meth:`~.computations.computeSrPowerDensity`.
colorScale (optional): 'lin' or 'log10' scale. Defaults to 'lin'.
"""
fig, ax = plt.subplots()
ax.set_xlabel("x position [mm]", fontsize=14)
ax.set_ylabel("y position [mm]", fontsize=14)
ax.tick_params(labelsize=14)
extent = np.array((pwrDensity["xax"][0], pwrDensity["xax"][-1], \
pwrDensity["yax"][0], pwrDensity["yax"][-1])) * 1e3
if colorScale == "lin":
im = ax.imshow(pwrDensity["data"], extent=extent,
origin="lower", cmap="jet")
elif colorScale == "log10":
im = ax.imshow(np.log10(pwrDensity["data"]), extent=extent,
origin="lower", cmap="jet")
else:
print("Color scale should be one of 'lin' or 'log10'!")
return
plt.tight_layout()
plt.show()
