def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_and_refpts_None(
):
e = elements.LongElement('e', 0.1)
f = elements.LongElement('e', 2.1)
print(get_s_pos([e, f], None))
numpy.testing.assert_equal(get_s_pos([e, f], None),
numpy.array([0, 0.1, 2.2]))
def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_using_bool_refpts(
):
e = elements.LongElement('e', 0.1)
f = elements.LongElement('e', 2.1)
lat = [e, f]
numpy.testing.assert_equal(get_s_pos(lat, numpy.ones(3, dtype=bool)),
numpy.array([0, 0.1, 2.2]))
def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_using_int_refpts(
):
e = elements.LongElement('e', 0.1)
f = elements.LongElement('e', 2.1)
lat = [e, f]
numpy.testing.assert_equal(get_s_pos(lat, range(len(lat) + 1)),
numpy.array([0, 0.1, 2.2]))
def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_using_bool_refpts(
):
e = elements.Element('e', 0.1)
f = elements.Element('e', 2.1)
lat = [e, f]
numpy.testing.assert_equal(
lattice.get_s_pos(lat, numpy.array((True, True, True))),
numpy.array([0, 0.1, 2.2]))
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 get_twiss(ring, dp=0.0, refpts=None, get_chrom=False, ddp=DDP):
"""
Determine Twiss parameters by first finding the transfer matrix.
"""
def twiss22(mat, ms):
"""
Calculate Twiss parameters from the standard 2x2 transfer matrix
(i.e. x or y).
"""
sin_mu_end = (numpy.sign(mat[0, 1]) *
math.sqrt(-mat[0, 1] * mat[1, 0] -
(mat[0, 0] - mat[1, 1])**2 / 4))
alpha0 = (mat[0, 0] - mat[1, 1]) / 2.0 / sin_mu_end
beta0 = mat[0, 1] / sin_mu_end
beta = ((ms[0, 0, :] * beta0 - ms[0, 1, :] * alpha0)**2 +
ms[0, 1, :]**2) / beta0
alpha = -((ms[0, 0, :] * beta0 - ms[0, 1, :] * alpha0) *
(ms[1, 0, :] * beta0 - ms[1, 1, :] * alpha0) +
ms[0, 1, :] * ms[1, 1, :]) / beta0
mu = numpy.arctan(ms[0, 1, :] /
(ms[0, 0, :] * beta0 - ms[0, 1, :] * alpha0))
mu = betatron_phase_unwrap(mu)
return alpha, beta, mu
chrom = None
refpts = lattice.uint32_refpts(refpts, len(ring))
nrefs = refpts.size
if refpts[-1] != len(ring):
refpts = numpy.append(refpts, [len(ring)])
orbit4, orbit = find_orbit4(ring, dp, refpts)
m44, mstack = find_m44(ring, dp, refpts, orbit4=orbit4)
ax, bx, mx = twiss22(m44[:2, :2], mstack[:2, :2, :])
ay, by, my = twiss22(m44[2:, 2:], mstack[2:, 2:, :])
tune = numpy.array((mx[-1], my[-1])) / (2 * numpy.pi)
twiss = numpy.zeros(nrefs, dtype=TWISS_DTYPE)
twiss['idx'] = refpts[:nrefs]
twiss['s_pos'] = lattice.get_s_pos(ring, refpts[:nrefs])
twiss['closed_orbit'] = numpy.rollaxis(orbit, -1)[:nrefs]
twiss['m44'] = numpy.rollaxis(mstack, -1)[:nrefs]
twiss['alpha'] = numpy.rollaxis(numpy.vstack((ax, ay)), -1)[:nrefs]
twiss['beta'] = numpy.rollaxis(numpy.vstack((bx, by)), -1)[:nrefs]
twiss['mu'] = numpy.rollaxis(numpy.vstack((mx, my)), -1)[:nrefs]
twiss['dispersion'] = numpy.NaN
# Calculate chromaticity by calling this function again at a slightly
# different momentum.
if get_chrom:
twissb, tuneb, _ = get_twiss(ring, dp + ddp, refpts[:nrefs])
chrom = (tuneb - tune) / ddp
twiss['dispersion'] = (twissb['closed_orbit'] -
twiss['closed_orbit']) / ddp
return twiss, tune, chrom
def get_mcf(ring, dp=0.0, ddp=DDP, keep_lattice=False):
"""Compute momentum compaction factor
PARAMETERS
ring lattice description
dp momentum deviation. Defaults to 0
KEYWORDS
keep_lattice Assume no lattice change since the previous tracking.
Defaults to False
ddp=1.0E-8 momentum deviation used for differentiation
"""
fp_a, _ = find_orbit4(ring, dp=dp - 0.5 * ddp, keep_lattice=keep_lattice)
fp_b, _ = find_orbit4(ring, dp=dp + 0.5 * ddp, keep_lattice=True)
fp = numpy.stack((fp_a, fp_b),
axis=0).T # generate a Fortran contiguous array
b = numpy.squeeze(lattice_pass(ring, fp, keep_lattice=True), axis=(2, 3))
ring_length = get_s_pos(ring, len(ring))
return (b[5, 1] - b[5, 0]) / ddp / ring_length[0]
def linopt6(ring,
refpts=None,
dp=None,
orbit=None,
cavpts=None,
twiss_in=None,
get_chrom=False,
get_w=False,
keep_lattice=False,
**kwargs):
"""Perform linear analysis of a fully coupled lattice
elemdata0, beamdata, elemdata = linopt6(lattice)
For circular machines, linopt6 analyses
the 4x4 1-turn transfer matrix if radiation is OFF, or
the 6x6 1-turn transfer matrix if radiation is ON.
For a transfer line, The "twiss_in" intput must contain either:
- a field 'R', as provided by ATLINOPT6, or
- the fields 'beta' and 'alpha', as provided by linopt and linopt6
PARAMETERS
ring lattice description.
refpts=None elements at which data is returned.
KEYWORDS
orbit avoids looking for the closed orbit if is already known
((6,) array)
keep_lattice Assume no lattice change since the previous tracking.
Defaults to False
XYStep=1.0e-8 transverse step for numerical computation
DPStep=1.0E-6 momentum deviation used for computation of
the closed orbit
twiss_in=None Initial conditions for transfer line optics. Record
array, as output by linopt or linopt6.
If present, the attribute 'R' will be used, otherwise
The attributes 'alpha' and 'beta' will be used. All
other attributes are ignored.
cavpts=None Cavity location for off-momentum tuning
OUTPUT
elemdata0 linear optics data at the entrance/end of the ring
beamdata lattice properties
elemdata linear optics at the points refered to by refpts, if
refpts is None an empty elemdata structure is returned.
elemdata is a record array with fields:
R R-matrices (3, 6, 6)
A A-matrix (6, 6)
M Transfer matrix from the entrance of the line (6, 6)
beta [betax, betay] vector
alpha [alphax, alphay] vector
dispersion (4,) dispersion vector
mu [mux, muy], betatron phases
s_pos longitudinal position [m]
closed_orbit (6,) closed orbit vector
W (2,) chromatic amplitude function
All values given at the entrance of each element specified in refpts.
Field values can be obtained with either
elemdata['beta'] or
elemdata.beta
beamdata is a record with fields:
tune Fractional tunes
damping_time Damping times [s]
chromaticity Chromaticities
REFERENCES
[1] Etienne Forest, Phys. Rev. E 58, 2481 – Published 1 August 1998
[2] Andrzej Wolski, Phys. Rev. ST Accel. Beams 9, 024001 –
Published 3 February 2006
[3] Brian W. Montague Report LEP Note 165, CERN, 1979
"""
def get_alphabeta(r):
beta = numpy.array([r[0, 0, 0], r[1, 2, 2]])
alpha = -numpy.array([r[0, 1, 0], r[1, 3, 2]])
return alpha, beta
# noinspection PyShadowingNames
def build_r(ring, dp, orbit, refpts=None, mxx=None, **kwargs):
""""""
if ring.radiation:
mt, ms = find_m66(ring, refpts=refpts, orbit=orbit, **kwargs)
else:
mt, ms = find_m44(ring, dp, refpts=refpts, orbit=orbit, **kwargs)
a0, vps = a_matrix(mt if mxx is None else mxx)
val0, vals = _r_analysis(a0, ms)
return ms, vps, val0, vals
def build_sigma(twin):
"""Build the initial distribution at entrance of the transfer line"""
try:
sigm = numpy.sum(twin.R, axis=0)
except AttributeError:
slices = [slice(2 * i, 2 * (i + 1)) for i in range(4)]
ab = numpy.stack((twin.alpha, twin.beta), axis=1)
sigm = numpy.zeros((4, 4))
for slc, (alpha, beta) in zip(slices, ab):
gamma = (1.0 + alpha * alpha) / beta
sigm[slc, slc] = numpy.array([[beta, -alpha], [-alpha, gamma]])
return sigm
def output6(r123, a, mu, *args):
"""Extract output parameters from Bk matrices"""
alpha, beta = get_alphabeta(r123)
dispersion = r123[2, :4, 4] / r123[2, 4, 4]
return (r123, a, alpha, beta, mu, dispersion) + args
def output4(r12, a, *args):
"""Extract output parameters from Bk matrices"""
alpha, beta = get_alphabeta(r12)
return (r12, a, alpha, beta) + args
def get_disp(ringup,
ringdn,
dpup,
dpdn,
refpts=None,
matpts=None,
keep_lattice=False,
**kwargs):
def off_momentum(rng, dp):
orb0, orbs = _find_orbit(rng,
dp,
refpts=refpts,
keep_lattice=keep_lattice,
**kwargs)
dp = orb0[4] # in 6D, dp comes out of find_orbit6
_, vps, el0, els = build_r(rng,
dp,
orb0,
refpts=matpts,
keep_lattice=True,
**kwargs)
tunes = numpy.mod(numpy.angle(vps) / 2.0 / pi, 1.0)
return dp, tunes, orb0, orbs, el0, els
def chromfunc(ddp, elup, eldn):
aup, bup = get_alphabeta(elup[0])
adn, bdn = get_alphabeta(eldn[0])
db = (bup - bdn) / ddp
mb = (bup + bdn) / 2
da = (aup - adn) / ddp
ma = (aup + adn) / 2
w = numpy.sqrt((da - ma / mb * db)**2 + (db / mb)**2)
return w
dpup, tunesup, o0up, orbup, el0up, elup = off_momentum(ringup, dpup)
dpdn, tunesdn, o0dn, orbdn, el0dn, eldn = off_momentum(ringdn, dpdn)
deltap = dpup - dpdn
chrom = (tunesup - tunesdn) / deltap
disp0 = (o0up - o0dn)[:4] / deltap
disp = ((oup - odn)[:4] / deltap for oup, odn in zip(orbup, orbdn))
w0 = chromfunc(deltap, el0up, el0dn)
w = (chromfunc(deltap, *v) for v in zip(elup, eldn))
return chrom, disp0, disp, w0, w
def unwrap(mu):
"""Remove the phase jumps"""
dmu = numpy.diff(numpy.concatenate((numpy.zeros((1, dms)), mu)),
axis=0)
jumps = dmu < -1.e-3
mu += numpy.cumsum(jumps, axis=0) * 2.0 * numpy.pi
dp_step = kwargs.get('DPStep', DConstant.DPStep)
if twiss_in is None: # Circular machine
mxx = None
else: # Transfer line
if orbit is None:
orbit = numpy.zeros((6, ))
sigma = build_sigma(twiss_in)
mxx = sigma.dot(jmat(sigma.shape[0] // 2))
orb0, orbs = _find_orbit(ring, dp, refpts=refpts, orbit=orbit, **kwargs)
dp = orb0[4]
ms, vps, el0, els = build_r(ring,
dp,
orb0,
refpts=refpts,
mxx=mxx,
keep_lattice=keep_lattice,
**kwargs)
dms = vps.size
length = ring.get_s_pos(len(ring))[0]
tunes = numpy.mod(numpy.angle(vps) / 2.0 / pi, 1.0)
damping_rates = -numpy.log(numpy.absolute(vps))
damping_times = length / clight / damping_rates
spos = get_s_pos(ring, refpts)
if dms >= 3: # 6D processsing
output = output6
dtype = W_DTYPE6
data0 = [*el0, numpy.identity(2 * dms), 0.0, orb0]
datas = [els, iter(ms), iter(spos), iter(orbs)]
if get_chrom or get_w:
f0 = get_rf_frequency(ring, cavpts=cavpts)
df = -dp_step * get_mcf(ring.radiation_off(copy=True)) * f0
rgup = set_rf_frequency(ring,
f0 + 0.5 * df,
cavpts=cavpts,
copy=True)
rgdn = set_rf_frequency(ring,
f0 - 0.5 * df,
cavpts=cavpts,
copy=True)
if get_w:
dtype = W_DTYPE6 + W_DTYPEW
chrom, _, _, w0, ws = get_disp(rgup,
rgdn,
None,
None,
matpts=refpts,
**kwargs)
data0.append(w0)
datas.append(ws)
else:
chrom, _, _, _, _ = get_disp(rgup, rgdn, None, None)
else:
chrom = numpy.NaN
else: # 4D processsing
output = output4
dtype = W_DTYPE4
data0 = [*el0]
datas = [els]
if get_w:
dtype = W_DTYPE4 + W_DTYPEW
chrom, d0, ds, w0, ws = get_disp(ring,
ring,
dp + 0.5 * dp_step,
dp - 0.5 * dp_step,
refpts=refpts,
matpts=refpts,
keep_lattice=True,
**kwargs)
data0 += [d0, numpy.identity(2 * dms), 0.0, orb0, w0]
datas += [ds, iter(ms), iter(spos), iter(orbs), ws]
elif get_chrom:
chrom, d0, ds, w0, ws = get_disp(ring,
ring,
dp + 0.5 * dp_step,
dp - 0.5 * dp_step,
refpts=refpts,
keep_lattice=True,
**kwargs)
data0 += [d0, numpy.identity(2 * dms), 0.0, orb0]
datas += [ds, iter(ms), iter(spos), iter(orbs)]
else:
chrom = numpy.NaN
data0 += [numpy.NaN, numpy.identity(2 * dms), 0.0, orb0]
datas += [repeat(numpy.NaN), iter(ms), iter(spos), iter(orbs)]
elemdata0 = numpy.array(output(*data0), dtype=dtype).view(numpy.recarray)
elemdata = numpy.fromiter((output(*el, *v) for el, *v in zip(*datas)),
dtype,
count=ring.refcount(refpts)).view(numpy.recarray)
beamdata = numpy.array(
(tunes, chrom, damping_times),
dtype=[('tune', numpy.float64, (dms, )),
('chromaticity', numpy.float64, (dms, )),
('damping_time', numpy.float64, (dms, ))]).view(numpy.recarray)
unwrap(elemdata.mu)
return elemdata0, beamdata, elemdata
def get_twiss(ring,
dp=0.0,
refpts=None,
get_chrom=False,
orbit=None,
keep_lattice=False,
ddp=DDP):
"""
Perform linear analysis of the NON-COUPLED lattices
twiss0, tune, chrom, twiss = get_twiss(ring, dp[, refpts])
PARAMETERS
ring lattice description.
dp=0.0 momentum deviation.
refpts=None 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.
KEYWORDS
orbit avoids looking for the closed orbit if is already known
((6,) array)
get_chrom=False compute dispersion and chromaticities. Needs computing
the optics at 2 different momentum deviations around
the central one.
keep_lattice Assume no lattice change since the previous tracking.
Defaults to False
ddp=1.0E-8 momentum deviation used for computation of
chromaticities and dispersion
OUTPUT
twiss0 linear optics data at the entrance/end of the ring
tune [tune_h, tune_v], fractional part of the linear tunes
chrom [ksi_h , ksi_v], chromaticities ksi = d(nu)/(dP/P).
Only computed if 'get_chrom' is True
twiss linear optics at the points refered to by refpts, if
refpts is None an empty twiss structure is returned.
twiss is a record array with fields:
idx element index in the ring
s_pos longitudinal position [m]
closed_orbit (6,) closed orbit vector
dispersion (4,) dispersion vector
Only computed if 'get_chrom' is True
m44 (4, 4) transfer matrix M from the beginning of ring
to the entrance of the element
mu (2,) betatron phase (modulo 2*pi)
beta [betax, betay] vector
alpha [alphax, alphay] vector
All values given at the entrance of each element specified in refpts.
Field values can be obtained with either
twiss['idx'] or
twiss.idx
See also linopt
"""
uintrefs = uint32_refpts(refpts, len(ring))
if orbit is None:
orbit, _ = find_orbit4(ring, dp, keep_lattice=keep_lattice)
keep_lattice = True
orbs = numpy.squeeze(lattice_pass(ring,
orbit.copy(order='K'),
refpts=uintrefs,
keep_lattice=keep_lattice),
axis=(1, 3)).T
m44, mstack = find_m44(ring, dp, uintrefs, orbit=orbit, keep_lattice=True)
nrefs = uintrefs.size
# Get initial twiss parameters
a0_x, b0_x, tune_x = _closure(m44[:2, :2])
a0_y, b0_y, tune_y = _closure(m44[2:, 2:])
tune = numpy.array([tune_x, tune_y])
# Calculate chromaticity by calling this function again at a slightly
# different momentum.
if get_chrom:
d0_up, tune_up, _, l_up = get_twiss(ring,
dp + 0.5 * ddp,
uintrefs,
keep_lattice=True)
d0_down, tune_down, _, l_down = get_twiss(ring,
dp - 0.5 * ddp,
uintrefs,
keep_lattice=True)
chrom = (tune_up - tune_down) / ddp
dispersion = (l_up['closed_orbit'] -
l_down['closed_orbit'])[:, :4] / ddp
disp0 = (d0_up['closed_orbit'] - d0_down['closed_orbit'])[:4] / ddp
else:
chrom = numpy.array([numpy.NaN, numpy.NaN])
dispersion = numpy.NaN
disp0 = numpy.NaN
twiss0 = numpy.rec.fromarrays((len(ring), get_s_pos(
ring, len(ring))[0], orbit, disp0, numpy.array(
[a0_x, a0_y]), numpy.array([b0_x, b0_y]), 2.0 * pi * tune, m44),
dtype=TWISS_DTYPE)
# Propagate to reference points
twiss = numpy.rec.array(numpy.zeros(nrefs, dtype=TWISS_DTYPE))
if nrefs > 0:
alpha_x, beta_x, mu_x = _twiss22(mstack[:, :2, :2], a0_x, b0_x)
alpha_z, beta_z, mu_z = _twiss22(mstack[:, 2:, 2:], a0_y, b0_y)
twiss['idx'] = uintrefs
twiss['s_pos'] = get_s_pos(ring, uintrefs[:nrefs])
twiss['closed_orbit'] = orbs
twiss['m44'] = mstack
twiss['alpha'] = numpy.stack((alpha_x, alpha_z), axis=1)
twiss['beta'] = numpy.stack((beta_x, beta_z), axis=1)
twiss['mu'] = numpy.stack((mu_x, mu_z), axis=1)
twiss['dispersion'] = dispersion
return twiss0, tune, chrom, twiss
def linopt(ring,
dp=0.0,
refpts=None,
get_chrom=False,
orbit=None,
keep_lattice=False,
ddp=DDP,
coupled=True):
"""
Perform linear analysis of a lattice
lindata0, tune, chrom, lindata = linopt(ring, dp[, refpts])
PARAMETERS
ring lattice description.
dp=0.0 momentum deviation.
refpts=None 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.
KEYWORDS
orbit avoids looking for the closed orbit if is already known
((6,) array)
get_chrom=False compute dispersion and chromaticities. Needs computing
the optics at 2 different momentum deviations around
the central one.
keep_lattice Assume no lattice change since the previous tracking.
Defaults to False
ddp=1.0E-8 momentum deviation used for computation of
chromaticities and dispersion
coupled=True if False, simplify the calculations by assuming
no H/V coupling
OUTPUT
lindata0 linear optics data at the entrance/end of the ring
tune [tune_A, tune_B], linear tunes for the two normal modes
of linear motion [1]
chrom [ksi_A , ksi_B], chromaticities ksi = d(nu)/(dP/P).
Only computed if 'get_chrom' is True
lindata linear optics at the points refered to by refpts, if
refpts is None an empty lindata structure is returned.
lindata is a record array with fields:
idx element index in the ring
s_pos longitudinal position [m]
closed_orbit (6,) closed orbit vector
dispersion (4,) dispersion vector
Only computed if 'get_chrom' is True
m44 (4, 4) transfer matrix M from the beginning of ring
to the entrance of the element [2]
A (2, 2) matrix A in [3]
B (2, 2) matrix B in [3]
C (2, 2) matrix C in [3]
gamma gamma parameter of the transformation to eigenmodes
mu [mux, muy], betatron phase (modulo 2*pi)
beta [betax, betay] vector
alpha [alphax, alphay] vector
All values given at the entrance of each element specified in refpts.
Field values can be obtained with either
lindata['idx'] or
lindata.idx
REFERENCES
[1] D.Edwars,L.Teng IEEE Trans.Nucl.Sci. NS-20, No.3, p.885-888, 1973
[2] E.Courant, H.Snyder
[3] D.Sagan, D.Rubin Phys.Rev.Spec.Top.-Accelerators and beams,
vol.2 (1999)
See also get_twiss
"""
# noinspection PyShadowingNames
def analyze(r44):
t44 = r44.reshape((4, 4))
mm = t44[:2, :2]
nn = t44[2:, 2:]
m = t44[:2, 2:]
n = t44[2:, :2]
gamma = sqrt(numpy.linalg.det(numpy.dot(n, C) + numpy.dot(G, nn)))
msa = (G.dot(mm) - m.dot(_jmt.dot(C.T.dot(_jmt.T)))) / gamma
msb = (numpy.dot(n, C) + numpy.dot(G, nn)) / gamma
cc = (numpy.dot(mm, C) + numpy.dot(G, m)).dot(
_jmt.dot(msb.T.dot(_jmt.T)))
return msa, msb, gamma, cc
uintrefs = uint32_refpts([] if refpts is None else refpts, len(ring))
if orbit is None:
orbit, _ = find_orbit4(ring, dp, keep_lattice=keep_lattice)
keep_lattice = True
orbs = numpy.squeeze(lattice_pass(ring,
orbit.copy(order='K'),
refpts=uintrefs,
keep_lattice=keep_lattice),
axis=(1, 3)).T
m44, mstack = find_m44(ring, dp, uintrefs, orbit=orbit, keep_lattice=True)
nrefs = uintrefs.size
M = m44[:2, :2]
N = m44[2:, 2:]
m = m44[:2, 2:]
n = m44[2:, :2]
if coupled:
# Calculate A, B, C, gamma at the first element
H = m + _jmt.dot(n.T.dot(_jmt.T))
t = numpy.trace(M - N)
t2 = t * t
t2h = t2 + 4.0 * numpy.linalg.det(H)
g = sqrt(1.0 + sqrt(t2 / t2h)) / sqrt(2.0)
G = numpy.diag((g, g))
C = -H * numpy.sign(t) / (g * sqrt(t2h))
A = G.dot(G.dot(M)) - numpy.dot(
G, (m.dot(_jmt.dot(C.T.dot(_jmt.T))) + C.dot(n))) + C.dot(
N.dot(_jmt.dot(C.T.dot(_jmt.T))))
B = G.dot(G.dot(N)) + numpy.dot(
G, (_jmt.dot(C.T.dot(_jmt.T.dot(m))) + n.dot(C))) + _jmt.dot(
C.T.dot(_jmt.T.dot(M.dot(C))))
else:
A = M
B = N
C = numpy.zeros((2, 2))
g = 1.0
# Get initial twiss parameters
a0_a, b0_a, tune_a = _closure(A)
a0_b, b0_b, tune_b = _closure(B)
tune = numpy.array([tune_a, tune_b])
if get_chrom:
d0_up, tune_up, _, l_up = linopt(ring,
dp + 0.5 * ddp,
uintrefs,
keep_lattice=True,
coupled=coupled)
d0_down, tune_down, _, l_down = linopt(ring,
dp - 0.5 * ddp,
uintrefs,
keep_lattice=True,
coupled=coupled)
chrom = (tune_up - tune_down) / ddp
dispersion = (l_up['closed_orbit'] -
l_down['closed_orbit'])[:, :4] / ddp
disp0 = (d0_up['closed_orbit'] - d0_down['closed_orbit'])[:4] / ddp
else:
chrom = numpy.array([numpy.NaN, numpy.NaN])
dispersion = numpy.NaN
disp0 = numpy.NaN
lindata0 = numpy.rec.fromarrays(
(len(ring), get_s_pos(ring, len(ring))[0], orbit, disp0,
numpy.array([a0_a, a0_b]), numpy.array(
[b0_a, b0_b]), 2.0 * pi * tune, m44, A, B, C, g),
dtype=LINDATA_DTYPE)
# Propagate to reference points
lindata = numpy.rec.array(numpy.zeros(nrefs, dtype=LINDATA_DTYPE))
if nrefs > 0:
if coupled:
MSA, MSB, gamma, CL = zip(*[analyze(ms44) for ms44 in mstack])
msa = numpy.stack(MSA, axis=0)
msb = numpy.stack(MSB, axis=0)
else:
msa = mstack[:, :2, :2]
msb = mstack[:, 2:, 2:]
gamma = 1.0
CL = numpy.zeros((1, 2, 2))
alpha_a, beta_a, mu_a = _twiss22(msa, a0_a, b0_a)
alpha_b, beta_b, mu_b = _twiss22(msb, a0_b, b0_b)
lindata['idx'] = uintrefs
lindata['s_pos'] = get_s_pos(ring, uintrefs)
lindata['closed_orbit'] = orbs
lindata['m44'] = mstack
lindata['alpha'] = numpy.stack((alpha_a, alpha_b), axis=1)
lindata['beta'] = numpy.stack((beta_a, beta_b), axis=1)
lindata['mu'] = numpy.stack((mu_a, mu_b), axis=1)
lindata['A'] = [
ms.dot(A.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msa
]
lindata['B'] = [
ms.dot(B.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msb
]
lindata['C'] = CL
lindata['gamma'] = gamma
lindata['dispersion'] = dispersion
return lindata0, tune, chrom, lindata
def linopt(ring,
dp=0.0,
refpts=None,
get_chrom=False,
orbit=None,
keep_lattice=False,
coupled=True,
twiss_in=None,
get_w=False,
**kwargs):
"""
Perform linear analysis of a lattice
lindata0, tune, chrom, lindata = linopt(ring, dp[, refpts])
PARAMETERS
ring lattice description.
dp=0.0 momentum deviation.
refpts=None 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.
KEYWORDS
orbit avoids looking for the closed orbit if is already known
((6,) array)
get_chrom=False compute dispersion and chromaticities. Needs computing
the tune and orbit at 2 different momentum deviations
around the central one.
keep_lattice Assume no lattice change since the previous tracking.
Defaults to False
XYStep=1.0e-8 transverse step for numerical computation
DPStep=1.0E-6 momentum deviation used for computation of
chromaticities and dispersion
coupled=True if False, simplify the calculations by assuming
no H/V coupling
twiss_in=None Initial twiss to compute transfer line optics of the
type lindata, the initial orbit in twiss_in is ignored,
only the beta and alpha are required other quatities
set to 0 if absent
get_w=False computes chromatic amplitude functions (W) [4], need to
compute the optics at 2 different momentum deviations
around the central one.
OUTPUT
lindata0 linear optics data at the entrance/end of the ring
tune [tune_A, tune_B], linear tunes for the two normal modes
of linear motion [1]
chrom [ksi_A , ksi_B], chromaticities ksi = d(nu)/(dP/P).
Only computed if 'get_chrom' is True
lindata linear optics at the points refered to by refpts, if
refpts is None an empty lindata structure is returned.
lindata is a record array with fields:
idx element index in the ring
s_pos longitudinal position [m]
closed_orbit (6,) closed orbit vector
dispersion (4,) dispersion vector
W (2,) chromatic amplitude function
Only computed if 'get_chrom' is True
m44 (4, 4) transfer matrix M from the beginning of ring
to the entrance of the element [2]
mu [mux, muy], betatron phase (modulo 2*pi)
beta [betax, betay] vector
alpha [alphax, alphay] vector
All values given at the entrance of each element specified in refpts.
In case coupled = True additional outputs are available:
A (2, 2) matrix A in [3]
B (2, 2) matrix B in [3]
C (2, 2) matrix C in [3]
gamma gamma parameter of the transformation to eigenmodes
Field values can be obtained with either
lindata['idx'] or
lindata.idx
REFERENCES
[1] D.Edwars,L.Teng IEEE Trans.Nucl.Sci. NS-20, No.3, p.885-888, 1973
[2] E.Courant, H.Snyder
[3] D.Sagan, D.Rubin Phys.Rev.Spec.Top.-Accelerators and beams,
vol.2 (1999)
[4] Brian W. Montague Report LEP Note 165, CERN, 1979
"""
# noinspection PyShadowingNames
def analyze(r44):
t44 = r44.reshape((4, 4))
mm = t44[:2, :2]
nn = t44[2:, 2:]
m = t44[:2, 2:]
n = t44[2:, :2]
gamma = sqrt(numpy.linalg.det(numpy.dot(n, C) + numpy.dot(G, nn)))
msa = (G.dot(mm) - m.dot(_jmt.dot(C.T.dot(_jmt.T)))) / gamma
msb = (numpy.dot(n, C) + numpy.dot(G, nn)) / gamma
cc = (numpy.dot(mm, C) + numpy.dot(G, m)).dot(
_jmt.dot(msb.T.dot(_jmt.T)))
return msa, msb, gamma, cc
xy_step = kwargs.pop('XYStep', DConstant.XYStep)
dp_step = kwargs.pop('DPStep', DConstant.DPStep)
uintrefs = uint32_refpts([] if refpts is None else refpts, len(ring))
# Get initial orbit
if twiss_in is None:
if orbit is None:
orbit, _ = find_orbit4(ring,
dp,
keep_lattice=keep_lattice,
XYStep=xy_step)
keep_lattice = True
disp0 = numpy.NaN
if get_chrom or get_w:
orbit_up, _ = find_orbit4(ring,
dp + 0.5 * dp_step,
XYStep=xy_step,
keep_lattice=keep_lattice)
orbit_down, _ = find_orbit4(ring,
dp - 0.5 * dp_step,
XYStep=xy_step,
keep_lattice=keep_lattice)
disp0 = numpy.array(orbit_up - orbit_down)[:4] / dp_step
else:
if orbit is None:
orbit = numpy.zeros((6, ))
disp0 = numpy.NaN
if get_chrom or get_w:
try:
disp0 = twiss_in['dispersion']
except KeyError:
print('Dispersion not found in twiss_in, setting to zero')
disp0 = numpy.zeros((4, ))
dorbit = numpy.hstack(
(0.5 * dp_step * disp0, numpy.array([0.5 * dp_step, 0])))
orbit_up = orbit + dorbit
orbit_down = orbit - dorbit
orbs = numpy.squeeze(lattice_pass(ring,
orbit.copy(order='K'),
refpts=uintrefs,
keep_lattice=keep_lattice),
axis=(1, 3)).T
m44, mstack = find_m44(ring,
dp,
uintrefs,
orbit=orbit,
keep_lattice=True,
XYStep=xy_step)
nrefs = uintrefs.size
M = m44[:2, :2]
N = m44[2:, 2:]
m = m44[:2, 2:]
n = m44[2:, :2]
# Calculate A, B, C, gamma at the first element
if coupled:
H = m + _jmt.dot(n.T.dot(_jmt.T))
t = numpy.trace(M - N)
t2 = t * t
t2h = t2 + 4.0 * numpy.linalg.det(H)
g = sqrt(1.0 + sqrt(t2 / t2h)) / sqrt(2.0)
G = numpy.diag((g, g))
C = -H * numpy.sign(t) / (g * sqrt(t2h))
A = G.dot(G.dot(M)) - numpy.dot(
G, (m.dot(_jmt.dot(C.T.dot(_jmt.T))) + C.dot(n))) + C.dot(
N.dot(_jmt.dot(C.T.dot(_jmt.T))))
B = G.dot(G.dot(N)) + numpy.dot(
G, (_jmt.dot(C.T.dot(_jmt.T.dot(m))) + n.dot(C))) + _jmt.dot(
C.T.dot(_jmt.T.dot(M.dot(C))))
else:
A = M
B = N
C = numpy.zeros((2, 2))
g = 1.0
# Get initial twiss parameters
if twiss_in is None:
a0_a, b0_a, tune_a = _closure(A)
a0_b, b0_b, tune_b = _closure(B)
tune = numpy.array([tune_a, tune_b])
else:
try:
a0_a, a0_b = twiss_in['alpha'][0], twiss_in['alpha'][1]
except KeyError:
raise ValueError('Initial alpha required for transfer line')
try:
b0_a, b0_b = twiss_in['beta'][0], twiss_in['beta'][1]
except KeyError:
raise ValueError('Initial beta required for transfer line')
try:
tune = numpy.array([twiss_in['mu'][0], twiss_in['mu'][1]
]) / (2 * pi)
except KeyError:
print('Mu not found in twiss_in, setting to zero')
tune = numpy.zeros((2, ))
# Get initial chromatic functions and dispersion
if get_w:
# noinspection PyUnboundLocalVariable
kwup = dict(orbit=orbit_up, twiss_in=twiss_in)
# noinspection PyUnboundLocalVariable
kwdown = dict(orbit=orbit_down, twiss_in=twiss_in)
param_up = linopt(ring,
dp=dp + 0.5 * dp_step,
refpts=uintrefs,
keep_lattice=True,
coupled=coupled,
XYStep=xy_step,
**kwup)
param_down = linopt(ring,
dp=dp - 0.5 * dp_step,
refpts=uintrefs,
keep_lattice=True,
coupled=coupled,
XYStep=xy_step,
**kwdown)
param_up_down = param_up + param_down
chrom, dispersion, w0, w = _chromfunc(dp_step, *param_up_down)
elif get_chrom:
# noinspection PyUnboundLocalVariable
kwup = dict(orbit=orbit_up, twiss_in=twiss_in)
# noinspection PyUnboundLocalVariable
kwdown = dict(orbit=orbit_down, twiss_in=twiss_in)
_, tune_up, _, _ = linopt(ring,
dp=dp + 0.5 * dp_step,
coupled=coupled,
keep_lattice=True,
XYStep=xy_step,
DPStep=dp_step,
**kwup)
orb_up = numpy.squeeze(lattice_pass(ring,
kwup['orbit'].copy(order='K'),
refpts=uintrefs,
keep_lattice=True),
axis=(1, 3)).T
_, tune_down, _, _ = linopt(ring,
dp=dp - 0.5 * dp_step,
coupled=coupled,
keep_lattice=True,
XYStep=xy_step,
DPStep=dp_step,
**kwdown)
orb_down = numpy.squeeze(lattice_pass(ring,
kwdown['orbit'].copy(order='K'),
refpts=uintrefs,
keep_lattice=True),
axis=(1, 3)).T
chrom = (tune_up - tune_down) / dp_step
dispersion = (orb_up - orb_down)[:, :4] / dp_step
w0 = numpy.array([numpy.NaN, numpy.NaN])
w = numpy.array([numpy.NaN, numpy.NaN])
else:
chrom = numpy.array([numpy.NaN, numpy.NaN])
dispersion = numpy.array([numpy.NaN, numpy.NaN, numpy.NaN, numpy.NaN])
w0 = numpy.array([numpy.NaN, numpy.NaN])
w = numpy.array([numpy.NaN, numpy.NaN])
lindata0 = numpy.rec.fromarrays(
(len(ring), get_s_pos(ring, len(ring))[0], orbit, disp0,
numpy.array([a0_a, a0_b]), numpy.array(
[b0_a, b0_b]), 2.0 * pi * tune, m44, A, B, C, g, w0),
dtype=LINDATA_DTYPE)
# Propagate to reference points
lindata = numpy.rec.array(numpy.zeros(nrefs, dtype=LINDATA_DTYPE))
if nrefs > 0:
if coupled:
MSA, MSB, gamma, CL = zip(*[analyze(ms44) for ms44 in mstack])
msa = numpy.stack(MSA, axis=0)
msb = numpy.stack(MSB, axis=0)
AL = [ms.dot(A.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msa]
BL = [ms.dot(B.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msb]
else:
msa = mstack[:, :2, :2]
msb = mstack[:, 2:, 2:]
AL = numpy.NaN
BL = numpy.NaN
CL = numpy.NaN
gamma = numpy.NaN
alpha_a, beta_a, mu_a = _twiss22(msa, a0_a, b0_a)
alpha_b, beta_b, mu_b = _twiss22(msb, a0_b, b0_b)
if twiss_in is not None:
qtmp = numpy.array([mu_a[-1], mu_b[-1]]) / (2 * numpy.pi)
qtmp -= numpy.floor(qtmp)
mu_a += tune[0] * 2 * pi
mu_b += tune[1] * 2 * pi
tune = qtmp
lindata['idx'] = uintrefs
lindata['s_pos'] = get_s_pos(ring, uintrefs)
lindata['closed_orbit'] = orbs
lindata['m44'] = mstack
lindata['alpha'] = numpy.stack((alpha_a, alpha_b), axis=1)
lindata['beta'] = numpy.stack((beta_a, beta_b), axis=1)
lindata['dispersion'] = dispersion
lindata['mu'] = numpy.stack((mu_a, mu_b), axis=1)
lindata['A'] = AL
lindata['B'] = BL
lindata['C'] = CL
lindata['gamma'] = gamma
lindata['W'] = w
return lindata0, tune, chrom, lindata
def test_get_s_pos_returns_length_for_lattice_with_one_element():
e = elements.LongElement('e', 0.1)
assert get_s_pos([e], [1]) == numpy.array([0.1])
def test_get_s_pos_returns_zero_for_empty_lattice():
numpy.testing.assert_equal(get_s_pos([], None), numpy.array((0,)))
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 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
