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)