def test_rfcavity(rin): rf = elements.RFCavity('rfcavity', 0.0, 187500, 3.5237e+8, 31, 6.e+9) lattice = [rf, rf, rf, rf] rin[4, 0] = 1e-6 rin[5, 0] = 1e-6 lattice_pass(lattice, rin, 1) expected = numpy.array([0., 0., 0., 0., 9.990769e-7, 1.e-6]).reshape(6, 1) numpy.testing.assert_allclose(rin, expected, atol=1e-12)
def test_aperture_outside_limits(rin): a = elements.Aperture('aperture', [-1e-3, 1e-3, -1e-4, 1e-4]) assert a.Length == 0 lattice = [a] rin[0, 0] = 1e-2 rin[2, 0] = -1e-2 lattice_pass(lattice, rin, 1) assert numpy.isnan(rin[0, 0]) assert rin[2, 0] == 0.0 # Only the 1st coordinate is nan, the rest is zero
def test_pyintegrator(hmba_lattice): params = { 'Length': 0, 'PassMethod': 'pyIdentityPass', } id_elem = Element('py_id', **params) pin = numpy.zeros(6) + 1.0e-6 pout1 = lattice_pass(hmba_lattice, pin.copy(), nturns=1) pin = numpy.zeros(6) + 1.0e-6 hmba_lattice = hmba_lattice + [id_elem] pout2 = lattice_pass(hmba_lattice, pin, nturns=1) numpy.testing.assert_equal(pout1, pout2)
def test_atpass(engine, lattices): py_lattice, ml_lattice, _ = lattices xy_step = 1.e-8 scaling = [xy_step, xy_step, xy_step, xy_step, xy_step, xy_step] ml_rin = engine.diag(matlab.double(scaling)) ml_rout = engine.atpass(ml_lattice, ml_rin, 1, 1) py_rin = numpy.asfortranarray(numpy.diag(scaling)) at.lattice_pass(py_lattice, py_rin, 1) assert_close(py_rin, ml_rout, rtol=0, atol=1.e-30)
def test_lattice_convert_to_list_if_incorrect_type(): lattice = numpy.array([elements.Drift('Drift', 1.0)]) rin = numpy.zeros((6, 2)) rin[0, 0] = 1e-6 r_original = numpy.copy(rin) r_out = lattice_pass(lattice, rin, 1) numpy.testing.assert_equal(r_original, r_out.reshape(6, 2))
def find_m44(ring, dp=0.0, refpts=None, orbit4=None, XYStep=XYDEFSTEP): """ Determine the transfer matrix for ring, by first finding the closed orbit. """ refpts = lattice.uint32_refpts(refpts, len(ring)) nrefs = refpts.size if refpts[-1] != len(ring): refpts = numpy.append(refpts, [len(ring)]) keeplattice = False if orbit4 is None: orbit4, _ = find_orbit4(ring, dp) keeplattice = True # Append zeros to closed 4-orbit. orbit6 = numpy.append(orbit4, (dp, 0.0)).reshape(6, 1) # Construct matrix of plus and minus deltas dmat = numpy.zeros((6, 8), order='F') for i in range(4): dmat[i, i] = 0.5 * XYStep dmat[i, i + 4] = -0.5 * XYStep # Add the deltas to multiple copies of the closed orbit in_mat = orbit6 + dmat out_mat = at.lattice_pass(ring, in_mat, refpts=refpts, keep_lattice=keeplattice) # out_mat: 8 particles at n refpts for one turn tmat3 = out_mat[:4, :, :, 0] # (x + d) - (x - d) / d mstack = (tmat3[:, :4, :] - tmat3[:, 4:8, :]) / XYStep m44 = mstack[:, :, -1] return m44, mstack[:, :, :nrefs]
def find_m66(ring, refpts=None, orbit=None, output_orbit=False, **kwargs): """find_m66 numerically finds the 6x6 transfer matrix of an accelerator lattice by differentiation of lattice_pass near the closed orbit. FINDM66 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 = find_m66(lattice) returns the full one-turn 6-by-6 matrix at the entrance of the first element m66, t = find_m66(lattice, refpts) returns 6x6 transfer matrices between the entrance of the first element and each element indexed by refpts. t is 6x6xNrefs array. ... = find_m66(lattice, ..., orbit=closed_orbit) - Does not search for closed orbit. Does not search for the closed orbit. Instead closed_orbit,a vector of initial conditions is used. This syntax is useful to specify the entrance orbit if lattice is not a ring or to avoid recomputing the closed orbit if it is already known. m66, t, orbit = find_m66(lattice, ..., output_orbit=True) Returns in addition the closed orbit at the entrance of the 1st element See also find_m44, find_orbit6 """ XYStep = kwargs.pop('XYStep', XYDEFSTEP) DPStep = kwargs.pop('DPStep', DPSTEP) keeplattice = False if orbit is None: orbit = find_orbit6(ring) keeplattice = True # Construct matrix of plus and minus deltas scaling = numpy.array([XYStep, XYStep, XYStep, XYStep, DPStep, DPStep]) dg = numpy.asfortranarray(0.5 * numpy.diag(scaling)) dmat = numpy.concatenate((dg, -dg, numpy.zeros((6, 1))), axis=1) in_mat = orbit.reshape(6, 1) + dmat refs = () if refpts is None else refpts out_mat = numpy.squeeze(at.lattice_pass(ring, in_mat, refpts=refs, keep_lattice=keeplattice), axis=3) # out_mat: 12 particles at n refpts for one turn # (x + d) - (x - d) / d m66 = (in_mat[:, :6] - in_mat[:, 6:-1]) / scaling.reshape((1, 6)) if refpts is not None: mstack = (out_mat[:, :6, :] - out_mat[:, 6:-1, :]) / XYStep if output_orbit: return m66, mstack, out_mat[:, -1, :] else: return m66, mstack else: return m66
def test_multiple_particles_lattice_pass(): lattice = [elements.Drift('Drift', 1.0)] rin = numpy.zeros((6, 2)) rin[0, 0] = 1e-6 # particle one offset in x rin[2, 1] = 1e-6 # particle two offset in y r_original = numpy.copy(rin) r_out = lattice_pass(lattice, rin, nturns=2) # particle position is not changed passing through the drift numpy.testing.assert_equal(r_original[:, 0], r_out[:, 0, 0, 0]) numpy.testing.assert_equal(r_original[:, 0], r_out[:, 0, 0, 1]) numpy.testing.assert_equal(r_original[:, 1], r_out[:, 1, 0, 0]) numpy.testing.assert_equal(r_original[:, 1], r_out[:, 1, 0, 1])
def test_linepass(engine, lattices): py_lattice, ml_lattice, _ = lattices xy_step = 1.e-8 scaling = [xy_step, xy_step, xy_step, xy_step, xy_step, xy_step] ml_rin = engine.diag(matlab.double(scaling)) ml_rout = numpy.array(engine.linepass(ml_lattice, ml_rin)) py_rin = numpy.asfortranarray(numpy.diag(scaling)) py_rout = numpy.squeeze(at.lattice_pass(py_lattice, py_rin, refpts=len(py_lattice))) assert_close(py_rout, ml_rout, rtol=0, atol=1.e-30)
def test_pydrift(): pydrift = [elements.Drift('drift', 1.0, PassMethod='pyDriftPass')] cdrift = [elements.Drift('drift', 1.0, PassMethod='DriftPass')] pyout = lattice_pass(pydrift, numpy.zeros(6) + 1.0e-6, nturns=1) cout = lattice_pass(cdrift, numpy.zeros(6) + 1.0e-6, nturns=1) numpy.testing.assert_equal(pyout, cout) set_shift(pydrift, [1.0e-3], [1.0e-3], relative=False) set_shift(cdrift, [1.0e-3], [1.0e-3], relative=False) pyout = lattice_pass(pydrift, numpy.zeros(6) + 1.0e-6, nturns=1) cout = lattice_pass(cdrift, numpy.zeros(6) + 1.0e-6, nturns=1) numpy.testing.assert_equal(pyout, cout) set_tilt(pydrift, [1.0e-3], relative=False) set_tilt(cdrift, [1.0e-3], relative=False) pyout = lattice_pass(pydrift, numpy.zeros(6) + 1.0e-6, nturns=1) cout = lattice_pass(cdrift, numpy.zeros(6) + 1.0e-6, nturns=1) numpy.testing.assert_equal(pyout, cout)
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 lattice_pass(dba_lattice, orbit, 1) assert_close(orbit[:4], orbit_copy[:4], atol=1e-12)
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. find_orbit6 starts the search from [0,,0, 0, 0, 0, 0], unless the third argument is specified: find_orbit6(ring,...,guess=initial_orbit) There exist 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; cod = find_orbit6(ring) cod: 6x1 vector - fixed point at the entrance of the 1-st element of the RING (x,px,y,py,delta,ct) cod, orbit = find_orbit6(ring, refpts) cod: 6x1 vector - fixed point at the entrance of the 1-st element of the RING (x,px,y,py,delta) orbit: 6xNrefs array - orbit at each location specified in refpts ... = find_orbit6(ring,...,guess=initial_orbit) sets the initial search to initial_orbit See also find_orbit4, find_sync_orbit. """ 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 = at.get_s_pos(ring, len(ring)) cavities = list(filter(at.checktype(at.RFCavity), ring)) if len(cavities) == 0: raise at.AtError('No cavity found in the lattice.') fRF = cavities[0].Frequency harm_number = cavities[0].HarmNumber theta = numpy.zeros((6, )) theta[5] = constants.speed_of_light * harm_number / fRF - 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 keeplattice = False while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = at.lattice_pass(ring, in_mat, refpts=[], keep_lattice=keeplattice) # 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 keeplattice = True if itercount == max_iterations: warnings.warn( at.AtWarning( 'Maximum number of iterations reached. Possible non-convergence' )) if refpts is None: output = ref_in else: all_points = numpy.squeeze( at.lattice_pass(ring, ref_in.copy(order='K'), refpts=refpts, keep_lattice=keeplattice)) output = (ref_in, all_points) return output
def find_sync_orbit(ring, dct=0.0, refpts=None, guess=None, **kwargs): """find_sync_orbit finds the closed orbit, synchronous with the RF cavity and momentum deviation dP (first 5 components of the phase space vector) % by numerically solving for a fixed point % of the one turn map M calculated with lattice_pass (X, PX, Y, PY, dP2, CT2 ) = M (X, PX, Y, PY, dP1, CT1) under constraints CT2 - CT1 = dCT = C(1/Frev - 1/Frev0) and dP2 = dP1 , where Frev0 = Frf0/HarmNumber is the design revolution frequency Frev = (Frf0 + dFrf)/HarmNumber is the imposed revolution frequency IMPORTANT!!! find_sync_orbit imposes a constraint (CT2 - CT1) and dP2 = dP1 but no constraint on the value of dP1, dP2 The algorithm assumes time-independent fixed-momentum ring to reduce the dimensionality of the problem. To impose this artificial constraint in find_sync_orbit PassMethod used for any element SHOULD NOT 1. change the longitudinal momentum dP (cavities , magnets with radiation) 2. have any time dependence (localized impedance, fast kickers etc). cod = find_sync_orbit(ring, dct) cod: 6x1 vector - fixed point at the entrance of the 1-st element of the ring (x,px,y,py,delta) cod ,orbit = find_sync_orbit(ring, dct, refpts) cod: 6x1 vector - fixed point at the entrance of the 1-st element of the ring (x,px,y,py,delta) orbit: 6xNrefs array - orbit at each location specified in refpts ... = find_sync_orbit(ring,...,guess=initial_orbit) sets the initial search to initial_orbit See also find_orbit4, find_orbit6. """ 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 delta_matrix = numpy.zeros((6, 6), order='F') for i in range(5): delta_matrix[i, i] = step_size theta5 = numpy.zeros((5, ), order='F') theta5[4] = dct id5 = numpy.zeros((5, 5), order='F') for i in range(4): id5[i, i] = 1.0 idx = numpy.array([0, 1, 2, 3, 5]) change = 1 itercount = 0 keeplattice = False while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = at.lattice_pass(ring, in_mat, refpts=[], keep_lattice=keeplattice) # the reference particle after one turn ref_out = in_mat[:, -1] # 5x5 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j5 = (in_mat[idx, :5] - in_mat[idx, 5:]) / step_size a = j5 - id5 # f'(r_n) - 1 b = ref_out[idx] - numpy.append(ref_in[:4], 0.0) - theta5 b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) r_next = ref_in - numpy.append(b_over_a, 0.0) # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keeplattice = True if itercount == max_iterations: warnings.warn( at.AtWarning( 'Maximum number of iterations reached. Possible non-convergence' )) if refpts is None: output = ref_in else: all_points = numpy.squeeze( at.lattice_pass(ring, ref_in.copy(order='K'), refpts=refpts, keep_lattice=keeplattice)) output = (ref_in, all_points) return output
def test_correct_dimensions_does_not_raise_error(rin): lattice_pass([], rin, 1) rin = numpy.zeros((6,)) lattice_pass([], rin, 1) rin = numpy.array(numpy.zeros((6, 2), order='F')) lattice_pass([], rin, 1)
def test_lattice_pass_raises_AssertionError_if_rin_incorrect_shape(input_dim): rin = numpy.zeros(input_dim) lattice = [] with pytest.raises(AssertionError): lattice_pass(lattice, rin)
def find_m44(ring, dp=0.0, refpts=None, orbit=None, output_orbit=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 = find_m44(lattice, dp=0.0) returns a full one-turn matrix at the entrance of the first element !!! With this syntax find_m44 assumes that the lattice is a ring and first finds the closed orbit m44, t = find_m44(lattice, dp, refpts) returns 4x4 transfer matrices between the entrance of the first element and each element indexed by refpts. t is a 4x4xNrefs array m44, t = find_m44(lattice, dp, refpts, full=True) Same as above except that matrixes returned in t are full 1-turn matrices at the entrance of each element indexed by refpts. ... = find_m44(lattice, ..., orbit=closed_orbit) Does not search for the closed orbit. Instead closed_orbit,a vector of initial conditions is used. This syntax is useful to specify the entrance orbit if lattice is not a ring or to avoid recomputing the closed orbit if it is already known. m44, t, orbit = find_m44(lattice, ..., output_orbit=True) Returns in addition the closed orbit at the entrance of the 1st element See also find_m66, find_orbit4 """ def mrotate(m): m = numpy.squeeze(m) return numpy.linalg.multi_dot([m, m44, _s4.T, m.T, _s4]) XYStep = kwargs.pop('XYStep', XYDEFSTEP) full = kwargs.pop('full', False) keeplattice = False if orbit is None: orbit = find_orbit4(ring, dp) keeplattice = True # Construct matrix of plus and minus deltas dg = numpy.asfortranarray(0.5 * numpy.diag([XYStep] * 6)[:, :4]) dmat = numpy.concatenate((dg, -dg, numpy.zeros((6, 1))), axis=1) # Add the deltas to multiple copies of the closed orbit in_mat = orbit.reshape(6, 1) + dmat refs = () if refpts is None else refpts out_mat = numpy.squeeze(at.lattice_pass(ring, in_mat, refpts=refs, keep_lattice=keeplattice), axis=3) # out_mat: 8 particles at n refpts for one turn # (x + d) - (x - d) / d m44 = (in_mat[:4, :4] - in_mat[:4, 4:-1]) / XYStep if refpts is not None: mstack = (out_mat[:4, :4, :] - out_mat[:4, 4:-1, :]) / XYStep if full: mstack = numpy.stack(map( mrotate, numpy.split(mstack, mstack.shape[2], axis=2)), axis=2) if output_orbit: return m44, mstack, out_mat[:, -1, :] else: return m44, mstack else: return m44
def find_orbit4(ring, dp=0.0, refpts=None): """findorbit4 finds the closed orbit in the 4-d transverse phase space by numerically solving for a fixed point of the one turn map M calculated with lattice_pass. (X, PX, Y, PY, dP2, CT2 ) = M (X, PX, Y, PY, dP1, CT1) under the CONSTANT MOMENTUM constraint, dP2 = dP1 = dP and there is NO constraint on the 6-th coordinate CT IMPORTANT!!! findorbit4 imposes a constraint on dP and relaxes the constraint on the revolution frequency. A physical storage ring does exactly the opposite: the momentum deviation of a particle on the closed orbit settles at the value such that the revolution is synchronous with the RF cavity HarmNumber*Frev = Frf To impose this artificial constraint in findorbit4, PassMethod used for any element SHOULD NOT 1. change the longitudinal momentum dP (cavities , magnets with radiation) 2. have any time dependence (localized impedance, fast kickers etc) findorbit4(RING, dP) is 4x1 vector - fixed point at the entrance of the 1-st element of the RING (x,px,y,py) findorbit4(RING, dP, refpts) is 4-by-Length(refpts) array of column vectors - fixed points (x,px,y,py) at the entrance of each element indexed refpts array. refpts is an array of increasing indexes that select elements from the range 0 to length(RING). See further explanation of refpts in the 'help' for FINDSPOS """ # We seek # - f(x) = x # - g(x) = f(x) - x = 0 # - g'(x) = f'(x) - 1 # Use a Newton-Raphson-type algorithm: # - r_n+1 = r_n - g(r_n) / g'(r_n) # - r_n+1 = r_n - (f(r_n) - r_n) / (f'(r_n) - 1) # # (f(r_n) - r_n) / (f'(r_n) - 1) can be seen as x = b/a where we use least # squares fitting to determine x when ax = b # f(r_n) - r_n is denoted b # f'(r_n) is the 4x4 jacobian, denoted j4 STEP_SIZE = 1e-6 MAX_ITERATIONS = 20 CONVERGENCE = 1e-12 r_in = numpy.zeros((6, ), order='F') r_in[4] = dp delta_matrix = numpy.zeros((6, 5), order='F') for i in range(4): delta_matrix[i, i] += STEP_SIZE change = 1 itercount = 0 keeplattice = False while (change > CONVERGENCE) and itercount < MAX_ITERATIONS: in_mat = r_in.reshape((6, 1)) + delta_matrix out_mat = at.lattice_pass(ring, in_mat, keep_lattice=keeplattice) # out_mat: 5 particles at one refpt and one turn out_mat = out_mat[:, :, 0, 0] # the reference particle after one turn ref_out = out_mat[:4, 4] # 4x4 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j4 = (out_mat[:4, :4] - ref_out.reshape((4, 1))) / STEP_SIZE a = j4 - numpy.identity(4) # f'(r_n) - 1 b = ref_out - r_in[:4] b_over_a, _, _, _ = numpy.linalg.lstsq(a, b) r_next = r_in - numpy.append(b_over_a, numpy.zeros((2, ))) # determine if we are close enough change = numpy.linalg.norm(r_next - r_in) itercount += 1 r_in = r_next keeplattice = True all_points = at.lattice_pass(ring, r_in, refpts=refpts, keep_lattice=keeplattice) # all_points: one particle at n refpts for one turn return r_in[:4], all_points[:4, 0, :, 0]
def find_orbit4(ring, dp=0.0, refpts=None, guess=None, **kwargs): """findorbit4 finds the closed orbit in the 4-d transverse phase space by numerically solving for a fixed point of the one turn map M calculated with lattice_pass. (X, PX, Y, PY, dP2, CT2 ) = M (X, PX, Y, PY, dP1, CT1) under the CONSTANT MOMENTUM constraint, dP2 = dP1 = dP and there is NO constraint on the 6-th coordinate CT IMPORTANT!!! findorbit4 imposes a constraint on dP and relaxes the constraint on the revolution frequency. A physical storage ring does exactly the opposite: the momentum deviation of a particle on the closed orbit settles at the value such that the revolution is synchronous with the RF cavity HarmNumber*Frev = Frf To impose this artificial constraint in find_orbit4, PassMethod used for any element SHOULD NOT 1. change the longitudinal momentum dP (cavities , magnets with radiation) 2. have any time dependence (localized impedance, fast kickers etc) cod = find_orbit4(ring, dP) cod: 6x1 vector - fixed point at the entrance of the 1-st element of the ring (x,px,y,py) cod ,orbit = find_orbit4(ring, dP, refpts) cod: 6x1 vector - fixed point at the entrance of the 1-st element of the ring (x,px,y,py) orbit: 6xNrefs array - orbit at each location specified in refpts ... = find_orbit4(ring,...,guess=initial_orbit) sets the initial search to initial_orbit See also find_sync_orbit, find_orbit6. """ # We seek # - f(x) = x # - g(x) = f(x) - x = 0 # - g'(x) = f'(x) - 1 # Use a Newton-Raphson-type algorithm: # - r_n+1 = r_n - g(r_n) / g'(r_n) # - r_n+1 = r_n - (f(r_n) - r_n) / (f'(r_n) - 1) # # (f(r_n) - r_n) / (f'(r_n) - 1) can be seen as x = b/a where we use least # squares fitting to determine x when ax = b # f(r_n) - r_n is denoted b # f'(r_n) is the 4x4 jacobian, denoted j4 convergence = kwargs.pop('convergence', CONVERGENCE) max_iterations = kwargs.pop('max_iterations', MAX_ITERATIONS) step_size = kwargs.pop('step_size', STEP_SIZE) if guess is None: ref_in = numpy.zeros((6, ), order='F') ref_in[4] = dp else: ref_in = guess delta_matrix = numpy.zeros((6, 5), order='F') for i in range(4): delta_matrix[i, i] = step_size id4 = numpy.asfortranarray(numpy.identity(4)) change = 1 itercount = 0 keeplattice = False while (change > convergence) and itercount < max_iterations: in_mat = ref_in.reshape((6, 1)) + delta_matrix _ = at.lattice_pass(ring, in_mat, refpts=[], keep_lattice=keeplattice) # the reference particle after one turn ref_out = in_mat[:, 4] # 4x4 jacobian matrix from numerical differentiation: # f(x+d) - f(x) / d j4 = (in_mat[:4, :4] - in_mat[:4, 4:]) / step_size a = j4 - id4 # f'(r_n) - 1 b = ref_out[:4] - ref_in[:4] b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1) r_next = ref_in - numpy.append(b_over_a, numpy.zeros((2, ))) # determine if we are close enough change = numpy.linalg.norm(r_next - ref_in) itercount += 1 ref_in = r_next keeplattice = True if itercount == max_iterations: warnings.warn( at.AtWarning( 'Maximum number of iterations reached. Possible non-convergence' )) if refpts is None: output = ref_in else: # We know that there is one particle and one turn, so select the # (6, nrefs) output. all_points = at.lattice_pass(ring, ref_in.copy(order='K'), refpts=refpts, keep_lattice=keeplattice)[:, 0, :, 0] output = (ref_in, all_points) return output