def test_find_m44_returns_same_answer_as_matlab(dba_lattice, refpts):
m44, mstack = physics.find_m44(dba_lattice, dp=DP, refpts=refpts)
assert_close(m44, M44_MATLAB, rtol=1e-5, atol=1e-7)
assert mstack.shape == (len(refpts), 4, 4)
m44, mstack = physics.find_m44(dba_lattice, dp=DP, refpts=refpts,
full=True)
assert_close(m44, M44_MATLAB, rtol=1e-5, atol=1e-7)
assert mstack.shape == (len(refpts), 4, 4)
def test_find_m44_no_refpts(dba_lattice):
m44 = physics.find_m44(dba_lattice, dp=DP)[0]
expected = numpy.array([[-0.66380, 2.23415, 0., 0.],
[-0.25037, -0.66380, 0., 0.],
[-1.45698e-31, -1.15008e-30, -0.99922, 0.26217],
[6.57748e-33, 8.75482e-32, -5.9497e-3, -0.99922]])
numpy.testing.assert_allclose(m44, expected, rtol=1e-5, atol=1e-7)
def test_find_m44_returns_same_answer_as_matlab(ring, refpts):
m44, mstack = physics.find_m44(ring, dp=DP, refpts=refpts)
numpy.testing.assert_allclose(m44[:4],
M44_MATLAB[:4],
rtol=1e-5,
atol=1e-7)
stack_size = 0 if refpts is None else len(refpts)
assert mstack.shape == (4, 4, stack_size)
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 test_find_m44(engine, ml_lattice, py_lattice, dp, refpts):
nelems = len(py_lattice)
refpts = range(nelems + 1) if refpts is None else refpts
nrefs = len(uint32_refpts(refpts, nelems))
# Python call
py_m44, py_mstack = physics.find_m44(py_lattice, dp, refpts)
# Matlab call
ml_m44, ml_mstack = engine.findm44(ml_lattice,
dp,
_ml_refs(refpts, nelems),
nargout=2)
ml_mstack = numpy.rollaxis(
numpy.asarray(ml_mstack).reshape((4, 4, nrefs)), -1)
# Comparison
numpy.testing.assert_almost_equal(py_m44, numpy.asarray(ml_m44), decimal=8)
numpy.testing.assert_almost_equal(py_mstack, ml_mstack, decimal=8)
def test_find_m44(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_m44, ml_mstack = engine.findm44(ml_lattice, dp, ml_refpts, nargout=2)
py_ml_m44 = numpy.asarray(ml_m44)
# Python call
py_m44, py_mstack = physics.find_m44(py_lattice, dp, refpts)
py_mstack = numpy.squeeze(py_mstack)
# Matches to 5 d.p.
numpy.testing.assert_almost_equal(py_ml_m44, py_m44.T, decimal=5)
assert py_mstack.T.shape == tuple(numpy.asarray(ml_mstack).shape)
# Matches to 5 d.p.
numpy.testing.assert_almost_equal(py_mstack.T,
numpy.asarray(ml_mstack),
decimal=5)
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