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