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