"""
This file is part of the openPMD-viewer.
This file contains diagnostics relevant for laser-plasma acceleration
Copyright 2015-2016, openPMD-viewer contributors
Authors: Soeren Jalas, Remi Lehe
License: 3-Clause-BSD-LBNL
"""
# Class that inherits from OpenPMDTimeSeries, and implements
# some standard diagnostics (emittance, etc.)
from openpmd_viewer import OpenPMDTimeSeries, FieldMetaInformation
import numpy as np
import scipy.constants as const
from scipy.optimize import curve_fit
from openpmd_viewer.openpmd_timeseries.utilities import sanitize_slicing
from openpmd_viewer.openpmd_timeseries.plotter import check_matplotlib
from scipy.signal import hilbert
try:
import matplotlib.pyplot as plt
except ImportError:
# Any error will be caught later by `check_matplotlib`
pass
[docs]
class LpaDiagnostics( OpenPMDTimeSeries ):
def __init__( self, path_to_dir, check_all_files=True, backend=None ):
"""
Initialize an OpenPMD time series with various methods to diagnose the
data
Parameters
----------
path_to_dir : string
The path to the directory where the openPMD files are.
For the moment, only HDF5 files are supported. There should be
one file per iteration, and the name of the files should end
with the iteration number, followed by '.h5' (e.g. data0005000.h5)
check_all_files: bool, optional
Check that all the files in the timeseries are consistent
(i.e. that they contain the same fields and particles,
with the same metadata)
For fast access to the files, this can be changed to False.
backend: string
Backend to be used for data reading. Can be `openpmd-api`
or `h5py`. If not provided will use `openpmd-api` if available
and `h5py` otherwise.
"""
OpenPMDTimeSeries.__init__( self, path_to_dir,
check_all_files=check_all_files, backend=backend )
[docs]
def get_energy_spread( self, t=None, iteration=None, species=None,
select=None, center='mean', width='std', property='energy' ):
"""
Calculate the central energy and energy spread according to the
particle weights
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
species : str
Particle species to use for calculations
select : dict, optional
Either None or a dictionary of rules
to select the particles, of the form
'x' : [-4., 10.] (Particles having x between -4 and 10 meters)
'z' : [0, 100] (Particles having z between 0 and 100 meters)
center : str
Method to find the central energy of the particle distribution
Should be one of
- 'mean'
- 'median'
width : str
Method to find the energy spread of the particle distribution
Should be one of
- 'std'
- 'mad'
property : str
Unit of energy. Should be one of
- 'energy' returns energy in MeV
- 'gamma' returns Lorentz factor
Returns
-------
A tuple of floats with:
- central energy
- energy spread
Returns NaN if particle selection is empty
"""
# Get particle data
ux, uy, uz, w, m = self.get_particle(
var_list=['ux', 'uy', 'uz', 'w', 'mass'], select=select,
species=species, t=t, iteration=iteration)
if len(w) == 0:
# Return NaN if no particles are found
return np.nan, np.nan
# Calculate Lorentz factor and energy for all particles
gamma = np.sqrt(1 + ux ** 2 + uy ** 2 + uz ** 2)
if property == 'energy':
prop = (gamma - 1) * m * const.c ** 2 / const.e * 1e-6
elif property == 'gamma':
prop = gamma
else:
raise ValueError('Invalid output property: %s'%property)
# Calculate weighted center
if center == 'mean':
p_center = w_ave(prop, w)
elif center == 'median':
p_center = w_median(prop, w)
else:
raise ValueError('Invalid center property: %s' % center)
# Calculate weighted width
if width == 'std':
p_width = w_std(prop, w)
elif width == 'mad':
p_width = w_mad(prop, w)
else:
raise ValueError('Invalid width property: %s' % width)
return p_center, p_width
[docs]
def get_mean_gamma( self, t=None, iteration=None, species=None,
select=None ):
"""
Calculate the mean gamma and standard deviation according to the
particle weights
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
species : string
Particle species to use for calculations
select : dict, optional
Either None or a dictionary of rules
to select the particles, of the form
'x' : [-4., 10.] (Particles having x between -4 and 10 meters)
'z' : [0, 100] (Particles having z between 0 and 100 meters)
Returns
-------
A tuple of floats with:
- mean weighted gamma
- weighted standard deviation of gamma
Returns NaN if particle selection is empty
"""
# Get particle data
ux, uy, uz, w = self.get_particle(
var_list=['ux', 'uy', 'uz', 'w'], select=select,
species=species, t=t, iteration=iteration )
if len(w) == 0:
# Return NaN if no particles are found
return np.nan, np.nan
# Calculate Lorentz factor for all particles
gamma = np.sqrt(1 + ux ** 2 + uy ** 2 + uz ** 2)
# Calculate weighted mean and average
try:
# Calculate mean gamma for selected particles
mean_gamma = np.average(gamma, weights=w)
except ZeroDivisionError:
# If selection is empty or all particles have weight zero,
# return NaN
mean_gamma = np.nan
std_gamma = w_std(gamma, w)
# Return the result
return mean_gamma, std_gamma
[docs]
def get_sigma_gamma_slice(self, dz, t=None, iteration=None, species=None,
select=None, plot=False, **kw):
"""
Calculate the standard deviation of gamma for particles in z-slices of
width dz
Parameters
----------
dz : float (in micrometers)
Width of slices in which to calculate sigma gamma
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
species : string
Particle species to use for calculations
select : dict, optional
Either None or a dictionary of rules
to select the particles, of the form
'x' : [-4., 10.] (Particles having x between -4 and 10 meters)
'z' : [0, 100] (Particles having z between 0 and 100 meters)
plot : bool, optional
Whether to plot the requested quantity
**kw : dict, otional
Additional options to be passed to matplotlib's `plot` method
Returns
-------
A tuple of arrays:
- Sigma gamma in each slice
- Central z position of each slice
"""
z, uz, ux, uy, w = self.get_particle(t=t, species=species,
select=select, var_list=['z', 'uz', 'ux', 'uy', 'w'],
iteration=iteration)
# Calculate gamma of each particle
gamma = np.sqrt(1 + (uz ** 2 + ux ** 2 + uy ** 2))
z0 = min(z)
zend = max(z)
N = int((zend - z0) / dz)
spreads = np.zeros(N + 1)
z_pos = np.linspace(z0, zend, N + 1)
zi = z0 + dz / 2.
i = 0
# Iterate over slices and calculate sigma gamma
while zi < zend:
z_filter = (z > zi - dz / 2.) & (z < zi + dz / 2.)
spreads[i] = w_std(gamma[z_filter], w[z_filter])
zi += dz
i += 1
# Plot the result if needed
if plot:
check_matplotlib()
iteration = self.iterations[ self._current_i ]
time_s = self.t[ self._current_i ]
plt.plot(z_pos, spreads, **kw)
plt.title("Slice energy spread at %.2e s (iteration %d)"
% (time_s, iteration), fontsize=self.plotter.fontsize)
plt.xlabel('$z \;(m)$', fontsize=self.plotter.fontsize)
plt.ylabel('$\sigma_\gamma (\Delta_z=%s m)$' % dz,
fontsize=self.plotter.fontsize)
return(spreads, z_pos)
[docs]
def get_charge( self, t=None, iteration=None, species=None, select=None ):
"""
Calculate the charge of the selcted particles.
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
species : string
Particle species to use for calculations
select : dict, optional
Either None or a dictionary of rules
to select the particles, of the form
'x' : [-4., 10.] (Particles having x between -4 and 10 meters)
'z' : [0, 100] (Particles having z between 0 and 100 meters)
Returns
-------
A float with the electric charge of the selected particles in Coulomb
"""
# Get particle data
w, q = self.get_particle( var_list=['w', 'charge'], species=species,
select=select, t=t, iteration=iteration )
# Calculate charge
charge = np.sum(w * q)
# Return the result
return( charge )
[docs]
def get_divergence( self, t=None, iteration=None, species=None,
select=None ):
"""
Calculate the divergence of the selected particles.
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
species : string
Particle species to use for calculations
select : dict, optional
Either None or a dictionary of rules
to select the particles, of the form
'x' : [-4., 10.] (Particles having x between -4 and 10 meters)
'z' : [0, 100] (Particles having z between 0 and 100 meters)
Returns
-------
A tuple with:
- divergence in x plane in rad
- divergence in y plane in rad
Returns NaN if particle selection is empty
"""
# Get particle data
ux, uy, uz, w = self.get_particle( var_list=['ux', 'uy', 'uz', 'w'],
t=t, iteration=iteration,
species=species, select=select )
if len(w) == 0:
# Return NaN if no particles are found
return np.nan, np.nan
# Calculate divergence
div_x = w_std( np.arctan2(ux, uz), w )
div_y = w_std( np.arctan2(uy, uz), w )
# Return the result
return div_x, div_y
[docs]
def get_emittance(self, t=None, iteration=None, species=None,
select=None, kind='normalized', description='projected',
nslices=0, beam_length=None):
"""
Calculate the RMS emittance.
(See K Floetmann: Some basic features of beam emittance. PRSTAB 2003)
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
species : string
Particle species to use for calculations
select : dict, optional
Either None or a dictionary of rules to select the particles,
of the form:
'x' : [-4., 10.] (Particles having x between -4 and 10 meters)
'z' : [0, 100] (Particles having z between 0 and 100 meters).
kind : string, optional
Kind of emittance to be computed. Can be 'normalized' or 'trace'.
description : string, optional
Type of emittance to be computed. Available options:
- 'projected' : projected emittance
- 'all-slices' : emittance within slices taken along the z
direction
- 'slice-averaged' : slice emittance averaged over all slices.
nslices : integer, optional
Number of slices to compute slice emittance. Required if
description='slice-average' or 'all-slices'.
beam_length : float (in meters), optional
Beam length, used to calculate slice positions when nslices>1.
By default, it is 4 times the standard deviation in z.
Returns
-------
If description='projected' or 'slice-averaged':
- beam emittance in the x plane (pi m rad)
- beam emittance in the y plane (pi m rad)
If description='all-slices':
- A 1d array with beam emittance in the x plane
(pi m rad) for each slice
- A 1d array with beam emittance in the y plane
(pi m rad) for each slice
- A 1d array with number of electrons in each slice
- A 1d array with slice centers
"""
if kind not in ['normalized', 'trace']:
raise ValueError('Argument `kind` not recognized.')
if description not in ['projected', 'slice-averaged', 'all-slices']:
raise ValueError('Argument `description` not recognized.')
# Wheter to compute slice emittance
do_slice_emittance = ( description in ['slice-averaged',
'all-slices'] )
if do_slice_emittance and not nslices > 0:
raise ValueError(
'nslices must be given if `description`=' + description + '.')
# Get particle data
x, y, z, ux, uy, uz, w = self.get_particle(
var_list=['x', 'y', 'z', 'ux', 'uy', 'uz', 'w'], t=t,
iteration=iteration, species=species, select=select )
# Normalized or trace-space emittance
if kind == 'normalized':
ux = ux
uy = uy
if kind == 'trace':
ux = ux / uz
uy = uy / uz
if do_slice_emittance:
# Get slice locations
zavg = w_ave(z, w)
z = z - zavg
if beam_length is None:
std = w_std(z, w)
beam_length = 4 * std
bins = np.linspace( -beam_length / 2, beam_length / 2, nslices )
binwidth = .5 * bins[1] - .5 * bins[0]
slice_centers = bins + .5 * binwidth
# Initialize slice emittance arrays
emit_slice_x = np.zeros(len(bins[:-1]))
emit_slice_y = np.zeros(len(bins[:-1]))
slice_weights = np.zeros(len(bins[:-1]))
# Loop over slices
for count, leftedge in enumerate(bins[:-1]):
# Get emittance in this slice
current_slice = (np.abs(z - slice_centers[count]) <= binwidth)
slice_weights[count] = np.sum(w[current_slice])
if slice_weights[count] > 0:
emit_x, emit_y = emittance_from_coord(
x[current_slice], y[current_slice],
ux[current_slice], uy[current_slice],
w[current_slice])
emit_slice_x[count] = emit_x
emit_slice_y[count] = emit_y
if description == 'all-slices':
return (emit_slice_x, emit_slice_y,
slice_weights, slice_centers)
else:
emit_x = w_ave(emit_slice_x, slice_weights)
emit_y = w_ave(emit_slice_y, slice_weights)
return emit_x, emit_y
else:
return emittance_from_coord(x, y, ux, uy, w)
[docs]
def get_current( self, t=None, iteration=None, species=None, select=None,
bins=100, plot=False, **kw ):
"""
Calculate the electric current along the z-axis for selected particles.
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user
species : string
Particle species to use for calculations
select : dict, optional
Either None or a dictionary of rules
to select the particles, of the form
'x' : [-4., 10.] (Particles having x between -4 and 10 meters)
'z' : [0, 100] (Particles having z between 0 and 100 meters)
bins : int, optional
Number of bins along the z-axis in which to calculate the current
plot : bool, optional
Whether to plot the requested quantity
**kw : dict, otional
Additional options to be passed to matplotlib's `plot` method
Returns
-------
A tuple of arrays containig
- The current in each bin in Ampere
- A FieldMetaInformation object
(see object's docstring for more details)
"""
# Get particle data
z, uz, uy, ux, w, q = self.get_particle(
var_list=['z', 'uz', 'uy', 'ux', 'w', 'charge'],
t=t, iteration=iteration,
species=species, select=select )
# Length to be seperated in bins
if w.size > 0:
min_z = np.min(z)
len_z = np.max(z) - min_z
# Calculate Lorentz factor for all particles
gamma = np.sqrt(1 + ux ** 2 + uy ** 2 + uz ** 2)
# Calculate particle velocities
vz = uz / gamma * const.c
# Length to be seperated in bins
len_z = np.max(z) - np.min(z)
vzq_sum, _ = np.histogram(z, bins=bins, weights=(vz * w * q))
# Calculate the current in each bin
current = np.abs(vzq_sum * bins / (len_z))
else:
current = np.zeros(bins)
len_z = 0
min_z = 0
# Info object with central position of the bins
info = FieldMetaInformation( {0: 'z'}, current.shape,
grid_spacing=(len_z / bins, ), grid_unitSI=1,
global_offset=(min_z + len_z / bins / 2,), position=(0,),
t=self.current_t, iteration=self.current_iteration)
# Plot the result if needed
if plot:
check_matplotlib()
iteration = self.iterations[ self._current_i ]
time_s = self.t[ self._current_i ]
plt.plot( info.z, current, **kw)
plt.title("Current at %.2e s (iteration %d)"
% (time_s, iteration ), fontsize=self.plotter.fontsize)
plt.xlabel('$z \;(m)$', fontsize=self.plotter.fontsize)
plt.ylabel('$I \;(A)$', fontsize=self.plotter.fontsize)
# Return the current and bin centers
return(current, info)
[docs]
def get_laser_envelope( self, t=None, iteration=None, pol=None,
laser_propagation='z',
m='all', theta=0, slice_across=None,
slice_relative_position=None, plot=False,
plot_range=[[None, None], [None, None]], **kw ):
"""
Calculate a laser field by filtering out high frequencies. Can either
return the envelope slice-wise or a full 2D envelope.
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
pol : string
Polarization of the field. Options are 'x', 'y', 'z'
laser_propagation : string, optional
Coordinate along which laser field propagates.
Default is 'z'.
m : int or str, optional
Only used for thetaMode geometry
Either 'all' (for the sum of all the modes)
or an integer (for the selection of a particular mode)
theta : float, optional
Only used for thetaMode geometry
The angle of the plane of observation, with respect to the x axis
slice_across : str or list of str, optional
Direction(s) across which the data should be sliced
+ In cartesian geometry, elements can be:
- 1d: 'z'
- 2d: 'x' and/or 'z'
- 3d: 'x' and/or 'y' and/or 'z'
+ In cylindrical geometry, elements can be 'r' and/or 'z'
Returned array is reduced by 1 dimension per slicing.
If slice_across is None, the full grid is returned.
Default is None.
slice_relative_position : float or list of float, optional
Number(s) between -1 and 1 that indicate where to slice the data,
along the directions in `slice_across`
-1 : lower edge of the simulation box
0 : middle of the simulation box
1 : upper edge of the simulation box
Default is 0.
plot : bool, optional
Whether to plot the requested quantity
plot_range : list of lists
A list containing 2 lists of 2 elements each
Indicates the values between which to clip the plot,
along the 1st axis (first list) and 2nd axis (second list)
Default: plots the full extent of the simulation box
**kw : dict, otional
Additional options to be passed to matplotlib's `plot`(1D) or
`imshow` (2D) method
Returns
-------
A tuple with:
- Envelope data (1D or 2D array)
- A FieldMetaInformation object
"""
# Check if polarization has been entered
if pol not in ['x', 'y', 'z']:
raise ValueError('The `pol` argument is missing or erroneous.')
# Prevent slicing across `laser_propagation`, when extracting the raw electric field
# (`laser_propagation` axis is needed for calculation of envelope)
# but record whether the user asked for slicing across `laser_propagation`,
# and whether a corresponding `slice_relative_position` coordinate
# along `laser_propagation` was given, so as to perform this slicing later in this function.
slicing_coord_laser = None
if slice_across is not None:
slice_across, slice_relative_position = \
sanitize_slicing(slice_across, slice_relative_position)
if laser_propagation in slice_across:
index_slicing_coord_laser = slice_across.index(laser_propagation)
slice_across.pop(index_slicing_coord_laser)
slicing_coord_laser = slice_relative_position.pop(index_slicing_coord_laser)
# Get field data, and perform Hilbert transform
field, info = self.get_field( t=t, iteration=iteration, field='E',
coord=pol, theta=theta, m=m,
slice_across=slice_across,
slice_relative_position=slice_relative_position )
inverted_axes_dict = {info.axes[key]: key for key in info.axes.keys()}
e_complx = hilbert(field, axis=inverted_axes_dict[laser_propagation])
envelope = np.abs(e_complx)
# If the user asked for slicing along `laser_propagation`, do it now
if slicing_coord_laser is not None:
slicing_index = inverted_axes_dict[laser_propagation]
coord_array = getattr( info, laser_propagation )
# Number of cells along the slicing direction
n_cells = len(coord_array)
# Index of the slice (prevent stepping out of the array)
i_cell = int( 0.5 * (slicing_coord_laser + 1.) * n_cells )
i_cell = max( i_cell, 0 )
i_cell = min( i_cell, n_cells - 1)
envelope = np.take( envelope, [i_cell], axis=slicing_index )
envelope = np.squeeze(envelope)
# Remove the sliced labels from the FieldMetaInformation
info._remove_axis(laser_propagation)
# Plot the result if needed
if plot:
geometry = self.fields_metadata['E']['geometry']
field_label = 'E%s (envelope)' %pol
if envelope.ndim == 1:
self.plotter.show_field_1d(envelope, info, field_label,
self._current_i, plot_range=plot_range, **kw)
elif envelope.ndim == 2:
self.plotter.show_field_2d(envelope, info, slice_across, m,
field_label, geometry, self._current_i,
plot_range=plot_range, **kw)
# Return the result
return( envelope, info )
[docs]
def get_main_frequency( self, t=None, iteration=None, pol=None, m='all',
method='max'):
"""
Calculate the angular frequency of a laser pulse.
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
pol : string
Polarization of the field. Options are 'x', 'y'
m : int or str, optional
Only used for thetaMode geometry
Either 'all' (for the sum of all the modes)
or an integer (for the selection of a particular mode)
method : string, optional
Method which is used to calculate the frequency of the pulse
'fit' : Fit a Gaussian curve to find central frequency
'max' : Take frequency with highest intensity in the spectrum
Returns
-------
A float with mean angular frequency
"""
# Extract the spectrum
spectrum, info = self.get_spectrum( t, iteration, pol, m )
# Calculate the main frequency
# Use the maximum
i_max = np.argmax( spectrum )
omega0 = info.omega[i_max]
if method == 'max':
return( omega0 )
# Gaussian fit
elif method == 'fit':
# Guess start values for fit
f0 = omega0
fmax = np.amax( spectrum )
fsigma = info.omegamax - info.omegamin
params, _ = curve_fit( gaussian_profile, info.omega,
spectrum, p0=[f0, fmax, fsigma])
return( params[0] )
else:
raise ValueError('Unknown method: {:s}'.format(method))
[docs]
def get_spectrum( self, t=None, iteration=None, pol=None,
m='all', plot=False, **kw ):
"""
Return the spectrum of the laser
(Absolute value of the Fourier transform of the fields.)
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
pol : string
Polarization of the field. Options are 'x', 'y'
m : int or str, optional
Only used for thetaMode geometry
Either 'all' (for the sum of all the modes)
or an integer (for the selection of a particular mode)
plot: bool, optional
Whether to plot the data
**kw : dict, otional
Additional options to be passed to matplotlib's `plot` method
Returns
-------
A tuple with:
- The 1D spectrum on axis
- A FieldMetaInformation object
"""
# Check if polarization has been entered
if pol not in ['x', 'y']:
raise ValueError('The `pol` argument is missing or erroneous.')
# Get a lineout along the 'z' axis,
slice_across = self._get_slicing_for_longitudinal_lineout()
# Get field data
field1d, info = self.get_field( t=t, iteration=iteration, field='E',
coord=pol, m=m, slice_across=slice_across )
# FFT of 1d data
dt = (info.z[1] - info.z[0]) / const.c # Integration step for the FFT
fft_field = np.fft.fft(field1d) * dt
# Take half of the data (positive frequencies only)
spectrum = abs( fft_field[ : int( len(fft_field) / 2 ) ] )
# Create a FieldMetaInformation object
T = (info.zmax - info.zmin) / const.c
spect_info = FieldMetaInformation( {0: 'omega'}, spectrum.shape,
grid_spacing=( 2 * np.pi / T, ), grid_unitSI=1,
global_offset=(0,), position=(0,),
t=self.current_t, iteration=self.current_iteration )
# Plot the field if required
if plot:
check_matplotlib()
iteration = self.iterations[ self._current_i ]
time_s = self.t[ self._current_i ]
plt.plot( spect_info.omega, spectrum, **kw )
plt.xlabel('$\omega \; (rad.s^{-1})$',
fontsize=self.plotter.fontsize )
plt.ylabel('Spectrum', fontsize=self.plotter.fontsize )
plt.title("Spectrum at %.2e s (iteration %d)"
% (time_s, iteration ), fontsize=self.plotter.fontsize)
return( spectrum, spect_info )
[docs]
def get_a0( self, t=None, iteration=None, pol=None ):
"""
Gives the laser strength a0 given by a0 = Emax * e / (me * c * omega)
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
pol : string
Polarization of the field. Options are 'x', 'y'
Returns
-------
Float with normalized vector potential a0
"""
# Get a lineout along the 'z' axis,
slice_across = self._get_slicing_for_longitudinal_lineout()
# Get the peak field from field envelope
Emax = np.amax(self.get_laser_envelope(t=t, iteration=iteration,
pol=pol, slice_across=slice_across)[0])
# Get mean frequency
omega = self.get_main_frequency(t=t, iteration=iteration, pol=pol)
# Calculate a0
a0 = Emax * const.e / (const.m_e * const.c * omega)
return( a0 )
[docs]
def get_ctau( self, t=None, iteration=None, pol=None, method='fit' ):
"""
Calculate the length of a (gaussian) laser pulse. Here 'length' means
the 'longitudinal waist' (i.e sqrt(2) * sigma_z).
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
pol : string
Polarization of the field. Options are 'x', 'y'
method : str, optional
The method which is used to compute ctau
'fit': Gaussian fit of the longitudinal profile
'rms': RMS radius, weighted by the longitudinal profile
('rms' tends to give more weight to the "wings" of the pulse)
Returns
-------
Float with ctau in meters
"""
# Get a lineout along the 'z' axis,
slice_across = self._get_slicing_for_longitudinal_lineout()
# Get the field envelope
E, info = self.get_laser_envelope(t=t, iteration=iteration,
pol=pol, slice_across=slice_across)
# Calculate ctau with RMS value
ctau = np.sqrt(2) * w_std(info.z, E)
if method == 'rms':
return( ctau )
# Calculate ctau from Gaussian fit
elif method == 'fit':
# Start value for E0 and x0 in fit
E0 = np.amax( E )
z0 = info.z[np.argmax( E )]
# Perform the fit
params, _ = curve_fit( gaussian_profile, info.z,
E, p0=[ z0, E0, ctau ])
return( params[2] )
else:
raise ValueError('Unknown method: {:s}'.format(method))
[docs]
def get_laser_waist( self, t=None, iteration=None, pol=None, theta=0,
laser_propagation='z', method='fit', profile_method='peak' ):
"""
Calculate the waist of a (gaussian) laser pulse. ( sqrt(2) * sigma_r)
In 3D, this function takes a slice across `y`, and thus computes the
waist in the `x-z` plane.
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
pol : string
Polarization of the field. Options are 'x', 'y'
theta : float, optional
Only used for thetaMode geometry
The angle of the plane of observation, with respect to the x axis
laser_propagation : string, optional
Coordinate along which laser field propagates.
Default is 'z'.
method : str, optional
The method which is used to compute the waist
'fit': Gaussian fit of the transverse profile
'rms': RMS radius, weighted by the transverse profile
('rms' tends to give more weight to the "wings" of the pulse)
profile_method : str, optional, default 'peak'
Method used to obtain the transverse profile. Options are: 'peak','projection'
Returns
-------
Float with laser waist in meters
"""
# In 3D, slice across 'y' by default
geometry = self.fields_metadata['E']['geometry']
if geometry == '3dcartesian':
slice_across = 'y'
else:
slice_across = None
# Get the field envelope (as 2D array)
field, info = self.get_laser_envelope(t=t, iteration=iteration,
pol=pol, laser_propagation=laser_propagation,
slice_across=slice_across, theta=theta)
assert field.ndim == 2
# Detect direction of laser propagation
inverted_axes_dict = {info.axes[key]: key for key in info.axes.keys()}
propagation_axis_index = inverted_axes_dict[laser_propagation]
# Get transverse positions
trans_pos = getattr(info, info.axes[(propagation_axis_index+1)%2])
if profile_method == 'peak':
# Find the indices of the maximum field
i_max = np.unravel_index( np.argmax( field ), field.shape )
# Pick the corresponding transverse slice
# (Transverse to laser propagation)
trans_profile = np.take( field, [i_max[propagation_axis_index]],
axis=propagation_axis_index ).flatten()
elif profile_method == 'projection':
# Project field along the propagation direction
trans_profile = np.sum( field, axis=propagation_axis_index )
else:
raise ValueError('Unknown profile_method: {:s}'.format(profile_method))
if trans_profile.sum() == 0:
# No laser
return( np.nan )
# Compute waist with RMS value
# (serves as initial guess when method=='fit')
w0 = np.sqrt(2) * w_std(trans_pos, trans_profile)
if method == 'rms':
return( w0 )
# Compute waist with Gaussian fit
elif method == 'fit':
# Get initial guess for the amplitude
E0 = trans_profile.max()
# Assume that the pulse is centered
x0 = 0
# Perform the fit
params, _ = curve_fit( gaussian_profile, trans_pos,
trans_profile, p0=[x0, E0, w0 ])
return( params[2] )
else:
raise ValueError('Unknown method: {:s}'.format(method))
[docs]
def get_spectrogram( self, t=None, iteration=None, pol=None,
plot=False, **kw ):
"""
Calculates the spectrogram of a laserpulse, by the FROG method.
Mathematically:
$$ s(\omega, \tau) = | \int_{-\infty}^{\infty} E(t) |E(t-\tau)|^2
\exp( -i\omega t) dt |^2 $$
See Trebino, R: Frequency Resolved Optical Gating: The measurements of
Ultrashort Laser Pulses: year 2000: formula 5.2
The time is centered around the laser pulse.
Parameters
----------
t : float (in seconds), optional
Time at which to obtain the data (if this does not correspond to
an existing iteration, the closest existing iteration will be used)
Either `t` or `iteration` should be given by the user.
iteration : int
The iteration at which to obtain the data
Either `t` or `iteration` should be given by the user.
pol : string
Polarization of the laser field. Options are 'x', 'y'
plot: bool, optional
Whether to plot the spectrogram
**kw : dict, otional
Additional options to be passed to matplotlib's `imshow` method
Returns
-------
- A 2d array with spectrogram
- info : a FieldMetaInformation object
(see the corresponding docstring)
"""
# Get a lineout along the 'z' axis
slice_across = self._get_slicing_for_longitudinal_lineout()
# Get the field envelope
env, _ = self.get_laser_envelope(t=t, iteration=iteration,
pol=pol, slice_across=slice_across)
# Get the field
E, info = self.get_field( t=t, iteration=iteration, field='E',
coord=pol, slice_across=slice_across)
Nz = len(E)
# Get time domain of the data
tmin = info.zmin / const.c
tmax = info.zmax / const.c
T = tmax - tmin
dt = T / Nz
# Normalize the Envelope
env /= np.sqrt(np.trapz(env ** 2, dx=dt))
# Allocate array for the gating function and the spectrogran
E_shift = np.zeros_like(E)
spectrogram = np.zeros((2 * Nz, Nz))
# Loop over the time variable of the spectrogram
for i in range( Nz * 2):
itau = i % Nz
# Shift the E field and fill the rest with zeros
if i < Nz:
E_shift[:itau] = env[ Nz - itau: Nz]
E_shift[itau:] = 0
else:
E_shift[itau:] = env[: Nz - itau]
E_shift[:itau] = 0
EE = E * E_shift ** 2
fft_EE = np.fft.fft(EE)
spectrogram[i, :] = np.abs(fft_EE) ** 2
# Rotate and flip array to have input form of imshow
spectrogram = np.flipud(np.rot90(spectrogram[:, int(Nz / 2):]))
# Find the time at which the wigner transform is the highest
maxi, maxj = np.unravel_index(spectrogram.argmax(), spectrogram.shape)
tmin = -(T - T / spectrogram.shape[1] * maxj)
info = FieldMetaInformation( {0: 'omega', 1: 't'}, spectrogram.shape,
grid_spacing=( 2 * np.pi / T, dt / 2. ), grid_unitSI=1,
global_offset=(0, tmin), position=(0, 0),
t=self.current_t, iteration=self.current_iteration )
# Plot the result if needed
if plot:
check_matplotlib()
iteration = self.iterations[ self._current_i ]
time_s = self.t[ self._current_i ]
plt.imshow( spectrogram, extent=info.imshow_extent, aspect='auto',
**kw)
plt.title("Spectrogram at %.2e s (iteration %d)"
% (time_s, iteration ), fontsize=self.plotter.fontsize)
plt.xlabel('$t \;(s)$', fontsize=self.plotter.fontsize )
plt.ylabel('$\omega \;(rad.s^{-1})$',
fontsize=self.plotter.fontsize )
return( spectrogram, info )
def _get_slicing_for_longitudinal_lineout(self):
"""
Return the `slice_across` argument which results in a 1D slice
along `z`, for the current geometry.
"""
geometry = self.fields_metadata['E']['geometry']
if geometry == "2dcartesian":
return 'x'
elif geometry == "3dcartesian":
return ['x', 'y']
elif geometry == "thetaMode":
return 'r'
else:
raise ValueError('Unknown geometry: %s' %geometry)
def w_ave( a, weights ):
"""
Calculate the weighted average of array `a`
Parameters
----------
a : 1d array
Calculate the weighted average for these a.
weights : 1d array
An array of weights for the values in a.
Returns
-------
Float with the weighted average
Returns nan if input array is empty
"""
# Check if input contains data
if not np.any(weights) and not np.any(a):
# If input is empty return NaN
return np.nan
else:
# Calculate the weighted average
average = np.average(a, weights=weights)
return( average )
def w_std( a, weights ):
"""
Calculate the weighted standard deviation.
Parameters
----------
a : array_like
Calculate the weighted standard deviation for these a.
weights : array_like
An array of weights for the values in a.
Returns
-------
Float with the weighted standard deviation.
Returns nan if input array is empty
"""
# Check if input contains data
if not np.any(weights) and not np.any(a):
# If input is empty return NaN
return np.nan
else:
# Calculate the weighted standard deviation
average = np.average(a, weights=weights)
variance = np.average((a - average) ** 2, weights=weights)
return( np.sqrt(variance) )
def w_median(a, weights):
"""
Compute the weighted median of a 1D numpy array.
Parameters
----------
a : ndarray
Input array (one dimension).
weights : ndarray
Array with the weights of the same size of `data`.
Returns
-------
median : float
The output value.
"""
quantile = .5
if not isinstance(a, np.matrix):
a = np.asarray(a)
if not isinstance(weights, np.matrix):
weights = np.asarray(weights)
if a.shape != weights.shape:
raise TypeError("the length of data and weights must be the same")
ind_sorted = np.argsort(a)
sorted_data = a[ind_sorted]
sorted_weights = weights[ind_sorted]
Sn = np.cumsum(sorted_weights)
# Center and normalize the cumsum (i.e. divide by the total sum)
Pn = (Sn - 0.5 * sorted_weights) / Sn[-1]
# Get the value of the weighted median
return np.interp(quantile, Pn, sorted_data)
def w_mad(a, w):
"""
Compute the weighted median absolute deviation of a 1D numpy array.
Parameters
----------
a : ndarray
Input array (one dimension).
weights : ndarray
Array with the weights of the same size of `data`.
Returns
-------
mad : float
The output value.
"""
med = w_median(a, w)
mad = w_median(np.abs(a - med), w)
return mad
def gaussian_profile( x, x0, E0, w0 ):
"""
Returns a Gaussian profile with amplitude E0 and waist w0.
(Used in order to fit the transverse laser profile and find the waist.)
Parameters
----------
x: 1darray of floats
An array of transverse positions (in meters)
x0: float
Position of the peak of the profile
E0: float
The amplitude at the peak of the profile
w0: float
The waist of the profile
Returns
-------
A 1darray of floats, of the same length as x
"""
return( E0 * np.exp( -(x - x0) ** 2 / w0 ** 2 ) )
def emittance_from_coord(x, y, ux, uy, w):
"""
Calculate emittance from arrays of particle coordinates.
Parameters
----------
x : arrays of floats
x position of particles
y : arrays of floats
y position of particles
ux : arrays of floats
ux normalized momentum of particles
uy : arrays of floats
uy normalized momentum of particles
w : arrays of floats
Particle weights
Returns
-------
emit_x : float
emittance in the x direction (m*rad)
emit_y : float
emittance in the y direction (m*rad)
"""
xm = w_ave( x, w )
xsq = w_ave( (x - xm) ** 2, w )
ym = w_ave( y, w )
ysq = w_ave( (y - ym) ** 2, w )
uxm = w_ave( ux, w )
uxsq = w_ave( (ux - uxm) ** 2, w )
uym = w_ave( uy, w )
uysq = w_ave( (uy - uym) ** 2, w )
xux = w_ave( (x - xm) * (ux - uxm), w )
yuy = w_ave( (y - ym) * (uy - uym), w )
emit_x = ( abs(xsq * uxsq - xux ** 2) )**.5
emit_y = ( abs(ysq * uysq - yuy ** 2) )**.5
return emit_x, emit_y