Commit 0fedba14 authored by Tong Zhang's avatar Tong Zhang
Browse files

ENH: Add support to draw ellipse and envelope, optimize code.

parent efaae4ce
......@@ -16,13 +16,14 @@ Show the available templates:
import sys
from phantasy_apps.allison_scanner.data import draw_beam_ellipse_with_params
from PyQt5.QtCore import pyqtSlot
from PyQt5.QtGui import QColor
from PyQt5.QtWidgets import QMainWindow
from PyQt5.QtWidgets import QMessageBox
from phantasy import MachinePortal
from phantasy_ui import BaseAppForm
from phantasy_apps.allison_scanner.data import draw_beam_ellipse_with_params
from .ui.ui_app import Ui_MainWindow
......@@ -30,6 +31,18 @@ MACH, SEGM = "ARIS", "F1"
ARIS_MP = MachinePortal(MACH, SEGM)
ARIS_LAT = ARIS_MP.work_lattice_conf
# key strings for Twiss X,Y parameters
TWISS_KEYS_X = [
i.format(u='x') for i in ('{u}_cen', '{u}p_cen', '{u}_rms', '{u}p_rms',
'emit_{u}', 'emitn_{u}', 'alpha_{u}', 'beta_{u}',
'gamma_{u}', 'total_intensity')
]
TWISS_KEYS_Y = [
i.format(u='y') for i in ('{u}_cen', '{u}p_cen', '{u}_rms', '{u}p_rms',
'emit_{u}', 'emitn_{u}', 'alpha_{u}', 'beta_{u}',
'gamma_{u}', 'total_intensity')
]
class MyAppWindow(BaseAppForm, Ui_MainWindow):
def __init__(self, version, **kws):
......@@ -63,6 +76,21 @@ class MyAppWindow(BaseAppForm, Ui_MainWindow):
# initialize quad1_name_cbb
self.quad1_name_cbb.currentTextChanged.emit(quad_name_list[0])
# initial vars for FLAME model
self.results = None
self.last_bs = None
self.fm = None
# envelope curves
o = self.envelope_plot
o.add_curve()
o.setLineID(0) # X
o.setLineColor(QColor('#0000FF'))
o.setLineLabel("$\sigma_x$")
o.setLineID(1) # Y
o.setLineColor(QColor('#FF0000'))
o.setLineLabel("$\sigma_y$")
@pyqtSlot('QString')
def on_quad1_name_changed(self, name: str) -> None:
"""When the current selected quad name is changed, do:
......@@ -77,8 +105,10 @@ class MyAppWindow(BaseAppForm, Ui_MainWindow):
self.quad_selected.current_setting('B2'))
except TypeError:
# current_settings('B2') is None --> most likely VA is not running
QMessageBox.critical(self, "ARIS Beam Ellipse",
"Cannot reach process variables, please either start virtual accelerator or ensure Channel Access is permittable.", QMessageBox.Ok, QMessageBox.Ok)
QMessageBox.critical(
self, "ARIS Beam Ellipse",
"Cannot reach process variables, please either start virtual accelerator or ensure Channel Access is permittable.",
QMessageBox.Ok, QMessageBox.Ok)
sys.exit(1)
self.quad1_grad_dsbox.valueChanged.connect(self.on_quad1_grad_changed)
......@@ -91,7 +121,14 @@ class MyAppWindow(BaseAppForm, Ui_MainWindow):
"""
q = self.quad_selected.name
print(f"'{q}' setting is: {grad} T/m")
# draw ellipse
# update simulation
ARIS_LAT.sync_settings()
_, fm = ARIS_LAT.run()
self.results, self.last_bs = fm.run(monitor='all')
self.fm = fm
# update drawing
self.update_drawing()
def update_drawing(self):
......@@ -99,59 +136,68 @@ class MyAppWindow(BaseAppForm, Ui_MainWindow):
Here I'm drawing the beam envelop along the entire beamline, try to
replace with your routine for beam ellipse drawing.
"""
try:
self.draw_ellipse()
except NotImplementedError:
self.draw_envelope()
self.draw_envelope()
self.draw_ellipse()
def draw_envelope(self):
"""Draw beam envelop onto the figure area.
"""
# online simulation
ARIS_LAT.sync_settings()
_, fm = ARIS_LAT.run()
r, _ = fm.run(monitor='all')
results_dict = fm.collect_data(r, 'pos', 'xrms', 'yrms')
results_dict = self.fm.collect_data(self.results, 'pos', 'xrms',
'yrms')
pos = results_dict['pos']
xrms = results_dict['xrms']
yrms = results_dict['yrms']
# update drawing
# Note: matplotlibbaseWidget is used here for generic drawing,
# for curve visualization, matplotlibCurveWidget is a better choice.
self.matplotlibbaseWidget.clear_figure()
self.matplotlibbaseWidget.axes.plot(
pos,
xrms,
'b-',
pos,
yrms,
'r-',
)
self.matplotlibbaseWidget.update_figure()
# x
self.envelope_plot.setLineID(0)
self.envelope_plot.update_curve(pos, xrms)
# y
self.envelope_plot.setLineID(1)
self.envelope_plot.update_curve(pos, yrms)
def draw_ellipse(self):
#raise NotImplementedError
"""Draw x and y beam ellipse onto the figure area.
"""
ARIS_LAT.sync_settings()
_, fm = ARIS_LAT.run()
r, s = fm.run(monitor='all')
keys_x = [i.format(u='x') for i in ('{u}_cen', '{u}p_cen', '{u}_rms', '{u}p_rms', 'emit_{u}', 'emitn_{u}','alpha_{u}', 'beta_{u}', 'gamma_{u}', 'total_intensity')]
vals_x = (s.xcen, s.xpcen, s.xrms, s.xprms, s.xemittance, s.xnemittance,
s.xtwiss_alpha, s.xtwiss_beta, (s.xtwiss_alpha**2+1)/s.xtwiss_beta, 1)
params_x = dict(zip(keys_x, vals_x))
self.matplotlibbaseWidget.clear_figure()
draw_beam_ellipse_with_params(params_x, color='b', factor=4, ax=self.matplotlibbaseWidget.axes, xoy='x', fill='g', anote=False)
self.matplotlibbaseWidget.update_figure()
keys_y = [i.format(u='y') for i in ('{u}_cen', '{u}p_cen', '{u}_rms', '{u}p_rms', 'emit_{u}', 'emitn_{u}','alpha_{u}', 'beta_{u}', 'gamma_{u}', 'total_intensity')]
vals_y = (s.ycen, s.ypcen, s.yrms, s.yprms, s.yemittance, s.ynemittance,
s.ytwiss_alpha, s.ytwiss_beta, (s.ytwiss_alpha**2+1)/s.ytwiss_beta, 1)
params_y = dict(zip(keys_y, vals_y))
self.matplotlibbaseWidget_1.clear_figure()
draw_beam_ellipse_with_params(params_y, color='b', factor=4, ax=self.matplotlibbaseWidget_1.axes, xoy='y', fill='g', anote=True)
self.matplotlibbaseWidget_1.update_figure()
#
s = self.last_bs
#
vals_x = (s.xcen, s.xpcen, s.xrms, s.xprms, s.xemittance,
s.xnemittance, s.xtwiss_alpha, s.xtwiss_beta,
(s.xtwiss_alpha**2 + 1) / s.xtwiss_beta, 1)
vals_y = (s.ycen, s.ypcen, s.yrms, s.yprms, s.yemittance,
s.ynemittance, s.ytwiss_alpha, s.ytwiss_beta,
(s.ytwiss_alpha**2 + 1) / s.ytwiss_beta, 1)
params_x = dict(zip(TWISS_KEYS_X, vals_x))
params_y = dict(zip(TWISS_KEYS_Y, vals_y))
self._plot_ellipse(self.x_ellipse_plot,
params_x,
color='b',
factor=4,
xoy='x',
fill='g',
anote=False)
self._plot_ellipse(self.y_ellipse_plot,
params_y,
color='r',
factor=4,
xoy='y',
fill='m',
anote=False)
def _plot_ellipse(self, figure_obj, params, **kws):
figure_obj.clear_figure()
draw_beam_ellipse_with_params(params,
ax=figure_obj.axes,
color=kws.get('color', 'b'),
factor=kws.get('factor', 4),
xoy=kws.get('xoy', 'x'),
fill=kws.get('fill', 'g'),
anote=kws.get('anote', False))
figure_obj.update_figure()
if __name__ == "__main__":
from PyQt5.QtWidgets import QApplication
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment