def test_tranport_twiss(self):
# transporting phase space coords:
f = 0.8
l = 1
FODO = Lattice([
QuadrupoleThin(2 * f),
Drift(l),
# expected beta minimum
QuadrupoleThin(-f),
Drift(l),
QuadrupoleThin(2 * f),
])
s, beta, alpha, gamma = FODO.twiss()
assert len(beta) == len(FODO) + 1
assert len(alpha) == len(FODO) + 1
assert len(gamma) == len(FODO) + 1
assert len(s) == len(FODO) + 1
# make sure the periodic solution is infact periodic
self.assertAlmostEqual(beta[0], beta[-1])
self.assertAlmostEqual(alpha[0], alpha[-1])
self.assertAlmostEqual(gamma[0], gamma[-1])
# make sure the beta is minimum in the right place
assert np.argmin(beta) == 2
# now in the v plane
s, beta, alpha, gamma = FODO.twiss(plane="v")
assert len(beta) == len(FODO) + 1
assert len(alpha) == len(FODO) + 1
assert len(gamma) == len(FODO) + 1
assert len(s) == len(FODO) + 1
# make sure the periodic solution is infact periodic
self.assertAlmostEqual(beta[0], beta[-1])
self.assertAlmostEqual(alpha[0], alpha[-1])
self.assertAlmostEqual(gamma[0], gamma[-1])
# make sure the beta is maximum in the right place
assert np.argmax(beta) == 2
def test_transport_error(self):
lat = Lattice([Drift(1)])
with self.assertRaises(ValueError):
lat.twiss()