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
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
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
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)))