From c24e3ca7b0943106bd6a1e767290e9178dab4935 Mon Sep 17 00:00:00 2001 From: Alexey Petrenko Date: Mon, 28 Mar 2022 12:24:16 +0500 Subject: [PATCH] First description of beamline and responses. MPL with scroll. --- aperture.json | 53 ++++++++++ beamline.json | 163 ++++++++++++++++++++++++++++++ plot_aperture.py | 85 ++++++++++++++++ plot_beamline.py | 86 ++++++++++++++++ plot_orbit_correction.py | 208 +++++++++++++++++++++++++++++++++++++++ plot_response.py | 82 +++++++++++++++ print_responses.py | 26 +++++ responses.json | 77 +++++++++++++++ 8 files changed, 780 insertions(+) create mode 100644 aperture.json create mode 100644 beamline.json create mode 100755 plot_aperture.py create mode 100755 plot_beamline.py create mode 100755 plot_orbit_correction.py create mode 100755 plot_response.py create mode 100755 print_responses.py create mode 100644 responses.json diff --git a/aperture.json b/aperture.json new file mode 100644 index 0000000..a284f1f --- /dev/null +++ b/aperture.json @@ -0,0 +1,53 @@ +{ + "Name": "Example beamline apertures", + "units": "m", + "Apertures": { + "DU50": { + "type": "elliptical", + "s": [0.0, 1.5], + "Xmax": [0.025, 0.025], + "Ymax": [0.025, 0.025] + }, + "kicker": { + "type": "rectangular", + "visible": false, + "s": [-0.4, -0.35, 0.35, 0.4], + "Xmax": [ 0.022, 0.018, 0.018, 0.022] + }, + "kicker 1": { + "definition": "kicker", + "s_shift": 1.0 + }, + "kicker 2": { + "definition": "kicker", + "s_shift": 6.0 + }, + "septum": { + "type": "rectangular", + "visible": false, + "s": [ -0.4, 0.4], + "Xmin": [-0.020,-0.020], + "Xmax": [ 0.012, 0.012], + "Ymax": [ 0.010, 0.010] + }, + "Septum 1": { + "definition": "septum", + "s_shift": 2.0 + }, + "bend": { + "type": "elliptical", + "visible": false, + "s": [-0.5, 0.5], + "Xmax": [0.030, 0.030], + "Ymax": [0.015, 0.015] + }, + "B10": { + "definition": "bend", + "s_shift": 3.0 + }, + "B20": { + "definition": "bend", + "s_shift": 5.0 + } + } +} diff --git a/beamline.json b/beamline.json new file mode 100644 index 0000000..4bfa972 --- /dev/null +++ b/beamline.json @@ -0,0 +1,163 @@ +{ + "Name": "Example Beamline", + "Beamline elements": { + "Beam Energy": { + "value": 100.0, + "value_units": "MeV" + }, + "BPM10H": { + "type": "BPM", + "plane": "X", + "s": 0.2, + "value": 5.0, + "value_units": "mm" + }, + "BPM10V": { + "type": "BPM", + "plane": "Y", + "s": 0.2, + "value": 3.0, + "value_units": "mm" + }, + "Q10": { + "type": "quad", + "s": 0.5, + "L": 0.2 + }, + "Screen 20": { + "type": "screen", + "s": 6.5 + }, + "S20H": { + "type": "BPM", + "plane": "X", + "s": 6.5, + "value": 20.0, + "value_units": "mm" + }, + "S20V": { + "type": "BPM", + "plane": "Y", + "s": 6.5, + "value": -2.0, + "value_units": "mm" + }, + "Cor10H": { + "type": "corrector", + "plane": "X", + "s": 0.5, + "kick": 5.0, + "min_kick": -10, + "max_kick": 10, + "kick_units": "mrad" + }, + "Cor10V": { + "type": "corrector", + "plane": "Y", + "s": 0.5, + "kick": -2.0, + "min_kick": -10, + "max_kick": 10, + "kick_units": "mrad" + }, + "Q20": { + "type": "quad", + "s": 1.5, + "L": 0.2 + }, + "Cor20H": { + "type": "corrector", + "plane": "X", + "s": 1.2, + "kick": 5.0, + "min_kick": -10, + "max_kick": 10, + "kick_units": "mrad" + }, + "Cor20V": { + "type": "corrector", + "plane": "Y", + "s": 1.2, + "kick": 2.0, + "min_kick": -10, + "max_kick": 10, + "kick_units": "mrad" + }, + "BPM20H": { + "type": "BPM", + "plane": "X", + "s": 1.8, + "value": -8.0, + "value_units": "mm" + }, + "BPM20V": { + "type": "BPM", + "plane": "Y", + "s": 1.8, + "value": -5.0, + "value_units": "mm" + }, + "B10": { + "type": "bend", + "s": 3.0, + "L": 1.0, + "Angle": 45 + }, + "B20": { + "type": "bend", + "s": 5.0, + "L": 1.0, + "Angle": 45 + }, + "Screen 10": { + "type": "screen", + "s": 3.8 + }, + "X-Corr after B10": { + "type": "corrector", + "plane": "X", + "s": 3.6, + "kick": 2.0, + "min_kick": -5, + "max_kick": 5, + "kick_units": "mrad" + }, + "Y-Corr after B10": { + "type": "corrector", + "plane": "Y", + "s": 3.6, + "kick": -4.0, + "min_kick": -5, + "max_kick": 5, + "kick_units": "mrad" + }, + "X-BPM after B10": { + "type": "BPM", + "plane": "X", + "s": 3.7, + "value": -1.0, + "value_units": "mm" + }, + "Y-BPM after B10": { + "type": "BPM", + "plane": "Y", + "s": 3.7, + "value": 2.5, + "value_units": "mm" + }, + "Q30": { + "type": "quad", + "s": 4.0, + "L": 0.2 + }, + "X-Corr after B20": { + "type": "corrector", + "plane": "X", + "s": 5.6, + "kick": 2.0, + "min_kick": -12, + "max_kick": 12, + "kick_units": "mrad" + } + } +} diff --git a/plot_aperture.py b/plot_aperture.py new file mode 100755 index 0000000..8bf2566 --- /dev/null +++ b/plot_aperture.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 + +import json, os, sys +import numpy as np +import matplotlib.pyplot as plt +from plot_beamline import plot_beamline + +def plot_aperture(cfg, ax, plane='X'): + + for name, itm in cfg['Apertures'].items(): + visible = True # visible by default + if "visible" in itm: + if not itm["visible"]: visible = False + + if visible: + s_shift = 0 + if "s_shift" in itm: + s_shift = itm['s_shift'] + if 'definition' in itm: + itm = cfg['Apertures'][itm['definition']] + + s = np.array(itm["s"]) + s_shift + + if plane == 'X': + if "Xmax" in itm: + Xmax = np.array(itm["Xmax"])*1e3 # mm + ax.plot(s, Xmax, lw=2, c='black', alpha=0.4) + if "Xmin" in itm: + Xmin = np.array(itm["Xmin"])*1e3 # mm + else: Xmin = -Xmax + ax.plot(s, Xmin, lw=2, c='black', alpha=0.4) + + if plane == 'Y': + if "Ymax" in itm: + Ymax = np.array(itm["Ymax"])*1e3 # mm + ax.plot(s, Ymax, lw=2, c='black', alpha=0.4) + if "Ymin" in itm: + Ymin = np.array(itm["Ymin"])*1e3 # mm + else: Ymin = -Ymax + ax.plot(s, Ymin, lw=2, c='black', alpha=0.4) + + +if __name__ == '__main__': + + aperture_file='aperture.json' + if len(sys.argv) > 1: + aperture_file = sys.argv[1] + + beamline_file='beamline.json' + if len(sys.argv) > 2: + beamline_file = sys.argv[2] + + if os.path.exists(aperture_file): + with open(aperture_file, "r") as f: + cfg = json.load(f) + + fig, (axx, axy) = plt.subplots(2, 1, sharex=True, sharey=True, figsize=(11,6)) + + title = aperture_file + if "Name" in cfg: + title = f'{cfg["Name"]} ({aperture_file})' + + axx.set_title(title) + + if os.path.exists(beamline_file): + with open(beamline_file, "r") as f: + beamline_cfg = json.load(f) + print(f"Beamline file: {beamline_file}") + + plot_beamline(beamline_cfg, axx, show_names=False, box_h=5.0) + plot_beamline(beamline_cfg, axy, show_names=False, box_h=5.0) + + plot_aperture(cfg, axx, plane="X") + plot_aperture(cfg, axy, plane="Y") + + axx.set_ylabel("$x$ (mm)") + + axy.set_xlabel("$s$ (m)") + axy.set_ylabel("$y$ (mm)") + + fig.tight_layout() + plt.show() + + else: + print(f"No {aperture_file} file to plot!") diff --git a/plot_beamline.py b/plot_beamline.py new file mode 100755 index 0000000..cdd8f7a --- /dev/null +++ b/plot_beamline.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 + +import json, os, sys +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.patches import Rectangle + +def plot_beamline(cfg, ax, show_names=True, box_h=0.1, alpha=0.5): + + s_names = {} + + for name, itm in cfg['Beamline elements'].items(): + if "s" in itm: + s = itm['s'] + + if s in s_names: + s_names[s].append(name) + else: + s_names[s] = [name] + + if show_names: + ax.plot([s,s], [0,-1], lw=1, color='blue', alpha=alpha) + + color = 'blue' + h = box_h + if 'type' in itm: + itm_type = itm['type'].upper() + if itm_type == "BEND": + color = 'green' + elif itm_type == "QUAD": + h = box_h*2 + color = 'red' + elif itm_type == "BPM": + color = 'black' + + if 'L' in itm: + L = itm['L'] + ax.add_patch(Rectangle((s-L/2, 0-h/2), L, h, facecolor=color, alpha=alpha)) + else: + ax.scatter(s, 0.0, c=color, alpha=alpha, s=20, marker='s') + + if show_names: + for s, itm in s_names.items(): + ax.text(s, -1, ", ".join(itm), verticalalignment='bottom', horizontalalignment='right', + color='blue', fontsize=9, rotation='vertical', alpha=alpha) + + s_values = np.array( list(s_names) ) + s_min = s_values.min() + s_max = s_values.max() + + ax.plot([s_min,s_max], [0,0], color='gray', alpha=alpha, lw=1) + + #plt.grid() + + #plt.ylim(-0.35,+0.35) + + #plt.ylabel("$U$ (V)") + + #plt.title("BPM" + BPM) + #plt.legend() + +if __name__ == '__main__': + json_file='beamline.json' + if len(sys.argv) > 1: + json_file = sys.argv[1] + + if os.path.exists(json_file): + with open(json_file, "r") as f: + cfg = json.load(f) + + fig, ax = plt.subplots(figsize=(11,4)) + + title = json_file + if "Name" in cfg: + title = f"{cfg['Name']} ({json_file})" + + ax.set_title(title) + + plot_beamline(cfg, ax) + ax.set_xlabel("$s$ (m)") + + fig.tight_layout() + plt.show() + + else: + print(f"No {json_file} file to plot!") diff --git a/plot_orbit_correction.py b/plot_orbit_correction.py new file mode 100755 index 0000000..95c3f51 --- /dev/null +++ b/plot_orbit_correction.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 + +import json, os, sys +import numpy as np +import matplotlib.pyplot as plt + +from plot_beamline import plot_beamline +from plot_aperture import plot_aperture + +def plot_orbit(beamline_cfg, ax, plane='X'): + s_values = [] + orbit_values = [] + BPM_names = [] + + for name, itm in beamline_cfg["Beamline elements"].items(): + if "type" in itm and itm["type"] == "BPM": + if "plane" in itm and itm["plane"] == plane: + if "s" in itm and "value" in itm: + s_values.append(itm["s"]) + orbit_values.append(itm["value"]) + BPM_names.append(name) + + s_values = np.array(s_values) + orbit_values = np.array(orbit_values) + BPM_names = np.array(BPM_names) + + ids = np.argsort(s_values) # indicies sorted by s + + line, = ax.plot(s_values[ids], orbit_values[ids], "--", marker="o", lw=3) + return line, BPM_names[ids] + +def plot_correctors(beamline_cfg, ax, plane='X'): + s_values=[] + kick_values=[] + corrector_names = [] + + for name, itm in beamline_cfg["Beamline elements"].items(): + if "type" in itm and itm["type"] == "corrector": + if "plane" in itm and itm["plane"] == plane: + if "s" in itm and "kick" in itm: + corrector_names.append(name) + s = itm["s"] + s_values.append(s) + kick_values.append(itm["kick"]) + if "max_kick" in itm: + ax.plot([s,s], [0,itm["max_kick"]], color="blue", alpha=0.4, lw=3) + if "min_kick" in itm: + ax.plot([s,s], [0,itm["min_kick"]], color="blue", alpha=0.4, lw=3) + + line, = ax.plot(s_values, kick_values, " ", marker="o", + markersize=8, markeredgewidth=1, markeredgecolor="black", + color="white", alpha=0.5) + return line, corrector_names + +def responses_to_corrector(element_names, corrector_name, responses_cfg): + slopes = [] + for name in element_names: + slope = 0 + response_id = f"{name} / {corrector_name}" + if response_id in responses_cfg["response data"]: + slope = responses_cfg["response data"][response_id]["slope"] + slopes.append(slope) + return np.array(slopes) + +def on_scroll(event): + global x_orbit_values, y_orbit_values, y_orbit_values + global x_cor_values, y_cor_values + + mouse_x = event.xdata + + if not mouse_x or not event.inaxes: return + + if event.inaxes is axcx: + line = cx_dots + names_data = x_correctors + cor_values = x_cor_values + if event.inaxes is axcy: + line = cy_dots + names_data = y_correctors + cor_values = y_cor_values + if event.inaxes is axx: + line = x_line + names_data = x_BPMs + if event.inaxes is axy: + line = y_line + names_data = y_BPMs + + s_data, value_data = line.get_data() + + i = np.argmin(np.abs(s_data-mouse_x)) + + #print('you pressed', event.button, event.xdata, event.ydata) + direction = -1.0 + if event.button == "up": direction = 1.0 + + v_min, v_max =event.inaxes.get_ylim() + step = direction*(v_max-v_min)/30.0 + + element_name = names_data[i] + element = beamline_cfg["Beamline elements"][element_name] + if element["type"] == "corrector": + #print(f"Apply correction: '{element_name}' change = {step:+.3f}") + x_resp = responses_to_corrector(x_BPMs, element_name, responses_cfg) + y_resp = responses_to_corrector(y_BPMs, element_name, responses_cfg) + + x_orbit_values = x_orbit_values + x_resp*step + y_orbit_values = y_orbit_values + y_resp*step + + cor_values[i]+= step + #value_data[i]+=step + #line.set_data(s_data, value_data) + + update_plot() + + #fig.canvas.draw() + #event.canvas.draw() + +def update_plot(): + + for line, values in zip([x_line, y_line, cx_dots, cy_dots], + [x_orbit_values, y_orbit_values, x_cor_values, y_cor_values]): + + s_data, _ = line.get_data() + line.set_data(s_data, values) + + fig.canvas.draw() + +if __name__ == '__main__': + + beamline_file='beamline.json' + if len(sys.argv) > 1: + beamline_file = sys.argv[1] + if not os.path.exists(beamline_file): + print(f"No {beamline_file} file to plot!") + exit() + + responses_file='responses.json' + if len(sys.argv) > 2: + responses_file = sys.argv[2] + if not os.path.exists(responses_file): + print(f"No {responses_file} file to use!") + exit() + + aperture_file='aperture.json' + if len(sys.argv) > 3: + aperture_file = sys.argv[3] + + with open(beamline_file, "r") as f: + beamline_cfg = json.load(f) + print(f"Beamline file: {beamline_file}") + + with open(responses_file, 'r') as f: + responses_cfg = json.load(f) + print(f"Response matrix file: {responses_file}") + + fig = plt.figure(figsize=(11,6)) + + axx = fig.add_subplot(411) + axy = fig.add_subplot(412, sharex=axx, sharey=axx) + axcx = fig.add_subplot(413, sharex=axx) + axcy = fig.add_subplot(414, sharex=axx, sharey=axcx) + + title = beamline_file + + axx.set_title(title) + + if os.path.exists(aperture_file): + with open(aperture_file, "r") as f: + aperture_cfg = json.load(f) + + plot_aperture(aperture_cfg, axx, plane="X") + plot_aperture(aperture_cfg, axy, plane="Y") + + plot_beamline(beamline_cfg, axx, show_names=False, box_h=5.0, alpha=0.2) + plot_beamline(beamline_cfg, axy, show_names=False, box_h=5.0, alpha=0.2) + plot_beamline(beamline_cfg, axcx, show_names=False, box_h=2.0, alpha=0.2) + plot_beamline(beamline_cfg, axcy, show_names=False, box_h=2.0, alpha=0.2) + + x_line, x_BPMs = plot_orbit(beamline_cfg, axx, plane='X') + y_line, y_BPMs = plot_orbit(beamline_cfg, axy, plane='Y') + + #print(x_BPMs,y_BPMs) + + x_orbit_values = x_line.get_data()[1] # mm + y_orbit_values = y_line.get_data()[1] # mm + + cx_dots, x_correctors = plot_correctors(beamline_cfg, axcx, plane='X') + cy_dots, y_correctors = plot_correctors(beamline_cfg, axcy, plane='Y') + + x_cor_values = cx_dots.get_data()[1] # mrad + y_cor_values = cy_dots.get_data()[1] # mrad + + #print(x_line.get_data()) + + axx.set_ylabel("$x$ (mm)") + + axy.set_ylabel("$y$ (mm)") + + axcx.set_ylabel("x-corr. (mrad)") + axcy.set_ylabel("y-corr. (mrad)") + axcy.set_xlabel("$s$ (m)") + + fig.subplots_adjust(left=0.06, bottom=0.08, right=0.99, top=0.96, hspace=0) + #cid = fig.canvas.mpl_connect('scroll_event', on_scroll) + cid = axcx.figure.canvas.mpl_connect('scroll_event', on_scroll) + + plt.show() + diff --git a/plot_response.py b/plot_response.py new file mode 100755 index 0000000..ea2b448 --- /dev/null +++ b/plot_response.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 + +import json, os, sys +import numpy as np +import matplotlib.pyplot as plt +from plot_beamline import plot_beamline + +def plot_response(resp_name, resp_data, beamline_cfg, ax, plane='X'): + s_values=[] + slope_values=[] + slope_error_values=[] + units = "" + for name, itm in resp_data.items(): + observed, varied = name.split(" / ") + if varied == resp_name and observed in beamline_cfg["Beamline elements"]: + element = beamline_cfg["Beamline elements"][observed] + if "plane" in element and element["plane"] == plane: + s_values.append(element["s"]) + units = itm["units"] + slope_values.append(itm["slope"]) + slope_error_values.append(itm["slope error"]) + ax.plot(s_values, slope_values, "--", marker="o", lw=3) + return s_values, slope_values, units + +if __name__ == '__main__': + + resp_name = None + if len(sys.argv) > 1: + resp_name = sys.argv[1] + + resp_file='responses.json' + if len(sys.argv) > 2: + resp_file = sys.argv[2] + + beamline_file='beamline.json' + if len(sys.argv) > 3: + beamline_file = sys.argv[3] + + if os.path.exists(resp_file) and os.path.exists(beamline_file): + with open(resp_file, "r") as f: + cfg = json.load(f) + if not resp_name: + resp_name = list(cfg["response data"].keys())[0] + observed, resp_name = resp_name.split(" / ") + print(f'Plot response to "{resp_name}"') + with open(beamline_file, "r") as f: + beamline_cfg = json.load(f) + print(f"Beamline file: {beamline_file}") + + fig, (axx, axy) = plt.subplots(2, 1, sharex=True, sharey=True, figsize=(11,6)) + + title = f'Response to "{resp_name}" ({resp_file})' + + axx.set_title(title) + + # plot beamline + s, x_slopes, units = \ + plot_response(resp_name, cfg["response data"], beamline_cfg, axx, plane="X") + axx.set_ylabel(f'x-response ({units})') + + s, y_slopes, units = \ + plot_response(resp_name, cfg["response data"], beamline_cfg, axy, plane="Y") + axy.set_ylabel(f'y-response ({units})') + + axy.set_xlabel("$s$ (m)") + + min_slope = np.min([np.min(x_slopes), np.min(y_slopes)]) + max_slope = np.max([np.max(x_slopes), np.max(y_slopes)]) + + slope_range = max_slope-min_slope + box_h = 0.1 + if slope_range > 0: + box_h = slope_range*0.1 + + plot_beamline(beamline_cfg, axx, show_names=False, box_h=box_h) + plot_beamline(beamline_cfg, axy, show_names=False, box_h=box_h) + + fig.tight_layout() + plt.show() + + else: + print(f"No files to plot!") diff --git a/print_responses.py b/print_responses.py new file mode 100755 index 0000000..ad0956b --- /dev/null +++ b/print_responses.py @@ -0,0 +1,26 @@ +#!/usr/bin/env python3 + +import json, os, sys +import numpy as np + +resp_file='responses.json' +if len(sys.argv) > 1: + resp_file = sys.argv[1] + +if os.path.exists(resp_file): + with open(resp_file, "r") as f: + cfg = json.load(f) + var_dict = {} + for name, itm in cfg["response data"].items(): + observed, varied = name.split(" / ") + if varied in var_dict: + var_dict[varied].append(observed) + else: + var_dict[varied] = [observed] + + for name, itm in var_dict.items(): + print(f"{name}: {itm}") +else: + print(f"No {resp_file} file to plot!") + +#input("Press Enter.") \ No newline at end of file diff --git a/responses.json b/responses.json new file mode 100644 index 0000000..44cd4c1 --- /dev/null +++ b/responses.json @@ -0,0 +1,77 @@ +{ + "Name": "Example response matrix", + "response data": { + "BPM10H / Beam Energy": { + "units": "mm/MeV", + "slope": 0.1, + "slope error": 0.005 + }, + "BPM10V / Beam Energy": { + "units": "mm/MeV", + "slope": 0.15, + "slope error": 0.006 + }, + "BPM20H / Beam Energy": { + "units": "mm/MeV", + "slope": 0.2, + "slope error": 0.01 + }, + "BPM20V / Beam Energy": { + "units": "mm/MeV", + "slope": -0.1, + "slope error": 0.002 + }, + "X-BPM after B10 / Beam Energy": { + "units": "mm/MeV", + "slope": -0.1, + "slope error": 0.01, + "Energy values (MeV)": [ 80, 90, 100, 110], + "Measured X-values (mm)": [-21, -9, 0, 11] + }, + "Y-BPM after B10 / Beam Energy": { + "units": "mm/MeV", + "slope": -0.2, + "slope error": 0.005 + }, + "BPM20H / Cor10H": { + "units": "mm/mrad", + "slope": 1.0, + "slope error": 0.06 + }, + "BPM20V / Cor10H": { + "units": "mm/mrad", + "slope": -1.0e-1, + "slope error": 0.03 + }, + "X-BPM after B10 / Cor10H": { + "units": "mm/mrad", + "slope": -1.5, + "slope error": 0.15 + }, + "Y-BPM after B10 / Cor10H": { + "units": "mm/mrad", + "slope": 1.7e-1, + "slope error": 0.05 + }, + "BPM20V / Cor10V": { + "units": "mm/mrad", + "slope": 1.2, + "slope error": 0.03 + }, + "S20H / X-Corr after B20": { + "units": "mm/mrad", + "slope": 4, + "slope error": 0.1 + }, + "X-BPM after B10 / Cor10V": { + "units": "mm/mrad", + "slope": 1.5e-1, + "slope error": 0.07 + }, + "Y-BPM after B10 / Cor10V": { + "units": "mm/mrad", + "slope": -1.3, + "slope error": 0.05 + } + } +}