209 lines
6.6 KiB
Python
Executable File
209 lines
6.6 KiB
Python
Executable File
#!/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()
|
|
|