from enum import Enum
from warnings import warn
from math import pi
from typing import Optional, Tuple
import numpy
from scipy.optimize import least_squares
from at.lattice import Lattice, Dipole, Wiggler, RFCavity, Refpts, EnergyLoss
from at.lattice import check_radiation, AtError, AtWarning
from at.lattice import QuantumDiffusion, Collective, SimpleQuantDiff
from at.lattice import get_bool_index, set_value_refpts
from at.constants import clight, Cgamma
from at.tracking import internal_lpass
__all__ = ['get_energy_loss', 'set_cavity_phase', 'ELossMethod',
'get_timelag_fromU0']
[docs]
class ELossMethod(Enum):
"""methods for the computation of energy losses"""
#: The losses are obtained from
#: :math:`E_{loss}=C_\gamma/2\pi . E^4 . I_2`.
#: Takes into account bending magnets and wigglers.
INTEGRAL = 1
#: The losses are obtained by tracking without cavities.
#: Needs radiation ON, takes into account all radiating elements
TRACKING = 2
[docs]
def get_energy_loss(ring: Lattice,
method: Optional[ELossMethod] = ELossMethod.INTEGRAL
) -> float:
"""Computes the energy loss per turn
Parameters:
ring: Lattice description
method: Method for energy loss computation.
See :py:class:`ELossMethod`.
Returns:
eloss (float): Energy loss per turn [eV]
"""
# noinspection PyShadowingNames
def integral(ring):
"""Losses = Cgamma / 2pi * EGeV^4 * i2
"""
def wiggler_i2(wiggler: Wiggler):
rhoinv = wiggler.Bmax / ring.BRho
coefh = wiggler.By[1, :]
coefv = wiggler.Bx[1, :]
return wiggler.Length * (numpy.sum(coefh * coefh) + numpy.sum(
coefv*coefv)) * rhoinv ** 2 / 2
def dipole_i2(dipole: Dipole):
return dipole.BendingAngle ** 2 / dipole.Length
def eloss_i2(eloss: EnergyLoss):
return eloss.EnergyLoss / coef
i2 = 0.0
coef = Cgamma / 2.0 / pi * ring.energy ** 4
for el in ring:
if isinstance(el, Dipole):
i2 += dipole_i2(el)
elif isinstance(el, Wiggler) and el.PassMethod != 'DriftPass':
i2 += wiggler_i2(el)
elif isinstance(el, EnergyLoss) and el.PassMethod != 'IdentityPass':
i2 += eloss_i2(el)
e_loss = coef * i2
return e_loss
# noinspection PyShadowingNames
@check_radiation(True)
def tracking(ring):
"""Losses from tracking
"""
ringtmp = ring.disable_6d(RFCavity, QuantumDiffusion, Collective,
SimpleQuantDiff, copy=True)
o6 = numpy.squeeze(internal_lpass(ringtmp, numpy.zeros(6),
refpts=len(ringtmp)))
if numpy.isnan(o6[0]):
dp = 0
for e in ringtmp:
ot = numpy.squeeze(internal_lpass([e], numpy.zeros(6)))
dp += -ot[4] * ring.energy
return dp
else:
return -o6[4] * ring.energy
if isinstance(method, str):
method = ELossMethod[method.upper()]
warn(FutureWarning('You should use {0!s}'.format(method)))
if method is ELossMethod.INTEGRAL:
return ring.periodicity * integral(ring)
elif method == ELossMethod.TRACKING:
return ring.periodicity * tracking(ring)
else:
raise AtError('Invalid method: {}'.format(method))
# noinspection PyPep8Naming
[docs]
def get_timelag_fromU0(ring: Lattice,
method: Optional[ELossMethod] = ELossMethod.TRACKING,
cavpts: Optional[Refpts] = None,
divider: Optional[int] = 4,
ts_tol: Optional[float] = 1.0e-9) -> Tuple[float, float]:
"""
Get the TimeLag attribute of RF cavities based on frequency,
voltage and energy loss per turn, so that the synchronous phase is zero.
An error occurs if all cavities do not have the same frequency.
Used in set_cavity_phase()
Parameters:
ring: Lattice description
method: Method for energy loss computation.
See :py:class:`ELossMethod`.
cavpts: Cavity location. If None, use all cavities.
This allows to ignore harmonic cavities.
divider: number of segments to search for ts
phis_tol: relative tolerance for ts calculation
Returns:
timelag (float): Timelag
ts (float): Time difference with the present value
"""
def singlev(values):
vals = numpy.unique(values)
if len(vals) > 1:
raise AtError('values not equal for all cavities')
return vals[0]
def eq(x, freq, rfv, tl0, u0):
omf = 2*numpy.pi*freq/clight
if u0 > 0.0:
eq1 = (numpy.sum(-rfv*numpy.sin(omf*(x-tl0)))-u0)/u0
else:
eq1 = numpy.sum(-rfv * numpy.sin(omf * (x - tl0)))
eq2 = numpy.sum(-omf*rfv*numpy.cos(omf*(x-tl0)))
if eq2 > 0:
return numpy.sqrt(eq1**2+eq2**2)
else:
return abs(eq1)
if cavpts is None:
cavpts = get_bool_index(ring, RFCavity)
u0 = get_energy_loss(ring, method=method) / ring.periodicity
freq = numpy.array([cav.Frequency for cav in ring.select(cavpts)])
rfv = numpy.array([cav.Voltage for cav in ring.select(cavpts)])
tl0 = numpy.array([cav.TimeLag for cav in ring.select(cavpts)])
try:
frf = singlev(freq)
tml = singlev(tl0)
except AtError:
ctmax = clight/numpy.amin(freq)/2
tt0 = tl0[numpy.argmin(freq)]
bounds = (-ctmax, ctmax)
args = (freq, rfv, tl0, u0)
r = []
for i in range(divider):
fact = (i+1)/divider
r.append(least_squares(eq, bounds[0]*fact+tt0,
args=args, bounds=bounds+tt0))
r.append(least_squares(eq, bounds[1]*fact+tt0,
args=args, bounds=bounds+tt0))
res = numpy.array([ri.fun[0] for ri in r])
ok = res < ts_tol
vals = numpy.array([abs(ri.x[0]).round(decimals=6) for ri in r])
if not numpy.any(ok):
raise AtError('No solution found for Phis, please check '
'RF settings')
if len(numpy.unique(vals[ok])) > 1:
warn(AtWarning('More than one solution found for Phis: use '
'best fit, please check RF settings'))
ts = -r[numpy.argmin(res)].x[0]
timelag = ts+tl0
else:
if u0 > numpy.sum(rfv):
raise AtError('Not enough RF voltage: unstable ring')
vrf = numpy.sum(rfv)
timelag = clight/(2*numpy.pi*frf)*numpy.arcsin(u0/vrf)
ts = timelag - tml
timelag *= numpy.ones(ring.refcount(cavpts))
return timelag, ts
[docs]
def set_cavity_phase(ring: Lattice,
method: ELossMethod = ELossMethod.TRACKING,
refpts: Optional[Refpts] = None,
cavpts: Optional[Refpts] = None,
copy: bool = False) -> None:
"""
Adjust the TimeLag attribute of RF cavities based on frequency,
voltage and energy loss per turn, so that the synchronous phase is zero.
An error occurs if all cavities do not have the same frequency.
.. Warning::
This function changes the time reference, this should be avoided
Parameters:
ring: Lattice description
method: Method for energy loss computation.
See :py:class:`ELossMethod`.
cavpts: Cavity location. If None, use all cavities.
This allows to ignore harmonic cavities.
copy: If True, returns a shallow copy of ring with new
cavity elements. Otherwise, modify ring in-place.
"""
# refpts is kept for backward compatibility
if cavpts is None and refpts is not None:
warn(FutureWarning('You should use "cavpts" instead of "refpts"'))
cavpts = refpts
elif cavpts is None:
cavpts = get_bool_index(ring, RFCavity)
timelag, _ = get_timelag_fromU0(ring, method=method, cavpts=cavpts)
set_value_refpts(ring, cavpts, 'TimeLag', timelag, copy=copy)
Lattice.get_energy_loss = get_energy_loss
Lattice.energy_loss = property(get_energy_loss)
Lattice.set_cavity_phase = set_cavity_phase
