def _get_rf_unique_attr(ring, attr, cavpts=None):
freq = numpy.unique(_get_rf_attr(ring, attr, cavpts=cavpts))
if len(freq) == 0:
raise AtError('No cavity found in the lattice')
elif len(freq) > 1:
raise AtError('{0} not equal for all cavities'.format(attr))
else:
return freq[0]
def get_timelag_fromU0(ring, method=ELossMethod.INTEGRAL, cavpts=None):
"""
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
KEYWORDS
method=ELossMethod.INTEGRAL
method for energy loss computation.
See "get_energy_loss".
cavpts=None Cavity location. If None, use all cavities.
This allows to ignore harmonic cavities.
RETURN
timelag
"""
if cavpts is None:
cavpts = get_cells(ring, checktype(RFCavity))
u0 = get_energy_loss(ring, method=method)
try:
rfv = ring.get_rf_voltage(cavpts=cavpts)
freq = ring.get_rf_frequency(cavpts=cavpts)
tl0 = ring.get_rf_timelag(cavpts=cavpts)
except AtError:
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)])
if u0 > numpy.sum(rfv):
raise AtError('Not enough RF voltage: unstable ring')
ctmax = 1 / numpy.amin(freq) * clight / 2
tt0 = tl0[numpy.argmin(freq)]
eq = lambda x: numpy.sum(rfv * numpy.sin(2 * pi * freq *
(x + tl0) / clight)) - u0
deq = lambda x: numpy.sum(rfv * numpy.cos(2 * pi * freq *
(x + tl0) / clight))
zero_diff = least_squares(deq,
-tt0 + ctmax / 2,
bounds=(-tt0, ctmax - tt0)).x[0]
if numpy.sign(deq(zero_diff - 1.0e-6)) > 0:
ts = least_squares(eq, (zero_diff - tt0) / 2,
bounds=(-tt0, zero_diff)).x[0]
else:
ts = least_squares(eq, (ctmax - tt0 + zero_diff) / 2,
bounds=(zero_diff, ctmax - tt0)).x[0]
timelag = ts + tl0
else:
if u0 > rfv:
raise AtError('Not enough RF voltage: unstable ring')
timelag = clight / (2 * pi * freq) * numpy.arcsin(u0 / rfv)
ts = timelag - tl0
timelag *= numpy.ones(refpts_len(ring, cavpts))
return timelag, ts
def reswall(self, wcomp, length, rvac, conduct, beta, yokoya_factor=1):
if wcomp is WakeComponent.Z:
raise AtError('Resitive wall not available '
'for WakeComponent: {}'.format(wcomp))
elif wcomp is (WakeComponent.DX or wcomp is WakeComponent.DY or
wcomp is WakeComponent.QX or wcomp is WakeComponent.QY):
return wake_functions.transverse_reswall(self._srange,
yokoya_factor, length,
rvac, conduct, beta)
else:
raise AtError('Invalid WakeComponent: {}'.format(wcomp))
def __init__(self, *args, **kwargs):
"""Lattice constructor"""
iterator = kwargs.pop('iterator', None)
scan = kwargs.pop('scan', False)
for key in list(k for k in kwargs.keys() if k.startswith('_')):
kwargs.pop(key)
if iterator is None:
elem_list, = args or [[]] # accept 0 or 1 argument
if isinstance(elem_list, Lattice):
elems = lattice_filter(kwargs, elem_list)
elif scan:
elems = params_filter(kwargs, type_filter, elem_list)
else:
elems = type_filter(kwargs, elem_list)
else:
elems = iterator(kwargs, *args)
super(Lattice, self).__init__(elems)
# set default values
kwargs.setdefault('name', '')
kwargs.setdefault('periodicity', 1)
if 'energy' not in kwargs:
raise AtError('Lattice energy is not defined')
# set attributes
self.update(kwargs)
def _orbit6(ring, cavpts=None, guess=None, keep_lattice=False, **kwargs):
"""Solver for 6D motion"""
convergence = kwargs.pop('convergence', DConstant.OrbConvergence)
max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter)
xy_step = kwargs.pop('XYStep', DConstant.XYStep)
dp_step = kwargs.pop('DPStep', DConstant.DPStep)
method = kwargs.pop('method', ELossMethod.TRACKING)
# Get revolution period
l0 = get_s_pos(ring, len(ring))
cavities = [elm for elm in ring if isinstance(elm, elements.RFCavity)]
if len(cavities) == 0:
raise AtError('No cavity found in the lattice.')
f_rf = cavities[0].Frequency
harm_number = cavities[0].HarmNumber
if guess is None:
_, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts)
ref_in = numpy.zeros((6,), order='F')
ref_in[5] = -dt
else:
ref_in = numpy.copy(guess)
theta = numpy.zeros((6,))
theta[5] = constants.speed_of_light * harm_number / f_rf - l0
scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + \
dp_step * numpy.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0])
delta_matrix = numpy.asfortranarray(
numpy.concatenate((numpy.diag(scaling), numpy.zeros((6, 1))), axis=1))
id6 = numpy.asfortranarray(numpy.identity(6))
change = 1
itercount = 0
while (change > convergence) and itercount < max_iterations:
in_mat = ref_in.reshape((6, 1)) + delta_matrix
_ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice)
# the reference particle after one turn
ref_out = in_mat[:, 6]
# 6x6 jacobian matrix from numerical differentiation:
# f(x+d) - f(x) / d
j6 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling
a = j6 - id6 # f'(r_n) - 1
b = ref_out[:] - ref_in[:] - theta
# b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1)
b_over_a = numpy.linalg.solve(a, b)
r_next = ref_in - b_over_a
# determine if we are close enough
change = numpy.linalg.norm(r_next - ref_in)
itercount += 1
ref_in = r_next
keep_lattice = True
if itercount == max_iterations:
warnings.warn(AtWarning('Maximum number of iterations reached. '
'Possible non-convergence'))
return ref_in
def a_matrix(tt):
"""A, eigval = a_matrix(T)
Find the A matrix from one turn map matrix T such that:
[Rotx 0 0 ]
inv(A)*T*A=[ 0 Rotz 0 ]
[ 0 0 Rots]
Order it so that it is close to the order of x,y,z
also ensure that positive modes are before negative so that
one has proper symplecticity
B. Nash July 18, 2013
INPUT
T (m, m) transfer matrix for 1 turn
OUTPUT
A (m, m) A-matrix
eigval (m,) vector of Eigen values of T
"""
nv = tt.shape[0]
dms = int(nv / 2)
jmt = jmatswap(dms)
select = numpy.arange(0, nv, 2)
rbase = numpy.stack((select, select), axis=1).flatten()
# noinspection PyTupleAssignmentBalance
lmbd, vv = eig(tt)
# Compute the norms
vp = numpy.dot(vv.conj().T, jmt)
n = -0.5j * numpy.sum(vp.T * vv, axis=0)
if any(abs(n) < 1.0E-12):
raise AtError('Unstable ring')
# Move positive before negatives
order = rbase + (n < 0)
vv = vv[:, order]
n = n[order]
lmbd = lmbd[order]
# Normalize vectors
vn = vv / numpy.sqrt(abs(n)).reshape((1, nv))
# find the vectors that project most onto x,y,z, and reorder
# nn will have structure
# n1x n1y n1z
# n2x n2y n2z
# n3x n3y n3z
nn = 0.5 * abs(numpy.sqrt(-1.j * vn.conj().T.dot(jmt).dot(_vxyz[dms - 1])))
rows = list(select)
order = []
for ixz in select:
ind = numpy.argmax(nn[rows, ixz])
order.append(rows[ind])
del rows[ind]
v_ordered = vn[:, order]
lmbd = lmbd[order]
aa = numpy.vstack((numpy.real(v_ordered), numpy.imag(v_ordered))).reshape(
(nv, nv), order='F')
return aa, lmbd
def set_cavity(ring,
Voltage=None,
Frequency=None,
HarmNumber=None,
TimeLag=None,
cavpts=None,
copy=False):
"""
PARAMETERS
ring lattice description
KEYWORDS
Frequency=None RF frequency. The special enum value Frf.NOMINAL
sets the frequency to the nominal value, given
ring length and harmonic number.
Voltage=None RF voltage for the full ring.
HarmNumber=None Harmonic number for the full ring.
TimeLag=None
cavpts=None Cavity location. If None, use all cavities.
This allows to ignore harmonic cavities.
copy=False If True, returns a shallow copy of ring with new
cavity elements. Otherwise, modify ring in-place
"""
def get_harm(hcell):
if hcell is not None:
return hcell * n_periods
else:
return ring.get_rf_harmonic_number(cavpts=cavpts)
if cavpts is None:
cavpts = get_cells(ring, checktype(elements.RFCavity))
n_cavities = ring.refcount(cavpts)
if n_cavities < 1:
raise AtError('No cavity found in the lattice')
n_periods = ring.periodicity
modif = {}
if Frequency is not None:
if Frequency is Frf.NOMINAL:
Frequency = (clight / ring.circumference) * get_harm(HarmNumber)
modif['Frequency'] = Frequency
if HarmNumber is not None:
modif['HarmNumber'] = HarmNumber / n_periods
if TimeLag is not None:
modif['TimeLag'] = TimeLag
if Voltage is not None:
modif['Voltage'] = Voltage / n_periods / n_cavities
# noinspection PyShadowingNames
@make_copy(copy)
def apply(ring, cavpts, modif):
for cavity in ring.select(cavpts):
cavity.update(modif)
return apply(ring, cavpts, modif)
def add(self, wtype, wcomp, *args, **kwargs):
if wtype is WakeType.FILE:
w = self.readwakefile(*args, **kwargs)
elif wtype is WakeType.TABLE:
w = self.resample(*args)
elif wtype is WakeType.RESONATOR:
w = self.resonator(wcomp, *args, **kwargs)
elif wtype is WakeType.RESWALL:
w = self.reswall(wcomp, *args, **kwargs)
else:
raise AtError('Invalid WakeType: {}'.format(wtype))
if self.components[wcomp] is None:
self.components[wcomp] = w
else:
self.components[wcomp] += w
def resonator(self,
wcomp,
frequency,
qfactor,
rshunt,
beta,
yokoya_factor=1):
if wcomp is WakeComponent.Z:
return wake_functions.long_resonator(self._srange, frequency,
qfactor, rshunt, beta)
elif (wcomp is WakeComponent.DX or wcomp is WakeComponent.DY
or wcomp is WakeComponent.QX or wcomp is WakeComponent.QY):
return wake_functions.transverse_resonator(self._srange, frequency,
qfactor, rshunt,
yokoya_factor, beta)
else:
raise AtError('Invalid WakeComponent: {}'.format(wcomp))
def matfile_generator(params, mat_file):
"""
run through matlab cells and generate AT elements
KEYWORDS
mat_file name of the .mat file
mat_key name of the Matlab variable containing the lattice.
Default: Matlab variable name if there is only one,
otherwise 'RING'
check=True if False, skip the coherence tests
quiet=False If True, suppress the warning for non-standard classes
"""
def mclean(data):
if data.dtype.type is numpy.str_:
# Convert strings in arrays back to strings.
return str(data[0])
elif data.size == 1:
v = data[0, 0]
if issubclass(v.dtype.type, numpy.void):
for f in v.dtype.fields:
v[f] = mclean(v[f])
# Return a scalar
return v
else:
# Remove any surplus dimensions in arrays.
return numpy.squeeze(data)
m = scipy.io.loadmat(params.setdefault('mat_file', mat_file))
matvars = [varname for varname in m if not varname.startswith('__')]
default_key = matvars[0] if (len(matvars) == 1) else 'RING'
key = params.setdefault('mat_key', default_key)
if key not in m.keys():
kok = [k for k in m.keys() if '__' not in k]
raise AtError('Selected mat_key does not exist, '
'please select in: {}'.format(kok))
check = params.pop('check', True)
quiet = params.pop('quiet', False)
cell_array = m[key].flat
for index, mat_elem in enumerate(cell_array):
elem = mat_elem[0, 0]
kwargs = {f: mclean(elem[f]) for f in elem.dtype.fields}
yield element_from_dict(kwargs, index=index, check=check, quiet=quiet)
def _orbit6(ring, cavpts=None, guess=None, keep_lattice=False, **kwargs):
"""Solver for 6D motion"""
def iscavity(elem):
return isinstance(elem, elements.RFCavity) and \
elem.PassMethod.endswith('CavityPass')
convergence = kwargs.pop('convergence', DConstant.OrbConvergence)
max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter)
xy_step = kwargs.pop('XYStep', DConstant.XYStep)
dp_step = kwargs.pop('DPStep', DConstant.DPStep)
method = kwargs.pop('method', ELossMethod.TRACKING)
l0 = get_s_pos(ring, len(ring))[0]
# Get the main RF frequency (the lowest)
try:
f_rf = min(elm.Frequency for elm in ring if iscavity(elm))
except ValueError:
raise AtError('No cavity found in the lattice.')
# gamma = self.energy / self.particle.mass
# beta = math.sqrt(1.0 - 1.0 / gamma / gamma)
# h = round(fmin*l0/beta/clight)
harm_number = round(f_rf * l0 / clight)
if guess is None:
_, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts)
# Getting timelag by tracking uses a different lattice,
# so we cannot now use the same one again.
if method is ELossMethod.TRACKING:
keep_lattice = False
ref_in = numpy.zeros((6, ), order='F')
ref_in[5] = -dt
else:
ref_in = numpy.copy(guess)
theta = numpy.zeros((6, ))
theta[5] = clight * harm_number / f_rf - l0
scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + \
dp_step * numpy.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0])
delta_matrix = numpy.asfortranarray(
numpy.concatenate((numpy.diag(scaling), numpy.zeros((6, 1))), axis=1))
id6 = numpy.asfortranarray(numpy.identity(6))
change = 1
itercount = 0
while (change > convergence) and itercount < max_iterations:
in_mat = ref_in.reshape((6, 1)) + delta_matrix
_ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice)
# the reference particle after one turn
ref_out = in_mat[:, 6]
# 6x6 jacobian matrix from numerical differentiation:
# f(x+d) - f(x) / d
j6 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling
a = j6 - id6 # f'(r_n) - 1
b = ref_out[:] - ref_in[:] - theta
# b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1)
b_over_a = numpy.linalg.solve(a, b)
r_next = ref_in - b_over_a
# determine if we are close enough
change = numpy.linalg.norm(r_next - ref_in)
itercount += 1
ref_in = r_next
keep_lattice = True
if itercount == max_iterations:
warnings.warn(
AtWarning('Maximum number of iterations reached. '
'Possible non-convergence'))
return ref_in
def get_energy_loss(ring, method=ELossMethod.INTEGRAL):
"""Compute the energy loss per turn [eV]
PARAMETERS
ring lattice description
KEYWORDS
method=ELossMethod.INTEGRAL method for energy loss computation
The enum class ELossMethod declares 2 values
INTEGRAL: The losses are obtained from
Losses = Cgamma / 2pi * EGeV^4 * i2
Takes into account bending magnets and wigglers.
TRACKING: The losses are obtained by tracking without cavities.
Needs radiation ON, takes into account all radiating elements.
"""
def integral(ring):
"""Losses = Cgamma / 2pi * EGeV^4 * i2
"""
def wiggler_i2(wiggler):
rhoinv = wiggler.Bmax / 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):
return dipole.BendingAngle**2 / dipole.Length
Brho = sqrt(ring.energy**2 - e_mass**2) / clight
i2 = 0.0
for el in ring:
if isinstance(el, Dipole):
i2 += dipole_i2(el)
elif isinstance(el, Wiggler) and el.PassMethod != 'DriftPass':
i2 += wiggler_i2(el)
e_loss = Cgamma / 2.0 / pi * ring.energy**4 * i2
return e_loss
@check_radiation(True)
def tracking(ring):
"""Losses from tracking
"""
ringtmp = ring.radiation_off(dipole_pass=None,
quadrupole_pass=None,
wiggler_pass=None,
sextupole_pass=None,
octupole_pass=None,
copy=True)
o0 = numpy.zeros(6)
o6 = numpy.squeeze(lattice_pass(ringtmp, o0, refpts=len(ringtmp)))
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))
def find_orbit6(ring, refpts=None, guess=None, **kwargs):
"""find_orbit6 finds the closed orbit in the full 6-D phase space
by numerically solving for a fixed point of the one turn
map M calculated with lattice_pass
(X, PX, Y, PY, DP, CT2 ) = M (X, PX, Y, PY, DP, CT1)
with constraint CT2 - CT1 = C*HarmNumber(1/Frf - 1/Frf0)
IMPORTANT!!! find_orbit6 is a realistic simulation
1. The Frf frequency in the RF cavities (may be different from Frf0)
imposes the synchronous condition
CT2 - CT1 = C*HarmNumber(1/Frf - 1/Frf0)
2. The algorithm numerically calculates
6-by-6 Jacobian matrix J6. In order for (J-E) matrix
to be non-singular it is NECESSARY to use a realistic
PassMethod for cavities with non-zero momentum kick
(such as ThinCavityPass).
3. find_orbit6 can find orbits with radiation.
In order for the solution to exist the cavity must supply
adequate energy compensation.
In the simplest case of a single cavity, it must have
'Voltage' field set so that Voltage > Erad - energy loss per turn
4. There is a family of solutions that correspond to different RF buckets
They differ in the 6-th coordinate by C*Nb/Frf. Nb = 1 .. HarmNum-1
5. The value of the 6-th coordinate found at the cavity gives
the equilibrium RF phase. If there is no radiation the phase is 0;
PARAMETERS
ring Sequence of AT elements
refpts elements at which data is returned. It can be:
1) an integer in the range [-len(ring), len(ring)-1]
selecting the element according to python indexing
rules. As a special case, len(ring) is allowed and
refers to the end of the last element,
2) an ordered list of such integers without duplicates,
3) a numpy array of booleans of maximum length
len(ring)+1, where selected elements are True.
Defaults to None, if refpts is None an empty array is
returned for orbit.
OUTPUT
orbit0 ((6,) closed orbit vector at the entrance of the
1-st element (x,px,y,py)
orbit (6, Nrefs) closed orbit vector at each location
specified in refpts
KEYWORDS
keep_lattice Assume no lattice change since the previous tracking.
Default: False
guess Initial value for the closed orbit. It may help
convergence. Default: (0, 0, 0, 0)
convergence Convergence criterion. Default: 1.e-12
max_iterations Maximum number of iterations. Default: 20
step_size Step size. Default: 1.e-6
See also find_orbit4, find_sync_orbit.
"""
keep_lattice = kwargs.pop('keep_lattice', False)
convergence = kwargs.pop('convergence', CONVERGENCE)
max_iterations = kwargs.pop('max_iterations', MAX_ITERATIONS)
step_size = kwargs.pop('step_size', STEP_SIZE)
ref_in = numpy.zeros((6, ), order='F') if guess is None else guess
# Get evolution period
l0 = get_s_pos(ring, len(ring))
cavities = [elem for elem in ring if isinstance(elem, elements.RFCavity)]
if len(cavities) == 0:
raise AtError('No cavity found in the lattice.')
f_rf = cavities[0].Frequency
harm_number = cavities[0].HarmNumber
theta = numpy.zeros((6, ))
theta[5] = constants.speed_of_light * harm_number / f_rf - l0
delta_matrix = numpy.zeros((6, 7), order='F')
for i in range(6):
delta_matrix[i, i] = step_size
id6 = numpy.asfortranarray(numpy.identity(6))
change = 1
itercount = 0
while (change > convergence) and itercount < max_iterations:
in_mat = ref_in.reshape((6, 1)) + delta_matrix
_ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice)
# the reference particle after one turn
ref_out = in_mat[:, 6]
# 6x6 jacobian matrix from numerical differentiation:
# f(x+d) - f(x) / d
j6 = (in_mat[:, :6] - in_mat[:, 6:]) / step_size
a = j6 - id6 # f'(r_n) - 1
b = ref_out[:] - ref_in[:] - theta
b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1)
r_next = ref_in - b_over_a
# determine if we are close enough
change = numpy.linalg.norm(r_next - ref_in)
itercount += 1
ref_in = r_next
keep_lattice = True
if itercount == max_iterations:
warnings.warn(
AtWarning('Maximum number of iterations reached. '
'Possible non-convergence'))
uint32refs = uint32_refpts(refpts, len(ring))
# bug in numpy < 1.13
all_points = numpy.empty(
(0, 6), dtype=float) if len(uint32refs) == 0 else numpy.squeeze(
lattice_pass(ring,
ref_in.copy(order='K'),
refpts=uint32refs,
keep_lattice=keep_lattice),
axis=(1, 3)).T
return ref_in, all_points
