def test_exact_hamiltonian_pass(rin):
drift = elements.Multipole('m1', 1, [0, 0, 0, 0], [0, 0, 0, 0])
drift.Type = 0
drift.PassMethod = 'ExactHamiltonianPass'
drift.BendingAngle = 0
l = Lattice([drift], name='lat', energy=3e9)
atpass(l, rin, 1)
def test_incorrect_dimensions_raises_value_error():
rin = numpy.array(numpy.zeros((7, 1)))
with pytest.raises(ValueError):
atpass([], rin, 1)
rin = numpy.array(numpy.zeros((1, 6)))
with pytest.raises(ValueError):
atpass([], rin, 1)
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
atpass(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_dipole(rin, dipole_class):
b = dipole_class('dipole', 1.0, 0.1, EntranceAngle=0.05, ExitAngle=0.05)
l = [b]
rin[0, 0] = 1e-6
rin_orig = numpy.copy(rin)
atpass(l, rin, 1)
rin_expected = numpy.array([1e-6, 0, 0, 0, 0, 1e-7]).reshape((6, 1))
numpy.testing.assert_almost_equal(rin_orig, rin_expected)
def test_monitor(rin):
mon = elements.Monitor('monitor')
assert mon.Length == 0
lattice = [mon]
rin = numpy.array(numpy.random.rand(*rin.shape), order='F')
rin_orig = numpy.array(rin, copy=True, order='F')
atpass(lattice, rin, 1)
numpy.testing.assert_equal(rin, rin_orig)
def test_marker(rin):
m = elements.Marker('marker')
assert m.Length == 0
lattice = [m]
rin = numpy.random.rand(*rin.shape)
rin_orig = numpy.array(rin, copy=True)
atpass(lattice, rin, 1)
numpy.testing.assert_equal(rin, rin_orig)
def test_drift_offset(rin):
d = elements.Drift('drift', 1)
lattice = [d]
rin[0, 0] = 1e-6
rin[2, 0] = 2e-6
rin_orig = numpy.array(rin, copy=True)
atpass(lattice, rin, 1)
numpy.testing.assert_equal(rin, rin_orig)
def test_ringparam(rin):
rp = elements.RingParam('ringparam', 1.e+09)
assert rp.Length == 0
lattice = [rp]
rin = numpy.array(numpy.random.rand(*rin.shape), order='F')
rin_orig = numpy.array(rin, copy=True, order='F')
atpass(lattice, rin, 1)
numpy.testing.assert_equal(rin, rin_orig)
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
atpass(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_aperture_inside_limits(rin):
a = elements.Aperture('aperture', [-1e-3, 1e-3, -1e-4, 1e-4])
assert a.Length == 0
lattice = [a]
rin[0, 0] = 1e-5
rin[2, 0] = -1e-5
rin_orig = numpy.array(rin, copy=True)
atpass(lattice, rin, 1)
numpy.testing.assert_equal(rin, rin_orig)
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
atpass(lattice, rin, 1)
assert numpy.isinf(rin[5, 0])
assert rin[2, 0] == -1e-2 # Only the 6th coordinate is marked as infinity
def test_one_particle_for_two_turns_with_no_refpts(rin):
d = elements.Drift('drift', 1.0)
lattice = [d]
rin[1][0] = 1e-6
rin[3][0] = -2e-6
atpass(lattice, rin, 2)
rout_expected = numpy.array([2e-6, 1e-6, -4e-6, -2e-6, 0, 5e-12]).reshape(6, 1)
# rin is changed in place
numpy.testing.assert_equal(rin, rout_expected)
def test_quad(rin):
q = elements.Quadrupole('quad', 0.4, k=1)
lattice = [q]
rin[0, 0] = 1e-6
atpass(lattice, rin, 1)
expected = numpy.array([
0.921060994002885, -0.389418342308651, 0, 0, 0, 0.000000010330489
]).reshape(6, 1) * 1e-6
numpy.testing.assert_allclose(rin, expected)
def test_drift_divergence(rin):
d = elements.Drift('drift', 1.0)
assert d.Length == 1
lattice = [d]
rin[1, 0] = 1e-6
rin[3, 0] = -2e-6
atpass(lattice, rin, 1)
# results from Matlab
rin_expected = numpy.array([1e-6, 1e-6, -2e-6, -2e-6, 0, 2.5e-12]).reshape(6, 1)
numpy.testing.assert_equal(rin, rin_expected)
def test_corrector(rin):
c = elements.Corrector('corrector', 0.0,
numpy.array([0.9, 0.5], dtype=numpy.float64))
assert c.Length == 0
lattice = [c]
rin[0, 0] = 1e-6
rin_orig = numpy.array(rin, copy=True)
rin_orig[1] = 0.9
rin_orig[3] = 0.5
atpass(lattice, rin, 1)
numpy.testing.assert_equal(rin, rin_orig)
def test_m66(rin, n):
m = numpy.random.rand(6, 6)
m66 = elements.M66('m66', m)
assert m66.Length == 0
lattice = [m66]
rin[n, 0] = 1e-6
atpass(lattice, rin, 1)
expected = numpy.array([
m[0, n], m[1, n], m[2, n], m[3, n], m[4, n], m[5, n]
]).reshape(6, 1) * 1e-6
numpy.testing.assert_equal(rin, expected)
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.atpass(py_lattice, py_rin, 1)
assert_close(py_rin, ml_rout, rtol=0, atol=1.e-30)
def test_dipole_bend_synonym(rin, dipole_class):
b = dipole_class('dipole', 1.0, 0.1, EntranceAngle=0.05, ExitAngle=0.05)
lat = [b]
rin[0, 0] = 1e-6
rin_orig = numpy.copy(rin)
atpass(lat, rin, 1)
rin_expected = numpy.array([1e-6, 0, 0, 0, 0, 1e-7]).reshape((6, 1))
numpy.testing.assert_almost_equal(rin_orig, rin_expected)
assert b.K == 0.0
b.PolynomB[1] = 0.2
assert b.K == 0.2
b.K = 0.1
assert b.PolynomB[1] == 0.1
def test_wiggler(rin):
period = 0.05
periods = 23
bmax = 1
by = numpy.array([1, 1, 0, 1, 1, 0], dtype=numpy.float64)
c = elements.Wiggler('wiggler', period * periods, period, bmax, 3e9, By=by)
assert abs(c.Length - 1.15) < 1e-10
lattice = [c]
# Expected value from Matlab AT.
expected = numpy.array(rin, copy=True)
expected[5] = 0.000000181809691064259
atpass(lattice, rin, 1)
numpy.testing.assert_allclose(rin, expected, atol=1e-12)
def test_exact_hamiltonian_pass_with_dls_dipole(rin):
bend = elements.Multipole('rb', 0.15, [0, 0, 0, 0],
[-0.0116333, 3.786786, 0, 0])
bend.Type = 1
bend.PassMethod = 'ExactHamiltonianPass'
bend.BendingAngle = -0.001745
bend.Energy = 3.5e9
bend.MaxOrder = 3
l = Lattice([bend], name='lat', energy=3.5e9)
atpass(l, rin, 1)
# Results from Matlab
expected = numpy.array([9.23965e-9, 1.22319e-5, 0, 0, 0,
-4.8100e-10]).reshape(6, 1)
numpy.testing.assert_allclose(rin, expected, rtol=1e-5, atol=1e-6)
def test_one_turn_for_demo_lattice(r_in, engine, ml_lattice, py_lattice):
for i in range(6):
# Change each item in r_in before calling.
r_in[i][0] = 1e-5
# Matlab call
r_out = engine.atpass(ml_lattice, r_in, 1, 1)
# Python setup
py_r_in = numpy.asfortranarray(r_in).reshape(6, 1)
py_r_out = numpy.asfortranarray(r_out).reshape(6, 1)
# Python call; py_r_in modified in place
at.atpass(py_lattice, py_r_in, 1)
numpy.testing.assert_almost_equal(py_r_in, py_r_out)
def test_drift_two_particles(rin):
d = elements.Drift('drift', 1.0)
assert d.Length == 1
lattice = [d]
two_rin = numpy.array(numpy.concatenate((rin, rin), axis=1), order='F')
# particle one is offset
two_rin[0, 0] = 1e-6
two_rin[2, 0] = 2e-6
# particle two has divergence
two_rin[1, 1] = 1e-6
two_rin[3, 1] = -2e-6
two_rin_orig = numpy.array(two_rin, copy=True)
atpass(lattice, two_rin, 1)
# results from Matlab
p1_expected = numpy.array(two_rin_orig[:, 0]).reshape(6, 1)
p2_expected = numpy.array([1e-6, 1e-6, -2e-6, -2e-6, 0, 2.5e-12]).reshape(6, 1)
two_rin_expected = numpy.concatenate((p1_expected, p2_expected), axis=1)
numpy.testing.assert_equal(two_rin, two_rin_expected)
def test_transposed_c_array_gives_same_result_as_fortran_array():
"""
This test explores the behaviour of C- and Fortran- aligned arrays.
Since we use the data from the numpy array directly, some of the
behaviour we see is not what we'd expect if we used numpy as intended.
"""
# Standard numpy array (6, 2) as used in pyAT.
rin_fortran = numpy.array(numpy.arange(12).reshape(6, 2), order='F') * 1e-5
# Taking an copy of the transpose and making sure it is C-aligned should
# give us the same layout of data in memory, but with a (2, 6) array.
rin_c = numpy.copy(rin_fortran.T, order='C')
lat = [elements.Drift('drift', 1.0)]
rout_fortran = atpass(lat, rin_fortran, 1)
# at.c does not accept (x, 6) arrays. This transpose allows rin_c
# to pass the dimension check, but does NOT change the layout in memory
# since in Python it returns a 'view' on the array.
# The AT C code would give the same result for rin_c without the
# transpose if the dimension check were not there.
rin_c_transposed = rin_c.T
rout_c = atpass(lat, rin_c_transposed, 1)
numpy.testing.assert_equal(rout_c, rout_fortran)
def test_incorrect_types_raises_value_error(rin):
with pytest.raises(TypeError):
atpass(1, rin, 1)
with pytest.raises(TypeError):
atpass([], 1, 1)
with pytest.raises(TypeError):
atpass([], rin, 'a')
def test_correct_dimensions_does_not_raise_error(rin):
l = []
atpass(l, rin, 1)
rin = numpy.zeros((6, ))
atpass(l, rin, 1)
rin = numpy.array(numpy.zeros((6, 2), order='F'))
atpass(l, rin, 1)
def test_incorrect_types_raises_value_error(rin):
l = []
with pytest.raises(ValueError):
atpass(1, rin, 1)
with pytest.raises(ValueError):
atpass(l, 1, 1)
with pytest.raises(ValueError):
atpass(l, rin, 'a')
def test_1d_particle():
# this is just a demonstration that while you can pass a 1d particle
# (shape (6,), you get back a 2d array (1, 6)
d = elements.Drift('drift', 1.0)
lattice = [d]
rin = numpy.zeros(6, )
rin[1] = 1e-6
# an empty refpts returns only the value at the end of the last turn
rout = atpass(lattice, rin, 1)
# the input may be 1d but the output is 2d
assert rout.shape == (6, 1)
rout_expected = numpy.array([1e-6, 1e-6, 0, 0, 0, 5e-13])
# rin is changed in place
numpy.testing.assert_equal(rin, rout_expected)
numpy.testing.assert_equal(rout, rout_expected.reshape(6, 1))
def test_one_particle_for_two_turns_with_empty_refpts(rin):
d = elements.Drift('drift', 1.0)
lattice = [d]
rin[1][0] = 1e-6
rin[3][0] = -2e-6
# an empty refpts returns only the value at the end of the last turn
rout = atpass(lattice,
rin,
2,
refpts=numpy.zeros((0, ), dtype=numpy.uint32))
assert rout.shape == (6, 1)
rout_expected = numpy.array([2e-6, 1e-6, -4e-6, -2e-6, 0,
5e-12]).reshape(6, 1)
# rin is changed in place
numpy.testing.assert_equal(rin, rout_expected)
numpy.testing.assert_equal(rout, rout_expected)
def test_1d_particle():
# This is just a demonstration that if you pass no refpts you get back
# a (6, *, 0, *) array. You may do this if you want only to operate
# on rin in-place.
lat = [elements.Drift('drift', 1.0)]
rin = numpy.zeros(6,)
rin[1] = 1e-6
# an empty refpts returns only the value at the end of the last turn
rout = atpass(lat, rin, 1)
# output shape: (dimensions, nparticles, refpts, nturns)
# if no refpts are supplied, the output is (6, 1, 0, 1) and contains
# no data
assert rout.shape == (6, 1, 0, 1)
rout_expected = numpy.array([1e-6, 1e-6, 0, 0, 0, 5e-13])
# rin is changed in place
numpy.testing.assert_equal(rin, rout_expected)
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)
