87 lines
2.2 KiB
Python
Executable File
87 lines
2.2 KiB
Python
Executable File
#!/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!")
|