Tracking integration

apace/tracking_integration.py

@@ -2,19 +2,24 @@
 import numpy as np
 from .classes import Drift, Dipole, Quadrupole, Sextupole
 
+# from pyprofilers import profile_by_line, simple_timer
 
-def runge_kutta_4(y0, t, h, element):
-    k1 = h * y_prime(y0, t, element)
-    k2 = h * y_prime(y0 + k1 / 2, t + h / 2, element)
-    k3 = h * y_prime(y0 + k2 / 2, t + h / 2, element)
-    k4 = h * y_prime(y0 + k3, t + h, element)
+
+# @profile_by_line
+def runge_kutta_4(t, y0, h, element):
+    k1 = h * y_prime(t, y0, element)
+    k2 = h * y_prime(t + h / 2, y0 + k1 / 2, element)
+    k3 = h * y_prime(t + h / 2, y0 + k2 / 2, element)
+    k4 = h * y_prime(t + h, y0 + k3, element)
     return t + h, y0 + (k1 + 2 * k2 + 2 * k3 + k4) / 6
 
 
-def y_prime(y, t, element):
+# @profile_by_line
+def y_prime(t, y, element):
     out = np.zeros(y.shape)
-    out[0] = np.copy(y[1])
-    out[2] = np.copy(y[3])
+    # TODO: is a copy needed here?
+    out[0] = y[1]
+    out[2] = y[3]
     if isinstance(element, Drift):
         out[1] = 0
         out[3] = 0
@@ -37,27 +42,42 @@ def y_prime(y, t, element):
     return out
 
 
-# TODO: this is still experimental and not well tested!!!
 class Tracking:
+    """TODO: this is still experimental and is not well tested!!!"""
+
     def __init__(self, lattice):
         self.lattice = lattice
 
-    def track(self, initial_distribution, step_size=0.01):
-        n_steps = math.ceil(self.lattice.length / step_size) + 10
+    def track(self, initial_distribution, max_step=0.01, n_turns=1, watchers=None):
         n_particles = initial_distribution.shape[1]
+        if watchers is None:
+            watchers = np.linspace(0, self.lattice.length)
+
+        n_watchers = len(watchers)
+        n_steps = n_watchers * n_turns
         s = np.empty(n_steps)
-        s[0] = 0
         trajectory = np.empty((n_steps, 6, n_particles))
-        trajectory[0] = initial_distribution
-
-        end = 0
-        i = 0
-        for element in self.lattice.arrangement:
-            end += element.length
-            while s[i] < end:
-                step_size_arg = min(step_size, end - s[0])
-                s[i + 1], trajectory[i + 1] = runge_kutta_4(
-                    trajectory[i], s[i], step_size_arg, element
-                )
-                i += 1
-        return s[: i + 1], trajectory[: i + 1]
+
+        dist = initial_distribution
+
+        for turn in range(n_turns):
+            # print(f"Turn {turn + 1}/{n_turns}")
+            pos = end = 0
+            watchers_iter = enumerate(iter(watchers))
+            i, watcher = next(watchers_iter)
+            for element in self.lattice:
+                end = round(end + element.length, 6)  # TODO: round to um see issue 69
+                while pos < end:
+                    watcher_dist = watcher - pos
+                    element_dist = end - pos
+                    step_size = min(watcher_dist, element_dist, max_step)
+                    pos, dist = runge_kutta_4(pos, dist, step_size, element)
+                    if watcher <= pos:
+                        if watcher == pos:
+                            j = turn * n_watchers + i
+                            trajectory[j] = dist
+                            s[j] = pos + turn * self.lattice.length
+
+                        i, watcher = next(watchers_iter, (None, math.inf))
+
+        return s, trajectory

docs/.gitignore

@@ -3,3 +3,4 @@ generated
 src
 examples
 auto_examples
+reference

docs/conf.py

@@ -63,17 +63,15 @@
     "examples_dirs": "../examples",  # path to your example scripts
     "gallery_dirs": "examples",  # path to where to save gallery generated output
     "backreferences_dir": False,
-    "filename_pattern": "/",  # plot all Python files
+    "filename_pattern": r".*\.py",  # plot all Python files
+    "ignore_pattern": r"/_.*.py",
 }
 
 # Intersphinx
 extensions.append("sphinx.ext.intersphinx")
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3/", None),
-    "numpy": (
-        "https://docs.scipy.org/doc/numpy/",
-        None,
-    ),  # TODO: check how sparse did it
+    "numpy": ("https://docs.scipy.org/doc/numpy/", None),  # TODO: look at sparse docs
     "scipy": ("https://docs.scipy.org/doc/scipy/reference/", None),
 }
 

examples/_tribs.py

@@ -0,0 +1,80 @@
+"""
+TRIBS - Transverse Resonance Island Buckets
+===========================================
+
+Plot x-x' phase space for a machine close to a tune resonace.
+"""
+
+# TODO: not working at all!
+
+
+#%%
+import matplotlib.pyplot as plt
+
+import apace as ap
+from apace.tracking_integration import Tracking
+from apace.plot import draw_lattice, twiss_plot
+
+lattice = ap.Lattice.from_file(
+    "/home/felix/Git/nobeam/lattices/b2_stduser_2019_05_07.json"
+)
+twiss = ap.Twiss(lattice)
+print(twiss.tune_x)
+
+quads = {"3": [], "4": [], "5": []}
+for element in lattice.elements:
+    if isinstance(element, ap.Quadrupole):
+        name1 = element.name[1]
+        if name1.isdigit() and int(name1) >= 3:
+            quads[element.name[1]].append(element)
+
+step = 0.001
+while (abs((twiss.tune_x % 1) - 1 / 3)) > 0.25:
+    # for q3 in quads["3"]:
+    #     q3.k1 += step
+    for q4 in quads["4"]:
+        q4.k1 += -step
+    # for q5 in quads["5"]:
+    #     q5.k1 += -step
+
+    print(twiss.tune_x)
+
+fig = twiss_plot(twiss)
+fig.savefig("test_twiss.pdf")
+
+#%%
+
+
+x_width = 0.001
+n_turns = 10
+dist = ap.distribution(10, x_dist="uniform", x_center=0.0005, x_width=0.001)
+tracking = Tracking(lattice)
+s, trajectroy = tracking.track(dist, n_turns=n_turns, max_step=0.01, watchers=[0])
+s_all, t_all = tracking.track(dist, n_turns=n_turns, max_step=0.01)
+tracking_mat = ap.MatrixTracking(lattice, dist, watch_points=[0], turns=n_turns)
+x_mat, x_dds_mat = tracking_mat.x, tracking_mat.x_dds
+
+#%%
+
+fig, (ax1, ax2) = plt.subplots(nrows=2)
+ax1 = plt.subplot(211)
+
+ax1.plot(s_all, t_all[:, 0])
+draw_lattice(
+    lattice,
+    draw_sub_lattices=False,
+    annotate_elements=False,
+    annotate_sub_lattices=False,
+)
+
+ax2 = plt.subplot(223)
+ax2.plot(trajectroy[:, 0], trajectroy[:, 1], "o")
+ax2.set_xlim(-2 * x_width, 2 * x_width)
+
+ax3 = plt.subplot(224)
+ax3.plot(x_mat, x_dds_mat, "o")
+
+plt.savefig("test.pdf")
+
+
+# %%

examples/_tune_bump.py

@@ -0,0 +1,83 @@
+"""
+Tune Bump
+=========
+
+This example shows how to shift the tune using quadrupoles.
+"""
+
+#%%
+# Imports
+import numpy as np
+from scipy.optimize import minimize
+import apace as ap
+import random
+
+#%%
+# OPTIONS
+url = "https://raw.githubusercontent.com/NoBeam/lattices/master/b2_stduser_2019_05_07.json"
+tune_x_target = 17.66
+method = "Nelder-Mead"
+max_value = 3
+
+#%%
+# SETUP
+lattice = ap.Lattice.from_file(url)
+twiss = ap.Twiss(lattice)
+fit_elements = [
+    element
+    for element in lattice.elements
+    if isinstance(element, ap.Quadrupole)
+    and element.name[1].isdigit()
+    and int(element.name[1]) >= 3
+]
+fit_elements = random.sample(fit_elements, 10)
+print(fit_elements)
+s = np.linspace(0, lattice.length, 1000)
+reference = [
+    np.interp(s, twiss.s, value) for value in (twiss.beta_x, twiss.beta_y, twiss.eta_x)
+]
+
+#%%
+# FITNESS FUNCTION
+def fitness_function(params):
+    for param, element in zip(params, fit_elements):
+        element.k1 = max(min(param, max_value), -max_value)
+
+    if not twiss.stable:
+        return float("nan")
+
+    values = twiss.beta_x, twiss.beta_x, twiss.eta_x
+    means = []
+    for value, value_ref in zip(values, reference):
+        beat = (np.interp(s, twiss.s, value) / value_ref) ** 4
+        np.clip(beat, 1, None, out=beat)
+        means.append(np.mean(beat))
+    return np.mean(means) + abs(twiss.tune_x - tune_x_target)
+
+
+#%%
+# MINIMIZE
+print(f"Tune before {twiss.tune_x}")
+n_iterations = 2
+for i in range(n_iterations):
+    start_values = np.array([element.k1 for element in fit_elements])
+    if i == 0:
+        initial_values = start_values
+
+    res = minimize(fitness_function, initial_values, method=method)
+    print(f"iteration: {i + 1}/{n_iterations}, result: {res.fun:.3f}")
+
+print(f"Tune after {twiss.tune_x}")
+
+# OUTPUT RESULTS
+#%%
+for element, initial_value in zip(fit_elements, initial_values):
+    print(f"{element.name+'.k1':>9} = {element.k1:+.3f} ({initial_value:+.3f})")
+print("\n")
+
+from apace.plot import twiss_plot
+
+ref_lattice = ap.Lattice.from_file(url)
+ref_twiss = ap.Twiss(ref_lattice)
+fig = twiss_plot(twiss, ref_twiss=ref_twiss)
+fig.savefig("tune_bump.pdf")

examples/fodo_twiss.py

@@ -5,7 +5,8 @@
 This example shows how to calulate and plot the Twiss parameter of a FOOD lattice.
 """
 
-#%% Create a fodo based ring
+#%%
+# Create a fodo based ring
 import numpy as np
 import apace as ap
 
@@ -20,7 +21,6 @@
 
 #%%
 # Output some info on the FODO lattice
-
 print(
     f"Overview of {fodo_ring.name}",
     f"Num of elements: {len(fodo_ring.arrangement)}",
@@ -31,7 +31,6 @@
 
 #%%
 # Create a new ``Twiss`` object to calculate the Twiss parameter
-
 twiss = ap.Twiss(fodo_ring)
 
 print(

tests/conftest.py

@@ -10,6 +10,17 @@
 FODO_CELL_JSON = json.loads((LATTICE_PATH / "fodo_cell.json").read_text())
 
 
+def pytest_addoption(parser):
+    parser.addoption(
+        "--plots", action="store_true", default=False, help="Plot test results"
+    )
+
+
+@pytest.fixture
+def plots(request):
+    return request.config.getoption("--plots")
+
+
 @pytest.fixture
 def base_path():
     return BASE_PATH

tests/test_plot.py

@@ -3,18 +3,20 @@
 import matplotlib.pyplot as plt
 
 
-def test_draw_lattice(fodo_ring, tmp_path):
+def test_draw_lattice(fodo_ring, tmp_path, plots):
     twiss = ap.Twiss(fodo_ring)
     _, ax = plt.subplots()
     plot_twiss(twiss, ax=ax)
     draw_lattice(fodo_ring)
-    plt.tight_layout()
-    plt.savefig(tmp_path / "test_draw_lattice.pdf")
+    if plots:
+        plt.tight_layout()
+        plt.savefig(tmp_path / "test_draw_lattice.pdf")
 
 
-def test_floor_plan(fodo_ring, tmp_path):
+def test_floor_plan(fodo_ring, tmp_path, plots):
     _, ax = plt.subplots()
     ax = floor_plan(fodo_ring, ax=ax)
     ax.invert_yaxis()
-    plt.tight_layout()
-    plt.savefig(tmp_path / "apace_test_floor_plan.pdf")
+    if plots:
+        plt.tight_layout()
+        plt.savefig(tmp_path / "apace_test_floor_plan.pdf")