def simple_ring():
ring = [elements.Drift('D1', 1, R1=numpy.eye(6), R2=numpy.eye(6)),
elements.Marker('M1', attr='a_value'), elements.M66('M66'),
elements.Drift('D2', 1, T1=numpy.zeros(6), T2=numpy.zeros(6)),
elements.Drift('D3', 1, R1=numpy.eye(6), R2=numpy.eye(6)),
elements.Drift('D4', 1, T1=numpy.zeros(6), T2=numpy.zeros(6))]
return ring
def test_drift_offset(rin):
d = elements.Drift('drift', 1)
rin[0, 0] = 1e-6
rin[2, 0] = 2e-6
rin_orig = rin.copy()
element_pass(d, rin)
numpy.testing.assert_equal(rin, rin_orig)
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 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_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_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_drift_divergence(rin):
d = elements.Drift('drift', 1.0)
assert d.Length == 1
rin[1, 0] = 1e-6
rin[3, 0] = -2e-6
element_pass(d, rin)
# 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_drift_divergence(rin):
d = elements.Drift('drift', 1.0)
assert d.Length == 1
lattice = [d]
rin[0][1] = 1e-6
rin[0][3] = -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(1, 6)
numpy.testing.assert_equal(rin, rin_expected)
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_lattice_string_ordering():
lat = Lattice([elements.Drift('D0', 1.0, attr1=numpy.array(0))],
name='lat', energy=5, periodicity=1, attr2=3)
latstr = str(lat)
assert latstr.startswith("Lattice(<1 elements>, name='lat', "
"energy=5, particle=Particle('electron'), "
"periodicity=1")
assert latstr.endswith("attr2=3)")
latrepr = repr(lat)
assert latrepr.startswith("Lattice([Drift('D0', 1.0, attr1=array(0))], "
"name='lat', "
"energy=5, particle=Particle('electron'), "
"periodicity=1")
assert latrepr.endswith("attr2=3)")
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_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_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_drift_two_particles(rin):
d = elements.Drift('drift', 1.0)
assert d.Length == 1
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 = two_rin.copy()
element_pass(d, two_rin)
# 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_divide_splits_attributes_correctly():
pre = elements.Drift('drift', 1)
post = pre.divide([0.2, 0.5, 0.3])
assert len(post) == 3
assert sum([e.Length for e in post]) == pre.Length
pre = elements.Dipole('dipole', 1, KickAngle=[0.5, -0.5], BendingAngle=0.2)
post = pre.divide([0.2, 0.5, 0.3])
assert len(post) == 3
assert sum([e.Length for e in post]) == pre.Length
assert sum([e.KickAngle[0] for e in post]) == pre.KickAngle[0]
assert sum([e.KickAngle[1] for e in post]) == pre.KickAngle[1]
assert sum([e.BendingAngle for e in post]) == pre.BendingAngle
pre = elements.RFCavity('rfc', 1, voltage=187500, frequency=3.5237e+8,
harmonic_number=31, energy=6.e+9)
post = pre.divide([0.2, 0.5, 0.3])
assert len(post) == 3
assert sum([e.Length for e in post]) == pre.Length
assert sum([e.Voltage for e in post]) == pre.Voltage
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 test_reuse_attributes(rin, reuse):
lat = [elements.Drift('drift', 1.0)]
rin[0, 0] = 1e-6
rin[1, 0] = 1e-6
rin_copy = numpy.copy(rin)
# two turns with original lattice
atpass(lat, rin, 2)
# one turn with original lattice
atpass(lat, rin_copy, 1)
# change an attribute
lat[0].Length = 2
# one turn with altered lattice
atpass(lat, rin_copy, 1, reuse=reuse)
if reuse:
numpy.testing.assert_equal(rin, rin_copy)
else:
with pytest.raises(AssertionError):
numpy.testing.assert_equal(rin, rin_copy)
def test_lattice_string_ordering():
lat = Lattice([elements.Drift('D0', 1.0, attr1=numpy.array(0))],
name='lat',
energy=5,
periodicity=1,
attr2=3)
# Default dictionary ordering is only in Python >= 3.6
if sys.version_info < (3, 6):
assert lat.__str__().startswith("Lattice(<1 elements>, ")
assert lat.__str__().endswith(", attr2=3)")
assert lat.__repr__().startswith("Lattice([Drift('D0', 1.0, "
"attr1=array(0))], ")
assert lat.__repr__().endswith(", attr2=3)")
else:
assert lat.__str__() == ("Lattice(<1 elements>, name='lat', energy=5,"
" periodicity=1, attr2=3)")
assert lat.__repr__() == ("Lattice([Drift('D0', 1.0, attr1=array(0))],"
" name='lat', energy=5, periodicity=1,"
" attr2=3)")
def test_drift_two_particles(rin):
d = elements.Drift('drift', 1.0)
assert d.Length == 1
lattice = [d]
two_rin = numpy.concatenate((rin, rin), axis=0)
# particle one is offset
two_rin[0][0] = 1e-6
two_rin[0][2] = 2e-6
# particle two has divergence
two_rin[1][1] = 1e-6
two_rin[1][3] = -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(1, 6)
p2_expected = numpy.array([1e-6, 1e-6, -2e-6, -2e-6, 0,
2.5e-12]).reshape(1, 6)
two_rin_expected = numpy.concatenate((p1_expected, p2_expected), axis=0)
numpy.testing.assert_equal(two_rin, two_rin_expected)
def test_insert_into_drift():
# Create elements
drift = elements.Drift('drift', 1)
monitor = elements.Monitor('bpm')
quad = elements.Quadrupole('quad', 0.3)
# Test None splitting behaviour
el_list = drift.insert([(0., None), (0.3, None), (0.7, None), (1., None)])
assert len(el_list) == 3
numpy.testing.assert_almost_equal([e.Length for e in el_list],
[0.3, 0.4, 0.3])
# Test normal insertion
el_list = drift.insert([(0.3, monitor), (0.7, quad)])
assert len(el_list) == 5
numpy.testing.assert_almost_equal([e.Length for e in el_list],
[0.3, 0.0, 0.25, 0.3, 0.15])
# Test insertion at either end produces -ve length drifts
el_list = drift.insert([(0.0, quad), (1.0, quad)])
assert len(el_list) == 5
numpy.testing.assert_almost_equal([e.Length for e in el_list],
[-0.15, 0.3, 0.7, 0.3, -0.15])
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)
# results from Matlab
rout1_expected = numpy.array([1e-6, 1e-6, -2e-6, -2e-6, 0,
2.5e-12]).reshape(6)
rout2_expected = numpy.array([2e-6, 1e-6, -4e-6, -2e-6, 0,
5e-12]).reshape(6)
# the second particle doesn't change
rout3_expected = numpy.zeros((6, ))
# the first two 6x1 columns are the two particles after the first turn
numpy.testing.assert_equal(rout[:, 0], rout1_expected)
numpy.testing.assert_equal(rout[:, 1], rout3_expected)
# the second two 6x1 columns are the two particles after the second turn
numpy.testing.assert_equal(rout[:, 2], rout2_expected)
numpy.testing.assert_equal(rout[:, 3], rout3_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_patpass_multiple_particles_and_turns():
nturns = 10
nparticles = 10
rin = numpy.zeros((6, nparticles))
d = elements.Drift('drift', 1.0)
assert d.Length == 1
lattice = [d]
rin[1, 0] = 1e-6
rin[3, 0] = -2e-6
rout = patpass(lattice, rin, nturns)
# results from Matlab
assert rout.shape == (6, nparticles * nturns)
rout_expected = numpy.array([1e-6, 1e-6, -2e-6, -2e-6, 0,
2.5e-12]).reshape(6, )
zeros = numpy.zeros((6, ))
numpy.testing.assert_equal(rout[:, 0], rout_expected)
# only the first particle is not all zeros
for i in range(nturns):
with pytest.raises(AssertionError):
numpy.testing.assert_equal(rout[:, i * nturns], zeros)
for j in range(1, nparticles):
numpy.testing.assert_equal(rout[:, i * nturns + j], zeros)
def test_patpass_multiple_particles_and_turns():
nturns = 10
nparticles = 10
rin = numpy.zeros((6, nparticles))
d = elements.Drift('drift', 1.0)
lattice = [d]
rin[1, 0] = 1e-6
rin[3, 0] = -2e-6
rout = patpass(lattice, rin, nturns)
# results from Matlab
assert rout.shape == (6, nparticles, 1, nturns)
rout_expected = numpy.array([1e-6, 1e-6, -2e-6, -2e-6, 0, 2.5e-12]).reshape(6,1)
zeros = numpy.zeros((6,))
# The fourth index is for nturns; the first element is after one turn.
numpy.testing.assert_equal(rout[:,0,:,0], rout_expected)
for i in range(nturns):
# The first particle is not all zeros.
with pytest.raises(AssertionError):
numpy.testing.assert_equal(rout[:,0,0,i], zeros)
# All other particles are all zeros.
for j in range(1, nparticles):
numpy.testing.assert_equal(rout[:,j,0,i], zeros)
def test_missing_length_raises_attribute_error(rin):
m = elements.Drift('drift', 1.0)
l = [m]
del m.Length
with pytest.raises(AttributeError):
atpass(l, rin, 1)
def test_missing_length_raises_attribute_error(rin):
lat = [elements.Drift('drift', 1.0)]
del lat[0].Length
with pytest.raises(AttributeError):
atpass(lat, rin, 1)
def test_element_string_ordering():
d = elements.Drift('D0', 1, attr=numpy.array(0))
assert d.__str__() == ("Drift:\nFamName : D0\nLength : 1.0\nPassMethod :"
" DriftPass\nattr : 0")
assert d.__repr__() == "Drift('D0', 1.0, attr=array(0))"