def compute(self, ring, *args, **kwargs): """Optics computation before evaluation of all constraints""" ld0, tune, chrom, ld = linopt(ring, refpts=self.refpts, get_chrom=self.get_chrom, **kwargs) return (ld[ref[self.refpts]] for ref in self.refs), (tune, chrom)
def test_linear_analysis(engine, lattices, dp): py_lattice, ml_lattice, _ = lattices nelems = len(py_lattice) fields = [('SPos', 's_pos'), ('ClosedOrbit', 'closed_orbit'), ('Dispersion', 'dispersion'), ('alpha', 'alpha'), ('beta', 'beta'), ('mu', 'mu'), ('M44', 'm44'), ('A', 'A'), ('B', 'B'), ('C', 'C'), ('gamma', 'gamma')] refpts = range(nelems + 1) py_data0, py_tune, py_chrom, py_data = physics.linopt(py_lattice, dp, refpts, True, ddp=1.E-6) # Matlab call ml_data, ml_tune, ml_chrom = engine.pyproxy('atlinopt', ml_lattice, dp, _ml_refs(refpts, nelems), nargout=3) ml_data0 = engine.pyproxy('atlinopt', ml_lattice, dp, _ml_refs(nelems, nelems), nargout=3)[0] # Comparison assert_close(py_tune, _py_data(ml_tune), atol=2.e-6, rtol=0) #1e-8 assert_close(py_chrom, _py_data(ml_chrom), atol=2.e-3, rtol=0) _compare_physdata(numpy.expand_dims(py_data0, 0), ml_data0, fields, rtol=1e-8) _compare_physdata(py_data, ml_data, fields, rtol=1e-8)
def test_linopt(dba_lattice, refpts): """Compare with Matlab results""" lindata0, tune, chrom, lindata = physics.linopt(dba_lattice, DP2, refpts, get_chrom=True) obs = lindata[-1] assert_close(tune, [0.355804633927360, 0.488487169156598], rtol=1e-8) assert_close(chrom, [-3.428312247995742, -1.597924047969101], rtol=2e-4) assert_close(obs['s_pos'], 56.209377216, atol=1e-9) assert_close(obs['closed_orbit'][:5], [0.000426438389644, -0.000000000287482, 0, 0, DP2], atol=1e-12) assert_close(obs['dispersion'], [0.156822576442091, -0.000000162902610, 0, 0], rtol=1e-7, atol=2e-10) expected = [[-0.616893565445970, 2.191800192047084, 0, 0], [-0.282617257709865, -0.616893094967279, 0, 0], [0, 0, -0.997384485665288, 0.466909288129080], [0, 0, -0.011187515624062, -0.997384713772755]] assert_close(obs['m44'], expected, rtol=1e-7, atol=1e-12) assert_close(obs['alpha'], [-2.988886505944512e-07, 1.578070569086581e-06], rtol=1e-8, atol=1e-8) assert_close(obs['beta'], [2.784841119739221, 6.460251554763623], rtol=5e-8, atol=1e-12) assert_close(obs['mu'], [2.235586452367286, 3.069255403991328], rtol=1e-8, atol=1e-12) assert_close(obs['gamma'], 1.0, rtol=1e-6, atol=1e-20) assert_close(obs['A'], [[-0.616892545020345, 2.191796566512299], [-0.282616790222590, -0.616892074542433]], rtol=1e-7, atol=1e-12) assert_close(obs['B'], [[-0.997384485665288, 0.466909288129080], [-0.011187515624062, -0.997384713772754]], rtol=1e-7, atol=1e-12) assert_close(obs['C'], [[0, 0], [0, 0]], rtol=1e-5, atol=1e-7)
def test_linopt_uncoupled(dba_lattice, refpts): """Compare with Matlab results""" lindata0, tune, chrom, lindata = physics.linopt(dba_lattice, DP2, refpts, coupled=False) obs = lindata[-1] assert_close(tune, [0.355804634603528, 0.488487169156732], rtol=1e-8) assert_close(obs['s_pos'], 56.209377216, atol=1e-9) assert_close(obs['closed_orbit'][:5], [0.000426438389644, -0.000000000287482, 0, 0, DP2], atol=1e-12) expected = [[-0.616893565445970, 2.191800192047084, 0, 0], [-0.282617257709865, -0.616893094967279, 0, 0], [0, 0, -0.997384485665288, 0.466909288129080], [0, 0, -0.011187515624062, -0.997384713772755]] assert_close(obs['m44'], expected, rtol=1e-7, atol=1e-12) assert_close(obs['alpha'], [-2.988885294797512e-07, 1.578069929495847e-06], rtol=1e-8, atol=1e-8) assert_close(obs['beta'], [2.784839991078270, 6.460248936505217], rtol=5e-8, atol=1e-12) assert_close(obs['mu'], [2.235586452367286, 3.069255403991328], rtol=1e-8, atol=1e-12)
def test_linopt_uncoupled(dba_lattice, refpts): lindata0, tune, chrom, lindata = physics.linopt(dba_lattice, DP, refpts, coupled=False) numpy.testing.assert_allclose(tune, [0.365529, 0.493713], rtol=1e-5) numpy.testing.assert_allclose(lindata['s_pos'][-1], 56.209377216, atol=1e-9) numpy.testing.assert_allclose( lindata['closed_orbit'][-1][:5], [1.091636e-7, 1.276757e-15, 4.238871e-33, 1.117703e-33, DP], atol=1e-12) expected_m44 = [[-0.663802, 2.234145, 0, 0], [-0.250372, -0.663802, 0, 0], [-1.456977e-31, -1.150075e-30, -0.99922, 0.262171], [6.577482e-33, 8.75482e-32, -5.949696e-3, -0.99922]] numpy.testing.assert_allclose(lindata['m44'][-1], expected_m44, rtol=1e-5, atol=1e-7) numpy.testing.assert_allclose(lindata['alpha'][-1], [-1.32787e-7, 1.85909e-7], rtol=1e-5, atol=1e-7) numpy.testing.assert_almost_equal(lindata['beta'][-1], [2.98719, 6.638115], decimal=4) numpy.testing.assert_almost_equal(lindata['mu'][-1], [2.296687, 3.102088], decimal=4)
def test_linopt(dba_lattice, refpts): lindata0, tune, chrom, lindata = physics.linopt(dba_lattice, DP, refpts, get_chrom=True) numpy.testing.assert_allclose(tune, [0.365529, 0.493713], rtol=1e-5) numpy.testing.assert_allclose(chrom, [-0.309037, -0.441859], rtol=1e-5) numpy.testing.assert_allclose(lindata['s_pos'][-1], 56.209377216, atol=1e-9) numpy.testing.assert_allclose( lindata['closed_orbit'][-1][:5], [1.091636e-7, 1.276757e-15, 4.238871e-33, 1.117703e-33, DP], atol=1e-12) numpy.testing.assert_allclose( lindata['dispersion'][-1], [1.107402e-2, 1.262031e-10, -2.139355e-25, 3.757804e-25], rtol=1e-5, atol=1e-7) expected = [[-0.663802, 2.234145, 0, 0], [-0.250372, -0.663802, 0, 0], [-1.456977e-31, -1.150075e-30, -0.99922, 0.262171], [6.577482e-33, 8.75482e-32, -5.949696e-3, -0.99922]] numpy.testing.assert_allclose(lindata['m44'][-1], expected, rtol=1e-5, atol=1e-7) numpy.testing.assert_allclose(lindata['alpha'][-1], [-1.32787e-7, 1.85909e-7], rtol=1e-5, atol=1e-7) numpy.testing.assert_almost_equal(lindata['beta'][-1], [2.98719, 6.638115], decimal=4) numpy.testing.assert_almost_equal(lindata['mu'][-1], [2.296687, 3.102088], decimal=4) numpy.testing.assert_almost_equal(lindata['gamma'][-1], 1) numpy.testing.assert_allclose(lindata['A'][-1], [[-0.6638, 2.23415], [-0.25037, -0.6638]], rtol=1e-5, atol=1e-7) numpy.testing.assert_allclose(lindata['B'][-1], [[-0.99922, 0.262171], [-0.00595, -0.99922]], rtol=1e-4, atol=1e-7) numpy.testing.assert_allclose( lindata['C'][-1], [[-9.87933e-32, -1.65044e-30], [-2.44501e-32, -2.91703e-31]], rtol=1e-5, atol=1e-7)
def tunes_vs_amp(ring, amp=None, dim=0, dp=0, window=1, nturns=512): """ Generates a range of tunes for varying x, y, or z amplitudes """ def _gen_part(ring, amp, dim, orbit, ld, nturns): 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]) ]]) part = numpy.array([ orbit, ] * len(amp)).T + 1.0e-6 part[dim, :] += amp part = lattice_pass(ring, part, nturns=nturns) sh = part.shape partx = numpy.reshape(part[0, :], (sh[1], sh[3])) partxp = numpy.reshape(part[1, :], (sh[1], sh[3])) party = numpy.reshape(part[2, :], (sh[1], sh[3])) partyp = numpy.reshape(part[3, :], (sh[1], sh[3])) px = numpy.array([ numpy.matmul(invx, [partx[i], partxp[i]]) for i in range(len(amp)) ]) py = numpy.array([ numpy.matmul(invy, [party[i], partyp[i]]) for i in range(len(amp)) ]) return px[:, 0, :] - 1j * px[:, 1, :], py[:, 0, :] - 1j * py[:, 1, :] l0, q0, _, _ = linopt(ring, dp=dp) orbit = l0['closed_orbit'] tunes = [] if amp is not None: partx, party = _gen_part(ring, amp, dim, orbit, l0, nturns) qx = get_tunes_harmonic(partx, 'laskar') qy = get_tunes_harmonic(party, 'laskar') tunes = numpy.vstack((qx, qy)).T return numpy.array(tunes)
def pldata_beta_disp(ring, refpts, **kwargs): """Generates data for plotting beta functions and dispersion""" # compute linear optics at the required locations data0, _, _, data = linopt(ring, refpts=refpts, get_chrom=True, **kwargs) # Extract the plot data s_pos = data['s_pos'] betax = data['beta'][:, 0] betaz = data['beta'][:, 1] dispersion = data['dispersion'][:, 0] # Left axis definition left = (r'$\beta$ [m]', s_pos, [betax, betaz], [r'$\beta_x$', r'$\beta_z$']) # Right axis definition right = ('dispersion [m]', s_pos, [dispersion], ['dispersion']) return 'Optical functions', left, right
def test_linear_analysis(engine, ml_lattice, py_lattice, dp, refpts, func_data): """N.B. a 'mu' comparison is left out for twiss data as the values for 'mu' returned by 'twissring' in Matlab are inconsistent with those from 'get_twiss' and 'linopt' in Python as well as those returned from 'atlinopt' in Matlab. """ nelems = len(py_lattice) refpts = range(nelems + 1) if refpts is None else refpts # Python call if func_data[0] == 'twissring': py_data0, py_tune, py_chrom, py_data = physics.get_twiss(py_lattice, dp, refpts, True, ddp=1.E-6) else: py_data0, py_tune, py_chrom, py_data = physics.linopt(py_lattice, dp, refpts, True, ddp=1.E-6) # Matlab call ml_data, ml_tune, ml_chrom = engine.pyproxy(func_data[0], ml_lattice, dp, _ml_refs(refpts, nelems), nargout=3) ml_data0 = engine.pyproxy(func_data[0], ml_lattice, dp, _ml_refs(nelems, nelems), nargout=3)[0] # Comparison numpy.testing.assert_almost_equal(py_tune, _py_data(ml_tune), decimal=6) numpy.testing.assert_almost_equal(py_chrom, _py_data(ml_chrom), decimal=4) _compare_physdata(numpy.expand_dims(py_data0, 0), ml_data0, func_data[1], decimal=5) _compare_physdata(py_data, ml_data, func_data[1], decimal=6)
def detuning(ring, xm=0.3e-4, ym=0.3e-4, npoints=3, dp=0, window=1, nturns=512): """ This function uses tunes_vs_amp to compute the tunes for the specified amplitudes. Then it fits this data and returns result for dQx/dx, dQy/dx, dQx/dy, dQy/dy """ lindata0, _, _, _ = linopt(ring, dp=dp) gamma = (1 + lindata0.alpha * lindata0.alpha) / lindata0.beta x = numpy.linspace(-xm, xm, npoints) y = numpy.linspace(-ym, ym, npoints) x2 = x * x y2 = y * y q_dx = tunes_vs_amp(ring, amp=x, dim=0, dp=dp, window=window, nturns=nturns) q_dy = tunes_vs_amp(ring, amp=y, dim=2, dp=dp, window=window, nturns=nturns) fx = numpy.polyfit(x2, q_dx, 1) fy = numpy.polyfit(y2, q_dy, 1) q0 = [[fx[1, 0], fx[1, 1]], [fy[1, 0], fy[1, 1]]] q1 = [[2 * fx[0, 0] / gamma[0], 2 * fx[0, 1] / gamma[0]], [2 * fy[0, 0] / gamma[1], 2 * fy[0, 1] / gamma[1]]] return numpy.array(q0), numpy.array(q1), x, q_dx, y, q_dy
def pldata_linear(ring, refpts, *keys, **kwargs): """data extraction function for plotting results of linopt""" class Lind(object): """helper class for lindata extraction""" lab2 = "xz" lab6 = ("x", "x'", "z", "z'", "l", "\\delta") id6 = "123456" params = dict(beta=(r'$\beta$ [m]', r'$\beta_{0}$', lab2), closed_orbit=('position [m]', r'${0}$', lab6), dispersion=('dispersion [m]', r'$\eta_{0}$', lab6), alpha=(r'$\alpha$', r'$\alpha_{0}$', lab2), mu=(r'Phase advance', r'$\mu_{0}$', lab2), gamma=('Gamma', 'Gamma', None), A=('A', r'$A_{{{0}{1}}}$', id6), B=('B', r'$B_{{{0}{1}}}$', id6), C=('C', r'$C_{{{0}{1}}}$', id6), m44=('Transfert', r'$T_{{{0},{1}}}$', id6)) @classmethod def extract(cls, lindata, key, *idx): def it(v): try: return iter(v) except TypeError: return iter([v]) axis_title, fmt, convert = cls.params[key] indices = list(zip(*(it(i) for i in idx))) print(indices) datay = (lindata[key][(slice(None), ) + ic] for ic in indices) labels = (fmt.format(*(convert[i] for i in ic)) for ic in indices) return axis_title, lindata['s_pos'], datay, labels title = kwargs.pop('title', 'Linear optics') data0, _, _, data = linopt(ring, refpts=refpts, get_chrom=True, **kwargs) return (title, ) + tuple(Lind.extract(data, *key) for key in keys)
def get_radiation_integrals(ring, dp=0.0, twiss=None): """ Compute the 5 radiation integrals for uncoupled lattices. No RF cavity or radiating element is allowed. PARAMETERS ring lattice description. dp=0.0 momentum deviation KEYWORDS twiss=None linear optics at all points (from linopt). If None, it will be computed. OUTPUT i1, i2, i3, i4, i5 """ def dipole_radiation(elem, vini, vend): """Analytically compute the radiation integrals in dipoles""" beta0 = vini.beta[0] alpha0 = vini.alpha[0] eta0 = vini.dispersion[0] etap0 = vini.dispersion[1] ll = elem.Length theta = elem.BendingAngle rho = ll / theta rho2 = rho * rho k2 = elem.K + 1.0 / rho2 eps1 = tan(elem.EntranceAngle) / rho eps2 = tan(elem.ExitAngle) / rho eta3 = vend.dispersion[0] alpha1 = alpha0 - beta0 * eps1 gamma1 = (1.0 + alpha1 * alpha1) / beta0 etap1 = etap0 + eta0 * eps1 etap2 = vend.dispersion[1] - eta3 * eps2 h0 = gamma1*eta0*eta0 + 2.0*alpha1*eta0*etap1 + beta0*etap1*etap1 if k2 != 0.0: if k2 > 0.0: # Focusing kl = ll * sqrt(k2) ss = sin(kl) / kl cc = cos(kl) else: # Defocusing kl = ll * sqrt(-k2) ss = sinh(kl) / kl cc = cosh(kl) eta_ave = (theta - (etap2 - etap1)) / k2 / ll bb = 2.0 * (alpha1 * eta0 + beta0 * etap1) * rho aa = -2.0 * (alpha1 * etap1 + gamma1 * eta0) * rho h_ave = h0 + (aa * (1.0 - ss) + bb * (1.0 - cc) / ll + gamma1 * (3.0 - 4.0 * ss + ss * cc) / 2.0 / k2 - alpha1 * (1.0 - cc) ** 2 / k2 / ll + beta0 * (1.0 - ss * cc) / 2.0 ) / k2 / rho2 else: eta_ave = 0.5 * (eta0 + eta3) - ll * ll / 12.0 / rho hp0 = 2.0 * (alpha1 * eta0 + beta0 * etap1) / rho h2p0 = 2.0 * (-alpha1 * etap1 + beta0 / rho - gamma1 * eta0) / rho h_ave = h0 + hp0 * ll / 2.0 + h2p0 * ll * ll / 6.0 \ - alpha1 * ll ** 3 / 4.0 / rho2 \ + gamma1 * ll ** 4 / 20.0 / rho2 di1 = eta_ave * ll / rho di2 = ll / rho2 di3 = ll / abs(rho) / rho2 di4 = eta_ave * ll * (2.0 * elem.K + 1.0 / rho2) / rho \ - (eta0 * eps1 + eta3 * eps2) / rho di5 = h_ave * ll / abs(rho) / rho2 return numpy.array([di1, di2, di3, di4, di5]) def wiggler_radiation(elem, dini): """Compute the radiation integrals in wigglers with the following approximations: - The wiggler is aligned with the closed orbit - The self-induced dispersion is neglected in I4 and I5, but is is used as a lower limit for the I5 contribution - I1, I2 are integrated analytically - I3 is integrated analytically for a single harmonic, numerically otherwise """ def b_on_axis(wig, s): """On-axis wiggler field""" def harm(coef, h, phi): return -Bmax * coef * numpy.cos(h*kws + phi) kw = 2 * pi / wig.Lw Bmax = wig.Bmax kws = kw * s zz = [numpy.zeros(kws.shape)] vh = zz + [harm(pb[1], pb[4], pb[5]) for pb in wig.By.T] vv = zz + [-harm(pb[1], pb[4], pb[5]) for pb in wig.Bx.T] bys = numpy.sum(numpy.stack(vh), axis=0) bxs = numpy.sum(numpy.stack(vv), axis=0) return bxs, bys le = elem.Length alphax0 = dini.alpha[0] betax0 = dini.beta[0] gammax0 = (alphax0 * alphax0 + 1) / betax0 eta0 = dini.dispersion[0] etap0 = dini.dispersion[1] H0 = gammax0*eta0*eta0 + 2*alphax0*eta0*etap0 + betax0*etap0*etap0 avebetax = betax0 + alphax0*le + gammax0*le*le/3 kw = 2 * pi / elem.Lw rhoinv = elem.Bmax / Brho coefh = elem.By[1, :] * rhoinv coefv = elem.Bx[1, :] * rhoinv coef2 = numpy.concatenate((coefh, coefv)) if len(coef2) == 1: di3 = le * coef2[0] ** 3 * 4 / 3 / pi else: bx, bz = b_on_axis(elem, numpy.linspace(0, elem.Lw, _NSTEP + 1)) rinv = numpy.sqrt(bx*bx + bz*bz) / Brho di3 = numpy.trapz(rinv ** 3) * le / _NSTEP di2 = le * (numpy.sum(coefh * coefh) + numpy.sum(coefv * coefv)) / 2 di1 = -di2 / kw / kw di4 = 0 if len(coefh) > 0: d5lim = 4 * avebetax * le * coefh[0] ** 5 / 15 / pi / kw / kw else: d5lim = 0 di5 = max(H0 * di3, d5lim) return numpy.array([di1, di2, di3, di4, di5]) Brho = sqrt(ring.energy**2 - e_mass**2) / clight integrals = numpy.zeros((5,)) if twiss is None: _, _, _, twiss = linopt(ring, dp, range(len(ring) + 1), get_chrom=True, coupled=False) elif len(twiss) != len(ring) + 1: raise ValueError('length of Twiss data should be {0}' .format(len(ring) + 1)) for (el, vini, vend) in zip(ring, twiss[:-1], twiss[1:]): if isinstance(el, elements.Dipole) and el.BendingAngle != 0.0: integrals += dipole_radiation(el, vini, vend) elif isinstance(el, elements.Wiggler) and el.PassMethod != 'DriftPass': integrals += wiggler_radiation(el, vini) return tuple(integrals)
def test_linopt_no_refpts(dba_lattice): lindata0, tune, chrom, lindata = physics.linopt(dba_lattice, DP, get_chrom=True) assert list(lindata) == [] assert len(physics.linopt(dba_lattice, DP, get_chrom=True)) == 4
def _get_chrom(ring, dp=0): _, _, chrom, _ = linopt(ring, dp=dp, get_chrom=True) return chrom
def _get_tune(ring, dp=0): _, tune, _, _ = linopt(ring, dp=dp) return tune
def get_radiation_integrals(ring, dp=0.0, twiss=None): """ Compute the 5 radiation integrals for uncoupled lattices. No RF cavity or radiating element is allowed. PARAMETERS ring lattice description. dp=0.0 momentum deviation KEYWORDS twiss=None linear optics at all points (from linopt). If None, it will be computed. OUTPUT i1, i2, i3, i4, i5 """ i1 = 0.0 i2 = 0.0 i3 = 0.0 i4 = 0.0 i5 = 0.0 if twiss is None: _, _, _, twiss = linopt(ring, dp, range(len(ring) + 1), get_chrom=True, coupled=False) elif len(twiss) != len(ring) + 1: raise ValueError( 'length of Twiss data should be {0}'.format(len(ring) + 1)) for (elem, vini, vend) in zip(ring, twiss[:-1], twiss[1:]): if isinstance(elem, elements.Dipole) and elem.BendingAngle != 0.0: beta0 = vini.beta[0] alpha0 = vini.alpha[0] gamma0 = (1.0 + alpha0 * alpha0) / beta0 eta0 = vini.dispersion[0] etap0 = vini.dispersion[1] ll = elem.Length theta = elem.BendingAngle rho = ll / theta rho2 = rho * rho k2 = elem.K + 1.0 / rho2 eps1 = tan(elem.EntranceAngle) / rho eps2 = tan(elem.ExitAngle) / rho eta3 = vend.dispersion[0] alpha1 = alpha0 - beta0 * eps1 etap1 = etap0 + eta0 * eps1 etap2 = vend.dispersion[1] - eta3 * eps2 h0 = gamma0 * eta0 * eta0 + 2.0 * alpha1 * eta0 * etap1 + beta0 * etap1 * etap1 if k2 != 0.0: if k2 > 0.0: # Focusing kl = ll * sqrt(k2) ss = sin(kl) / kl cc = cos(kl) else: # Defocusing kl = ll * sqrt(-k2) ss = sinh(kl) / kl cc = cosh(kl) eta_ave = (theta - (etap2 - etap1)) / k2 / ll bb = 2.0 * (alpha1 * eta0 + beta0 * etap1) * rho aa = -2.0 * (alpha1 * etap1 + gamma0 * eta0) * rho h_ave = h0 + (aa * (1.0 - ss) + bb * (1.0 - cc) / ll + gamma0 * (3.0 - 4.0 * ss + ss * cc) / 2.0 / k2 - alpha1 * (1.0 - cc)**2 / k2 / ll + beta0 * (1.0 - ss * cc) / 2.0) / k2 / rho2 else: eta_ave = 0.5 * (eta0 + eta3) - ll * ll / 12.0 / rho hp0 = 2.0 * (alpha1 * eta0 + beta0 * etap1) / rho h2p0 = 2.0 * (-alpha1 * etap1 + beta0 / rho - gamma0 * eta0) / rho h_ave = h0 + hp0*ll/2.0 + h2p0*ll*ll/6.0 \ - alpha1*ll**3/4.0/rho2 \ + gamma0*ll**4/20.0/rho2 i1 += eta_ave * ll / rho i2 += ll / rho2 i3 += ll / abs(rho) / rho2 i4 += eta_ave * ll * (2.0*elem.K+1.0/rho2) / rho \ - (eta0*eps1 + eta3*eps2)/rho i5 += h_ave * ll / abs(rho) / rho2 return i1, i2, i3, i4, i5