Source code for pysrw.magnets

"""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