"""Observer and wavefront objects."""
from copy import deepcopy
import numpy as np
from . import srwlib
from . import configuration
[docs]
class Observer():
[docs]
def __init__(self, centerCoord,
obsXextension=1e-3, obsYextension=1e-3,
obsXres=10e-6, obsYres=10e-6,
nx=None, ny=None):
"""Mesh at the observation plane.
Args:
centerCoord: the [x, y, z] position of the observer center.
obsXextension (optional): the width of the observation plane.
Defaults to 1e-3.
obsYextension (optional): the height of the observation plane.
Defaults to 1e-3.
obsXres (optional): the horizontal resolution of the mesh.
Defaults to 10e-6.
obsYres (optional): the vertical resolution of the mesh.
Defaults to 10e-6.
nx (optional): number of horizontal bins.
If passed, `obsXres` is ignored.
ny (optional): number of vertical bins.
If passed, `obsYres` is ignored.
Note that either the resolution or the number of bins must be passed
for both directions.
"""
self.centerCoord = centerCoord
self.obsXextension = obsXextension
self.obsYextension = obsYextension
self._computeParams()
if nx is None:
self.obsXres = obsXres
self.nx = int(self.obsXextension / self.obsXres)
else:
self.obsXres = self.obsXextension / nx
self.nx = nx
if ny is None:
self.obsYres = obsYres
self.ny = int(self.obsYextension / self.obsYres)
else:
self.obsYres = self.obsYextension / ny
self.ny = ny
self._computeAxes()
def _computeParams(self):
self.zPos = self.centerCoord[2]
self.xStart = self.centerCoord[0] - self.obsXextension / 2
self.xFin = self.centerCoord[0] + self.obsXextension / 2
self.yStart = self.centerCoord[1] - self.obsYextension / 2
self.yFin = self.centerCoord[1] + self.obsYextension / 2
def _computeAxes(self):
self.xax = np.linspace(self.xStart, self.xFin, self.nx)
self.yax = np.linspace(self.yStart, self.yFin, self.ny)
[docs]
class Stokes(srwlib.SRWLStokes):
[docs]
def __init__(self, observer, wavelength):
"""Class to represent the Stokes parameters of the radiation field.
Args:
observer: an observation mesh instance of :py:class:`~Observer`.
wavelength: the observation wavelength in nanometer.
"""
super().__init__()
self.allocate(1, observer.nx, observer.ny)
self.mesh.zStart = observer.zPos
self.mesh.eStart = configuration.HC_CONST / wavelength
self.mesh.eFin = configuration.HC_CONST / wavelength
self.mesh.xStart = observer.xStart
self.mesh.xFin = observer.xFin
self.mesh.yStart = observer.yStart
self.mesh.yFin = observer.yFin
self.wavelength = wavelength
[docs]
def getPwrDensity(self):
"""Compute the power density associated with the radiation Stokes
parameters.
Returns:
dict: a dictionary containing the two coordinate axes as long as
the matrix of the power density data in W/mm^2.
"""
nx = self.mesh.nx
ny = self.mesh.ny
data = np.reshape(self.arS[:nx*ny], (ny, nx))
xax = np.linspace(self.mesh.xStart, self.mesh.xFin, nx)
yax = np.linspace(self.mesh.yStart, self.mesh.yFin, ny)
intensity = {"xax": xax, "yax": yax, "data": data}
return intensity
[docs]
class Wavefront(srwlib.SRWLWfr):
[docs]
def __init__(self, observer, wavelength):
"""Class to represent the spatial distribution of the radiation field.
Args:
observer: an observation mesh instance of :py:class:`~Observer`.
wavelength: the observation wavelength in nanometer.
"""
super().__init__()
self.allocate(1, observer.nx, observer.ny)
self.mesh.zStart = observer.zPos
self.mesh.eStart = configuration.HC_CONST / wavelength
self.mesh.eFin = configuration.HC_CONST / wavelength
self.mesh.xStart = observer.xStart
self.mesh.xFin = observer.xFin
self.mesh.yStart = observer.yStart
self.mesh.yFin = observer.yFin
self.wavelength = wavelength
[docs]
def getWfrI(self, pol="TOT", practicalUnits=True):
"""Compute the optical of the selected polarization component.
Args:
pol (optional): one of the :py:const:`~pysrw.configuration.POL`
polarization states or an angle in degrees for any arbitrary
linear polarization.
Defaults to "TOT".
practicalUnits (optional): if True, the intensity is computed in
ph/(s mm^2 nm) otherwise the default ph/(s mm^2 0.1%BW) are
used. Defaults to True.
Returns:
dict: a dictionary containing the two coordinate axes as long as
the matrix of the intensity data.
"""
if pol in configuration.POL.keys():
data = srwlib.array("f", [0] * self.mesh.nx * self.mesh.ny)
polIdx = configuration.POL[pol]
srwlib.srwl.CalcIntFromElecField(data, self, polIdx,
0, # single-electron flux
3, # compute vs x&y coordinate
self.mesh.eStart,
0, 0) # reference x and y position
data = np.reshape(data, (self.mesh.ny, self.mesh.nx))
elif isinstance(pol, int):
polAngleRad = pol * np.pi / 180
ExR, ExI, EyR, EyI = self.getWfrField()
data = np.reshape(
(ExR*np.cos(polAngleRad) + EyR*np.sin(polAngleRad))**2 +\
(ExI*np.cos(polAngleRad) + EyI*np.sin(polAngleRad))**2,
(self.mesh.ny, self.mesh.nx))
xax = np.linspace(self.mesh.xStart, self.mesh.xFin, self.mesh.nx)
yax = np.linspace(self.mesh.yStart, self.mesh.yFin, self.mesh.ny)
data *= configuration.CONFIG.PART_CHARGE**2
if practicalUnits:
data *= 1e3 / self.wavelength
intensity = {"xax": xax, "yax": yax, "data": data}
return intensity
[docs]
def getWfrField(self):
"""Extract the array of each field component from the wavefront.
Use :py:meth:`~.Wavefront.getWfrFieldMat` if you want to extract the
field components as two-dimensional matrices.
Returns:
list: list of Re{Ex}, Im{Ex}, Re{Ey}, Re{Ey} components of the
field as arrays.
"""
ExR = np.array(self.arEx[::2])
ExI = np.array(self.arEx[1::2])
EyR = np.array(self.arEy[::2])
EyI = np.array(self.arEy[1::2])
return ExR, ExI, EyR, EyI
[docs]
def getWfrFieldMat(self):
"""Extract the matrix of each field component from the wavefront.
Returns:
list: list of Re{Ex}, Im{Ex}, Re{Ey}, Re{Ey} components of the
field as two-dimensional matrices.
"""
ExR, ExI, EyR, EyI = self.getWfrField()
ExR_mat = np.reshape(deepcopy(ExR), (self.mesh.ny, self.mesh.nx))
ExI_mat = np.reshape(deepcopy(ExI), (self.mesh.ny, self.mesh.nx))
EyR_mat = np.reshape(deepcopy(EyR), (self.mesh.ny, self.mesh.nx))
EyI_mat = np.reshape(deepcopy(EyI), (self.mesh.ny, self.mesh.nx))
return ExR_mat, ExI_mat, EyR_mat, EyI_mat
[docs]
def replaceWfrFieldWithFieldMat(self, ExR, ExI, EyR, EyI):
"""Replace the field of this wavefront with the passed matrices.
The shape of the passed matrices must be consistent with the current
mesh of the wavefront.
Args:
ExR: matrix of the real-horizontal field component.
ExI: matrix of the imaginary-horizontal field component.
EyR: matrix of the real-vertical field component.
EyI: matrix of the imaginary-vertical field component.
"""
if ((self.mesh.ny, self.mesh.nx) == np.shape(ExR)) and \
((self.mesh.ny, self.mesh.nx) == np.shape(ExI)) and \
((self.mesh.ny, self.mesh.nx) == np.shape(EyR)) and \
((self.mesh.ny, self.mesh.nx) == np.shape(EyI)):
self.arEx = srwlib.array('f')
self.arEy = srwlib.array('f')
arEx = np.zeros(2 * self.mesh.ny * self.mesh.nx)
arEy = np.zeros(2 * self.mesh.ny * self.mesh.nx)
arEx[::2] = np.reshape(ExR, (1, self.mesh.ny * self.mesh.nx))
arEx[1::2] = np.reshape(ExI, (1, self.mesh.ny * self.mesh.nx))
arEy[::2] = np.reshape(EyR, (1, self.mesh.ny * self.mesh.nx))
arEy[1::2] = np.reshape(EyI, (1, self.mesh.ny * self.mesh.nx))
self.arEx = srwlib.array('f', arEx)
self.arEy = srwlib.array('f', arEy)
else:
print("New field matrices not consistent with wavefront shape!")
[docs]
def getWfrPhase(self, pol="H"):
"""Extract the phase of wavefront.
Args:
pol (optional): the 'H' or 'V' polarization component.
Defaults to "H".
Returns:
dict: an intensity-like dictionary with the phase as 'data'.
"""
ExR, ExI, EyR, EyI = self.getWfrFieldMat()
phase = self.getWfrI() # get intensity and replace data with phase
if pol == "H":
phase["data"] = np.angle(ExR + 1j * ExI)
elif pol == "V":
phase["data"] = np.angle(EyR + 1j * EyI)
else:
raise Exception(f"invalid value {pol} for 'pol'")
return phase
[docs]
class WavefrontMultiWavelength(srwlib.SRWLWfr):
"""Class to represent the spectral distribution of the radiation field.
Notice that the points of the spectrum are uniformly sampled by SRW in the
frequency (or photon energy) axis. These points are therefore not
uniformely distributed in wavelength.
Args:
observer: an observation mesh instance of :py:class:`~Observer`.
The spectrum is computed at the centre of the observer.
The extension and the resolution of the observer are ignored.
fromWavelength: lower wavelength in nanometers.
toWavelength: upper wavelength in nanometers.
numPoints: number of wavelengths to simulate.
"""
[docs]
def __init__(self, observer, fromWavelength, toWavelength, numPoints):
super().__init__()
self.allocate(numPoints, 1, 1)
self.mesh.zStart = observer.zPos
self.mesh.eStart = configuration.HC_CONST / toWavelength
self.mesh.eFin = configuration.HC_CONST / fromWavelength
self.mesh.xStart = observer.centerCoord[0]
self.mesh.xFin = observer.centerCoord[0]
self.mesh.yStart = observer.centerCoord[1]
self.mesh.yFin = observer.centerCoord[1]
self.presFT = configuration.DOMAIN["frequency"]
self.unitElFld = configuration.SR_FIELD_UNITS["sqr_phflux"]
self.eax = np.linspace(self.mesh.eStart, self.mesh.eFin, self.mesh.ne)
self.wlax = configuration.HC_CONST / self.eax
[docs]
def getWfrField(self):
"""Extract the array of each field component from the wavefront.
Returns:
list: list of Re{Ex}, Im{Ex}, Re{Ey}, Re{Ey} components of the
field as arrays.
"""
ExR = np.array(self.arEx[::2])
ExI = np.array(self.arEx[1::2])
EyR = np.array(self.arEy[::2])
EyI = np.array(self.arEy[1::2])
return ExR, ExI, EyR, EyI
[docs]
def getSpectrumI(self, pol="TOT", practicalUnits=True):
"""Compute the spectral distribution of the selected polarization
component.
Args:
pol (optional): one of the :py:const:`~pysrw.configuration.POL`
polarization states or an angle in degrees for any arbitrary
linear polarization.
Defaults to "TOT".
practicalUnits (optional): if True, the intensity is computed in
ph/(s mm^2 nm) otherwise the default ph/(s mm^2 0.1%BW) are
used. Defaults to True.
Returns:
dict: a dictionary containing the photon energy axis
(uniform sampling), the corresponding wavelength axis and the
intensity data.
"""
if pol in configuration.POL.keys():
data = srwlib.array("f", [0] * self.mesh.ne)
polIdx = configuration.POL[pol]
srwlib.srwl.CalcIntFromElecField(data, self, polIdx,
0, # single-electron flux
0, # compute vs photon energy
self.mesh.eStart,
self.mesh.xStart, self.mesh.yStart)
elif isinstance(pol, int):
polAngleRad = pol * np.pi / 180
ExR, ExI, EyR, EyI = self.getWfrField()
data = (ExR*np.cos(polAngleRad) + EyR*np.sin(polAngleRad))**2 +\
(ExI*np.cos(polAngleRad) + EyI*np.sin(polAngleRad))**2
data = np.array(data) * configuration.CONFIG.PART_CHARGE**2
if practicalUnits:
data *= 1e3 / self.wlax
spectrum = {"wlax": self.wlax, "eax": self.eax, "data": data}
return spectrum
