def test_orbit_maxiter_warnings(hmba_lattice):
with pytest.warns(AtWarning):
physics.find_orbit4(hmba_lattice, max_iterations=1)
with pytest.warns(AtWarning):
physics.find_sync_orbit(hmba_lattice, max_iterations=1)
with pytest.warns(AtWarning):
physics.find_orbit6(hmba_lattice, max_iterations=1)
def test_orbit_maxiter_warnings(hmba_lattice):
hmba_lattice_rad = hmba_lattice.radiation_on(copy=True)
with pytest.warns(AtWarning):
physics.find_orbit4(hmba_lattice, max_iterations=1)
with pytest.warns(AtWarning):
physics.find_sync_orbit(hmba_lattice, max_iterations=1)
with pytest.warns(AtWarning):
physics.find_orbit6(hmba_lattice_rad, max_iterations=1)
def test_find_orbit4_with_two_refpts_with_and_without_guess(dba_lattice):
expected = numpy.array(
[[8.148212e-6, 1.0993354e-5, 0, 0, DP, 2.963929e-6],
[3.0422808e-8, 9.1635269e-8, 0, 0, DP, 5.9280346e-6]])
_, all_points = physics.find_orbit4(dba_lattice, DP, [49, 99])
numpy.testing.assert_allclose(all_points, expected, atol=1e-12)
_, all_points = physics.find_orbit4(dba_lattice, DP, [49, 99],
numpy.array([0., 0., 0., 0., DP, 0.]))
numpy.testing.assert_allclose(all_points, expected, atol=1e-12)
def _find_orbit(ring, dp=None, refpts=None, orbit=None, ct=None, **kwargs):
""""""
if ring.radiation:
if dp is not None:
warn(AtWarning('In 6D, "dp" and "ct" are ignored'))
if orbit is None:
orbit, _ = find_orbit6(ring, **kwargs)
else:
if dp is None:
dp = 0.0
if orbit is None:
if ct is not None:
orbit, _ = find_sync_orbit(ring, ct, **kwargs)
else:
orbit, _ = find_orbit4(ring, dp, **kwargs)
if refpts is None:
orbs = []
else:
orbs = numpy.squeeze(lattice_pass(ring,
orbit.copy(order='K'),
refpts=refpts,
keep_lattice=True),
axis=(1, 3)).T
return orbit, orbs
def gen_detuning_elem(ring, orbit=None):
"""
Generates an element that for detuning with amplitude
"""
if orbit is None:
orbit, _ = find_orbit4(ring)
[lindata0, tunes, xsi, lindata] = ring.linopt(dp=0,
get_chrom=True,
orbit=orbit)
r0, r1, x, q_dx, y, q_dy = detuning(ring,
xm=1.0e-4,
ym=1.0e-4,
npoints=3,
dp=0)
rp = ring.periodicity
nonlin_elem = Element('NonLinear',
PassMethod='DeltaQPass',
Betax=lindata0.beta[0],
Betay=lindata0.beta[1],
Alphax=lindata0.alpha[0],
Alphay=lindata0.alpha[1],
Qpx=xsi[0],
Qpy=xsi[1],
A1=r1[0][0] * rp,
A2=r1[0][1] * rp,
A3=r1[1][1] * rp,
T1=-orbit,
T2=orbit)
return nonlin_elem
def gen_centroid(ring, ampl, nturns, dp, remove_dc):
orbit, _ = find_orbit4(ring, dp)
ld, _, _, _ = linopt(ring, dp, orbit=orbit)
invx = numpy.array([[1 / numpy.sqrt(ld['beta'][0]), 0],
[
ld['alpha'][0] / numpy.sqrt(ld['beta'][0]),
numpy.sqrt(ld['beta'][0])
]])
invy = numpy.array([[1 / numpy.sqrt(ld['beta'][1]), 0],
[
ld['alpha'][1] / numpy.sqrt(ld['beta'][1]),
numpy.sqrt(ld['beta'][1])
]])
p0 = numpy.array([
orbit,
] * 2).T
p0[0, 0] += ampl
p0[2, 1] += ampl
p1 = lattice_pass(ring, p0, refpts=len(ring), nturns=nturns)
cent_x = p1[0, 0, 0, :]
cent_xp = p1[1, 0, 0, :]
cent_y = p1[2, 1, 0, :]
cent_yp = p1[3, 1, 0, :]
if remove_dc:
cent_x -= numpy.mean(cent_x)
cent_y -= numpy.mean(cent_y)
cent_xp -= numpy.mean(cent_xp)
cent_yp -= numpy.mean(cent_yp)
cent_x, cent_xp = numpy.matmul(invx, [cent_x, cent_xp])
cent_y, cent_yp = numpy.matmul(invy, [cent_y, cent_yp])
return (cent_x - 1j * cent_xp, cent_y - 1j * cent_yp)
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 test_find_orbit4(engine, ml_lattice, py_lattice, dp, refpts):
# Matlab call
ml_refpts = (matlab.double([]) if refpts is None else matlab.double(
list(r + 1 for r in refpts)))
ml_orbit4 = engine.findorbit4(ml_lattice, dp, ml_refpts)
py_ml_orbit4 = numpy.asarray(ml_orbit4)
# Python call
py_orbit4 = physics.find_orbit4(py_lattice, dp, refpts)
numpy.testing.assert_almost_equal(py_ml_orbit4, py_orbit4.T)
def test_find_orbit4(engine, ml_lattice, py_lattice, dp, refpts):
nelems = len(py_lattice)
refpts = range(nelems + 1) if refpts is None else refpts
# Python call
py_orb4, py_orbit4 = physics.find_orbit4(py_lattice, dp, refpts)
# Matlab call
ml_orbit4, ml_orb4 = engine.findorbit4(ml_lattice,
dp,
_ml_refs(refpts, nelems),
nargout=2)
ml_orbit4 = numpy.rollaxis(numpy.asarray(ml_orbit4), -1)
# Comparison
numpy.testing.assert_almost_equal(py_orb4, _py_data(ml_orb4), decimal=8)
numpy.testing.assert_almost_equal(py_orbit4[:, :4], ml_orbit4, decimal=8)
def test_find_orbit4_with_two_refpts(ring):
_, all_points = physics.find_orbit4(ring, DP, [49, 99])
expected = numpy.array(
[[8.148212e-6, 1.0993354e-5, 0, 0, DP, 2.963929e-6],
[3.0422808e-8, 9.1635269e-8, 0, 0, DP, 5.9280346e-6]]).T
numpy.testing.assert_allclose(all_points, expected, atol=1e-12)
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 find_m66(ring, refpts=None, orbit=None, keep_lattice=False, **kwargs):
"""find_m66 numerically finds the 6x6 transfer matrix of an accelerator
lattice by differentiation of lattice_pass near the closed orbit.
find_m66 uses find_orbit6 to search for the closed orbit in 6-D
In order for this to work the ring MUST have a CAVITY element
m66, t = find_m66(lattice, refpts)
m66: full one-turn 6-by-6 matrix at the entrance of the
first element.
t: 6x6 transfer matrices between the entrance of the first
element and each element indexed by refpts (nrefs, 6, 6) array.
PARAMETERS
ring lattice description
dp momentum deviation. Defaults to 0
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 mstack.
KEYWORDS
keep_lattice=False When True, assume no lattice change since the
previous tracking.
orbit=None Avoids looking for the closed orbit if is already
known (6,) array
XYStep=1.e-8 transverse step for numerical computation
See also find_m44, find_orbit6
"""
xy_step = kwargs.pop('XYStep', DConstant.XYStep)
dp_step = kwargs.pop('DPStep', DConstant.DPStep)
if orbit is None:
if ring.radiation:
orbit, _ = find_orbit6(ring, keep_lattice=keep_lattice,
XYStep=xy_step, DPStep=dp_step, **kwargs)
else:
orbit, _ = find_orbit4(ring, keep_lattice=keep_lattice,
XYStep=xy_step, **kwargs)
keep_lattice = True
# Construct matrix of plus and minus deltas
# scaling = 2*xy_step*numpy.array([1.0, 0.1, 1.0, 0.1, 1.0, 1.0])
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])
dg = numpy.asfortranarray(0.5 * numpy.diag(scaling))
dmat = numpy.concatenate((dg, -dg), axis=1)
in_mat = orbit.reshape(6, 1) + dmat
refs = uint32_refpts(refpts, len(ring))
out_mat = numpy.rollaxis(
numpy.squeeze(lattice_pass(ring, in_mat, refpts=refs,
keep_lattice=keep_lattice), axis=3), -1
)
# out_mat: 12 particles at n refpts for one turn
# (x + d) - (x - d) / d
m66 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling
if len(refs) > 0:
mstack = (out_mat[:, :, :6] - out_mat[:, :, 6:]) / scaling
else:
mstack = numpy.empty((0, 6, 6), dtype=float)
return m66, mstack
def test_find_orbit4_result_unchanged_by_atpass(dba_lattice):
orbit, _ = physics.find_orbit4(dba_lattice, DP)
orbit_copy = numpy.copy(orbit)
orbit[4] = DP
atpass(dba_lattice, orbit, 1)
assert_close(orbit[:4], orbit_copy[:4], atol=1e-12)
def test_find_orbit4_finds_zeros_if_dp_zero(dba_lattice):
orbit4, _ = physics.find_orbit4(dba_lattice, 0)
expected = numpy.zeros((6,))
assert_close(orbit4, expected, atol=1e-7)
def test_find_orbit4(dba_lattice):
orbit4, _ = physics.find_orbit4(dba_lattice, DP)
expected = numpy.array([1.091636e-7, 1.276747e-15, 0, 0, DP, 0])
assert_close(orbit4, expected, atol=1e-12)
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_find_orbit4_result_unchanged_by_atpass(ring):
orbit4, _ = physics.find_orbit4(ring, DP)
orbit6 = numpy.append(orbit4, numpy.zeros((1, 2)))
orbit6[4] = DP
orbit6_pass = atpass(ring, orbit6, 1)[:, 0]
numpy.testing.assert_allclose(orbit4, orbit6_pass[:4], atol=1e-12)
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 test_find_orbit4_finds_zeros_if_dp_zero(ring):
orbit4 = physics.find_orbit4(ring, 0)
expected = numpy.zeros((6, ))
numpy.testing.assert_allclose(orbit4, expected)
def test_find_orbit4(ring):
orbit4 = physics.find_orbit4(ring, DP)
expected = numpy.array([1.091636e-7, 1.276747e-15, 0, 0, DP, 0])
numpy.testing.assert_allclose(orbit4, expected, atol=1e-12)
def find_m44(ring,
dp=0.0,
refpts=None,
orbit=None,
keep_lattice=False,
**kwargs):
"""find_m44 numerically finds the 4x4 transfer matrix of an accelerator
lattice for a particle with relative momentum deviation DP
IMPORTANT!!! find_m44 assumes constant momentum deviation.
PassMethod used for any element in the lattice SHOULD NOT
1. change the longitudinal momentum dP
(cavities , magnets with radiation, ...)
2. have any time dependence (localized impedance, fast kickers, ...)
m44, t = find_m44(lattice, dp=0.0, refpts)
return 4x4 transfer matrices between the entrance of the first element
and each element indexed by refpts.
m44: full one-turn matrix at the entrance of the first element
t: 4x4 transfer matrices between the entrance of the first
element and each element indexed by refpts:
(Nrefs, 4, 4) array
Unless an input orbit is introduced, find_m44 assumes that the lattice is
a ring and first finds the closed orbit.
PARAMETERS
ring lattice description
dp momentum deviation. Defaults to 0
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 mstack.
KEYWORDS
keep_lattice=False When True, assume no lattice change since the
previous tracking.
full=False When True, matrices are full 1-turn matrices at
the entrance of each
element indexed by refpts.
orbit=None Avoids looking for the closed orbit if is already
known (6,) array
XYStep=6.055e-6 transverse step for numerical computation
See also find_m66, find_orbit4
"""
def mrotate(m):
m = numpy.squeeze(m)
return m.dot(m44.dot(_jmt.T.dot(m.T.dot(_jmt))))
xy_step = kwargs.pop('XYStep', XYDEFSTEP)
full = kwargs.pop('full', False)
if orbit is None:
orbit, _ = find_orbit4(ring, dp, keep_lattice=keep_lattice)
keep_lattice = True
# Construct matrix of plus and minus deltas
dg = numpy.asfortranarray(0.5 * numpy.diag([xy_step] * 6)[:, :4])
dmat = numpy.concatenate((dg, -dg), axis=1)
# Add the deltas to multiple copies of the closed orbit
in_mat = orbit.reshape(6, 1) + dmat
refs = uint32_refpts(refpts, len(ring))
out_mat = numpy.rollaxis(
numpy.squeeze(lattice_pass(ring,
in_mat,
refpts=refs,
keep_lattice=keep_lattice),
axis=3), -1)
# out_mat: 8 particles at n refpts for one turn
# (x + d) - (x - d) / d
m44 = (in_mat[:4, :4] - in_mat[:4, 4:]) / xy_step
if len(refs) > 0:
mstack = (out_mat[:, :4, :4] - out_mat[:, :4, 4:]) / xy_step
if full:
mstack = numpy.stack([mrotate(mat) for mat in mstack], axis=0)
else:
mstack = numpy.empty((0, 4, 4), dtype=float)
return m44, mstack
def compute(self, ring, *args, **kwargs):
"""Orbit computation before evaluation of all constraints"""
orbit0, orbit = find_orbit4(ring, refpts=self.refpts, **kwargs)
return (orbit[ref[self.refpts]].T for ref in self.refs), ()
def test_find_orbit4_result_unchanged_by_atpass(ring):
orbit = physics.find_orbit4(ring, DP)
orbit_copy = numpy.copy(orbit)
orbit[4] = DP
atpass(ring, orbit, 1)
numpy.testing.assert_allclose(orbit[:4], orbit_copy[:4], atol=1e-12)
def test_find_orbit4_produces_same_result_with_keep_lattice_True(dba_lattice):
orbit0, _ = physics.find_orbit4(dba_lattice)
orbit1, _ = physics.find_orbit4(dba_lattice, keep_lattice=True)
assert_close(orbit0, orbit1, rtol=0, atol=1e-12)
