from matplotlib.colors import LogNorm
import matplotlib.ticker
import numpy as np
import matplotlib.pyplot as plt
from math import sqrt
[docs]def plot_trajectory_position(trajectory, show=True, ofile='', axis='Z', figsize=[18, 4.5], ret=False):
"""Plot the trajectory position. You can optionally change the axis, output to a file, show or not, and return or not
:param trajectory: Particle trajectory
:type trajectory: list
:param show: to show the plot or not
:type show: bool
:param ofile: output file name
:type ofile: str
:param axis: which axis to plot trajectory on
:type axis: str
:param figsize: dimension of the figure
:type figsize: list
:param ret: to return the plot or not
:type ret: bool
"""
# Get coordinate lists
X = [item[1][0] for item in trajectory]
Y = [item[1][1] for item in trajectory]
Z = [item[1][2] for item in trajectory]
if axis is 'X':
X1Label = 'X [m]'
X2Label = 'Y [m]'
X3Label = 'Z [m]'
X1 = X
X2 = Y
X3 = Z
elif axis is 'Y':
X1Label = 'Y [m]'
X2Label = 'Z [m]'
X3Label = 'X [m]'
X1 = Y
X2 = Z
X3 = X
elif axis is 'Z':
X1Label = 'Z [m]'
X2Label = 'X [m]'
X3Label = 'Y [m]'
X1 = Z
X2 = X
X3 = Y
# Plot X and Y vs. Z
plt.figure(1, figsize=figsize)
plt.subplot(131)
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.plot(X1, X2)
plt.xlabel(X1Label)
plt.ylabel(X2Label)
plt.title('Particle Trajectory')
plt.subplot(132)
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.plot(X1, X3)
plt.xlabel(X1Label)
plt.ylabel(X3Label)
plt.title('Particle Trajectory')
plt.subplot(133)
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.plot(X2, X3)
plt.xlabel(X2Label)
plt.ylabel(X3Label)
plt.title('Particle Trajectory')
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret is True:
return plt
return
[docs]def plot_trajectory_velocity(trajectory, show=True, ofile='', figsize=[18, 4.5], ret=False):
"""Plot the trajectory velocity. You can optionally change the axis, output to a file, show or not, and return or not
:param trajectory: Particle trajectory
:type trajectory: list
:param show: to show the plot or not
:type show: bool
:param ofile: output file name
:type ofile: str
:param axis: which axis to plot trajectory on
:type axis: str
:param figsize: dimension of the figure
:type figsize: list
:param ret: to return the plot or not
:type ret: bool
"""
# Get coordinate lists
T = [item[0] for item in trajectory]
VX = [item[2][0] for item in trajectory]
VY = [item[2][1] for item in trajectory]
VZ = [item[2][2] for item in trajectory]
# Plot VX, VY, VZ vs. T
plt.figure(1, figsize=figsize)
plt.subplot(131)
plt.plot(T, VX)
plt.xlabel('T [s]')
plt.ylabel('$\\beta_x$')
plt.title('Particle $\\beta_x$')
plt.subplot(132)
plt.plot(T, VY)
plt.xlabel('T [s]')
plt.ylabel('$\\beta_y$')
plt.title('Particle $\\beta_y$')
plt.subplot(133)
plt.plot(T, VZ)
plt.xlabel('T [s]')
plt.ylabel('$\\beta_z$')
plt.title('Particle $\\beta_z$')
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_power_density(V, title=None, xlabel='X1 Axis [$m$]', ylabel='X2 Axis [$m$]', show=True, ofile='', figsize=None, ret=False, x1=None, x2=None):
"""Plot a 2D histogram with equal spacing"""
if x1 is not None or x2 is not None:
return plot_power_density_2d1d(V, x1=x1, x2=x2, title=title, show=show, ofile=ofile, figsize=figsize, ret=ret)
X = [item[0][0] for item in V]
Y = [item[0][1] for item in V]
P = [item[1] for item in V]
NX = len(np.unique(X))
NY = len(np.unique(Y))
plt.figure(1, figsize=figsize)
plt.hist2d(X, Y, bins=[NX, NY], weights=P)
plt.colorbar(format='%.0e')
#cb.formatter.set_scientific(True)
#cb.formatter.set_powerlimits((0, 0))
#cb.ax.yaxis.set_offset_position('right')
#cb.update_ticks()
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
if title is None:
title = 'Power Density [$W / mm^2$]'
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.title(title)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_power_density_2d1d(V, x1=None, x2=None, title=None, xlabel='[$m$]', ylabel='[$W / mm^2$]', xlim=None, ylim=None, show=True, ofile='', figsize=None, ret=False):
"""Plot a 2D histogram with equal spacing"""
if x1 is not None and x2 is not None:
raise ValueError('x1 and x2 cannot both be defined')
if x1 is None and x2 is None:
raise ValueError('x1 or x2 must be defined')
X = [item[0][0] for item in V]
Y = [item[0][1] for item in V]
P = [item[1] for item in V]
XValues = np.unique(X).tolist()
YValues = np.unique(Y).tolist()
NX = len(XValues)
NY = len(YValues)
x1_index0 = -1
x1_index1 = -1
x1_frac0 = 0.5
x2_index0 = -1
x2_index1 = -1
x2_frac0 = 0.5
XC = []
YP = []
title_position = ''
if x1 is not None:
if xlabel is None:
xlabel = 'x1 [m]'
title_position = 'x1 = ' + str(round(x1, 4))
if x1 in XValues:
x1_index0 = XValues.index(x1)
x1_index1 = x1_index0
else:
x1_index1 = next(x[0] for x in enumerate(XValues) if x[1] > x1)
x1_index0 = x1_index1 - 1
if x1_index0 < 0:
raise ValueError('x1 is not in range')
for iy in range(NY):
v0 = P[NY*x1_index0 + iy] * x1_frac0
v1 = P[NY*x1_index1 + iy] * (1 - x1_frac0)
YP.append(Y[NY*x1_index0 + iy])
XC.append(v0 + v1)
elif x2 is not None:
if xlabel is None:
xlabel = 'x2 [m]'
title_position = 'x2 = ' + str(round(x2, 4))
if x2 in YValues:
x2_index0 = YValues.index(x2)
x2_index1 = x2_index0
else:
x2_index1 = next(x[0] for x in enumerate(YValues) if x[1] > x2)
x2_index0 = x2_index1 - 1
if x2_index0 < 0:
raise ValueError('x2 is not in range')
for ix in range(NX):
v0 = P[NY*ix + x2_index0] * x2_frac0
v1 = P[NY*ix + x2_index1] * (1 - x2_frac0)
YP.append(X[NY*ix + x2_index0])
XC.append(v0 + v1)
plt.figure(1, figsize=figsize)
if ylim is not None: plt.ylim(ylim[0], ylim[1])
if xlim is not None: plt.xlim(xlim[0], xlim[1])
plt.plot(YP, XC)
plt.xlabel(xlabel)
plt.ylabel(ylabel)
if title is None:
title = 'Power Density [$W / mm^2$]'
plt.title(title + ' at ' + title_position)
else:
plt.title(title)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_power_density_1d(V, title='Power Density [$W / mm^2$]', xlabel='[$m$]', ylabel='[$W / mm^2$]', xlim=None, ylim=None, show=True, ofile='', figsize=None, ret=False):
"""Plot a 1D power density"""
X = [item[0][0] for item in V]
P = [item[1] for item in V]
plt.figure(1, figsize=figsize)
plt.plot(X, P)
if ylim is not None: plt.ylim(ylim[0], ylim[1])
if xlim is not None: plt.xlim(xlim[0], xlim[1])
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.title(title)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_flux(V, title='Flux [$\gamma / mm^2 / 0.1\%bw / s$]', xlabel='X1 Axis [$m$]', ylabel='X2 Axis [$m$]', clim=None, show=True, ofile='', figsize=None, ylim=None, xlim=None, colorbar=True, ret=False, nticks_cb=None):
"""Plot a 2D histogram with equal spacing"""
X = [item[0][0] for item in V]
Y = [item[0][1] for item in V]
P = [item[1] for item in V]
NX = len(np.unique(X))
NY = len(np.unique(Y))
# Size and limits
plt.figure(1, figsize=figsize)
if ylim is not None: plt.ylim(ylim[0], ylim[1])
if xlim is not None: plt.xlim(xlim[0], xlim[1])
plt.hist2d(X, Y, bins=[NX, NY], weights=P)
#cb.formatter.set_scientific(True)
#cb.formatter.set_powerlimits((0, 0))
#cb.ax.yaxis.set_offset_position('right')
#cb.update_ticks()
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
if colorbar is True:
cb = plt.colorbar(format='%.2e')
if nticks_cb is not None:
tick_locator = matplotlib.ticker.MaxNLocator(nbins=nticks_cb)
cb.locator = tick_locator
cb.update_ticks()
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.title(title)
if clim is not None:
plt.clim(clim)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
plt.close()
return
[docs]def plot_spectrum(S, log=False, show=True, ofile='', title='Spectrum', xlabel='Energy [eV]', ylabel='$\gamma / mm^2 / 0.1\%bw / s$', figsize=None, ylim=None, xlim=None, transparent=True, ret=False, xticks=None, **kwargs):
"""Plot the spectrum"""
# Size and limits
plt.figure(1, figsize=figsize)
if ylim is not None: plt.ylim(ylim[0], ylim[1])
if xlim is not None: plt.xlim(xlim[0], xlim[1])
X = [item[0] for item in S]
Y = [item[1] for item in S]
plt.plot(X, Y, **kwargs)
if log:
plt.yscale('log')
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.title(title)
if xticks is not None:
plt.xticks(xticks)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight', transparent=transparent)
if show == True:
plt.show()
if ret is True:
return plt
return
def plot_spectra(spectra, label=None, show=True, ofile='', title='', loc=None, log=False, loglog=False, xlabel='Energy [eV]', ylabel='[$\gamma / mm^2 / 0.1\%bw / s$]', figsize=None, ylim=None, xlim=None, ret=False, axis=None, transparent=True, xticks=None, xvlines=None, **kwargs):
# Size and limits
plt.figure(1, figsize=figsize)
if loglog:
plt.loglog()
if ylim is not None: plt.ylim(ylim[0], ylim[1])
if xlim is not None: plt.xlim(xlim[0], xlim[1])
for i in range(len(spectra)):
s = spectra[i]
X = [item[0] for item in s]
Y = [item[1] for item in s]
if label is not None:
if label[i] is not None:
plt.plot(X, Y, label=label[i], **kwargs)
else:
plt.plot(X, Y, **kwargs)
else:
plt.plot(X, Y, **kwargs)
if xvlines is not None:
for xvline in xvlines:
plt.axvline(x=xvline, color='red', linestyle='dashed', linewidth=1)
if log:
plt.yscale('log')
plt.legend(loc=loc)
if xlabel is not None:
plt.xlabel(xlabel)
if ylabel is not None:
plt.ylabel(ylabel)
if title == None:
title='Spectrum'
plt.title(title)
if axis is not None:
plt.axis(axis)
if xticks is not None:
plt.xticks(xticks)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight', transparent=transparent)
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_bfield(osr, mymin=-1, mymax=1, ylim=None, show=True, ofile='', axis='Z', npoints=20000, between_two_points=None, ret=False):
"""Plot the magnetic field as a function of Z"""
P = []
Bx = []
By = []
Bz = []
if between_two_points is not None:
p0 = between_two_points[0]
p1 = between_two_points[1]
step = [ (p1[0] - p0[0]) / float(npoints - 1), (p1[1] - p0[1]) / float(npoints - 1), (p1[2] - p0[2]) / float(npoints - 1)]
distance = sqrt( pow(p1[0]-p0[0], 2) + pow(p1[1]-p0[1], 2) + pow(p1[2]-p0[2], 2) )
P = np.linspace(0, distance, npoints)
for i in range(npoints):
x = p0[0] + step[0] * float(i)
y = p0[1] + step[1] * float(i)
z = p0[2] + step[2] * float(i)
Bx.append(osr.get_bfield([x, y, z])[0])
By.append(osr.get_bfield([x, y, z])[1])
Bz.append(osr.get_bfield([x, y, z])[2])
axis = 'Position'
else:
P = np.linspace(mymin, mymax, npoints)
if axis is 'X':
Bx = [osr.get_bfield([p, 0, 0])[0] for p in P]
By = [osr.get_bfield([p, 0, 0])[1] for p in P]
Bz = [osr.get_bfield([p, 0, 0])[2] for p in P]
elif axis is 'Y':
Bx = [osr.get_bfield([0, p, 0])[0] for p in P]
By = [osr.get_bfield([0, p, 0])[1] for p in P]
Bz = [osr.get_bfield([0, p, 0])[2] for p in P]
elif axis is 'Z':
Bx = [osr.get_bfield([0, 0, p])[0] for p in P]
By = [osr.get_bfield([0, 0, p])[1] for p in P]
Bz = [osr.get_bfield([0, 0, p])[2] for p in P]
else:
raise
plt.figure(1, figsize=(18, 4.5))
plt.subplot(131)
plt.plot(P, Bx)
plt.xlabel(axis + ' [m]')
plt.ylabel('Bx [T]')
plt.ylim(ylim)
plt.subplot(132)
plt.plot(P, By)
plt.xlabel(axis + ' [m]')
plt.ylabel('By [T]')
plt.ylim(ylim)
plt.subplot(133)
plt.plot(P, Bz)
plt.xlabel(axis + ' [m]')
plt.ylabel('Bz [T]')
plt.ylim(ylim)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_efield(osr, mymin=-1, mymax=1, ylim=None, show=True, ofile='', axis='Z', npoints=20000, between_two_points=None, ret=False):
"""Plot the electric field as a function of Z"""
P = []
Bx = []
By = []
Bz = []
if between_two_points is not None:
p0 = between_two_points[0]
p1 = between_two_points[1]
step = [ (p1[0] - p0[0]) / float(npoints - 1), (p1[1] - p0[1]) / float(npoints - 1), (p1[2] - p0[2]) / float(npoints - 1)]
distance = sqrt( pow(p1[0]-p0[0], 2) + pow(p1[1]-p0[1], 2) + pow(p1[2]-p0[2], 2) )
P = np.linspace(0, distance, npoints)
for i in range(npoints):
x = p0[0] + step[0] * float(i)
y = p0[1] + step[1] * float(i)
z = p0[2] + step[2] * float(i)
Bx.append(osr.get_efield([x, y, z])[0])
By.append(osr.get_efield([x, y, z])[1])
Bz.append(osr.get_efield([x, y, z])[2])
axis = 'Position'
else:
P = np.linspace(mymin, mymax, npoints)
if axis is 'X':
Bx = [osr.get_efield([p, 0, 0])[0] for p in P]
By = [osr.get_efield([p, 0, 0])[1] for p in P]
Bz = [osr.get_efield([p, 0, 0])[2] for p in P]
elif axis is 'Y':
Bx = [osr.get_efield([0, p, 0])[0] for p in P]
By = [osr.get_efield([0, p, 0])[1] for p in P]
Bz = [osr.get_efield([0, p, 0])[2] for p in P]
elif axis is 'Z':
Bx = [osr.get_efield([0, 0, p])[0] for p in P]
By = [osr.get_efield([0, 0, p])[1] for p in P]
Bz = [osr.get_efield([0, 0, p])[2] for p in P]
else:
raise
plt.figure(1, figsize=(18, 4.5))
plt.subplot(131)
plt.plot(P, Bx)
plt.xlabel(axis + ' [m]')
plt.ylabel('Ex [V/m]')
plt.ylim(ylim)
plt.subplot(132)
plt.plot(P, By)
plt.xlabel(axis + ' [m]')
plt.ylabel('Ey [V/m]')
plt.ylim(ylim)
plt.subplot(133)
plt.plot(P, Bz)
plt.xlabel(axis + ' [m]')
plt.ylabel('Ez [V/m]')
plt.ylim(ylim)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
return
[docs]def total_power(pd):
"""Calculate the total power in a uniform grid.
This will not work for a non-uniform grid. Different NX and NY are ok."""
X = [item[0][0] for item in pd]
Y = [item[0][1] for item in pd]
P = [item[1] for item in pd]
NX = len(np.unique(X))
NY = len(np.unique(Y))
dx = (max(X) - min(X)) / (NX - 1)
dy = (max(Y) - min(Y)) / (NY - 1)
return dx * dy * sum(P) * 1e6
[docs]def plot_electric_field_vs_time(efield, show=True, ofile='', ret=False):
"""Plot the electric field as a function of time"""
T = [item[0] for item in efield]
Ex = [item[1][0] for item in efield]
Ey = [item[1][1] for item in efield]
Ez = [item[1][2] for item in efield]
plt.figure(1, figsize=(18, 9))
plt.subplot(131)
plt.plot(T, Ex)
plt.xlabel('T [s]')
plt.ylabel('Ex')
plt.title('Electric Field (Ex)')
plt.figure(1, figsize=(18, 9))
plt.subplot(132)
plt.plot(T, Ey)
plt.xlabel('T [s]')
plt.ylabel('Ey')
plt.title('Electric Field (Ey)')
plt.figure(1, figsize=(18, 9))
plt.subplot(133)
plt.plot(T, Ez)
plt.xlabel('T [s]')
plt.ylabel('Ez')
plt.title('Electric Field (Ez)')
if ret:
return plt
return
[docs]def plot_undulator_flux_onaxis(oth, period, nperiods, harmonics, minimum=0, bfield=None, K=None, show=True, ret=False, title='Flux On-Axis', figsize=None, ofile=None):
'''Plot the on-axis flux of an undulator. More docstring needed'''
if bfield is None and K is None:
raise ValueError('bfield or K must be defined')
if bfield is not None and K is not None:
raise ValueError('bfield and K cannot both be defined. pick one or the other')
if bfield is not None:
if len(bfield) is not 2:
raise ValueError('bfield must be: [min, max]')
else:
K = [oth.undulator_K(bfield[0], period), oth.undulator_K(bfield[1], period)]
# Size and limits
plt.figure(1, figsize=figsize)
plt.loglog()
# Loop over all harmonics
for i in harmonics:
R = []
for k in np.linspace(K[1], K[0], 300):
ev_flux = oth.undulator_flux_onaxis(K=k,
period=period,
nperiods=nperiods,
harmonic=i
)
if ev_flux[1] >= minimum:
R.append(ev_flux)
X = []
Y = []
for j in range(len(R)):
X.append(R[j][0])
Y.append(R[j][1])
plt.plot(X, Y, label=str(i))
plt.legend(title='Harmonics')
plt.xlabel('Energy [eV]')
plt.ylabel('[$\gamma / mrad^2 / 0.1\%bw / s$]')
plt.title(title)
if ofile is not None:
plt.savefig(ofile, bbox_inches='tight', transparent=transparent)
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_undulator_brightness(oth, period, nperiods, harmonics, minimum=0, bfield=None, K=None, show=True, ret=False, title='Brightness', figsize=None, ofile=None):
'''Plot the brightness of an undulator.'''
if bfield is None and K is None:
raise ValueError('bfield or K must be defined')
if bfield is not None and K is not None:
raise ValueError('bfield and K cannot both be defined. pick one or the other')
if bfield is not None:
if len(bfield) is not 2:
raise ValueError('bfield must be: [min, max]')
else:
K = [oth.undulator_K(bfield[0], period), oth.undulator_K(bfield[1], period)]
# Size and limits
plt.figure(1, figsize=figsize)
plt.loglog()
# Loop over all harmonics
for i in harmonics:
R = []
for k in np.linspace(K[1], K[0], 300):
ev_brightness = oth.undulator_brightness(K=k,
period=period,
nperiods=nperiods,
harmonic=i
)
if ev_brightness[1] >= minimum:
R.append(ev_brightness)
X = []
Y = []
for j in range(len(R)):
X.append(R[j][0])
Y.append(R[j][1])
plt.plot(X, Y, label=str(i))
plt.legend(title='Harmonics')
plt.xlabel('Energy [eV]')
plt.ylabel('[$\gamma / mm^2 / mrad^2 / 0.1\%bw / s$]')
plt.title(title)
if ofile is not None:
plt.savefig(ofile, bbox_inches='tight', transparent=True)
if show == True:
plt.show()
if ret:
return plt
return
[docs]def plot_flux_spectrum(F, S, energy=None, title='Flux [$\gamma / mm^2 / 0.1\%bw / s]$', xlabel='X1 Axis [$m$]', ylabel='X2 Axis [$m$]', show=True, ofile='', figsize=[10, 3], ylim=None, xlim=None, colorbar=True, ret=False):
"""Plot a 2D histogram with equal spacing"""
X = [item[0][0] for item in F]
Y = [item[0][1] for item in F]
P = [item[1] for item in F]
NX = len(np.unique(X))
NY = len(np.unique(Y))
# Size and limits
plt.figure(1, figsize=figsize)
plt.subplot(121)
if ylim is not None: plt.ylim(ylim[0], ylim[1])
if xlim is not None: plt.xlim(xlim[0], xlim[1])
plt.hist2d(X, Y, bins=[NX, NY], weights=P)
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
if colorbar is True:
plt.colorbar(format='%.0e')
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.title(title)
plt.subplot(122)
X2 = [item[0] for item in S]
Y2 = [item[1] for item in S]
plt.plot(X2, Y2)
plt.xlabel('Energy [eV]')
plt.ylabel('$\gamma / mm^2 / 0.1\%bw / s$')
plt.title('Spectrum')
if energy is not None:
plt.axvline(x=energy, color='red')
plt.figure(1)
if ofile != '':
plt.savefig(ofile, bbox_inches='tight')
if show == True:
plt.show()
if ret:
return plt
return