"""Magnetic structures."""
import numpy as np
from scipy.constants import c
from . import srwlib
from . import configuration
PREDEFINED_DIPOLES = {
"SPSarc": [743.08, 6.22, 210e-3],
"LHCD3": [6013.00, 9.45, 86e-3],
"FCCarc": [10760.00, 23.94, 107e-3],
"CLEAR": [1.12, 0.471, 3e-3]
}
"""Data of known dipole magnets arranged as ['bendingR', 'coreL', 'edgeL']."""
PREDEFINED_UNDULATORS = {
"LHC": ["v", -1, 0, 0.28, 1, 4.5]
}
"""Data of known undulators arranged as ['undPlane', 'symmetry',
'initialPhase', 'undPeriod', 'numPer', 'magField']."""
[docs]
class Magnet():
[docs]
def __init__(self, energy=None, centerCoord=[0, 0, 0]):
self.energy = energy
if energy is not None:
self.gamma = energy * \
configuration.CONFIG.PART_CHARGE / configuration.CONFIG.PART_MASS
else:
self.gamma = None
self.centerCoord = centerCoord
def _calcExtremes(self, totalLength):
self.minZ = self.centerCoord[2] - 0.5 * (totalLength)
self.maxZ = self.centerCoord[2] + 0.5 * (totalLength)
@staticmethod
def _scaleMagField(magField):
return - magField * configuration.CONFIG.B_SCALING
[docs]
class Dipole(Magnet):
[docs]
def __init__(self, energy, coreL, edgeL,
magField=None, bendingR=None, centerCoord=[0, 0, 0]):
"""A dipole magnet with a decaying fringe field.
Args:
energy: beam energy in GeV.
coreL: dipole magnetic length.
edgeL: characteristic length of the fringe field.
If not known, pass yoke diameter as rough estimate.
magField (optional): dipole constant field in T.
Ignored if `bendingR` is passed as well.
bendingR (optional): particle bending radius.
If provided, the magField is computed from the bending radius.
centerCoord (optional): magnet centre position.
Defaults to [0, 0, 0].
"""
super().__init__(energy=energy, centerCoord=centerCoord)
if bendingR is not None:
self.bendingR = bendingR
self.magField = energy * 1e9 / (np.abs(bendingR) * c)
elif magField is not None:
self.magField = magField
self.bendingR = energy *1e9 / (c * magField)
else:
raise Exception("either bendingR or magField required!")
self.magFieldScaled = self._scaleMagField(self.magField)
self.coreL = coreL
self.edgeL = edgeL
fieldParam = 1 # dipolar field
normalOrSkew = "n"
self.srwElement = srwlib.SRWLMagFldM(self.magFieldScaled,
fieldParam, normalOrSkew,
coreL, edgeL)
self._calcExtremes(coreL + 2 * edgeL)
[docs]
def getLambdaCritical(self):
"""Get the dipole critical wavelength.
Returns:
float: :math:`\lambda_c` in nanometers.
"""
return 4 * np.pi / 3 * self.bendingR / self.gamma**3 * 1e9
[docs]
class Quadrupole(Magnet):
[docs]
def __init__(self, coreL, edgeL, magGradient, centerCoord=[0, 0, 0]):
"""A (normal) quadrupole magnet.
Args:
coreL: the magnetic length.
edgeL: characteristic length of the fringe field.
If not known, pass yoke diameter as rough estimate.
magGradient: the magnetic gradient in T/m.
centerCoord (optional): magnet centre position.
Defaults to [0, 0, 0].
"""
# TODO check the gradient definition / normalization
super().__init__(centerCoord=centerCoord)
self.coreL = coreL
self.edgeL = edgeL
self.magGradientScaled = self._scaleMagField(magGradient)
fieldParam = 2 # quadrupolar field
normalOrSkew = "n"
bendingR = 0
self.srwElement = srwlib.SRWLMagFldM(self.magGradientScaled,
fieldParam, normalOrSkew,
coreL, edgeL, bendingR)
self._calcExtremes(coreL + 2 * edgeL)
[docs]
class UndulatorPlanar(Magnet):
[docs]
def __init__(self, energy, undPeriod, numPer, magField,
undPlane="v", symmetry=-1, initialPhase=0,
centerCoord=[0,0,0]):
"""A planar undulator with sinusoidal magnetic field which smoothly
decays approaching the edge.
Args:
energy: particle energy in GeV.
undPeriod: periodicity of the magnetic field.
numPer: number of periods.
magField: peak magnetic field in the gap.
undPlane (optional): 'h' or 'v' to indicate the direction of the
magnetic field. Defaults to 'v'.
symmetry (optional): 1 for symmetric (cos like) or -1 for
anti-symmetric (sin like) field profile. Defaults to -1.
initialPhase (optional): initial phase of the field oscillations.
Defaults to 0.
centerCoord (optional): undulator centre position.
Defaults to [0, 0, 0].
"""
super().__init__(energy=energy, centerCoord=centerCoord)
self.undPlane = undPlane
self.undPeriod = undPeriod
self.numPer = numPer
self.magField = magField
self.magFieldScaled = self._scaleMagField(magField)
self.initialPhase = initialPhase
self.symmetry = symmetry
harmonicNumber = 1 # doesn't matter for single harmonic undulator
transFieldCoeff = 1 #
# create an undulator with a single harmonic component
harmonic = srwlib.SRWLMagFldH(harmonicNumber, undPlane,
self.magFieldScaled, initialPhase, symmetry,
transFieldCoeff)
self.srwElement = srwlib.SRWLMagFldU([harmonic],
self.undPeriod, self.numPer)
self._calcExtremes(self.numPer * self.undPeriod)
[docs]
def getK(self):
"""Compute the deflection parameter of the undulator.
Returns:
float: the K parameter.
"""
return self.srwElement.get_K() * \
configuration.CONFIG.PART_CHARGE / configuration.CONFIG.PART_MASS
[docs]
def getLambda1(self):
"""Compute the wavelength of the on-axis fundamental harmonic.
Returns:
float: value of :math:`\lambda_1` in nanometers.
"""
K = self.getK()
lambda1 = self.undPeriod / (2 * self.gamma**2) * (1 + K**2/2) * 1e9
return lambda1
[docs]
class CustomMagnet(Magnet):
[docs]
def __init__(self, magField, ranges,
numPer=1, centerCoord=[0, 0, 0], interp="bicubic"):
"""A magnetic device with an arbitrary field.
Args:
magField: a list with [Bx, By, Bz]. Each element is a numpy
tensor which determines the 3D distribution of the
corresponding field component.
ranges: a list with [rx, ry, rz] which determines the absolute
extension of the magnetic region in meters.
numPer (optional): number of repetition of the custom field,
useful to create periodic devices with custom fields.
Defaults to 1.
centerCoord (optional): magnet centre position.
Defaults to [0, 0, 0].
interp (optional): one of the :py:data:`.configuration.MAG_INTERP`
interpolation strategies to adapt the `magField` to the mesh of
the device. Defaults to 'bicubic'.
"""
super().__init__(centerCoord=centerCoord)
# TODO check consistency of sign convention
self.magField = np.array(magField)
self.magFieldScaled = - self._scaleMagField(self.magField)
Bx = self.magFieldScaled[0]
By = self.magFieldScaled[1]
Bz = self.magFieldScaled[2]
rx, ry, rz = ranges
nx, ny, nz = np.shape(Bx)
n = nx * ny * nz
arBx = srwlib.array("d", np.reshape(Bx, n, order="F"))
arBy = srwlib.array("d", np.reshape(By, n, order="F"))
arBz = srwlib.array("d", np.reshape(Bz, n, order="F"))
_interp = configuration.MAG_INTERP[interp]
self.srwElement = srwlib.SRWLMagFld3D(_arBx=arBx,_arBy=arBy,_arBz=arBz,
_nx=nx, _ny=ny, _nz=nz,
_rx=rx, _ry=ry, _rz=rz,
_nRep=numPer, _interp=_interp)
self._calcExtremes(numPer * rz)
[docs]
class MagnetsContainer():
[docs]
def __init__(self, magnets, lengthPadFraction=1.0, lengthResolution=50e-6,
forceMinZ=None, forceMaxZ=None):
"""Container of the magnetic elements. This class groups the magnetic
devices and defines the integration parameters for the trajectory
simulation.
Args:
magnets: list of magnetic elements, instances of the classes
available in :py:mod:`~.magnets`.
lengthPadFraction (optional): extra distance as a fraction of the
minimum container extension. Defaults to 1.0.
lengthResolution (optional): discretization of the longitudinal
coordinate for the computation of the trajectory, in meters.
Defaults to 50e-6.
forceMinZ (optional): user-defined position to start the trajectory
computation. If None, the start is assumed at the
entrance of the container, adding the extra margin specified
by `lengthPadFraction`. Defaults to None.
forceMaxZ (optional): same as `forceMinZ` for the integration end.
Defaults to None.
"""
# TODO implement (and test!) also tilt and rotation of elements
self.lengthPadFraction = lengthPadFraction
self.forceMinZ = forceMinZ
self.forceMaxZ = forceMaxZ
self.lengthResolution = lengthResolution
self.sortedMagnets = self._sortMagnets(magnets)
srwlMagnets = [magnet.srwElement for magnet in self.sortedMagnets]
xPosList = [magnet.centerCoord[0] for magnet in self.sortedMagnets]
yPosList = [magnet.centerCoord[1] for magnet in self.sortedMagnets]
zPosList = [magnet.centerCoord[2] for magnet in self.sortedMagnets]
self.srwCnt = srwlib.SRWLMagFldC(srwlMagnets,
srwlib.array("d", xPosList),
srwlib.array("d", yPosList),
srwlib.array("d", zPosList))
[docs]
def getIntegrationParams(self):
"""Get the start, end and discretization step of the longitudinal
coordinate used to compute the particle trajectory.
Returns:
list: parameters for the trajectory computation.
"""
if (self.forceMinZ is None) or (self.forceMaxZ is None):
minZ = min([magnet.minZ for magnet in self.sortedMagnets])
maxZ = max([magnet.maxZ for magnet in self.sortedMagnets])
lengthPad = (maxZ - minZ) * self.lengthPadFraction / 2
minZ = minZ - lengthPad
maxZ = maxZ + lengthPad
else:
minZ = self.forceMinZ
maxZ = self.forceMaxZ
nPts = int((maxZ - minZ) / self.lengthResolution)
return minZ, maxZ, nPts
def _sortMagnets(self, magnets):
return sorted(magnets, key=lambda x: x.centerCoord[2])
# ============================ PREDEFINED MAGNETS =============================
[docs]
def predefinedDipole(dipoleType, energy, centerCoord=[0,0,0]):
"""Create a dipole from the list of known devices
:py:data:`PREDEFINED_DIPOLES`.
Args:
dipoleType: the name of the known dipole.
energy: the beam energy, needed to retrieve the field from the bending
radius.
centerCoord (optional): the [x,y,z] of the dipole centre.
Defaults to [0,0,0].
Returns:
Dipole: the predefined dipole.
"""
args = {
"bendingR": PREDEFINED_DIPOLES[dipoleType][0],
"coreL": PREDEFINED_DIPOLES[dipoleType][1],
"edgeL": PREDEFINED_DIPOLES[dipoleType][2],
"energy": energy,
"centerCoord": centerCoord
}
return Dipole(**args)
[docs]
def predefinedUndulator(undType, energy, centerCoord=[0,0,0]):
"""Create a (planar) undulator from the list of known devices
:py:data:`PREDEFINED_UNDULATORS`.
Args:
undType: the name of the known undulator.
energy: the beam energy in GeV.
centerCoord (optional): the [x,y,z] of the dipole centre.
Defaults to [0,0,0].
Returns:
UndulatorPlanar: the predefined (planar) undulator.
"""
args = {
"undPlane": PREDEFINED_UNDULATORS[undType][0],
"symmetry": PREDEFINED_UNDULATORS[undType][1],
"initialPhase": PREDEFINED_UNDULATORS[undType][2],
"undPeriod": PREDEFINED_UNDULATORS[undType][3],
"numPer": PREDEFINED_UNDULATORS[undType][4],
"magField": PREDEFINED_UNDULATORS[undType][5],
"energy": energy,
"centerCoord": centerCoord
}
return UndulatorPlanar(**args)
[docs]
def predefinedLHCsource(energy, sourceFlag="D3undD3"):
"""Create the magnet container including a user-defined set of LHC
source devices. The Beam 1 is taken as reference since the magnetic
structure is perfectly symmetric for Beam 2. The LHC source consists of
three devices, two D3-type dipoles and one undulator. The undulator and the
D3 dipole at the right side of Pt4 (D3R) are the main sources at
injection and flat-top energy respectively. The D3 dipole at the left side
of Pt4 (D3L) is negligible in most practical cases but still influences
the light emitted at flat-top energy.
Args:
energy: the beam energy in GeV.
sourceFlag (optional): the set of source devices to include in the
container. Allowed values are 'D3undD3' for the complete source,
'und' for the undulator only, 'D3L' or 'D3R' for either two
dipoles, 'undD3' for the undulator and the D3R and 'D3D3' for the
two dipoles without undulator.
Defaults to 'D3undD3'.
Returns:
MagnetsContainer: the container with the selected set of source
devices.
"""
und = predefinedUndulator(undType="LHC", energy=energy,
centerCoord=[0,0,0])
D3L = predefinedDipole(dipoleType="LHCD3", energy=energy,
centerCoord=[0, 0, -97.15])
D3R = predefinedDipole(dipoleType="LHCD3", energy=energy,
centerCoord=[0, 0, 5.85])
if sourceFlag == "und":
magnets = [und]
params = [-3, 3, 50e-6]
elif sourceFlag == "undD3":
magnets = [und, D3R]
params = [-3, 15, 50e-6]
elif sourceFlag == "D3R":
magnets = [D3R]
params = [-3, 15, 50e-6]
elif sourceFlag == "D3L":
magnets = [D3L]
params = [-103, -91, 50e-6]
elif sourceFlag == "D3D3":
magnets = [D3L, D3R]
params = [-103, 15, 100e-6]
elif sourceFlag == "D3undD3":
magnets = [D3L, und, D3R]
params = [-103, 15, 100e-6]
else:
raise Exception(f"unrecognized 'sourceFlag' {sourceFlag}")
minZ, maxZ, res = params
magCnt = MagnetsContainer(magnets,
forceMinZ=minZ, forceMaxZ=maxZ,
lengthResolution=res)
return magCnt