diff --git a/apace/tracking_integration.py b/apace/tracking_integration.py index d91cfcf..7cb78f8 100644 --- a/apace/tracking_integration.py +++ b/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 diff --git a/docs/.gitignore b/docs/.gitignore index 1079822..fdc49a9 100644 --- a/docs/.gitignore +++ b/docs/.gitignore @@ -3,3 +3,4 @@ generated src examples auto_examples +reference diff --git a/docs/conf.py b/docs/conf.py index 7db49e8..a186221 100644 --- a/docs/conf.py +++ b/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), } diff --git a/examples/_tribs.py b/examples/_tribs.py new file mode 100644 index 0000000..c607ec2 --- /dev/null +++ b/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") + + +# %% diff --git a/examples/_tune_bump.py b/examples/_tune_bump.py new file mode 100644 index 0000000..ec255e8 --- /dev/null +++ b/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") diff --git a/examples/fodo_twiss.py b/examples/fodo_twiss.py index e410ac0..8963767 100644 --- a/examples/fodo_twiss.py +++ b/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( diff --git a/tests/conftest.py b/tests/conftest.py index 5f69024..a29d10c 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_plot.py b/tests/test_plot.py index 812a913..ea463b4 100644 --- a/tests/test_plot.py +++ b/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") diff --git a/tests/test_tracking_integration.py b/tests/test_tracking_integration.py index 2e13a16..c2e7668 100644 --- a/tests/test_tracking_integration.py +++ b/tests/test_tracking_integration.py @@ -2,25 +2,47 @@ import numpy as np import apace as ap from apace.tracking_integration import Tracking +import math from apace.plot import draw_lattice import matplotlib.pyplot as plt -def test_quadrupole(): +def test_quadrupole(tmp_path, plots): 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") + s, trajectory = tracking.track(distribution, n_turns=2) + + if plots: + _, ax = plt.subplots(figsize=(20, 5)) + ax.plot(s, trajectory[:, 0]) + draw_lattice(lattice, draw_sub_lattices=False) + plt.savefig(tmp_path / "test_quadrupole.pdf") + + +def test_fodo_ring(tmp_path, fodo_ring, plots): + b1 = fodo_ring["B1"] + b1.e1 = 0 + b1.e2 = 0 + dist = ap.distribution(2, x_dist="uniform", x_width=0.001) + tracking_mat = ap.MatrixTracking(fodo_ring, dist) + tracking_int = Tracking(fodo_ring) + s_mat, x_mat = tracking_mat.orbit_position, tracking_mat.x + result = tracking_int.track(dist, max_step=math.inf) + s_int, x_int = result[0], result[1][:, 0] + assert np.allclose(np.interp(s_int, s_mat, x_mat[:, 0]), x_int[:, 0], rtol=0.01) + + if plots: + plt.plot(s_mat, x_mat, "--", label="Matrix Method", alpha=0.5) + plt.gca().set_prop_cycle(None) + plt.plot(s_int, x_int, "x", label="Integration of EOM") + ap.plot.draw_lattice(fodo_ring) + plt.legend(loc="lower right") + plt.savefig("/tmp/test_fodo_ring.pdf")