Esempio n. 1
0
    def test_tranport_phasespace(self):
        # transporting phase space coords:
        lat = Lattice([Drift(1)])
        s, x, x_prime, y, y_prime, dp = lat.transport([1, 0, 0, 0, 0])
        assert len(x) == 2
        assert len(x_prime) == 2
        assert len(y) == 2
        assert len(y_prime) == 2
        assert len(s) == 2
        assert np.allclose(x, np.array([1, 1]))
        assert np.allclose(x_prime, np.array([0, 0]))
        assert np.allclose(y, np.array([0, 0]))
        assert np.allclose(y_prime, np.array([0, 0]))
        assert np.allclose(s, np.array([0, 1]))
        assert dp.shape == x.shape

        s, x, x_prime, y, y_prime, dp = lat.transport([1, 1, 0, 0, 0])
        assert len(x) == 2
        assert len(x_prime) == 2
        assert len(y) == 2
        assert len(y_prime) == 2
        assert len(s) == 2
        assert np.allclose(x, np.array([1, 2]))
        assert np.allclose(x_prime, np.array([1, 1]))
        assert np.allclose(y, np.array([0, 0]))
        assert np.allclose(y_prime, np.array([0, 0]))
        assert np.allclose(s, np.array([0, 1]))
        assert dp.shape == x.shape

        # make sure the transport is consistent:
        lat = Lattice([Drift(2)])
        s_1, x_1, x_prime_1, y_1, y_prime_1, dp = lat.transport(
            [1, 1, 0, 0, 0])
        lat = Lattice([Drift(1), Drift(1)])
        s_2, x_2, x_prime_2, y_2, y_prime_2, dp = lat.transport(
            [1, 1, 0, 0, 0])
        assert len(x_2) == 3
        assert len(x_prime_2) == 3
        assert len(y_2) == 3
        assert len(y_prime_2) == 3
        assert len(s_2) == 3
        assert x_1[-1] == x_2[-1]
        assert x_prime_1[-1] == x_prime_2[-1]
        assert y_1[-1] == y_2[-1]
        assert y_prime_1[-1] == y_prime_2[-1]
        assert s_1[-1] == s_2[-1]
        assert dp.shape == x_2.shape
Esempio n. 2
0
 def test_transport_distribution(self):
     f = 0.8
     l = 1
     FODO = Lattice([
         QuadrupoleThin(2 * f),
         Drift(l),
         QuadrupoleThin(f),
         Drift(l),
         QuadrupoleThin(2 * f),
     ])
     n_particles = 10
     beam = Beam(n_particles=n_particles)
     s, x, x_prime, y, y_prime, dp = FODO.transport(
         beam.match(FODO.twiss_solution()))
     assert x.shape[-1] == len(FODO) + 1
     assert x_prime.shape[-1] == len(FODO) + 1
     assert y.shape[-1] == len(FODO) + 1
     assert y_prime.shape[-1] == len(FODO) + 1
     assert len(s) == len(FODO) + 1
     assert x.shape[0] == n_particles
     assert x_prime.shape[0] == n_particles
     assert y.shape[0] == n_particles
     assert y_prime.shape[0] == n_particles
Esempio n. 3
0
 def test_transport_ellipse(self):
     # transporting phase space coords:
     f = 0.8
     l = 1
     n_angles = 100
     FODO = Lattice([
         QuadrupoleThin(2 * f),
         Drift(l),
         QuadrupoleThin(-f),
         Drift(l),
         QuadrupoleThin(2 * f),
     ])
     beam = Beam()
     s, x, x_prime, y, y_prime, dp = FODO.transport(
         beam.ellipse(FODO.twiss_solution(), n_angles=n_angles))
     assert x.shape[-1] == len(FODO) + 1
     assert x_prime.shape[-1] == len(FODO) + 1
     assert y.shape[-1] == len(FODO) + 1
     assert y_prime.shape[-1] == len(FODO) + 1
     assert len(s) == len(FODO) + 1
     assert x.shape[0] == n_angles
     assert x_prime.shape[0] == n_angles
     assert y.shape[0] == n_angles
     assert y_prime.shape[0] == n_angles
Esempio n. 4
0
 def test_transport_input(self):
     # just checking it doesn't raise ValueError
     lat = Lattice([Drift(1)])
     # make sure arrays also work
     lat.transport(np.array([1, 2, 1, 1, 1]))
     lat.transport(
         np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]))
     # make sure lists work
     lat.transport([1, 2, 3, 1, 2])
     lat.transport([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]])
     # tuples ?
     lat.transport((1, 2, 3, 1, 2))
     lat.transport(((1, 2, 3), (1, 2, 3), (1, 2, 3), (1, 2, 3), (1, 2, 3)))