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