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_orbit6_produces_same_result_with_keep_lattice_True(hmba_lattice): hmba_lattice = hmba_lattice.radiation_on(quadrupole_pass=None, copy=True) orbit0, _ = physics.find_orbit6(hmba_lattice) # Technicality - the default arguments to find_orbit6 mean that # keep_lattice argument is always false. orbit1, _ = physics.find_orbit6(hmba_lattice, keep_lattice=True) # With INTEGRAL keep_lattice does take effect. orbit2, _ = physics.find_orbit6( hmba_lattice, keep_lattice=True, method=physics.ELossMethod.INTEGRAL ) assert_close(orbit0, orbit1, rtol=0, atol=1e-12) assert_close(orbit0, orbit2, rtol=0, atol=1e-12)
def tapering(ring, multipoles=True, niter=1, **kwargs): """ Scales magnet strength with local energy to cancel the closed orbit and optics errors due to synchrotron radiations. PolynomB is used for dipoles such that the machine geometry is maintained. This is the ideal tapering scheme where magnets and multipoles components (PolynomB and PolynomA) are scaled individually. !!! WARNING: This method works only for lattices without errors and corrections: if not all corrections and field errors will also be scaled !!! tapering(ring) or ring.tapering() PARAMETERS ring lattice description. KEYWORDS multipoles=True scale all multipoles niter=1 number of iteration XYStep=1.0e-8 transverse step for numerical computation DPStep=1.0E-6 momentum deviation used for computation of orbit6 """ xy_step = kwargs.pop('XYStep', DConstant.XYStep) dp_step = kwargs.pop('DPStep', DConstant.DPStep) dipoles = get_refpts(ring, Dipole) b0 = get_value_refpts(ring, dipoles, 'BendingAngle') k0 = get_value_refpts(ring, dipoles, 'PolynomB', index=0) ld = get_value_refpts(ring, dipoles, 'Length') for i in range(niter): _, o6 = find_orbit6(ring, refpts=range(len(ring) + 1), XYStep=xy_step, DPStep=dp_step) dpps = (o6[dipoles, 4] + o6[dipoles + 1, 4]) / 2 set_value_refpts(ring, dipoles, 'PolynomB', b0 / ld * dpps + k0 * (1 + dpps), index=0) if multipoles: mults = get_refpts(ring, Multipole) k0 = get_value_refpts(ring, dipoles, 'PolynomB', index=0) _, o6 = find_orbit6(ring, refpts=range(len(ring) + 1), XYStep=xy_step, DPStep=dp_step) dpps = (o6[mults, 4] + o6[mults + 1, 4]) / 2 for dpp, el in zip(dpps, ring[mults]): el.PolynomB *= 1 + dpp el.PolynomA *= 1 + dpp set_value_refpts(ring, dipoles, 'PolynomB', k0, index=0)
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 _dmatr(ring, orbit=None, keep_lattice=False): """ compute the cumulative diffusion and orbit matrices over the ring """ nelems = len(ring) energy = ring.energy allrefs = uint32_refpts(range(nelems + 1), nelems) if orbit is None: orbit, _ = find_orbit6(ring, keep_lattice=keep_lattice) keep_lattice = True orbs = numpy.squeeze(lattice_pass(ring, orbit.copy(order='K'), refpts=allrefs, keep_lattice=keep_lattice), axis=(1, 3)).T b0 = numpy.zeros((6, 6)) bb = [ find_mpole_raddiff_matrix(elem, elemorb, energy) if elem.PassMethod.endswith('RadPass') else b0 for elem, elemorb in zip(ring, orbs) ] bbcum = numpy.stack(list(_cumulb(zip(ring, orbs, bb))), axis=0) return bbcum, orbs
def test_find_orbit6_raises_AtError_if_there_is_no_cavity(dba_lattice): with pytest.raises(at.lattice.utils.AtError): physics.find_orbit6(dba_lattice)
def test_find_orbit6(hmba_lattice): hmba_lattice = hmba_lattice.radiation_on(copy=True) refpts = numpy.ones(len(hmba_lattice), dtype=bool) orbit6, all_points = physics.find_orbit6(hmba_lattice, refpts) assert_close(orbit6, orbit6_MATLAB, rtol=0, atol=1e-12)
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=6.055e-6 transverse step for numerical computation DPStep=6.055e-6 longitudinal step for numerical computation See also find_m44, find_orbit6 """ xy_step = kwargs.pop('XYStep', XYDEFSTEP) dp_step = kwargs.pop('DPStep', DPSTEP) if orbit is None: orbit, _ = find_orbit6(ring, keep_lattice=keep_lattice) keep_lattice = True # Construct matrix of plus and minus deltas scaling = numpy.array( [xy_step, xy_step, xy_step, xy_step, dp_step, dp_step]) 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.reshape((1, 6)) if len(refs) > 0: mstack = (out_mat[:, :, :6] - out_mat[:, :, 6:]) / xy_step else: mstack = numpy.empty((0, 6, 6), dtype=float) return m66, mstack
def test_find_orbit6(hmba_lattice): expected = numpy.zeros((len(hmba_lattice), 6)) refpts = numpy.ones(len(hmba_lattice), dtype=bool) _, all_points = physics.find_orbit6(hmba_lattice, refpts) numpy.testing.assert_allclose(all_points, expected, atol=1e-12)
def ohmi_envelope(ring, refpts=None, orbit=None, keep_lattice=False): """ Calculate the equilibrium beam envelope in a circular accelerator using Ohmi's beam envelope formalism [1] emit0, beamdata, emit = ohmi_envelope(ring[, refpts]) PARAMETERS ring Lattice object. 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=None Avoids looking for the closed orbit if it is already known (6,) array) keep_lattice=False Assume no lattice change since the previous tracking OUTPUT emit0 emittance data at the start/end of the ring beamdata beam parameters at the start of the ring emit emittance data at the points refered to by refpts, if refpts is None an empty structure is returned. emit is a record array with fields: r66 (6, 6) equilibrium envelope matrix R r44 (4, 4) betatron emittance matrix (dpp = 0) m66 (6, 6) transfer matrix from the start of the ring orbit6 (6,) closed orbit emitXY (2,) betatron emittance projected on xxp and yyp emitXYZ (3,) 6x6 emittance projected on xxp, yyp, ldp beamdata is a record array with fields: tunes tunes of the 3 normal modes damping_rates damping rates of the 3 normal modes mode_matrices R-matrices of the 3 normal modes mode_emittances equilibrium emittances of the 3 normal modes Field values can be obtained with either emit['r66'] or emit.r66 REFERENCES [1] K.Ohmi et al. Phys.Rev.E. Vol.49. (1994) """ def cumulb(it): """accumulate diffusion matrices""" cumul = numpy.zeros((6, 6)) yield cumul for el, orbin, b in it: m = find_elem_m66(el, orbin) cumul = m.dot(cumul).dot(m.T) + b yield cumul def process(r66): # projections on xx', zz', ldp emit3sq = numpy.array([det(r66[s, s]) for s in _submat]) # Prevent from unrealistic negative values of the determinant emit3 = numpy.sqrt(numpy.maximum(emit3sq, 0.0)) # Emittance cut for dpp=0 if emit3[0] < 1.E-13: # No equilibrium emittance r44 = numpy.nan * numpy.ones((4, 4)) elif emit3[1] < 1.E-13: # Uncoupled machine minv = inv(r66[[0, 1, 4, 5], :][:, [0, 1, 4, 5]]) r44 = numpy.zeros((4, 4)) r44[:2, :2] = inv(minv[:2, :2]) else: # Coupled machine minv = inv(r66) r44 = inv(minv[:4, :4]) # betatron emittances (dpp=0) emit2sq = numpy.array( [det(r44[s, s], check_finite=False) for s in _submat[:2]]) # Prevent from unrealistic negative values of the determinant emit2 = numpy.sqrt(numpy.maximum(emit2sq, 0.0)) return r44, emit2, emit3 def propag(m, cumb, orbit6): """Propagate the beam matrix to refpts""" sigmatrix = m.dot(rr).dot(m.T) + cumb m44, emit2, emit3 = process(sigmatrix) return sigmatrix, m44, m, orbit6, emit2, emit3 nelems = len(ring) uint32refs = uint32_refpts(refpts, nelems) allrefs = uint32_refpts(range(nelems + 1), nelems) energy = ring.energy if orbit is None: orbit, _ = find_orbit6(ring, keep_lattice=keep_lattice) keep_lattice = True orbs = numpy.squeeze(lattice_pass(ring, orbit.copy(order='K'), refpts=allrefs, keep_lattice=keep_lattice), axis=(1, 3)).T mring, ms = find_m66(ring, uint32refs, orbit=orbit, keep_lattice=True) b0 = numpy.zeros((6, 6)) bb = [ find_mpole_raddiff_matrix(elem, orbit, energy) if elem.PassMethod.endswith('RadPass') else b0 for elem in ring ] bbcum = numpy.stack(list(cumulb(zip(ring, orbs, bb))), axis=0) # ------------------------------------------------------------------------ # Equation for the moment matrix R is # R = MRING*R*MRING' + BCUM; # We rewrite it in the form of Lyapunov-Sylvester equation to use scipy's # solve_sylvester function # A*R + R*B = Q # where # A = inv(MRING) # B = -MRING' # Q = inv(MRING)*BCUM # ------------------------------------------------------------------------ aa = inv(mring) bb = -mring.T qq = numpy.dot(aa, bbcum[-1]) rr = solve_sylvester(aa, bb, qq) rr = 0.5 * (rr + rr.T) rr4, emitxy, emitxyz = process(rr) r66data = get_tunes_damp(mring, rr) data0 = numpy.rec.fromarrays((rr, rr4, mring, orbit, emitxy, emitxyz), dtype=ENVELOPE_DTYPE) if uint32refs.shape == (0, ): data = numpy.recarray((0, ), dtype=ENVELOPE_DTYPE) else: data = numpy.rec.fromrecords(list( map(propag, ms, bbcum[uint32refs], orbs[uint32refs, :])), dtype=ENVELOPE_DTYPE) return data0, r66data, data