def test_length_changed():
drift = ap.Drift("Drift", length=1)
cell_1 = ap.Lattice("Cell1", [drift, drift])
cell_2 = ap.Lattice("Cell2", [cell_1, cell_1])
cell_3 = ap.Lattice("Cell3", [cell_2, cell_2])
initial_length = cell_3.length
for i in range(2, 10):
drift.length += 1
print("TEST:", drift.length, cell_3.length)
assert i * initial_length == cell_3.length
def test_unique_names():
drift_1 = ap.Drift("Drift", length=1)
drift_2 = ap.Drift("Drift", length=2)
with pytest.raises(ap.AmbiguousNameError):
ap.Lattice("Lattice", [drift_1, drift_2])
cell_1 = ap.Lattice("cell", [])
cell_2 = ap.Lattice("cell", [])
with pytest.raises(ap.AmbiguousNameError):
ap.Lattice("Lattice", [cell_1, cell_2])
def test_indices():
e0 = ap.Element("e0", length=1)
e1 = ap.Element("e1", length=1)
e2 = ap.Element("e2", length=1)
l0 = ap.Lattice("l0", (e0, e1, e2))
assert [0] == l0.indices[e0]
assert [1] == l0.indices[e1]
assert [2] == l0.indices[e2]
l1 = ap.Lattice("l1", (e0, l0, e1, l0, e2))
# e0, e0, e1, e2, e1, e0, e1, e2, e2
# 0, 1, 2, 3, 4, 5, 6, 7, 8
assert [0, 1, 5] == l1.indices[e0]
assert [2, 4, 6] == l1.indices[e1]
assert [3, 7, 8] == l1.indices[e2]
def test_attributes(fodo_cell):
fodo_ring = ap.Lattice("fodo-ring", 8 * [fodo_cell])
d1 = fodo_ring["D1"]
assert isinstance(d1, ap.Drift)
assert 0.55 == fodo_ring["D1"].length
b1: ap.Dipole = fodo_ring["B1"]
assert isinstance(b1, ap.Dipole)
assert 0.392701 == fodo_ring["B1"].angle
assert 0.1963505 == fodo_ring["B1"].e1
assert 0.1963505 == fodo_ring["B1"].e2
assert 1.5 == fodo_ring["B1"].length
q1 = fodo_ring["Q1"]
assert isinstance(q1, ap.Quadrupole)
assert 0.2 == fodo_ring["Q1"].length
assert 1.2 == fodo_ring["Q1"].k1
q2 = fodo_ring["Q2"]
assert isinstance(q2, ap.Quadrupole)
assert 0.4 == fodo_ring["Q2"].length
assert -1.2 == fodo_ring["Q2"].k1
fodo_cell = fodo_ring["FODO"]
assert isinstance(fodo_cell, ap.Lattice)
assert 6.0 == fodo_cell.length
assert isinstance(fodo_ring, ap.Lattice)
assert 48 == fodo_ring.length
def test_quadrupole():
d1 = ap.Drift("D1", length=5)
d2 = ap.Drift("D2", length=0.2)
q = ap.Quadrupole("Q1", length=0.2, k1=1.5)
s = ap.Sextupole("S1", length=0.4, k2=1000)
# s = ap.Drift("S1", length=0.4)
lattice = ap.Lattice("Lattice", [d2, q, d2, s, d1])
distribution = ap.distribution(5,
x_dist="uniform",
x_width=0.0002,
delta_dist="uniform",
delta_width=0.2)
tracking = Tracking(lattice)
s, trajectory = tracking.track(distribution)
_, ax = plt.subplots(figsize=(20, 5))
ax.plot(s, trajectory[:, 0])
plt.ylim(-0.0002, 0.0002)
draw_lattice(lattice, draw_sub_lattices=False)
plt.savefig("/tmp/test_quadrupole.pdf")
Twiss parameter of a FODO lattice
=================================
This example shows how to calulate and plot the Twiss parameter of a FOOD lattice.
"""
#%% Create a fodo based ring
import numpy as np
import apace as ap
D1 = ap.Drift("D1", length=0.55)
d1 = ap.Drift("D1", length=0.55)
b1 = ap.Dipole("B1", length=1.5, angle=0.392701, e1=0.1963505, e2=0.1963505)
q1 = ap.Quadrupole("Q1", length=0.2, k1=1.2)
q2 = ap.Quadrupole("Q2", length=0.4, k1=-1.2)
fodo_cell = ap.Lattice("FODO_CELL", [q1, d1, b1, d1, q2, d1, b1, d1, q1])
fodo_ring = ap.Lattice("FODO_RING", [fodo_cell] * 8)
#%%
# Output some info on the FODO lattice
print(
f"Overview of {fodo_ring.name}",
f"Num of elements: {len(fodo_ring.arrangement)}",
f"Lattice Length : {fodo_ring.length}",
f"Cell Length : {fodo_cell.length}",
sep="\n",
)
#%%
# Create a new ``Twiss`` object to calculate the Twiss parameter
def fodo_ring(fodo_cell):
return ap.Lattice("fodo-ring", 8 * [fodo_cell])
import numpy as np
from random import randrange
import apace as ap
# FODO circular accelerator from Klaus Wille Chapter 3.13.3
D1 = ap.Drift("D1", length=0.55)
Q1 = ap.Quadrupole("Q1", length=0.2, k1=1.2)
B1 = ap.Dipole("B1", length=1.5, angle=0.392701, e1=0.1963505, e2=0.1963505)
Q2 = ap.Quadrupole("Q2", length=0.4, k1=-1.2)
fodo = ap.Lattice("fodo-lattice", [Q1, D1, B1, D1, Q2, D1, B1, D1, Q1])
ring = ap.Lattice("fodo-ring", [fodo] * 8)
twiss = ap.Twiss(ring)
# a different start_idx should make no difference!
twiss.start_idx = randrange(twiss.n_kicks)
def test_beta():
beta_x_start = twiss.beta_x[0]
beta_x_end = twiss.beta_x[-1]
beta_y_start = twiss.beta_y[0]
beta_y_end = twiss.beta_y[-1]
assert twiss.stable
assert 9.8176 == round(beta_x_start, 4)
assert 9.8176 == round(beta_x_end, 4)
assert 1.2376 == round(beta_y_start, 4)
assert 1.2376 == round(beta_y_end, 4)
def test_dispersion():
def test_print_tree(fodo_ring):
nested1 = ap.Lattice("nested1", [fodo_ring])
nested2 = ap.Lattice("nested2", [nested1])
nested3 = ap.Lattice("nested3", [nested2])
nested4 = ap.Lattice("nested4", 2 * [nested3])
nested4.print_tree()
