def __init__(self, segment, basename, line_view_config):
self.hook = HookCollection(
plot=None,
plot_ax=None,
destroy=None,
)
# create figure
self.figure = figure = FigurePair()
self.segment = segment
self.config = line_view_config
self.title = line_view_config['title'][basename]
self.sname = sname = 's'
self.xname = xname = basename + 'x'
self.yname = yname = basename + 'y'
axx, axy = figure.axx, figure.axy
axx.twiss_name, axx.twiss_conj = xname, yname
axy.twiss_name, axy.twiss_conj = yname, xname
self.axes = {xname: axx, yname: axy}
# plot style
self._label = line_view_config['label']
unit_names = line_view_config['unit']
self.unit = {col: getattr(units, unit_names[col])
for col in [sname, xname, yname]}
# subscribe for updates
TwissCurveSegment(segment, self)
DrawLineElements(self, self.config['element_style'])
self.segment.hook.update.connect(self.update)
def __init__(self, session, sequence, range, beam, twiss_args,
show_element_indicators):
"""
:param Session session:
:param str sequence:
:param tuple range:
"""
self.hook = HookCollection(
update=None,
remove=None,
show_element_indicators=None,
)
self.session = session
self.sequence = session.madx.sequences[sequence]
self.start, self.stop = self.parse_range(range)
self.range = (normalize_range_name(self.start.name),
normalize_range_name(self.stop.name))
self._beam = beam
self._twiss_args = twiss_args
self._show_element_indicators = show_element_indicators
self._use_beam(beam)
raw_elements = self.sequence.elements
# TODO: provide uncached version of elements with units:
self.elements = list(map(
session.utool.dict_add_unit, raw_elements))
# TODO: self.hook.create(self)
self.twiss()
def __init__(self, app, show=True):
"""
Create notebook frame.
Extends wx.Frame.__init__.
"""
self.hook = HookCollection(
init='madgui.core.mainframe.init',
menu='madgui.core.mainframe.menu',
reset=None,
)
super(MainFrame, self).__init__(
None, -1,
title='MadGUI',
size=wx.Size(800, 600))
self.views = []
self.app = app
self.env = {
'frame': self,
'views': self.views,
'session': None,
}
self.madx_units = unit.UnitConverter(
unit.from_config_dict(self.app.conf['madx_units']))
self.CreateControls()
self.Show(show)
def __init__(self, segment, rules):
self.hook = HookCollection(
stop=None,
add_constraint=None,
remove_constraint=None,
clear_constraints=None)
self.segment = segment
self.constraints = {}
self._elements = segment.elements
self._rules = rules
self._variable_parameters = {}
def __init__(self, popup, segment, element_name):
"""Start to manage the popup window."""
self.hook = HookCollection(
set_element=None,
close=None)
self.segment = segment
self.popup = popup
self._closed = False
segment.hook.update.connect(self.update)
self.popup.Bind(wx.EVT_CLOSE, self.OnClose)
# this comes last, as it implies an update
self.element_name = element_name
def __init__(self, parent, view, **kwargs):
"""
Initialize panel and connect the view.
Extends wx.App.__init__.
"""
self.hook = HookCollection(
init='madgui.widget.figure.init',
capture_mouse=None)
super(FigurePanel, self).__init__(parent, **kwargs)
self.capturing = False
self.view = view
# couple figure to canvas
self.canvas = Canvas(self, -1, view.figure.figure)
view.canvas = self.canvas
# create a toolbar
self.toolbar = Toolbar(self.canvas)
self.hook.init(self)
self.toolbar.Realize()
# put elements into sizer
sizer = wx.BoxSizer(wx.VERTICAL)
sizer.Add(self.canvas, 1, wx.EXPAND)
sizer.Add(self.toolbar, 0 , wx.LEFT | wx.EXPAND)
self.SetSizer(sizer)
# setup mouse capturing
self.hook.capture_mouse.connect(self.on_capture_mouse)
self.toolbar.Bind(wx.EVT_TOOL,
self.on_zoom_or_pan,
id=self.get_zoom_id())
self.toolbar.Bind(wx.EVT_TOOL,
self.on_zoom_or_pan,
id=self.get_pan_id())
# get notified when frame is destroyed
self.Bind(wx.EVT_WINDOW_DESTROY, self.on_destroy)
def __init__(self, args=None, conf=None):
"""
Create an application instance.
:param dict args: preprocessed command line parameters
"""
self.hook = HookCollection(
init='madgui.core.app.init')
self.args = args
self.conf = conf
self.dist = working_set.find(Requirement.parse('madgui'))
self.add_entry_points(self.entry_points)
# Add all entry point maps (strings like `App.entry_point` above) that
# are registered under 'madgui.entry_points'. This indirection renders
# the plugin mechanism more dynamic and allows plugins to be defined
# more easily by eliminating the need to execute 'setup.py' each time
# an entrypoint is added, changed or removed. Instead, their setup
# step only needs to create a single entrypoint which is less likely
# to change.
for ep in iter_entry_points('madgui.entry_points'):
self.add_entry_points(ep.load())
super(App, self).__init__(redirect=False)
def __init__(self, panel):
"""Add toolbar tool to panel and subscribe to capture events."""
self.hook = HookCollection(
start='madgui.component.matching.start')
self.cid = None
self.segment = panel.view.segment
self.panel = panel
self.view = panel.view
self.matcher = None
# toolbar tool
res = PackageResource('madgui.data')
with res.open('cursor.xpm') as xpm:
img = wx.ImageFromStream(xpm, wx.BITMAP_TYPE_XPM)
bmp = wx.BitmapFromImage(img)
self.toolbar = panel.toolbar
self.tool = panel.toolbar.AddCheckTool(
wx.ID_ANY,
bitmap=bmp,
shortHelp='Beam matching',
longHelp='Match by specifying constraints for envelope x(s), y(s).')
panel.Bind(wx.EVT_TOOL, self.OnMatchClick, self.tool)
# setup mouse capture
panel.hook.capture_mouse.connect(self.stop_match)
class MatchTool(object):
"""
Controller that performs matching when clicking on an element.
"""
def __init__(self, panel):
"""Add toolbar tool to panel and subscribe to capture events."""
self.hook = HookCollection(
start='madgui.component.matching.start')
self.cid = None
self.segment = panel.view.segment
self.panel = panel
self.view = panel.view
self.matcher = None
# toolbar tool
res = PackageResource('madgui.data')
with res.open('cursor.xpm') as xpm:
img = wx.ImageFromStream(xpm, wx.BITMAP_TYPE_XPM)
bmp = wx.BitmapFromImage(img)
self.toolbar = panel.toolbar
self.tool = panel.toolbar.AddCheckTool(
wx.ID_ANY,
bitmap=bmp,
shortHelp='Beam matching',
longHelp='Match by specifying constraints for envelope x(s), y(s).')
panel.Bind(wx.EVT_TOOL, self.OnMatchClick, self.tool)
# setup mouse capture
panel.hook.capture_mouse.connect(self.stop_match)
def OnMatchClick(self, event):
"""Invoked when user clicks Match-Button"""
if event.IsChecked():
self.start_match()
else:
self.stop_match()
def start_match(self):
"""Start matching mode."""
self.panel.hook.capture_mouse()
self.cid = self.view.figure.canvas.mpl_connect(
'button_press_event',
self.on_match)
app = self.panel.GetTopLevelParent().app
self.matcher = Matching(self.segment, app.conf['matching'])
self.hook.start(self.matcher, self.view)
def stop_match(self):
"""Stop matching mode."""
if self.cid is not None:
self.view.figure.canvas.mpl_disconnect(self.cid)
self.cid = None
self.toolbar.ToggleTool(self.tool.Id, False)
self.matcher.stop()
def on_match(self, event):
"""
Draw new constraint and perform matching.
Invoked after the user clicks in matching mode.
"""
axes = event.inaxes
if axes is None:
return
name = axes.twiss_name
conj = axes.twiss_conj
elem = self.segment.element_by_position(
event.xdata * self.view.unit['s'])
if elem is None or 'name' not in elem:
return
if event.button == 2:
self.matcher.remove_constraint(name, elem)
self.matcher.remove_constraint(conj, elem)
return
elif event.button != 1:
return
orig_cursor = self.panel.GetCursor()
wait_cursor = wx.StockCursor(wx.CURSOR_WAIT)
self.panel.SetCursor(wait_cursor)
# By default, the list of constraints will be reset. The shift/alt
# keys are used to add more constraints.
pressed_keys = event.key or ''
add_keys = ['shift', 'control']
if not any(add_key in pressed_keys for add_key in add_keys):
self.matcher.clear_constraints()
# add the clicked constraint
envelope = event.ydata * self.view.unit[name]
self.matcher.add_constraint(name, elem, envelope)
# add another constraint to hold the orthogonal axis constant
orth_env = self.segment.get_twiss(elem, conj)
self.matcher.add_constraint(conj, elem, orth_env)
self.matcher.match()
self.panel.SetCursor(orig_cursor)
class Matching(object):
def __init__(self, segment, rules):
self.hook = HookCollection(
stop=None,
add_constraint=None,
remove_constraint=None,
clear_constraints=None)
self.segment = segment
self.constraints = {}
self._elements = segment.elements
self._rules = rules
self._variable_parameters = {}
def stop(self):
self.clear_constraints()
self.hook.stop()
def _allvars(self, axis):
try:
allvars = self._variable_parameters[axis]
except KeyError:
# filter element list for usable types:
param_spec = self._rules.get(axis, {})
allvars = [(elem, param_spec[elem['type']])
for elem in self._elements
if elem['type'] in param_spec]
self._variable_parameters[axis] = allvars
return allvars
def match(self):
"""Perform matching according to current constraints."""
segment = self.segment
simul = self.segment.session
trans = MatchTransform(segment)
# transform constraints (envx => betx, etc)
trans_constr = {}
for axis, constr in self.constraints.items():
for elem, value in constr:
trans_name, trans_value = getattr(trans, axis)(value)
this_constr = trans_constr.setdefault(trans_name, [])
this_constr.append((elem, trans_value))
# The following uses a greedy algorithm to select all elements that
# can be used for varying. This means that for advanced matching it
# will most probably not work.
# Copy all needed variable lists (for later modification):
allvars = {axis: self._allvars(axis)[:]
for axis in trans_constr}
vary = []
for axis, constr in trans_constr.items():
for elem, envelope in constr:
at = elem['at']
allowed = [v for v in allvars[axis] if v[0]['at'] < at]
if not allowed:
# No variable in range found! Ok.
continue
v = max(allowed, key=lambda v: v[0]['at'])
expr = _get_any_elem_param(v[0], v[1])
if expr is None:
allvars[axis].remove(v)
else:
vary.append(expr)
for c in allvars.values():
try:
c.remove(v)
except ValueError:
pass
# create constraints list to be passed to Madx.match
constraints = []
for name, constr in trans_constr.items():
for elem, val in constr:
constraints.append({
'range': elem['name'],
name: simul.utool.strip_unit(name, val)})
twiss_args = simul.utool.dict_strip_unit(segment.twiss_args)
simul.madx.match(sequence=segment.sequence.name,
vary=vary,
constraints=constraints,
twiss_init=twiss_args)
segment.twiss()
def _gconstr(self, axis):
return self.constraints.get(axis, [])
def _sconstr(self, axis):
return self.constraints.setdefault(axis, [])
def find_constraint(self, axis, elem):
"""Find and return the constraint for the specified element."""
return [c for c in self._gconstr(axis) if c[0] == elem]
def add_constraint(self, axis, elem, envelope):
"""Add constraint and perform matching."""
self.remove_constraint(axis, elem)
self._sconstr(axis).append( (elem, envelope) )
self.hook.add_constraint()
def remove_constraint(self, axis, elem):
"""Remove the constraint for elem."""
try:
orig = self.constraints[axis]
except KeyError:
return
filtered = [c for c in orig if c[0]['name'] != elem['name']]
if filtered:
self.constraints[axis] = filtered
else:
del self.constraints[axis]
if len(filtered) < len(orig):
self.hook.remove_constraint()
def clear_constraints(self):
"""Remove all constraints."""
self.constraints = {}
self.hook.clear_constraints()
class TwissView(object):
"""Instanciate an FigurePair + XYCurve(Envelope)."""
@classmethod
def create(cls, session, frame, basename):
"""Create a new view panel as a page in the notebook frame."""
if not session.segment:
return
view = cls(session.segment, basename, frame.app.conf['line_view'])
panel = frame.AddView(view, view.title)
return view
def __init__(self, segment, basename, line_view_config):
self.hook = HookCollection(
plot=None,
plot_ax=None,
destroy=None,
)
# create figure
self.figure = figure = FigurePair()
self.segment = segment
self.config = line_view_config
self.title = line_view_config['title'][basename]
self.sname = sname = 's'
self.xname = xname = basename + 'x'
self.yname = yname = basename + 'y'
axx, axy = figure.axx, figure.axy
axx.twiss_name, axx.twiss_conj = xname, yname
axy.twiss_name, axy.twiss_conj = yname, xname
self.axes = {xname: axx, yname: axy}
# plot style
self._label = line_view_config['label']
unit_names = line_view_config['unit']
self.unit = {col: getattr(units, unit_names[col])
for col in [sname, xname, yname]}
# subscribe for updates
TwissCurveSegment(segment, self)
DrawLineElements(self, self.config['element_style'])
self.segment.hook.update.connect(self.update)
def destroy(self):
self.segment.hook.update.disconnect(self.update)
self.hook.destroy()
def update(self):
self.figure.draw()
def get_label(self, name):
return self._label[name] + ' ' + get_unit_label(self.unit[name])
def plot(self):
fig = self.figure
axx = fig.axx
axy = fig.axy
sname, xname, yname = self.sname, self.xname, self.yname
# start new plot
fig.start_plot()
axx.set_ylabel(self.get_label(xname))
axy.set_ylabel(self.get_label(yname))
fig.set_slabel(self.get_label(sname))
# invoke plot hooks
self.hook.plot_ax(axx, xname)
self.hook.plot_ax(axy, yname)
self.hook.plot()
# finish and draw:
fig.draw()
def get_axes_name(self, axes):
return axes.twiss_name
def get_conjugate(self, name):
return axes.twiss_conj
class App(wx.App):
"""
Core application class.
Use App.main() to run the application.
:ivar args: command line arguments
"""
version = _version
usage = _usage
entry_points = """
[madgui.core.app.init]
mainframe = madgui.core.mainframe:MainFrame
[madgui.core.mainframe.menu]
online_control = madgui.online.control:Control
[madgui.widget.figure.init]
matchtool = madgui.component.matchtool:MatchTool
selecttool = madgui.component.selecttool:SelectTool
comparetool = madgui.component.comparetool:CompareTool
statusbar = madgui.component.lineview:UpdateStatusBar.create
[madgui.component.matching.start]
drawconstraints = madgui.component.lineview:DrawConstraints
"""
@classmethod
def main(cls, argv=None):
"""
Create an application instance and run the MainLoop.
:param list argv: command line parameters
"""
from docopt import docopt
args = docopt(cls.usage, argv, version=cls.version)
conf = load_config(args['--config'])
cls(args, conf).MainLoop()
def __init__(self, args=None, conf=None):
"""
Create an application instance.
:param dict args: preprocessed command line parameters
"""
self.hook = HookCollection(
init='madgui.core.app.init')
self.args = args
self.conf = conf
self.dist = working_set.find(Requirement.parse('madgui'))
self.add_entry_points(self.entry_points)
# Add all entry point maps (strings like `App.entry_point` above) that
# are registered under 'madgui.entry_points'. This indirection renders
# the plugin mechanism more dynamic and allows plugins to be defined
# more easily by eliminating the need to execute 'setup.py' each time
# an entrypoint is added, changed or removed. Instead, their setup
# step only needs to create a single entrypoint which is less likely
# to change.
for ep in iter_entry_points('madgui.entry_points'):
self.add_entry_points(ep.load())
super(App, self).__init__(redirect=False)
def OnInit(self):
"""Initialize the application and create main window."""
# allow plugin components to create stuff (frame!)
self.hook.init(self)
# signal wxwidgets to enter the main loop
return True
def add_entry_points(self, entry_map_section):
"""Add entry points."""
recursive_merge(
self.dist.get_entry_map(),
EntryPoint.parse_map(entry_map_section, self.dist))
class FigurePanel(wx.Panel):
"""
Display panel for a matplotlib figure.
"""
def __init__(self, parent, view, **kwargs):
"""
Initialize panel and connect the view.
Extends wx.App.__init__.
"""
self.hook = HookCollection(
init='madgui.widget.figure.init',
capture_mouse=None)
super(FigurePanel, self).__init__(parent, **kwargs)
self.capturing = False
self.view = view
# couple figure to canvas
self.canvas = Canvas(self, -1, view.figure.figure)
view.canvas = self.canvas
# create a toolbar
self.toolbar = Toolbar(self.canvas)
self.hook.init(self)
self.toolbar.Realize()
# put elements into sizer
sizer = wx.BoxSizer(wx.VERTICAL)
sizer.Add(self.canvas, 1, wx.EXPAND)
sizer.Add(self.toolbar, 0 , wx.LEFT | wx.EXPAND)
self.SetSizer(sizer)
# setup mouse capturing
self.hook.capture_mouse.connect(self.on_capture_mouse)
self.toolbar.Bind(wx.EVT_TOOL,
self.on_zoom_or_pan,
id=self.get_zoom_id())
self.toolbar.Bind(wx.EVT_TOOL,
self.on_zoom_or_pan,
id=self.get_pan_id())
# get notified when frame is destroyed
self.Bind(wx.EVT_WINDOW_DESTROY, self.on_destroy)
def on_destroy(self, event):
"""Invoked when C++ window is destroyed."""
self.view.destroy()
def on_zoom_or_pan(self, event):
"""Capture mouse, after Zoom/Pan tools were clicked."""
if event.IsChecked():
self.capturing = True
self.hook.capture_mouse()
self.capturing = False
event.Skip()
def on_capture_mouse(self):
"""Disable Zoom/Pan tools when someone captures the mouse."""
if self.capturing:
return
zoom_id = self.get_zoom_id()
if self.toolbar.GetToolState(zoom_id):
self.toolbar.zoom()
self.toolbar.ToggleTool(zoom_id, False)
pan_id = self.get_pan_id()
if self.toolbar.GetToolState(pan_id):
self.toolbar.pan()
self.toolbar.ToggleTool(pan_id, False)
def get_pan_id(self):
try:
return self.toolbar.wx_ids['Pan']
except AttributeError:
return self.toolbar._NTB2_PAN
def get_zoom_id(self):
try:
return self.toolbar.wx_ids['Zoom']
except AttributeError:
return self.toolbar._NTB2_ZOOM
class Segment(object):
"""
Simulate one fixed segment, i.e. sequence + range.
:ivar Madx madx:
:ivar list elements:
:ivar dict twiss_args:
"""
_columns = [
'name', 'l', 'angle', 'k1l',
's',
'x', 'y',
'betx','bety',
'alfx', 'alfy',
]
# TODO: extend list of merge-columns
_mixin_columns = [
'x', 'y',
'betx','bety',
'alfx', 'alfy',
]
def __init__(self, session, sequence, range, beam, twiss_args,
show_element_indicators):
"""
:param Session session:
:param str sequence:
:param tuple range:
"""
self.hook = HookCollection(
update=None,
remove=None,
show_element_indicators=None,
)
self.session = session
self.sequence = session.madx.sequences[sequence]
self.start, self.stop = self.parse_range(range)
self.range = (normalize_range_name(self.start.name),
normalize_range_name(self.stop.name))
self._beam = beam
self._twiss_args = twiss_args
self._show_element_indicators = show_element_indicators
self._use_beam(beam)
raw_elements = self.sequence.elements
# TODO: provide uncached version of elements with units:
self.elements = list(map(
session.utool.dict_add_unit, raw_elements))
# TODO: self.hook.create(self)
self.twiss()
@property
def madx(self):
return self.session.madx
@property
def utool(self):
return self.session.utool
@property
def data(self):
return {
'sequence': self.sequence.name,
'range': self.range,
'beam': self.beam,
'twiss': self.twiss_args,
}
def get_element_info(self, element):
"""Get :class:`ElementInfo` from element name or index."""
if isinstance(element, ElementInfo):
return element
if isinstance(element, (basestring, dict)):
element = self.sequence.elements.index(element)
element_data = self.session.utool.dict_add_unit(
self.sequence.elements[element])
if element < 0:
element += len(self.sequence.elements)
return ElementInfo(element_data['name'], element, element_data['at'])
def parse_range(self, range):
"""Convert a range str/tuple to a tuple of :class:`ElementInfo`."""
if isinstance(range, basestring):
range = range.split('/')
start_name, stop_name = range
return (self.get_element_info(start_name),
self.get_element_info(stop_name))
def destroy(self):
self.session.segment = None
self.hook.remove()
@property
def show_element_indicators(self):
return self._show_element_indicators
@show_element_indicators.setter
def show_element_indicators(self, show):
if show == self._show_element_indicators:
return
self._show_element_indicators = show
self.hook.show_element_indicators()
@property
def twiss_args(self):
return self._twiss_args
@twiss_args.setter
def twiss_args(self, twiss_args):
self._twiss_args = twiss_args
self.twiss()
@property
def beam(self):
"""Get the beam parameter dictionary."""
return self._beam
@beam.setter
def beam(self, beam):
"""Set beam from a parameter dictionary."""
self._beam = beam
self._use_beam(beam)
self.twiss()
def _use_beam(self, beam):
beam = self.utool.dict_strip_unit(beam)
beam = dict(beam, sequence=self.sequence.name)
self.madx.command.beam(**beam)
def element_by_position(self, pos):
"""Find optics element by longitudinal position."""
if pos is None:
return None
for elem in self.elements:
at, L = elem['at'], elem['l']
if pos >= at and pos <= at+L:
return elem
return None
def get_element_index(self, elem):
"""Get element index by it name."""
return self.sequence.elements.index(elem)
def get_twiss(self, elem, name):
"""Return beam envelope at element."""
element = self.get_element_info(elem)
if not self.contains(element):
return None
return self.tw[name][element.index - self.start.index]
def contains(self, element):
return (self.start.index <= element.index and
self.stop.index >= element.index)
def twiss(self):
"""Recalculate TWISS parameters."""
results = self.raw_twiss()
# Update TWISS results
self.tw = self.utool.dict_add_unit(results)
self.summary = self.utool.dict_add_unit(results.summary)
# data post processing
self.tw['s'] += self.start.at
self.pos = self.tw['s']
self.tw['envx'] = (self.tw['betx'] * self.summary['ex'])**0.5
self.tw['envy'] = (self.tw['bety'] * self.summary['ey'])**0.5
# Create aliases for x,y that have non-empty common prefix. The goal
# is to make the config file entries less awkward that hold this
# prefix:
self.tw['posx'] = self.tw['x']
self.tw['posy'] = self.tw['y']
self.hook.update()
def _get_twiss_args(self, **kwargs):
twiss_init = self.utool.dict_strip_unit(self.twiss_args)
twiss_args = {
'sequence': self.sequence.name,
'range': self.range,
'columns': self._columns,
'twiss_init': twiss_init,
}
twiss_args.update(kwargs)
return twiss_args
def raw_twiss(self, **kwargs):
return self.madx.twiss(**self._get_twiss_args(**kwargs))
def get_transfer_map(self, beg_elem, end_elem):
"""
Get the transfer matrix R(i,j) between the two elements.
This requires a full twiss call, so don't do it too often.
"""
info = self.get_element_info
twiss_args = self._get_twiss_args()
twiss_args['range_'] = (info(beg_elem).name, info(end_elem).name)
twiss_args['tw_range'] = twiss_args.pop('range')
return self.madx.get_transfer_map_7d(**twiss_args)
class ElementView(object):
"""
Control class for filling a TableDialog with beam line element details.
"""
def __init__(self, popup, segment, element_name):
"""Start to manage the popup window."""
self.hook = HookCollection(
set_element=None,
close=None)
self.segment = segment
self.popup = popup
self._closed = False
segment.hook.update.connect(self.update)
self.popup.Bind(wx.EVT_CLOSE, self.OnClose)
# this comes last, as it implies an update
self.element_name = element_name
def OnClose(self, event):
"""Disconnect the manager, after the popup window was closed."""
self.segment.hook.update.disconnect(self.update)
self._closed = True
self.hook.close()
event.Skip()
def __nonzero__(self):
return not self._closed
__bool__ = __nonzero__
@property
def element_name(self):
return self._element_name
@element_name.setter
def element_name(self, name):
self._element_name = name
self.hook.set_element()
self.update()
@property
def element(self):
elements = self.segment.session.madx.active_sequence.elements
raw_element = elements[self.element_name]
return self.segment.session.utool.dict_add_unit(raw_element)
def update(self):
"""
Update the contents of the managed popup window.
"""
el = self.element
rows = list(el.items())
# convert to title case:
rows = [(k.title(),v) for (k,v) in rows]
# presort alphanumerically:
# (with some luck the order on the elements with equal key in the
# subsequent sort will be left invariant)
rows = sorted(rows)
# sort preferred elements to top:
order = ['Name',
'Type',
'At',
'L',
'Ksl',
'Knl',
]
order = dict(zip(order, range(-len(order), 0)))
rows = sorted(rows, key=lambda row: order.get(row[0], len(order)))
rows = filter(lambda row: row[0] != 'Vary', rows)
# update view:
self.popup.rows = rows
class MainFrame(MDIParentFrame):
"""
Notebook window class for MadGUI (main window).
"""
def __init__(self, app, show=True):
"""
Create notebook frame.
Extends wx.Frame.__init__.
"""
self.hook = HookCollection(
init='madgui.core.mainframe.init',
menu='madgui.core.mainframe.menu',
reset=None,
)
super(MainFrame, self).__init__(
None, -1,
title='MadGUI',
size=wx.Size(800, 600))
self.views = []
self.app = app
self.env = {
'frame': self,
'views': self.views,
'session': None,
}
self.madx_units = unit.UnitConverter(
unit.from_config_dict(self.app.conf['madx_units']))
self.CreateControls()
self.Show(show)
def CreateControls(self):
# create notebook
self.Bind(wx.EVT_CLOSE, self.OnClose)
statusbar = self.CreateStatusBar()
statusbar.SetFont(monospace(10))
# create menubar and listen to events:
self.SetMenuBar(self._CreateMenu())
self.session = None
self._ResetSession()
def _ResetSession(self, session=None):
"""Associate a new Session with this frame."""
# start new session if necessary
if session is None:
session = Session(self.madx_units)
# remove existing associations
if self.session:
self.session.close()
CloseMDIChildren(self)
# add new associations
self._NewLogTab()
self.session = session
self.env['session'] = session
self.env['madx'] = session.madx
self.env['libmadx'] = session.libmadx
self.env.pop('segment', None)
self.env.pop('sequence', None)
self.env.pop('elements', None)
self.env.pop('twiss', None)
threading.Thread(target=self._read_stream,
args=(session.remote_process.stdout,)).start()
self.hook.reset()
@Cancellable
def _LoadFile(self, event=None):
self._ConfirmResetSession()
wildcards = [("Model files", "*.cpymad.yml"),
("MAD-X files", "*.madx", "*.str", "*.seq"),
("All files", "*")]
with OpenDialog(self, "Open model", wildcards) as dlg:
dlg.Directory = self.app.conf.get('model_path', '.')
ShowModal(dlg)
name = dlg.Filename
repo = FileResource(dlg.Directory)
session = Session.load(self.madx_units, repo, name)
self._ResetSession(session)
if not session.madx.sequences:
return
with Dialog(self) as dialog:
widget = ModelWidget(dialog, session)
data = widget.Query(session.data)
session.init_segment(data)
segment = session.segment
self.env['segment'] = segment
self.env['sequence'] = segment.sequence
self.env['elements'] = segment.elements
self.env['twiss'] = segment.sequence.twiss_table
TwissView.create(session, self, basename='env')
@Cancellable
def _SaveModel(self, event=None):
data = self.session.data
pathes = data.get('init-files', [])
if pathes:
folder = os.path.dirname(pathes[0])
else:
folder = self.app.conf.get('model_path', '.')
wildcards = [("cpymad model files", "*.cpymad.yml"),
("All files", "*")]
with SaveDialog(self, 'Save model', wildcards) as dlg:
dlg.Directory = folder
ShowModal(dlg)
path = dlg.Path
self.session.save(path)
def _CanSaveModel(self, event):
event.Enable(bool(self.session.segment))
@Cancellable
def _EditTwiss(self, event=None):
segment = self.GetActiveFigurePanel().view.segment
utool = self.madx_units
with Dialog(self) as dialog:
widget = TwissWidget(dialog, session=self.session)
segment.twiss_args = widget.Query(segment.twiss_args)
@Cancellable
def _SetBeam(self, event=None):
segment = self.GetActiveFigurePanel().view.segment
with Dialog(self) as dialog:
widget = BeamWidget(dialog, session=self.session)
segment.beam = widget.Query(segment.beam)
def _ShowIndicators(self, event):
panel = self.GetActiveFigurePanel()
segment = panel.view.segment
segment.show_element_indicators = event.Checked()
def _UpdateShowIndicators(self, event):
segment = self.GetActiveFigurePanel().view.segment
event.Check(bool(segment.show_element_indicators))
def _ConfirmResetSession(self):
"""Prompt the user to confirm resetting the current session."""
if not self.session.segment:
return False
question = 'Open new MAD-X session? Unsaved changes will be lost.'
answer = wx.MessageBox(
question, 'Reset session',
wx.OK | wx.CANCEL | wx.ICON_QUESTION,
parent=self)
if answer == wx.OK:
return True
raise CancelAction
def _CreateMenu(self):
"""Create a menubar."""
# TODO: this needs to be done more dynamically. E.g. use resource
# files and/or a plugin system to add/enable/disable menu elements.
MenuItem = menu.Item
Menu = menu.Menu
Separator = menu.Separator
menubar = self.menubar = wx.MenuBar()
menu.extend(self, menubar, [
Menu('&Session', [
MenuItem('&New session window\tCtrl+N',
'Open a new session window',
self.OnNewWindow),
MenuItem('&Python shell\tCtrl+P',
'Open a tab with a python shell',
self._NewCommandTab),
Separator,
MenuItem('&Open model\tCtrl+O',
'Load model or open new model from a MAD-X file.',
self._LoadFile),
MenuItem('&Save model as..\tCtrl+S',
'Save the current model (beam + twiss) to a file',
self._SaveModel,
self._CanSaveModel),
Separator,
MenuItem('&Reset session',
'Clear the MAD-X session state.',
lambda _: self._ResetSession()),
Separator,
MenuItem('&Close',
'Close window',
self.OnQuit,
id=wx.ID_CLOSE),
]),
Menu('&View', [
MenuItem('&Envelope',
'Open new tab with beam envelopes.',
lambda _: TwissView.create(self.session,
self, basename='env')),
MenuItem('&Position',
'Open new tab with beam position.',
lambda _: TwissView.create(self.session,
self, basename='pos')),
]),
Menu('&Manage', [
MenuItem('&Initial conditions',
'Add/remove/edit TWISS initial conditions.',
self._EditTwiss),
MenuItem('&Beam',
'Set beam.',
self._SetBeam),
Separator,
MenuItem('Show &element indicators',
'Show indicators for beam line elements.',
self._ShowIndicators,
self._UpdateShowIndicators,
wx.ITEM_CHECK),
]),
Menu('&Help', [
MenuItem('&About',
'Show about dialog.',
lambda _: show_about_dialog(self)),
]),
])
# Create menu items
self.hook.menu(self, menubar)
self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateMenu, menubar)
self._IsEnabledTop = {self.ViewMenuIndex: True,
self.TabMenuIndex: True}
return menubar
@property
def ViewMenuIndex(self):
return 1
@property
def TabMenuIndex(self):
return 2
def OnNewWindow(self, event):
"""Open a new frame."""
self.__class__(self.app)
def AddView(self, view, title):
"""Add new notebook tab for the view."""
# TODO: remove this method in favor of a event based approach?
child = MDIChildFrame(self, -1, title)
panel = FigurePanel(child, view)
sizer = wx.BoxSizer(wx.VERTICAL)
sizer.Add(panel, 1, wx.EXPAND)
child.SetSizer(sizer)
def OnPageClose(event):
self.views.remove(view)
event.Skip()
child.Bind(wx.EVT_CLOSE, OnPageClose)
ShowMDIChildFrame(child)
view.plot()
self.views.append(view)
return panel
def GetActivePanel(self):
"""Return the Panel which is currently active."""
if self.GetActiveChild():
return self.GetActiveChild().GetChildren()[0]
def GetActiveFigurePanel(self):
"""Return the FigurePanel which is currently active or None."""
panel = self.GetActivePanel()
if isinstance(panel, FigurePanel):
return panel
return None
def OnClose(self, event):
# We want to terminate the remote session, otherwise _read_stream
# may hang:
try:
self.session.close()
except IOError:
# The connection may already be terminated in case MAD-X crashed.
pass
CloseMDIChildren(self)
event.Skip()
def OnLogTabClose(self, event):
"""Prevent the command tab from closing, if other tabs are open."""
if self.views:
event.Veto()
else:
self.Close()
def OnQuit(self, event):
"""Close the window."""
self.Close()
def OnUpdateMenu(self, event):
if not self.session.madx:
return
enable_view = bool(self.session.madx.sequences)
# we only want to call EnableTop() if the state is actually
# different from before, since otherwise this will cause very
# irritating flickering on windows. Because menubar.IsEnabledTop is
# bugged on windows, we need to keep track ourself:
# if enable != self.menubar.IsEnabledTop(idx):
view_menu_index = self.ViewMenuIndex
if enable_view != self._IsEnabledTop[view_menu_index]:
self.menubar.EnableTop(view_menu_index, enable_view)
self._IsEnabledTop[view_menu_index] = enable_view
# Enable/Disable &Tab menu
enable_tab = bool(self.GetActiveFigurePanel())
tab_menu_index = self.TabMenuIndex
if enable_tab != self._IsEnabledTop[tab_menu_index]:
self.menubar.EnableTop(tab_menu_index, enable_tab)
self._IsEnabledTop[tab_menu_index] = enable_tab
event.Skip()
def _NewCommandTab(self, event=None):
"""Open a new command tab."""
child = MDIChildFrame(self, -1, "Command")
crust = Shell(child, locals=self.env)
sizer = wx.BoxSizer(wx.VERTICAL)
sizer.Add(crust, 1, wx.EXPAND)
child.SetSizer(sizer)
ShowMDIChildFrame(child)
def _NewLogTab(self):
child = MDIChildFrame(self, -1, "Log")
# Create a tab for logging
textctrl = wx.TextCtrl(child, wx.ID_ANY,
style=wx.TE_MULTILINE|wx.TE_READONLY)
textctrl.SetFont(monospace(10))
sizer = wx.BoxSizer(wx.VERTICAL)
sizer.Add(textctrl, 1, wx.EXPAND)
child.SetSizer(sizer)
child.Bind(wx.EVT_CLOSE, self.OnLogTabClose)
ShowMDIChildFrame(child)
self._log_ctrl = textctrl
self._basicConfig(logging.INFO,
'%(asctime)s %(levelname)s %(name)s: %(message)s',
'%H:%M:%S')
def _basicConfig(self, level, fmt, datefmt=None):
"""Configure logging."""
stream = TextCtrlStream(self._log_ctrl)
root = logging.RootLogger(level)
manager = logging.Manager(root)
formatter = logging.Formatter(fmt, datefmt)
handler = logging.StreamHandler(stream)
handler.setFormatter(formatter)
root.addHandler(handler)
# store member variables:
self._log_stream = stream
self._log_manager = manager
def getLogger(self, name='root'):
return self._log_manager.getLogger(name)
def _read_stream(self, stream):
# The file iterator seems to be buffered:
for line in iter(stream.readline, b''):
try:
self._log_stream.write(line)
except:
break
