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