def test_compute_twiss_solution(self):
transfer_matrix = np.array([[1, 1, 0], [0, 1, 0], [0, 0, 1]])
with self.assertRaises(ValueError):
utils.compute_twiss_solution(transfer_matrix)
# TODO: make sure this is correct
transfer_matrix_fodo = Lattice([
QuadrupoleThin(2 * 0.8),
Drift(1),
QuadrupoleThin(-0.8),
Drift(1),
QuadrupoleThin(2 * 0.8),
]).m.h
expected = np.array([[3.33066560e00], [1.11528141e-16],
[3.00240288e-01]])
assert np.allclose(utils.compute_twiss_solution(transfer_matrix_fodo),
expected)
def test_init(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8)])
cons = Constraints(lat)
assert cons.targets == []
assert cons.free_parameters == []
assert cons._lattice == lat
def test_repr(self):
lat = Lattice([Drift(1)])
repr(lat.constraints)
def test_match(self):
lat = Lattice([Drift(1)])
with self.assertRaises(ValueError):
lat.constraints.match()
lat.constraints.add_free_parameter("drift", "l")
with self.assertRaises(ValueError):
# missing target
lat.constraints.match()
lat.constraints.clear()
target = TargetPhasespace("drift", x=2)
lat.constraints.add_target(target)
with self.assertRaises(ValueError):
# missing free parameter
lat.constraints.match()
# Compute drift length to reach a x coord of 10 meters:
lat = Lattice([Drift(1)])
lat.constraints.add_free_parameter(element="drift", attribute="l")
target = TargetPhasespace("drift", x=10, initial=[0, 1, 0, 0, 0])
lat.constraints.add_target(target)
matched_lat, res = lat.constraints.match()
assert res.success
self.assertAlmostEqual(matched_lat[0].l, 10)
# make sure the original did not change.
assert lat[0].l == 1
# Compute drift length to reach a x coord of 5 meters after the first
# Drift:
drift_0 = Drift(1)
drift_1 = Drift(1)
lat = Lattice([drift_0, drift_1])
lat.constraints.add_free_parameter(drift_0, "l")
target = TargetPhasespace(drift_0, x=5, initial=[0, 1, 0, 0, 0])
lat.constraints.add_target(target)
matched_lat, res = lat.constraints.match()
assert res.success
self.assertAlmostEqual(matched_lat[0].l, 5)
self.assertAlmostEqual(matched_lat[1].l, 1)
# make sure the original did not change.
assert lat[0].l == 1
assert lat[1].l == 1
# Compute drift length to reach a x coord of 5 meters after the second
# Drift with equal lengths of both Drifts:
drift_0 = Drift(1)
drift_1 = Drift(1)
lat = Lattice([drift_0, drift_1])
lat.constraints.add_free_parameter("drift", "l")
target = TargetPhasespace(drift_1, x=5, initial=[0, 1, 0, 0, 0])
lat.constraints.add_target(target)
matched_lat, res = lat.constraints.match()
assert res.success
self.assertAlmostEqual(matched_lat[0].l, 2.5)
self.assertAlmostEqual(matched_lat[1].l, 2.5)
# make sure the original did not change.
assert lat[0].l == 1
assert lat[1].l == 1
# Compute the quad strengths of a fodo to get a beta min of 0.5 meters
lat = Lattice(
[
QuadrupoleThin(1.6, name="quad_f"),
Drift(1),
QuadrupoleThin(-0.8, name="quad_d"),
Drift(1),
QuadrupoleThin(1.6, name="quad_f"),
]
)
lat.constraints.add_free_parameter("quad_f", "f")
lat.constraints.add_free_parameter("quad_d", "f")
target = TargetTwiss("quad_d", beta=0.5, plane="h")
lat.constraints.add_target(target)
matched, opt_res = lat.constraints.match()
s, beta, *_ = matched.twiss()
self.assertAlmostEqual(min(beta), 0.5)
assert opt_res.success
# same thing but now with constraints such that the magnet strengths are
# equal
matched, opt_res = lat.constraints.match(
constraints=({"type": "eq", "fun": lambda x: x[0] + 2 * x[1]})
)
s, beta, *_ = matched.twiss()
self.assertAlmostEqual(min(beta), 0.5)
assert matched[0].f == -2 * matched[2].f
def test_repr(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8)])
cons = Constraints(lat)
repr(cons)
def test_add_free_parameter(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8)])
cons = Constraints(lat)
cons.add_free_parameter("drift", "l")
assert cons.free_parameters[0].element == "drift"
assert cons.free_parameters[0].attribute == "l"
def test_add_target(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8)])
cons = Constraints(lat)
target = TargetPhasespace("quadrupole", x=10, x_prime=1, y=0, y_prime=0, dp=0)
cons.add_target(target)
assert cons.targets[0] == target
def test_repr(self):
lat = Lattice([Drift(1)])
plotter = Plotter(lat)
repr(plotter)
def test_some_list_methods(self):
lat = Lattice()
assert len(lat) == 0
with self.assertRaises(TypeError):
lat.m
lat.append(Drift(1))
assert len(lat) == 1
assert np.allclose(lat.m, Drift(1).m)
lat.append(Drift(1))
assert len(lat) == 2
assert np.allclose(lat.m, Drift(2).m)
drift = Drift(1)
quad = QuadrupoleThin(0.8)
lat = Lattice([drift, quad])
assert lat[0] == drift
assert lat[1] == quad
lat.reverse()
assert lat[0] == quad
assert lat[1] == drift
lat = Lattice([Drift(1)])
assert len(lat) == 1
assert np.allclose(lat.m, Drift(1).m)
lat = lat * 2
assert len(lat) == 2
assert np.allclose(lat.m, Drift(2).m)
lat = Lattice([Drift(1)])
lat = lat + lat
assert len(lat) == 2
assert np.allclose(lat.m, Drift(2).m)
lat = Lattice([Drift(1)])
lat.extend([Drift(1)])
assert len(lat) == 2
assert np.allclose(lat.m, Drift(2).m)
lat = Lattice([Drift(1)])
lat.insert(0, QuadrupoleThin(0.8))
assert len(lat) == 2
assert np.allclose(lat.m, Drift(1).m @ QuadrupoleThin(0.8).m)
lat = Lattice([Drift(1)])
popped = lat.pop(0)
assert len(lat) == 0
assert isinstance(popped, Drift)
drift = Drift(1)
lat = Lattice([drift])
lat.remove(drift)
assert len(lat) == 0
lat = Lattice([Drift(1)])
lat.clear()
assert len(lat) == 0
def test_repr_(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8), Dipole(1, 1)])
repr(lat)
def test_search(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8), Dipole(1, 1)])
assert lat.search("drift") == [0]
assert lat.search("quadrupole") == [1]
assert lat.search("dipole") == [2]
lat = Lattice([
QuadrupoleThin(0.8, name="quad_f"),
QuadrupoleThin(-0.8, name="quad_d")
])
assert lat.search("quad_f") == [0]
assert lat.search("quad_d") == [1]
assert lat.search("quad_[fd]") == [0, 1]
with self.assertRaises(ValueError):
lat.search("drift")
def test_transfer_matrixes(self):
lat = Lattice([Drift(1), Drift(1)])
assert np.allclose(lat.m, Drift(2).m)
lat = Lattice([Drift(1), QuadrupoleThin(0.8)])
assert np.allclose(lat.m, QuadrupoleThin(0.8).m @ Drift(1).m)
def test_plot(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8)])
lat.plot()
def test_transport_error(self):
lat = Lattice([Drift(1)])
with self.assertRaises(ValueError):
lat.twiss()
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)))
def test_init(self):
lat = Lattice()
assert len(lat) == 0
lat = Lattice([Drift(1)])
assert len(lat) == 1