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
