def test_find_m44(engine, ml_lattice, py_lattice, dp, refpts):
nelems = len(py_lattice)
refpts = range(nelems + 1) if refpts is None else refpts
nrefs = len(uint32_refpts(refpts, nelems))
# Python call
py_m44, py_mstack = physics.find_m44(py_lattice, dp, refpts)
# Matlab call
ml_m44, ml_mstack = engine.findm44(ml_lattice,
dp,
_ml_refs(refpts, nelems),
nargout=2)
ml_mstack = numpy.rollaxis(
numpy.asarray(ml_mstack).reshape((4, 4, nrefs)), -1)
# Comparison
numpy.testing.assert_almost_equal(py_m44, numpy.asarray(ml_m44), decimal=8)
numpy.testing.assert_almost_equal(py_mstack, ml_mstack, decimal=8)
def test_two_particles_for_two_turns():
rin = numpy.asfortranarray(numpy.zeros((6, 2)))
d = elements.Drift('drift', 1.0)
lattice = [d]
rin[1][0] = 1e-6
rin[3][0] = -2e-6
rout = atpass(lattice, rin, 2, refpts=uint32_refpts([1], 1))
# results from Matlab
rout_particle1_turn1 = numpy.array([1e-6, 1e-6, -2e-6, -2e-6, 0, 2.5e-12]).reshape(6,1)
rout_particle1_turn2 = numpy.array([2e-6, 1e-6, -4e-6, -2e-6, 0, 5e-12]).reshape(6,1)
# the second particle doesn't change
rout_particle2 = numpy.zeros((6,1))
# the second index is particle number
numpy.testing.assert_equal(rout[:,0,:,0], rout_particle1_turn1)
numpy.testing.assert_equal(rout[:,1,:,0], rout_particle2)
# the fourth index is turn number
numpy.testing.assert_equal(rout[:,0,:,1], rout_particle1_turn2)
numpy.testing.assert_equal(rout[:,1,:,1], rout_particle2)
def _ml_refs(refpts, nelems):
"""Convert refpoints to Matlab"""
uintrefs = uint32_refpts(refpts, nelems)
return matlab.double([ref + 1 for ref in uintrefs])
def _ml_refs(refpts, nelems):
"""Convert refpoints to Matlab"""
uintrefs = at.uint32_refpts(refpts, nelems)
# noinspection PyUnresolvedReferences
return matlab.double([ref+1 for ref in uintrefs])
def get_twiss(ring, dp=0.0, refpts=None, get_chrom=False, ddp=DDP):
"""Determine Twiss parameters by first finding the transfer matrix.
The Twiss structured array has nrefs elements, so:
* twiss['idx'].shape is (nrefs,)
* twiss['closed_orbit'].shape is (nrefs, 6).
Returns:
twiss - structured array
tune - numpy array of shape (2,)
chrom - numpy array of shape (2,)
"""
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 = at.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, orbit=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]
# Use rollaxis to get the arrays in the correct shape for the twiss
# structured array - that is, with nrefs as the first dimension.
twiss['s_pos'] = at.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'])[:, :4] / ddp
return twiss, tune, chrom