def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_and_refpts_None( ): e = elements.LongElement('e', 0.1) f = elements.LongElement('e', 2.1) print(get_s_pos([e, f], None)) numpy.testing.assert_equal(get_s_pos([e, f], None), numpy.array([0, 0.1, 2.2]))
def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_using_bool_refpts( ): e = elements.LongElement('e', 0.1) f = elements.LongElement('e', 2.1) lat = [e, f] numpy.testing.assert_equal(get_s_pos(lat, numpy.ones(3, dtype=bool)), numpy.array([0, 0.1, 2.2]))
def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_using_int_refpts( ): e = elements.LongElement('e', 0.1) f = elements.LongElement('e', 2.1) lat = [e, f] numpy.testing.assert_equal(get_s_pos(lat, range(len(lat) + 1)), numpy.array([0, 0.1, 2.2]))
def test_get_s_pos_returns_all_points_for_lattice_with_two_elements_using_bool_refpts( ): e = elements.Element('e', 0.1) f = elements.Element('e', 2.1) lat = [e, f] numpy.testing.assert_equal( lattice.get_s_pos(lat, numpy.array((True, True, True))), numpy.array([0, 0.1, 2.2]))
def _orbit6(ring, cavpts=None, guess=None, keep_lattice=False, **kwargs): """Solver for 6D motion""" convergence = kwargs.pop('convergence', DConstant.OrbConvergence) max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter) xy_step = kwargs.pop('XYStep', DConstant.XYStep) dp_step = kwargs.pop('DPStep', DConstant.DPStep) method = kwargs.pop('method', ELossMethod.TRACKING) # Get revolution period l0 = get_s_pos(ring, len(ring)) cavities = [elm for elm in ring if isinstance(elm, elements.RFCavity)] if len(cavities) == 0: raise AtError('No cavity found in the lattice.') f_rf = cavities[0].Frequency harm_number = cavities[0].HarmNumber if guess is None: _, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts) ref_in = numpy.zeros((6,), order='F') ref_in[5] = -dt else: ref_in = numpy.copy(guess) theta = numpy.zeros((6,)) theta[5] = constants.speed_of_light * harm_number / f_rf - l0 scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + \ dp_step * numpy.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]) delta_matrix = numpy.asfortranarray( numpy.concatenate((numpy.diag(scaling), numpy.zeros((6, 1))), axis=1)) id6 = numpy.asfortranarray(numpy.identity(6)) change = 1 itercount = 0 while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice) # the reference particle after one turn ref_out = in_mat[:, 6] # 6x6 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j6 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling a = j6 - id6 # f'(r_n) - 1 b = ref_out[:] - ref_in[:] - theta # b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) b_over_a = numpy.linalg.solve(a, b) r_next = ref_in - b_over_a # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keep_lattice = True if itercount == max_iterations: warnings.warn(AtWarning('Maximum number of iterations reached. ' 'Possible non-convergence')) return ref_in
def get_twiss(ring, dp=0.0, refpts=None, get_chrom=False, ddp=DDP): """ Determine Twiss parameters by first finding the transfer matrix. """ def twiss22(mat, ms): """ Calculate Twiss parameters from the standard 2x2 transfer matrix (i.e. x or y). """ sin_mu_end = (numpy.sign(mat[0, 1]) * math.sqrt(-mat[0, 1] * mat[1, 0] - (mat[0, 0] - mat[1, 1])**2 / 4)) alpha0 = (mat[0, 0] - mat[1, 1]) / 2.0 / sin_mu_end beta0 = mat[0, 1] / sin_mu_end beta = ((ms[0, 0, :] * beta0 - ms[0, 1, :] * alpha0)**2 + ms[0, 1, :]**2) / beta0 alpha = -((ms[0, 0, :] * beta0 - ms[0, 1, :] * alpha0) * (ms[1, 0, :] * beta0 - ms[1, 1, :] * alpha0) + ms[0, 1, :] * ms[1, 1, :]) / beta0 mu = numpy.arctan(ms[0, 1, :] / (ms[0, 0, :] * beta0 - ms[0, 1, :] * alpha0)) mu = betatron_phase_unwrap(mu) return alpha, beta, mu chrom = None refpts = lattice.uint32_refpts(refpts, len(ring)) nrefs = refpts.size if refpts[-1] != len(ring): refpts = numpy.append(refpts, [len(ring)]) orbit4, orbit = find_orbit4(ring, dp, refpts) m44, mstack = find_m44(ring, dp, refpts, orbit4=orbit4) ax, bx, mx = twiss22(m44[:2, :2], mstack[:2, :2, :]) ay, by, my = twiss22(m44[2:, 2:], mstack[2:, 2:, :]) tune = numpy.array((mx[-1], my[-1])) / (2 * numpy.pi) twiss = numpy.zeros(nrefs, dtype=TWISS_DTYPE) twiss['idx'] = refpts[:nrefs] twiss['s_pos'] = lattice.get_s_pos(ring, refpts[:nrefs]) twiss['closed_orbit'] = numpy.rollaxis(orbit, -1)[:nrefs] twiss['m44'] = numpy.rollaxis(mstack, -1)[:nrefs] twiss['alpha'] = numpy.rollaxis(numpy.vstack((ax, ay)), -1)[:nrefs] twiss['beta'] = numpy.rollaxis(numpy.vstack((bx, by)), -1)[:nrefs] twiss['mu'] = numpy.rollaxis(numpy.vstack((mx, my)), -1)[:nrefs] twiss['dispersion'] = numpy.NaN # Calculate chromaticity by calling this function again at a slightly # different momentum. if get_chrom: twissb, tuneb, _ = get_twiss(ring, dp + ddp, refpts[:nrefs]) chrom = (tuneb - tune) / ddp twiss['dispersion'] = (twissb['closed_orbit'] - twiss['closed_orbit']) / ddp return twiss, tune, chrom
def get_mcf(ring, dp=0.0, ddp=DDP, keep_lattice=False): """Compute momentum compaction factor PARAMETERS ring lattice description dp momentum deviation. Defaults to 0 KEYWORDS keep_lattice Assume no lattice change since the previous tracking. Defaults to False ddp=1.0E-8 momentum deviation used for differentiation """ fp_a, _ = find_orbit4(ring, dp=dp - 0.5 * ddp, keep_lattice=keep_lattice) fp_b, _ = find_orbit4(ring, dp=dp + 0.5 * ddp, keep_lattice=True) fp = numpy.stack((fp_a, fp_b), axis=0).T # generate a Fortran contiguous array b = numpy.squeeze(lattice_pass(ring, fp, keep_lattice=True), axis=(2, 3)) ring_length = get_s_pos(ring, len(ring)) return (b[5, 1] - b[5, 0]) / ddp / ring_length[0]
def linopt6(ring, refpts=None, dp=None, orbit=None, cavpts=None, twiss_in=None, get_chrom=False, get_w=False, keep_lattice=False, **kwargs): """Perform linear analysis of a fully coupled lattice elemdata0, beamdata, elemdata = linopt6(lattice) For circular machines, linopt6 analyses the 4x4 1-turn transfer matrix if radiation is OFF, or the 6x6 1-turn transfer matrix if radiation is ON. For a transfer line, The "twiss_in" intput must contain either: - a field 'R', as provided by ATLINOPT6, or - the fields 'beta' and 'alpha', as provided by linopt and linopt6 PARAMETERS ring lattice description. refpts=None elements at which data is returned. KEYWORDS orbit avoids looking for the closed orbit if is already known ((6,) array) keep_lattice Assume no lattice change since the previous tracking. Defaults to False XYStep=1.0e-8 transverse step for numerical computation DPStep=1.0E-6 momentum deviation used for computation of the closed orbit twiss_in=None Initial conditions for transfer line optics. Record array, as output by linopt or linopt6. If present, the attribute 'R' will be used, otherwise The attributes 'alpha' and 'beta' will be used. All other attributes are ignored. cavpts=None Cavity location for off-momentum tuning OUTPUT elemdata0 linear optics data at the entrance/end of the ring beamdata lattice properties elemdata linear optics at the points refered to by refpts, if refpts is None an empty elemdata structure is returned. elemdata is a record array with fields: R R-matrices (3, 6, 6) A A-matrix (6, 6) M Transfer matrix from the entrance of the line (6, 6) beta [betax, betay] vector alpha [alphax, alphay] vector dispersion (4,) dispersion vector mu [mux, muy], betatron phases s_pos longitudinal position [m] closed_orbit (6,) closed orbit vector W (2,) chromatic amplitude function All values given at the entrance of each element specified in refpts. Field values can be obtained with either elemdata['beta'] or elemdata.beta beamdata is a record with fields: tune Fractional tunes damping_time Damping times [s] chromaticity Chromaticities REFERENCES [1] Etienne Forest, Phys. Rev. E 58, 2481 – Published 1 August 1998 [2] Andrzej Wolski, Phys. Rev. ST Accel. Beams 9, 024001 – Published 3 February 2006 [3] Brian W. Montague Report LEP Note 165, CERN, 1979 """ def get_alphabeta(r): beta = numpy.array([r[0, 0, 0], r[1, 2, 2]]) alpha = -numpy.array([r[0, 1, 0], r[1, 3, 2]]) return alpha, beta # noinspection PyShadowingNames def build_r(ring, dp, orbit, refpts=None, mxx=None, **kwargs): """""" if ring.radiation: mt, ms = find_m66(ring, refpts=refpts, orbit=orbit, **kwargs) else: mt, ms = find_m44(ring, dp, refpts=refpts, orbit=orbit, **kwargs) a0, vps = a_matrix(mt if mxx is None else mxx) val0, vals = _r_analysis(a0, ms) return ms, vps, val0, vals def build_sigma(twin): """Build the initial distribution at entrance of the transfer line""" try: sigm = numpy.sum(twin.R, axis=0) except AttributeError: slices = [slice(2 * i, 2 * (i + 1)) for i in range(4)] ab = numpy.stack((twin.alpha, twin.beta), axis=1) sigm = numpy.zeros((4, 4)) for slc, (alpha, beta) in zip(slices, ab): gamma = (1.0 + alpha * alpha) / beta sigm[slc, slc] = numpy.array([[beta, -alpha], [-alpha, gamma]]) return sigm def output6(r123, a, mu, *args): """Extract output parameters from Bk matrices""" alpha, beta = get_alphabeta(r123) dispersion = r123[2, :4, 4] / r123[2, 4, 4] return (r123, a, alpha, beta, mu, dispersion) + args def output4(r12, a, *args): """Extract output parameters from Bk matrices""" alpha, beta = get_alphabeta(r12) return (r12, a, alpha, beta) + args def get_disp(ringup, ringdn, dpup, dpdn, refpts=None, matpts=None, keep_lattice=False, **kwargs): def off_momentum(rng, dp): orb0, orbs = _find_orbit(rng, dp, refpts=refpts, keep_lattice=keep_lattice, **kwargs) dp = orb0[4] # in 6D, dp comes out of find_orbit6 _, vps, el0, els = build_r(rng, dp, orb0, refpts=matpts, keep_lattice=True, **kwargs) tunes = numpy.mod(numpy.angle(vps) / 2.0 / pi, 1.0) return dp, tunes, orb0, orbs, el0, els def chromfunc(ddp, elup, eldn): aup, bup = get_alphabeta(elup[0]) adn, bdn = get_alphabeta(eldn[0]) db = (bup - bdn) / ddp mb = (bup + bdn) / 2 da = (aup - adn) / ddp ma = (aup + adn) / 2 w = numpy.sqrt((da - ma / mb * db)**2 + (db / mb)**2) return w dpup, tunesup, o0up, orbup, el0up, elup = off_momentum(ringup, dpup) dpdn, tunesdn, o0dn, orbdn, el0dn, eldn = off_momentum(ringdn, dpdn) deltap = dpup - dpdn chrom = (tunesup - tunesdn) / deltap disp0 = (o0up - o0dn)[:4] / deltap disp = ((oup - odn)[:4] / deltap for oup, odn in zip(orbup, orbdn)) w0 = chromfunc(deltap, el0up, el0dn) w = (chromfunc(deltap, *v) for v in zip(elup, eldn)) return chrom, disp0, disp, w0, w def unwrap(mu): """Remove the phase jumps""" dmu = numpy.diff(numpy.concatenate((numpy.zeros((1, dms)), mu)), axis=0) jumps = dmu < -1.e-3 mu += numpy.cumsum(jumps, axis=0) * 2.0 * numpy.pi dp_step = kwargs.get('DPStep', DConstant.DPStep) if twiss_in is None: # Circular machine mxx = None else: # Transfer line if orbit is None: orbit = numpy.zeros((6, )) sigma = build_sigma(twiss_in) mxx = sigma.dot(jmat(sigma.shape[0] // 2)) orb0, orbs = _find_orbit(ring, dp, refpts=refpts, orbit=orbit, **kwargs) dp = orb0[4] ms, vps, el0, els = build_r(ring, dp, orb0, refpts=refpts, mxx=mxx, keep_lattice=keep_lattice, **kwargs) dms = vps.size length = ring.get_s_pos(len(ring))[0] tunes = numpy.mod(numpy.angle(vps) / 2.0 / pi, 1.0) damping_rates = -numpy.log(numpy.absolute(vps)) damping_times = length / clight / damping_rates spos = get_s_pos(ring, refpts) if dms >= 3: # 6D processsing output = output6 dtype = W_DTYPE6 data0 = [*el0, numpy.identity(2 * dms), 0.0, orb0] datas = [els, iter(ms), iter(spos), iter(orbs)] if get_chrom or get_w: f0 = get_rf_frequency(ring, cavpts=cavpts) df = -dp_step * get_mcf(ring.radiation_off(copy=True)) * f0 rgup = set_rf_frequency(ring, f0 + 0.5 * df, cavpts=cavpts, copy=True) rgdn = set_rf_frequency(ring, f0 - 0.5 * df, cavpts=cavpts, copy=True) if get_w: dtype = W_DTYPE6 + W_DTYPEW chrom, _, _, w0, ws = get_disp(rgup, rgdn, None, None, matpts=refpts, **kwargs) data0.append(w0) datas.append(ws) else: chrom, _, _, _, _ = get_disp(rgup, rgdn, None, None) else: chrom = numpy.NaN else: # 4D processsing output = output4 dtype = W_DTYPE4 data0 = [*el0] datas = [els] if get_w: dtype = W_DTYPE4 + W_DTYPEW chrom, d0, ds, w0, ws = get_disp(ring, ring, dp + 0.5 * dp_step, dp - 0.5 * dp_step, refpts=refpts, matpts=refpts, keep_lattice=True, **kwargs) data0 += [d0, numpy.identity(2 * dms), 0.0, orb0, w0] datas += [ds, iter(ms), iter(spos), iter(orbs), ws] elif get_chrom: chrom, d0, ds, w0, ws = get_disp(ring, ring, dp + 0.5 * dp_step, dp - 0.5 * dp_step, refpts=refpts, keep_lattice=True, **kwargs) data0 += [d0, numpy.identity(2 * dms), 0.0, orb0] datas += [ds, iter(ms), iter(spos), iter(orbs)] else: chrom = numpy.NaN data0 += [numpy.NaN, numpy.identity(2 * dms), 0.0, orb0] datas += [repeat(numpy.NaN), iter(ms), iter(spos), iter(orbs)] elemdata0 = numpy.array(output(*data0), dtype=dtype).view(numpy.recarray) elemdata = numpy.fromiter((output(*el, *v) for el, *v in zip(*datas)), dtype, count=ring.refcount(refpts)).view(numpy.recarray) beamdata = numpy.array( (tunes, chrom, damping_times), dtype=[('tune', numpy.float64, (dms, )), ('chromaticity', numpy.float64, (dms, )), ('damping_time', numpy.float64, (dms, ))]).view(numpy.recarray) unwrap(elemdata.mu) return elemdata0, beamdata, elemdata
def get_twiss(ring, dp=0.0, refpts=None, get_chrom=False, orbit=None, keep_lattice=False, ddp=DDP): """ Perform linear analysis of the NON-COUPLED lattices twiss0, tune, chrom, twiss = get_twiss(ring, dp[, refpts]) PARAMETERS ring lattice description. dp=0.0 momentum deviation. refpts=None elements at which data is returned. It can be: 1) an integer in the range [-len(ring), len(ring)-1] selecting the element according to python indexing rules. As a special case, len(ring) is allowed and refers to the end of the last element, 2) an ordered list of such integers without duplicates, 3) a numpy array of booleans of maximum length len(ring)+1, where selected elements are True. KEYWORDS orbit avoids looking for the closed orbit if is already known ((6,) array) get_chrom=False compute dispersion and chromaticities. Needs computing the optics at 2 different momentum deviations around the central one. keep_lattice Assume no lattice change since the previous tracking. Defaults to False ddp=1.0E-8 momentum deviation used for computation of chromaticities and dispersion OUTPUT twiss0 linear optics data at the entrance/end of the ring tune [tune_h, tune_v], fractional part of the linear tunes chrom [ksi_h , ksi_v], chromaticities ksi = d(nu)/(dP/P). Only computed if 'get_chrom' is True twiss linear optics at the points refered to by refpts, if refpts is None an empty twiss structure is returned. twiss is a record array with fields: idx element index in the ring s_pos longitudinal position [m] closed_orbit (6,) closed orbit vector dispersion (4,) dispersion vector Only computed if 'get_chrom' is True m44 (4, 4) transfer matrix M from the beginning of ring to the entrance of the element mu (2,) betatron phase (modulo 2*pi) beta [betax, betay] vector alpha [alphax, alphay] vector All values given at the entrance of each element specified in refpts. Field values can be obtained with either twiss['idx'] or twiss.idx See also linopt """ uintrefs = uint32_refpts(refpts, len(ring)) if orbit is None: orbit, _ = find_orbit4(ring, dp, keep_lattice=keep_lattice) keep_lattice = True orbs = numpy.squeeze(lattice_pass(ring, orbit.copy(order='K'), refpts=uintrefs, keep_lattice=keep_lattice), axis=(1, 3)).T m44, mstack = find_m44(ring, dp, uintrefs, orbit=orbit, keep_lattice=True) nrefs = uintrefs.size # Get initial twiss parameters a0_x, b0_x, tune_x = _closure(m44[:2, :2]) a0_y, b0_y, tune_y = _closure(m44[2:, 2:]) tune = numpy.array([tune_x, tune_y]) # Calculate chromaticity by calling this function again at a slightly # different momentum. if get_chrom: d0_up, tune_up, _, l_up = get_twiss(ring, dp + 0.5 * ddp, uintrefs, keep_lattice=True) d0_down, tune_down, _, l_down = get_twiss(ring, dp - 0.5 * ddp, uintrefs, keep_lattice=True) chrom = (tune_up - tune_down) / ddp dispersion = (l_up['closed_orbit'] - l_down['closed_orbit'])[:, :4] / ddp disp0 = (d0_up['closed_orbit'] - d0_down['closed_orbit'])[:4] / ddp else: chrom = numpy.array([numpy.NaN, numpy.NaN]) dispersion = numpy.NaN disp0 = numpy.NaN twiss0 = numpy.rec.fromarrays((len(ring), get_s_pos( ring, len(ring))[0], orbit, disp0, numpy.array( [a0_x, a0_y]), numpy.array([b0_x, b0_y]), 2.0 * pi * tune, m44), dtype=TWISS_DTYPE) # Propagate to reference points twiss = numpy.rec.array(numpy.zeros(nrefs, dtype=TWISS_DTYPE)) if nrefs > 0: alpha_x, beta_x, mu_x = _twiss22(mstack[:, :2, :2], a0_x, b0_x) alpha_z, beta_z, mu_z = _twiss22(mstack[:, 2:, 2:], a0_y, b0_y) twiss['idx'] = uintrefs twiss['s_pos'] = get_s_pos(ring, uintrefs[:nrefs]) twiss['closed_orbit'] = orbs twiss['m44'] = mstack twiss['alpha'] = numpy.stack((alpha_x, alpha_z), axis=1) twiss['beta'] = numpy.stack((beta_x, beta_z), axis=1) twiss['mu'] = numpy.stack((mu_x, mu_z), axis=1) twiss['dispersion'] = dispersion return twiss0, tune, chrom, twiss
def linopt(ring, dp=0.0, refpts=None, get_chrom=False, orbit=None, keep_lattice=False, ddp=DDP, coupled=True): """ Perform linear analysis of a lattice lindata0, tune, chrom, lindata = linopt(ring, dp[, refpts]) PARAMETERS ring lattice description. dp=0.0 momentum deviation. refpts=None elements at which data is returned. It can be: 1) an integer in the range [-len(ring), len(ring)-1] selecting the element according to python indexing rules. As a special case, len(ring) is allowed and refers to the end of the last element, 2) an ordered list of such integers without duplicates, 3) a numpy array of booleans of maximum length len(ring)+1, where selected elements are True. KEYWORDS orbit avoids looking for the closed orbit if is already known ((6,) array) get_chrom=False compute dispersion and chromaticities. Needs computing the optics at 2 different momentum deviations around the central one. keep_lattice Assume no lattice change since the previous tracking. Defaults to False ddp=1.0E-8 momentum deviation used for computation of chromaticities and dispersion coupled=True if False, simplify the calculations by assuming no H/V coupling OUTPUT lindata0 linear optics data at the entrance/end of the ring tune [tune_A, tune_B], linear tunes for the two normal modes of linear motion [1] chrom [ksi_A , ksi_B], chromaticities ksi = d(nu)/(dP/P). Only computed if 'get_chrom' is True lindata linear optics at the points refered to by refpts, if refpts is None an empty lindata structure is returned. lindata is a record array with fields: idx element index in the ring s_pos longitudinal position [m] closed_orbit (6,) closed orbit vector dispersion (4,) dispersion vector Only computed if 'get_chrom' is True m44 (4, 4) transfer matrix M from the beginning of ring to the entrance of the element [2] A (2, 2) matrix A in [3] B (2, 2) matrix B in [3] C (2, 2) matrix C in [3] gamma gamma parameter of the transformation to eigenmodes mu [mux, muy], betatron phase (modulo 2*pi) beta [betax, betay] vector alpha [alphax, alphay] vector All values given at the entrance of each element specified in refpts. Field values can be obtained with either lindata['idx'] or lindata.idx REFERENCES [1] D.Edwars,L.Teng IEEE Trans.Nucl.Sci. NS-20, No.3, p.885-888, 1973 [2] E.Courant, H.Snyder [3] D.Sagan, D.Rubin Phys.Rev.Spec.Top.-Accelerators and beams, vol.2 (1999) See also get_twiss """ # noinspection PyShadowingNames def analyze(r44): t44 = r44.reshape((4, 4)) mm = t44[:2, :2] nn = t44[2:, 2:] m = t44[:2, 2:] n = t44[2:, :2] gamma = sqrt(numpy.linalg.det(numpy.dot(n, C) + numpy.dot(G, nn))) msa = (G.dot(mm) - m.dot(_jmt.dot(C.T.dot(_jmt.T)))) / gamma msb = (numpy.dot(n, C) + numpy.dot(G, nn)) / gamma cc = (numpy.dot(mm, C) + numpy.dot(G, m)).dot( _jmt.dot(msb.T.dot(_jmt.T))) return msa, msb, gamma, cc uintrefs = uint32_refpts([] if refpts is None else refpts, len(ring)) if orbit is None: orbit, _ = find_orbit4(ring, dp, keep_lattice=keep_lattice) keep_lattice = True orbs = numpy.squeeze(lattice_pass(ring, orbit.copy(order='K'), refpts=uintrefs, keep_lattice=keep_lattice), axis=(1, 3)).T m44, mstack = find_m44(ring, dp, uintrefs, orbit=orbit, keep_lattice=True) nrefs = uintrefs.size M = m44[:2, :2] N = m44[2:, 2:] m = m44[:2, 2:] n = m44[2:, :2] if coupled: # Calculate A, B, C, gamma at the first element H = m + _jmt.dot(n.T.dot(_jmt.T)) t = numpy.trace(M - N) t2 = t * t t2h = t2 + 4.0 * numpy.linalg.det(H) g = sqrt(1.0 + sqrt(t2 / t2h)) / sqrt(2.0) G = numpy.diag((g, g)) C = -H * numpy.sign(t) / (g * sqrt(t2h)) A = G.dot(G.dot(M)) - numpy.dot( G, (m.dot(_jmt.dot(C.T.dot(_jmt.T))) + C.dot(n))) + C.dot( N.dot(_jmt.dot(C.T.dot(_jmt.T)))) B = G.dot(G.dot(N)) + numpy.dot( G, (_jmt.dot(C.T.dot(_jmt.T.dot(m))) + n.dot(C))) + _jmt.dot( C.T.dot(_jmt.T.dot(M.dot(C)))) else: A = M B = N C = numpy.zeros((2, 2)) g = 1.0 # Get initial twiss parameters a0_a, b0_a, tune_a = _closure(A) a0_b, b0_b, tune_b = _closure(B) tune = numpy.array([tune_a, tune_b]) if get_chrom: d0_up, tune_up, _, l_up = linopt(ring, dp + 0.5 * ddp, uintrefs, keep_lattice=True, coupled=coupled) d0_down, tune_down, _, l_down = linopt(ring, dp - 0.5 * ddp, uintrefs, keep_lattice=True, coupled=coupled) chrom = (tune_up - tune_down) / ddp dispersion = (l_up['closed_orbit'] - l_down['closed_orbit'])[:, :4] / ddp disp0 = (d0_up['closed_orbit'] - d0_down['closed_orbit'])[:4] / ddp else: chrom = numpy.array([numpy.NaN, numpy.NaN]) dispersion = numpy.NaN disp0 = numpy.NaN lindata0 = numpy.rec.fromarrays( (len(ring), get_s_pos(ring, len(ring))[0], orbit, disp0, numpy.array([a0_a, a0_b]), numpy.array( [b0_a, b0_b]), 2.0 * pi * tune, m44, A, B, C, g), dtype=LINDATA_DTYPE) # Propagate to reference points lindata = numpy.rec.array(numpy.zeros(nrefs, dtype=LINDATA_DTYPE)) if nrefs > 0: if coupled: MSA, MSB, gamma, CL = zip(*[analyze(ms44) for ms44 in mstack]) msa = numpy.stack(MSA, axis=0) msb = numpy.stack(MSB, axis=0) else: msa = mstack[:, :2, :2] msb = mstack[:, 2:, 2:] gamma = 1.0 CL = numpy.zeros((1, 2, 2)) alpha_a, beta_a, mu_a = _twiss22(msa, a0_a, b0_a) alpha_b, beta_b, mu_b = _twiss22(msb, a0_b, b0_b) lindata['idx'] = uintrefs lindata['s_pos'] = get_s_pos(ring, uintrefs) lindata['closed_orbit'] = orbs lindata['m44'] = mstack lindata['alpha'] = numpy.stack((alpha_a, alpha_b), axis=1) lindata['beta'] = numpy.stack((beta_a, beta_b), axis=1) lindata['mu'] = numpy.stack((mu_a, mu_b), axis=1) lindata['A'] = [ ms.dot(A.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msa ] lindata['B'] = [ ms.dot(B.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msb ] lindata['C'] = CL lindata['gamma'] = gamma lindata['dispersion'] = dispersion return lindata0, tune, chrom, lindata
def linopt(ring, dp=0.0, refpts=None, get_chrom=False, orbit=None, keep_lattice=False, coupled=True, twiss_in=None, get_w=False, **kwargs): """ Perform linear analysis of a lattice lindata0, tune, chrom, lindata = linopt(ring, dp[, refpts]) PARAMETERS ring lattice description. dp=0.0 momentum deviation. refpts=None elements at which data is returned. It can be: 1) an integer in the range [-len(ring), len(ring)-1] selecting the element according to python indexing rules. As a special case, len(ring) is allowed and refers to the end of the last element, 2) an ordered list of such integers without duplicates, 3) a numpy array of booleans of maximum length len(ring)+1, where selected elements are True. KEYWORDS orbit avoids looking for the closed orbit if is already known ((6,) array) get_chrom=False compute dispersion and chromaticities. Needs computing the tune and orbit at 2 different momentum deviations around the central one. keep_lattice Assume no lattice change since the previous tracking. Defaults to False XYStep=1.0e-8 transverse step for numerical computation DPStep=1.0E-6 momentum deviation used for computation of chromaticities and dispersion coupled=True if False, simplify the calculations by assuming no H/V coupling twiss_in=None Initial twiss to compute transfer line optics of the type lindata, the initial orbit in twiss_in is ignored, only the beta and alpha are required other quatities set to 0 if absent get_w=False computes chromatic amplitude functions (W) [4], need to compute the optics at 2 different momentum deviations around the central one. OUTPUT lindata0 linear optics data at the entrance/end of the ring tune [tune_A, tune_B], linear tunes for the two normal modes of linear motion [1] chrom [ksi_A , ksi_B], chromaticities ksi = d(nu)/(dP/P). Only computed if 'get_chrom' is True lindata linear optics at the points refered to by refpts, if refpts is None an empty lindata structure is returned. lindata is a record array with fields: idx element index in the ring s_pos longitudinal position [m] closed_orbit (6,) closed orbit vector dispersion (4,) dispersion vector W (2,) chromatic amplitude function Only computed if 'get_chrom' is True m44 (4, 4) transfer matrix M from the beginning of ring to the entrance of the element [2] mu [mux, muy], betatron phase (modulo 2*pi) beta [betax, betay] vector alpha [alphax, alphay] vector All values given at the entrance of each element specified in refpts. In case coupled = True additional outputs are available: A (2, 2) matrix A in [3] B (2, 2) matrix B in [3] C (2, 2) matrix C in [3] gamma gamma parameter of the transformation to eigenmodes Field values can be obtained with either lindata['idx'] or lindata.idx REFERENCES [1] D.Edwars,L.Teng IEEE Trans.Nucl.Sci. NS-20, No.3, p.885-888, 1973 [2] E.Courant, H.Snyder [3] D.Sagan, D.Rubin Phys.Rev.Spec.Top.-Accelerators and beams, vol.2 (1999) [4] Brian W. Montague Report LEP Note 165, CERN, 1979 """ # noinspection PyShadowingNames def analyze(r44): t44 = r44.reshape((4, 4)) mm = t44[:2, :2] nn = t44[2:, 2:] m = t44[:2, 2:] n = t44[2:, :2] gamma = sqrt(numpy.linalg.det(numpy.dot(n, C) + numpy.dot(G, nn))) msa = (G.dot(mm) - m.dot(_jmt.dot(C.T.dot(_jmt.T)))) / gamma msb = (numpy.dot(n, C) + numpy.dot(G, nn)) / gamma cc = (numpy.dot(mm, C) + numpy.dot(G, m)).dot( _jmt.dot(msb.T.dot(_jmt.T))) return msa, msb, gamma, cc xy_step = kwargs.pop('XYStep', DConstant.XYStep) dp_step = kwargs.pop('DPStep', DConstant.DPStep) uintrefs = uint32_refpts([] if refpts is None else refpts, len(ring)) # Get initial orbit if twiss_in is None: if orbit is None: orbit, _ = find_orbit4(ring, dp, keep_lattice=keep_lattice, XYStep=xy_step) keep_lattice = True disp0 = numpy.NaN if get_chrom or get_w: orbit_up, _ = find_orbit4(ring, dp + 0.5 * dp_step, XYStep=xy_step, keep_lattice=keep_lattice) orbit_down, _ = find_orbit4(ring, dp - 0.5 * dp_step, XYStep=xy_step, keep_lattice=keep_lattice) disp0 = numpy.array(orbit_up - orbit_down)[:4] / dp_step else: if orbit is None: orbit = numpy.zeros((6, )) disp0 = numpy.NaN if get_chrom or get_w: try: disp0 = twiss_in['dispersion'] except KeyError: print('Dispersion not found in twiss_in, setting to zero') disp0 = numpy.zeros((4, )) dorbit = numpy.hstack( (0.5 * dp_step * disp0, numpy.array([0.5 * dp_step, 0]))) orbit_up = orbit + dorbit orbit_down = orbit - dorbit orbs = numpy.squeeze(lattice_pass(ring, orbit.copy(order='K'), refpts=uintrefs, keep_lattice=keep_lattice), axis=(1, 3)).T m44, mstack = find_m44(ring, dp, uintrefs, orbit=orbit, keep_lattice=True, XYStep=xy_step) nrefs = uintrefs.size M = m44[:2, :2] N = m44[2:, 2:] m = m44[:2, 2:] n = m44[2:, :2] # Calculate A, B, C, gamma at the first element if coupled: H = m + _jmt.dot(n.T.dot(_jmt.T)) t = numpy.trace(M - N) t2 = t * t t2h = t2 + 4.0 * numpy.linalg.det(H) g = sqrt(1.0 + sqrt(t2 / t2h)) / sqrt(2.0) G = numpy.diag((g, g)) C = -H * numpy.sign(t) / (g * sqrt(t2h)) A = G.dot(G.dot(M)) - numpy.dot( G, (m.dot(_jmt.dot(C.T.dot(_jmt.T))) + C.dot(n))) + C.dot( N.dot(_jmt.dot(C.T.dot(_jmt.T)))) B = G.dot(G.dot(N)) + numpy.dot( G, (_jmt.dot(C.T.dot(_jmt.T.dot(m))) + n.dot(C))) + _jmt.dot( C.T.dot(_jmt.T.dot(M.dot(C)))) else: A = M B = N C = numpy.zeros((2, 2)) g = 1.0 # Get initial twiss parameters if twiss_in is None: a0_a, b0_a, tune_a = _closure(A) a0_b, b0_b, tune_b = _closure(B) tune = numpy.array([tune_a, tune_b]) else: try: a0_a, a0_b = twiss_in['alpha'][0], twiss_in['alpha'][1] except KeyError: raise ValueError('Initial alpha required for transfer line') try: b0_a, b0_b = twiss_in['beta'][0], twiss_in['beta'][1] except KeyError: raise ValueError('Initial beta required for transfer line') try: tune = numpy.array([twiss_in['mu'][0], twiss_in['mu'][1] ]) / (2 * pi) except KeyError: print('Mu not found in twiss_in, setting to zero') tune = numpy.zeros((2, )) # Get initial chromatic functions and dispersion if get_w: # noinspection PyUnboundLocalVariable kwup = dict(orbit=orbit_up, twiss_in=twiss_in) # noinspection PyUnboundLocalVariable kwdown = dict(orbit=orbit_down, twiss_in=twiss_in) param_up = linopt(ring, dp=dp + 0.5 * dp_step, refpts=uintrefs, keep_lattice=True, coupled=coupled, XYStep=xy_step, **kwup) param_down = linopt(ring, dp=dp - 0.5 * dp_step, refpts=uintrefs, keep_lattice=True, coupled=coupled, XYStep=xy_step, **kwdown) param_up_down = param_up + param_down chrom, dispersion, w0, w = _chromfunc(dp_step, *param_up_down) elif get_chrom: # noinspection PyUnboundLocalVariable kwup = dict(orbit=orbit_up, twiss_in=twiss_in) # noinspection PyUnboundLocalVariable kwdown = dict(orbit=orbit_down, twiss_in=twiss_in) _, tune_up, _, _ = linopt(ring, dp=dp + 0.5 * dp_step, coupled=coupled, keep_lattice=True, XYStep=xy_step, DPStep=dp_step, **kwup) orb_up = numpy.squeeze(lattice_pass(ring, kwup['orbit'].copy(order='K'), refpts=uintrefs, keep_lattice=True), axis=(1, 3)).T _, tune_down, _, _ = linopt(ring, dp=dp - 0.5 * dp_step, coupled=coupled, keep_lattice=True, XYStep=xy_step, DPStep=dp_step, **kwdown) orb_down = numpy.squeeze(lattice_pass(ring, kwdown['orbit'].copy(order='K'), refpts=uintrefs, keep_lattice=True), axis=(1, 3)).T chrom = (tune_up - tune_down) / dp_step dispersion = (orb_up - orb_down)[:, :4] / dp_step w0 = numpy.array([numpy.NaN, numpy.NaN]) w = numpy.array([numpy.NaN, numpy.NaN]) else: chrom = numpy.array([numpy.NaN, numpy.NaN]) dispersion = numpy.array([numpy.NaN, numpy.NaN, numpy.NaN, numpy.NaN]) w0 = numpy.array([numpy.NaN, numpy.NaN]) w = numpy.array([numpy.NaN, numpy.NaN]) lindata0 = numpy.rec.fromarrays( (len(ring), get_s_pos(ring, len(ring))[0], orbit, disp0, numpy.array([a0_a, a0_b]), numpy.array( [b0_a, b0_b]), 2.0 * pi * tune, m44, A, B, C, g, w0), dtype=LINDATA_DTYPE) # Propagate to reference points lindata = numpy.rec.array(numpy.zeros(nrefs, dtype=LINDATA_DTYPE)) if nrefs > 0: if coupled: MSA, MSB, gamma, CL = zip(*[analyze(ms44) for ms44 in mstack]) msa = numpy.stack(MSA, axis=0) msb = numpy.stack(MSB, axis=0) AL = [ms.dot(A.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msa] BL = [ms.dot(B.dot(_jmt.dot(ms.T.dot(_jmt.T)))) for ms in msb] else: msa = mstack[:, :2, :2] msb = mstack[:, 2:, 2:] AL = numpy.NaN BL = numpy.NaN CL = numpy.NaN gamma = numpy.NaN alpha_a, beta_a, mu_a = _twiss22(msa, a0_a, b0_a) alpha_b, beta_b, mu_b = _twiss22(msb, a0_b, b0_b) if twiss_in is not None: qtmp = numpy.array([mu_a[-1], mu_b[-1]]) / (2 * numpy.pi) qtmp -= numpy.floor(qtmp) mu_a += tune[0] * 2 * pi mu_b += tune[1] * 2 * pi tune = qtmp lindata['idx'] = uintrefs lindata['s_pos'] = get_s_pos(ring, uintrefs) lindata['closed_orbit'] = orbs lindata['m44'] = mstack lindata['alpha'] = numpy.stack((alpha_a, alpha_b), axis=1) lindata['beta'] = numpy.stack((beta_a, beta_b), axis=1) lindata['dispersion'] = dispersion lindata['mu'] = numpy.stack((mu_a, mu_b), axis=1) lindata['A'] = AL lindata['B'] = BL lindata['C'] = CL lindata['gamma'] = gamma lindata['W'] = w return lindata0, tune, chrom, lindata
def test_get_s_pos_returns_length_for_lattice_with_one_element(): e = elements.LongElement('e', 0.1) assert get_s_pos([e], [1]) == numpy.array([0.1])
def test_get_s_pos_returns_zero_for_empty_lattice(): numpy.testing.assert_equal(get_s_pos([], None), numpy.array((0,)))
def _orbit6(ring, cavpts=None, guess=None, keep_lattice=False, **kwargs): """Solver for 6D motion""" def iscavity(elem): return isinstance(elem, elements.RFCavity) and \ elem.PassMethod.endswith('CavityPass') convergence = kwargs.pop('convergence', DConstant.OrbConvergence) max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter) xy_step = kwargs.pop('XYStep', DConstant.XYStep) dp_step = kwargs.pop('DPStep', DConstant.DPStep) method = kwargs.pop('method', ELossMethod.TRACKING) l0 = get_s_pos(ring, len(ring))[0] # Get the main RF frequency (the lowest) try: f_rf = min(elm.Frequency for elm in ring if iscavity(elm)) except ValueError: raise AtError('No cavity found in the lattice.') # gamma = self.energy / self.particle.mass # beta = math.sqrt(1.0 - 1.0 / gamma / gamma) # h = round(fmin*l0/beta/clight) harm_number = round(f_rf * l0 / clight) if guess is None: _, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts) # Getting timelag by tracking uses a different lattice, # so we cannot now use the same one again. if method is ELossMethod.TRACKING: keep_lattice = False ref_in = numpy.zeros((6, ), order='F') ref_in[5] = -dt else: ref_in = numpy.copy(guess) theta = numpy.zeros((6, )) theta[5] = clight * harm_number / f_rf - l0 scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + \ dp_step * numpy.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]) delta_matrix = numpy.asfortranarray( numpy.concatenate((numpy.diag(scaling), numpy.zeros((6, 1))), axis=1)) id6 = numpy.asfortranarray(numpy.identity(6)) change = 1 itercount = 0 while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice) # the reference particle after one turn ref_out = in_mat[:, 6] # 6x6 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j6 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling a = j6 - id6 # f'(r_n) - 1 b = ref_out[:] - ref_in[:] - theta # b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) b_over_a = numpy.linalg.solve(a, b) r_next = ref_in - b_over_a # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keep_lattice = True if itercount == max_iterations: warnings.warn( AtWarning('Maximum number of iterations reached. ' 'Possible non-convergence')) return ref_in
def find_orbit6(ring, refpts=None, guess=None, **kwargs): """find_orbit6 finds the closed orbit in the full 6-D phase space by numerically solving for a fixed point of the one turn map M calculated with lattice_pass (X, PX, Y, PY, DP, CT2 ) = M (X, PX, Y, PY, DP, CT1) with constraint CT2 - CT1 = C*HarmNumber(1/Frf - 1/Frf0) IMPORTANT!!! find_orbit6 is a realistic simulation 1. The Frf frequency in the RF cavities (may be different from Frf0) imposes the synchronous condition CT2 - CT1 = C*HarmNumber(1/Frf - 1/Frf0) 2. The algorithm numerically calculates 6-by-6 Jacobian matrix J6. In order for (J-E) matrix to be non-singular it is NECESSARY to use a realistic PassMethod for cavities with non-zero momentum kick (such as ThinCavityPass). 3. find_orbit6 can find orbits with radiation. In order for the solution to exist the cavity must supply adequate energy compensation. In the simplest case of a single cavity, it must have 'Voltage' field set so that Voltage > Erad - energy loss per turn 4. There is a family of solutions that correspond to different RF buckets They differ in the 6-th coordinate by C*Nb/Frf. Nb = 1 .. HarmNum-1 5. The value of the 6-th coordinate found at the cavity gives the equilibrium RF phase. If there is no radiation the phase is 0; PARAMETERS ring Sequence of AT elements refpts elements at which data is returned. It can be: 1) an integer in the range [-len(ring), len(ring)-1] selecting the element according to python indexing rules. As a special case, len(ring) is allowed and refers to the end of the last element, 2) an ordered list of such integers without duplicates, 3) a numpy array of booleans of maximum length len(ring)+1, where selected elements are True. Defaults to None, if refpts is None an empty array is returned for orbit. OUTPUT orbit0 ((6,) closed orbit vector at the entrance of the 1-st element (x,px,y,py) orbit (6, Nrefs) closed orbit vector at each location specified in refpts KEYWORDS keep_lattice Assume no lattice change since the previous tracking. Default: False guess Initial value for the closed orbit. It may help convergence. Default: (0, 0, 0, 0) convergence Convergence criterion. Default: 1.e-12 max_iterations Maximum number of iterations. Default: 20 step_size Step size. Default: 1.e-6 See also find_orbit4, find_sync_orbit. """ keep_lattice = kwargs.pop('keep_lattice', False) convergence = kwargs.pop('convergence', CONVERGENCE) max_iterations = kwargs.pop('max_iterations', MAX_ITERATIONS) step_size = kwargs.pop('step_size', STEP_SIZE) ref_in = numpy.zeros((6, ), order='F') if guess is None else guess # Get evolution period l0 = get_s_pos(ring, len(ring)) cavities = [elem for elem in ring if isinstance(elem, elements.RFCavity)] if len(cavities) == 0: raise AtError('No cavity found in the lattice.') f_rf = cavities[0].Frequency harm_number = cavities[0].HarmNumber theta = numpy.zeros((6, )) theta[5] = constants.speed_of_light * harm_number / f_rf - l0 delta_matrix = numpy.zeros((6, 7), order='F') for i in range(6): delta_matrix[i, i] = step_size id6 = numpy.asfortranarray(numpy.identity(6)) change = 1 itercount = 0 while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = lattice_pass(ring, in_mat, refpts=[], keep_lattice=keep_lattice) # the reference particle after one turn ref_out = in_mat[:, 6] # 6x6 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j6 = (in_mat[:, :6] - in_mat[:, 6:]) / step_size a = j6 - id6 # f'(r_n) - 1 b = ref_out[:] - ref_in[:] - theta b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) r_next = ref_in - b_over_a # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keep_lattice = True if itercount == max_iterations: warnings.warn( AtWarning('Maximum number of iterations reached. ' 'Possible non-convergence')) uint32refs = uint32_refpts(refpts, len(ring)) # bug in numpy < 1.13 all_points = numpy.empty( (0, 6), dtype=float) if len(uint32refs) == 0 else numpy.squeeze( lattice_pass(ring, ref_in.copy(order='K'), refpts=uint32refs, keep_lattice=keep_lattice), axis=(1, 3)).T return ref_in, all_points