"""Functions involved in the propagation of a wavefront."""
import numpy as np
from . import srwlib
[docs]
def groupOpticalSystem(opticalSystem, groupSeq):
"""Inspect the elements of an optical system and group a given sequence of
elements. Useful if the user is not interested in the intensity on some
intermediate planes (for instance just after a lens). The computation of a
sequence is more efficicent because it is performed at the SRW level and
requires also less memory cause the wavefront is overridden.
Args:
opticalSystem: the optical system dictionary.
groupSeq: list of element types to group.
Returns:
list: list of optical containers where elements are grouped according
to the specified sequence.
"""
seqLen = len(groupSeq)
optElTypes = [optEl["type"] for optEl in opticalSystem.values()]
optEls = [{key: opticalSystem[key]} for key in opticalSystem.keys()]
if len(optElTypes) < len(groupSeq):
# nothing to group
optCnts = [optEls]
else:
optCnts = []
i = 0
while i < len(optEls):
if seqLen != 0 and optElTypes[i:i + seqLen] == groupSeq:
# sequence detected, group elements in one container
optCnts.append([optEl for optEl in optEls[i:i + seqLen]])
i += seqLen
else:
# spare element, add to list of optical containers
optCnts.append([optEls[i]])
i += 1
return optCnts
[docs]
def cropOpticalSystem(opticalSystem, startPropAt=None, stopPropAt=None):
"""Crop an optical system to propagate a wavefront only within a
sub-sequence of the initial system.
Args:
opticalSystem: the dictionary of the whole optical system.
startPropAt (optional): the name of the optical element to start from.
If None, the first element is picked as start.
Defaults to None.
stopPropAt (optional): the name of the optical element to end with.
If None, the last element is picked as stop.
Defaults to None.
Returns:
dict: the dictionary of the cropped optical system.
"""
optElKeys = list(opticalSystem.keys())
startPropAt = optElKeys[0] if startPropAt is None else startPropAt
stopPropAt = optElKeys[-1] if stopPropAt is None else stopPropAt
startIdx = optElKeys.index(startPropAt)
stopIdx = optElKeys.index(stopPropAt) + 1
opticalSystemCropped = {
key: opticalSystem[key] for key in optElKeys[startIdx:stopIdx]
}
return opticalSystemCropped
# =============================== PARSE TO SRW ================================
[docs]
def parsePropPar(autoResizeBefore=False, autoResizeAfter=False,
autoResizePrecision=1.0,
quadPhaseSemiAnal=False,
resizeInFourier=False,
hRangeChange=1.0, hResChange=1.0,
vRangeChange=1.0, vResChange=1.0):
"""Get the list of parameters for the propagation through an optical
element.
Args:
autoResizeBefore (optional): if True, auto-resize the wfr before the
propagation through the element.
Defaults to False.
autoResizeBefore (optional): if True, auto-resize the wfr after the
propagation through the element.
Defaults to False.
autoResizePrecision (optional): relative precision for propagation with
auto-resizing.
Defaults to 1.0.
quadPhaseSemiAnal (optional): use semi-analytical treatment of the
quadratic (leading) phase.
Defaults to False.
resizeInFourier (optional): resize in frequency domain using FTT.
Defaults to False.
hRangeChange (optional): resize horizontal extension.
Defaults to 1.0.
hResChange (optional): resize horizontal resolution.
Defaults to 1.0.
vRangeChange (optional): resize vertical extension.
Defaults to 1.0.
vResChange (optional): resize vertical resolution.
Defaults to 1.0.
Returns:
list: [description]
"""
SRWpropPar = [int(autoResizeBefore), int(autoResizeAfter),
float(autoResizePrecision),
int(quadPhaseSemiAnal),
int(resizeInFourier),
hRangeChange, hResChange,
vRangeChange, vResChange]
return SRWpropPar
[docs]
def parseToSRWelement(optEl):
"""Get the SRW element and propagation parameter from the PySRW definition.
Args:
optEl: dictionary with the definition of a valid PySRW optical element.
Returns:
list: a list with the SRW element and the propagation parameters.
"""
optElType = optEl["type"]
if optElType == "drift":
SRWoptEl, SRWpropPar = _parseDrift(optEl)
elif optElType in ["rectAp", "rectOb", "circAp", "circOb"]:
SRWoptEl, SRWpropPar = _parseApOrObs(optEl)
elif optElType == "lens":
SRWoptEl, SRWpropPar = _parseLens(optEl)
elif optElType == "transmission":
SRWoptEl, SRWpropPar = _parseTransmission(optEl)
else:
raise Exception(f"optical element type '{optElType}' not recognized")
return SRWoptEl, SRWpropPar
def _parseDrift(optEl):
# convert PySRW drift element to SRWLOptD
SRWoptEl = srwlib.SRWLOptD(optEl["length"])
if "propParam" in optEl.keys():
SRWpropPar = parsePropPar(**optEl["propParam"])
else:
SRWpropPar = parsePropPar() # leave default values, mesh not adapted
return SRWoptEl, SRWpropPar
def _parseApOrObs(optEl):
# convert PySRW aperture/obstacle element to SRWLOptA
elType = optEl["type"].lower()
shape = "c" if "circ" in elType else "r"
apOrObs = "a" if "ap" in elType else "o"
SRWoptEl = srwlib.SRWLOptA(shape, apOrObs,
_Dx=optEl["extension"][0],
_Dy=optEl["extension"][1],
_x=optEl["centre"][0],
_y=optEl["centre"][1])
if "propParam" in optEl.keys():
SRWpropPar = parsePropPar(**optEl["propParam"])
else:
SRWpropPar = parsePropPar()
return SRWoptEl, SRWpropPar
def _parseLens(optEl):
# convert PySRW lens element to SRWLOptL
if isinstance(optEl["f"], (list, tuple)):
# cylindrical lens
focal = optEl["f"]
else:
# axial lens
focal = [optEl["f"], optEl["f"]]
SRWoptEl = srwlib.SRWLOptL(_Fx=focal[0], _Fy=focal[1],
_x=optEl["centre"][0], _y=optEl["centre"][1])
if "propParam" in optEl.keys():
SRWpropPar = parsePropPar(**optEl["propParam"])
else:
SRWpropPar = parsePropPar()
return SRWoptEl, SRWpropPar
def _parseTransmission(optEl):
# convert PySRW transmission element to SRWLOptT
opd = optEl["opd"] * 1e-9
ampl = optEl["amp"]
extension = optEl["extension"]
centre = optEl["centre"]
nx, ny = np.shape(opd)
Tr = np.zeros(2*nx*ny)
Tr[::2] = np.reshape(ampl, (nx*ny))
Tr[1::2] = np.reshape(opd, (nx*ny))
arTr = srwlib.array("d", Tr)
SRWoptEl = srwlib.SRWLOptT(_nx= nx, _ny=ny,
_rx= extension[0], _ry=extension[1],
_arTr=arTr, _x=centre[0], _y=centre[1])
if "propParam" in optEl.keys():
SRWpropPar = parsePropPar(**optEl["propParam"])
else:
SRWpropPar = parsePropPar()
return SRWoptEl, SRWpropPar
