def __setitem__(self, key, values):
try:
super(Lattice, self).__setitem__(key, values)
except TypeError:
key = uint32_refpts(key, len(self))
for i, v in zip(*numpy.broadcast_arrays(key, values)):
super(Lattice, self).__setitem__(i, v)
def __delitem__(self, key):
try:
super(Lattice, self).__delitem__(key)
except TypeError:
key = uint32_refpts(key, len(self))
for i in reversed(key):
super(Lattice, self).__delitem__(i)
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 _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 patpass(ring, r_in, nturns, refpts=None, reuse=True, pool_size=None):
"""
Simple parallel implementation of atpass(). If more than one particle
is supplied, use multiprocessing to run each particle in a separate
process.
INPUT:
ring lattice description
r_in: 6xN array: input coordinates of N particles
nturns: number of passes through the lattice line
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, meaning no refpts, equivelent to
passing an empty array for calculation purposes.
"""
if not reuse:
raise ValueError('patpass does not support altering lattices')
if refpts is None:
refpts = len(ring)
refs = uint32_refpts(refpts, len(ring))
if pool_size is None:
pool_size = multiprocessing.cpu_count()
return _patpass(ring, r_in, nturns, refs, pool_size)
def patpass(ring,
r_in,
nturns=1,
refpts=None,
losses=False,
pool_size=None,
globvar=True,
**kwargs):
"""
Simple parallel implementation of atpass(). If more than one particle
is supplied, use multiprocessing to run each particle in a separate
process. In case a single particle is provided or the ring contains
ImpedanceTablePass element, atpass is returned
INPUT:
ring lattice description
r_in: 6xN array: input coordinates of N particles
nturns: number of passes through the lattice line
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, meaning no refpts, equivelent to
passing an empty array for calculation purposes.
losses Activate loss maps
pool_size number of processes, if None the min(npart,nproc) is used
globvar For linux machines speed-up is achieved by defining a global
ring variable, this can be disabled using globvar=False
OUTPUT:
(6, N, R, T) array containing output coordinates of N particles
at R reference points for T turns.
If losses ==True: {islost,turn,elem,coord} dictionnary containing
flag for particles lost (True -> particle lost), turn, element and
coordinates at which the particle is lost. Set to zero for particles
that survived
"""
if refpts is None:
refpts = len(ring)
refpts = uint32_refpts(refpts, len(ring))
pm_ok = [e.PassMethod in elements._collective for e in ring]
if len(numpy.atleast_1d(r_in[0])) > 1 and not any(pm_ok):
if pool_size is None:
pool_size = min(len(r_in[0]), multiprocessing.cpu_count())
return _atpass(ring,
r_in,
pool_size,
globvar,
nturns=nturns,
refpts=refpts,
losses=losses)
else:
if any(pm_ok):
warn(AtWarning('Collective PassMethod found: use single process'))
return atpass(ring, r_in, nturns=nturns, refpts=refpts, losses=losses)
def iterator(self, key):
"""Iterates over the indices selected by a slice or an array"""
if isinstance(key, slice):
indices = key.indices(len(self))
rg = range(*indices)
else:
rg = uint32_refpts(key, len(self))
return (super(Lattice, self).__getitem__(i) for i in rg)
def i_range(self):
"""Range of elements inside the range of interest"""
try:
i_range = self._i_range
except AttributeError:
self.s_range = None
i_range = self._i_range
return uint32_refpts(i_range, len(self))
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 lattice_pass(lattice, r_in, nturns=1, refpts=None, keep_lattice=False):
"""lattice_pass tracks particles through each element of a lattice
calling the element-specific tracking function specified in the
lattice[i].PassMethod field.
Note:
* lattice_pass(lattice, r_in, refpts=len(line)) is the same as
lattice_pass(lattice, r_in) since the reference point len(line) is the
exit of the last element
* lattice_pass(lattice, r_in, refpts=0) is a copy of r_in since the
reference point 0 is the entrance of the first element
PARAMETERS
lattice: iterable of AT elements
r_in: (6, N) array: input coordinates of N particles.
r_in is modified in-place and reports the coordinates at
the end of the tracking. For the the best efficiency, r_in
should be given as F_CONTIGUOUS numpy array.
nturns: number of passes through the lattice line
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, meaning no refpts, equivelent to
passing an empty array for calculation purposes.
keep_lattice: use elements persisted from a previous call to at.atpass.
If True, assume that the lattice has not changed since
that previous call.
OUTPUT
(6, N, R, T) array containing output coordinates of N particles
at R reference points for T turns.
"""
assert r_in.shape[0] == 6 and r_in.ndim in (1, 2), DIMENSION_ERROR
if not isinstance(lattice, list):
lattice = list(lattice)
nelems = len(lattice)
if refpts is None:
refpts = nelems
refs = uint32_refpts(refpts, nelems)
# atpass returns 6xAxBxC array where n = x*y*z;
# * A is number of particles;
# * B is number of refpts
# * C is the number of turns
if r_in.flags.f_contiguous:
return atpass(lattice, r_in, nturns, refs, int(keep_lattice))
else:
r_fin = numpy.asfortranarray(r_in)
r_out = atpass(lattice, r_fin, nturns, refs, int(keep_lattice))
r_in[:] = r_fin[:]
return r_out
def test_uint32_refpts_throws_ValueError_correctly(ref_in):
with pytest.raises(ValueError):
uint32_refpts(ref_in, 2)
def test_uint32_refpts_converts_numerical_inputs_correctly(ref_in, expected):
numpy.testing.assert_equal(uint32_refpts(ref_in, 5), expected)
ref_in2 = numpy.asarray(ref_in).astype(numpy.float64)
numpy.testing.assert_equal(uint32_refpts(ref_in2, 5), expected)
ref_in2 = numpy.asarray(ref_in).astype(numpy.int64)
numpy.testing.assert_equal(uint32_refpts(ref_in2, 5), expected)
def test_uint32_refpts_converts_other_input_types_correctly(ref_in, expected):
numpy.testing.assert_equal(uint32_refpts(ref_in, 5), expected)
def find_orbit6(ring,
refpts=None,
orbit=None,
dp=None,
dct=None,
keep_lattice=False,
**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 lattice description (radiation must be ON)
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
orbit=None avoids looking for initial the closed orbit if is
already known ((6,) array). find_orbit6 propagates it
to the specified refpts.
guess Initial value for the closed orbit. It may help
convergence. The default is computed from the energy
loss of the ring
keep_lattice Assume no lattice change since the previous tracking.
Default: False
method Method for energy loss computation (see get_energy_loss)
default: ELossMethod.TRACKING
cavpts=None Cavity location. If None, use all cavities. This is used
to compute the initial synchronous phase.
convergence Convergence criterion. Default: 1.e-12
max_iterations Maximum number of iterations. Default: 20
XYStep Step size. Default: DConstant.XYStep
DPStep Step size. Default: DConstant.DPStep
See also find_orbit4, find_sync_orbit.
"""
if not (dp is None and dct is None):
warnings.warn(AtWarning('In 6D, "dp" and "dct" are ignored'))
if orbit is None:
orbit = _orbit6(ring, keep_lattice=keep_lattice, **kwargs)
keep_lattice = True
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,
orbit.copy(order='K'),
refpts=uint32refs,
keep_lattice=keep_lattice),
axis=(1, 3)).T
return orbit, all_points
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, dP, CT2 ) = M (X, PX, Y, PY, dP, CT1)
under the constraint dCT = CT2 - CT1 = C/Frev - C/Frev0, 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).
PARAMETERS
ring Sequence of AT elements
dct Path lehgth deviation. Default: 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 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_orbit6.
"""
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
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
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[:, -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
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
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, dP, CT2 ) = M (X, PX, Y, PY, dP, CT1)
under the CONSTANT MOMENTUM constraint dP and with 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)
PARAMETERS
ring Sequence of AT elements
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 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 (6,) initial value for the closed orbit. It may help
convergence. Default: (0, 0, 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_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
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
ref_in[4] = dp
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
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[:, 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
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
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
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 find_m44(ring,
dp=0.0,
refpts=None,
orbit=None,
keep_lattice=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, t = find_m44(lattice, dp=0.0, refpts)
return 4x4 transfer matrices between the entrance of the first element
and each element indexed by refpts.
m44: full one-turn matrix at the entrance of the first element
t: 4x4 transfer matrices between the entrance of the first
element and each element indexed by refpts:
(Nrefs, 4, 4) array
Unless an input orbit is introduced, find_m44 assumes that the lattice is
a ring and first finds the closed orbit.
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.
full=False When True, matrices are full 1-turn matrices at
the entrance of each
element indexed by refpts.
orbit=None Avoids looking for the closed orbit if is already
known (6,) array
XYStep=6.055e-6 transverse step for numerical computation
See also find_m66, find_orbit4
"""
def mrotate(m):
m = numpy.squeeze(m)
return m.dot(m44.dot(_jmt.T.dot(m.T.dot(_jmt))))
xy_step = kwargs.pop('XYStep', XYDEFSTEP)
full = kwargs.pop('full', False)
if orbit is None:
orbit, _ = find_orbit4(ring, dp, keep_lattice=keep_lattice)
keep_lattice = True
# Construct matrix of plus and minus deltas
dg = numpy.asfortranarray(0.5 * numpy.diag([xy_step] * 6)[:, :4])
dmat = numpy.concatenate((dg, -dg), axis=1)
# Add the deltas to multiple copies of the closed orbit
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: 8 particles at n refpts for one turn
# (x + d) - (x - d) / d
m44 = (in_mat[:4, :4] - in_mat[:4, 4:]) / xy_step
if len(refs) > 0:
mstack = (out_mat[:, :4, :4] - out_mat[:, :4, 4:]) / xy_step
if full:
mstack = numpy.stack([mrotate(mat) for mat in mstack], axis=0)
else:
mstack = numpy.empty((0, 4, 4), dtype=float)
return m44, mstack
def match(ring,
variables,
constraints,
verbose=2,
max_nfev=1000,
diff_step=1.0e-10,
method=None):
"""Perform matching of constraints by varying variables
PARAMETERS
ring ring lattice or transfer line
variables sequence of Variable objects
constraints sequance of Constraints objects
KEYWORDS
verbose=2 See scipy.optimize.least_squares
max_nfev=1000 "
diff_step=1.0e-10 "
"""
def fun(vals):
for value, variable in zip(vals, variables):
variable.set(ring1, value)
variable.set(ring2, value)
c1 = [cons.evaluate(ring1) for cons in cst1]
c2 = [cons.evaluate(ring2) for cons in cst2]
return np.concatenate(c1 + c2, axis=None)
cst1 = [cons for cons in constraints if not cons.rad]
cst2 = [cons for cons in constraints if cons.rad]
ring1 = ring.copy() # Make a shallow copy of ring
for var in variables: # Make a deep copy of varying elements
for ref in uint32_refpts(var.refpts, len(ring1)):
ring1[ref] = ring1[ref].deepcopy()
ring2 = ring1.radiation_on(copy=True)
aaa = [(var.get(ring1), var.bounds) for var in variables]
vini, bounds = zip(*aaa)
bounds = np.array(bounds).T
cini1 = [cst.values(ring1) for cst in cst1]
cini2 = [cst.values(ring2) for cst in cst2]
ntargets = sum(np.size(a) for a in chain.from_iterable(cini1 + cini2))
if method is None:
if np.all(abs(bounds) == np.inf) and ntargets >= len(variables):
method = 'lm'
else:
method = 'trf'
if verbose >= 1:
print('\n{} constraints, {} variables, using method {}\n'.format(
ntargets, len(variables), method))
least_squares(fun,
vini,
bounds=bounds,
verbose=verbose,
max_nfev=max_nfev,
method=method,
diff_step=diff_step)
if verbose >= 1:
print(Constraints.header())
for cst, ini in zip(cst1, cini1):
print(cst.status(ring1, initial=ini))
for cst, ini in zip(cst2, cini2):
print(cst.status(ring2, initial=ini))
print(Variable.header())
for var, vini in zip(variables, vini):
print(var.status(ring1, vini=vini))
return ring1
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_uint32_refpts_throws_ValueError_if_input_invalid(input):
with pytest.raises(ValueError):
r = lattice.uint32_refpts(input, 2)
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_uint32_refpts_handles_array_as_input():
expected = numpy.array([1, 2, 3], dtype=numpy.uint32)
requested = numpy.array([1, 2, 3], dtype=numpy.uint32)
numpy.testing.assert_equal(lattice.uint32_refpts(requested, 5), expected)
requested = numpy.array([1, 2, 3], dtype=numpy.float64)
numpy.testing.assert_equal(lattice.uint32_refpts(requested, 5), expected)
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 find_orbit4(ring,
dp=0.0,
refpts=None,
dct=None,
orbit=None,
keep_lattice=False,
**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, dP, CT2 ) = M (X, PX, Y, PY, dP, CT1)
under the CONSTANT MOMENTUM constraint dP and with 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)
PARAMETERS
ring lattice description (radiation must be OFF)
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 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
dct=None path lengthening. If specified, dp is ignored and
the off-momentum is deduced from the path lengthening.
orbit=None avoids looking for initial the closed orbit if is
already known ((6,) array). find_orbit4 propagates it
to the specified refpts.
guess (6,) initial value for the closed orbit. It may help
convergence. Default: (0, 0, 0, 0, 0, 0)
keep_lattice Assume no lattice change since the previous tracking.
Default: False
convergence Convergence criterion. Default: 1.e-12
max_iterations Maximum number of iterations. Default: 20
XYStep Step size. Default: DConstant.XYStep
See also find_sync_orbit, find_orbit6.
"""
if orbit is None:
if dct is not None:
orbit = _orbit_dct(ring, dct, keep_lattice=keep_lattice, **kwargs)
else:
orbit = _orbit_dp(ring, dp, keep_lattice=keep_lattice, **kwargs)
keep_lattice = True
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,
orbit.copy(order='K'),
refpts=uint32refs,
keep_lattice=keep_lattice),
axis=(1, 3)).T
return orbit, all_points
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 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)
bbcum, orbs = _dmatr(ring, orbit=orbit, keep_lattice=keep_lattice)
mring, ms = find_m66(ring, uint32refs, orbit=orbs[0], keep_lattice=True)
# ------------------------------------------------------------------------
# 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, orbs[0], 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
def find_sync_orbit(ring,
dct=0.0,
refpts=None,
dp=None,
orbit=None,
keep_lattice=False,
**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, dP, CT2 ) = M (X, PX, Y, PY, dP, CT1)
under the constraint dCT = CT2 - CT1 = C/Frev - C/Frev0, 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).
PARAMETERS
ring lattice description (radiation must be OFF)
dct Path length deviation. Default: 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 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
orbit=None avoids looking for initial the closed orbit if is
already known ((6,) array). find_sync_orbit propagates
it to the specified refpts.
guess (6,) initial value for the closed orbit. It may help
convergence. Default: (0, 0, 0, 0, 0, 0)
keep_lattice Assume no lattice change since the previous tracking.
Default: False
convergence Convergence criterion. Default: 1.e-12
max_iterations Maximum number of iterations. Default: 20
XYStep Step size. Default: DConstant.XYStep
See also find_orbit4, find_orbit6.
"""
if orbit is None:
if dp is not None:
orbit = _orbit_dp(ring, dp, keep_lattice=keep_lattice, **kwargs)
else:
orbit = _orbit_dct(ring, dct, keep_lattice=keep_lattice, **kwargs)
keep_lattice = True
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,
orbit.copy(order='K'),
refpts=uint32refs,
keep_lattice=keep_lattice),
axis=(1, 3)).T
return orbit, all_points