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)