def test_transfer_matrix(self):
dipole = Dipole(rho_dipole, angle_dipole)
expected_transfer_matrix = np.array([
[0.5, 1.37832224e02, 0, 0, 7.95774715e01],
[-5.44139809e-03, 0.5, 0, 0, 8.66025404e-01],
[0, 0, 1, l_dipole, 0],
[0, 0, 0, 1, 0],
[0, 0, 0, 0, 1],
])
m = dipole._get_transfer_matrix()
assert np.allclose(m, expected_transfer_matrix)
assert np.allclose(dipole.m, m)
def test_plotter(self):
lat = Lattice([
Drift(1),
Quadrupole(1 / 100, 1),
Quadrupole(-1 / 100, 1),
Quadrupole(0, 1),
QuadrupoleThin(0.6),
QuadrupoleThin(-0.6),
QuadrupoleThin(0),
Dipole(1, np.pi / 2),
DipoleThin(np.pi / 16),
SextupoleThin(1),
CustomThin(np.identity(5)),
])
plotter = Plotter(lat)
fig, axes = plotter.top_down()
assert isinstance(fig, plt.Figure)
assert isinstance(axes, plt.Axes)
fig, axes = plotter.layout()
assert isinstance(fig, plt.Figure)
assert isinstance(axes, plt.Axes)
fig, axes = plotter()
assert isinstance(fig, plt.Figure)
assert isinstance(axes, plt.Axes)
# testing the margins
lat = Lattice([
Drift(0.01),
])
plotter = Plotter(lat)
fig, axes = plotter.top_down()
assert isinstance(fig, plt.Figure)
assert isinstance(axes, plt.Axes)
def test_slice(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8)])
lat_sliced = lat.slice(Drift, 10)
assert len(lat_sliced) == 11
assert all([isinstance(element, Drift) for element in lat_sliced[:10]])
assert all([element.length == 0.1 for element in lat_sliced[:10]])
assert np.allclose(lat.m, lat_sliced.m)
lat = Lattice([Drift(1), Dipole(np.pi / 4, 100)])
lat = lat.slice(Drift, 10)
assert len(lat) == 11
lat = lat.slice(Dipole, 10)
assert len(lat) == 20
def test_copy(self):
dipole = Dipole(rho_dipole, angle_dipole)
copy = dipole.copy()
assert copy.rho == dipole.rho
assert copy.theta == dipole.theta
assert copy.name == dipole.name
# make sure that if the instance's attribute is changed
# copying takes the new values.
dipole = Dipole(rho_dipole, angle_dipole)
dipole.rho = 2 * dipole.rho
dipole.theta = 2 * dipole.theta
copy = dipole.copy()
assert copy.rho == dipole.rho
assert copy.theta == dipole.theta
assert copy.name == dipole.name
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_serialize(self):
dipole = Dipole(rho_dipole, angle_dipole)
dic = dipole._serialize()
assert dic["element"] == "Dipole"
assert dic["rho"] == rho_dipole
assert dic["theta"] == angle_dipole
assert dic["name"] == dipole.name
# make sure that if the instance's attribute is changed
# the serialization takes the new values.
dipole = Dipole(rho_dipole, angle_dipole)
dipole.rho = 2 * dipole.rho
dipole.theta = 2 * dipole.theta
dic = dipole._serialize()
assert dic["element"] == "Dipole"
assert dic["rho"] == dipole.rho
assert dic["theta"] == dipole.theta
assert dic["name"] == dipole.name
def test_copy(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8), Dipole(1, 1)])
lat_copy = lat.copy()
assert len(lat) == len(lat_copy)
assert lat[0].l == lat_copy[0].l
assert lat[0].name == lat_copy[0].name
assert lat[1].f == lat_copy[1].f
assert lat[1].name == lat_copy[1].name
assert lat[2].rho == lat_copy[2].rho
assert lat[2].theta == lat_copy[2].theta
assert lat[2].name == lat_copy[2].name
# new element instances are created
assert not any(
[id(orig) == id(copy) for orig, copy in zip(lat, lat_copy)])
lat_shallow_copy = lat.copy(deep=False)
# the same instances are in both lattices
assert all([
id(orig) == id(copy) for orig, copy in zip(lat, lat_shallow_copy)
])
def test_dxztheta_ds(self):
dipole = Dipole(rho_dipole, angle_dipole)
# TODO: do the math to check this.
dipole._dxztheta_ds(0, l_dipole)
def test_repr(self):
repr(Dipole(rho_dipole, angle_dipole))
def test_slice(self):
dipole = Dipole(rho_dipole, angle_dipole)
assert len(dipole.slice(10)) == 10
assert isinstance(dipole.slice(10), Lattice)
assert np.allclose(dipole.slice(10).m, dipole.m)
assert dipole.slice(10)[0].name == dipole.name + "_slice_0"
def test_init(self):
dipole = Dipole(rho_dipole, angle_dipole)
assert dipole.length == l_dipole
assert dipole.name.startswith("dipole_")
dipole = Dipole(rho_dipole, angle_dipole, name="some_name")
assert dipole.name == "some_name"
def test_repr_(self):
lat = Lattice([Drift(1), QuadrupoleThin(0.8), Dipole(1, 1)])
repr(lat)