def __init__(self, parent=None):
QtGui.QWidget.__init__(self)
self.ui = Ui_ELE()
self.ui.setupUi(self)
self.parent = parent
self.fileExtension = '.ele'
self._setup_widgets()
self.update_widget_state()
self.process = QtCore.QProcess(self)
self.process.readyReadStandardOutput.connect(self._process_stdout)
self.process.readyReadStandardError.connect(self._process_stderr)
self.process.started.connect(self._process_started)
self.process.finished.connect(self._process_finished)
# states: 'summary' or 'full'
self.status_mode = 'summary'
self.summary_html = ''
self.container = self
class RbEle(QtGui.QWidget):
acceptsFileTypes = []
defaultTitle = 'Elegant'
task = 'Run an Elegant simulation'
category = 'simulations'
ERROR_FILE_NAME = 'elegant_errors.txt'
OUTPUT_FILE_NAME = 'elegant_output.txt'
ELEGANT_BASE_NAME = 'elegantSimulation'
ELEGANT_TEMPLATE = '''
&run_setup
lattice = "{latticeFileName}",
use_beamline = {beamlineName},
default_order = 2,
p_central_mev = {momentum},
output = %s.out,
centroid = %s.cen,
sigma = %s.sig,
parameters = %s.param,
random_number_seed = 987654321,
combine_bunch_statistics = 0,
concat_order = 2,
tracking_updates = 1,
echo_lattice = 0,
print_statistics = 1,
&end
&run_control
n_steps = 1,
reset_rf_for_each_step = 1,
&end
&twiss_output
matched = 0,
concat_order = 3,
beta_x = 5, alpha_x = 0,
beta_y = 5, alpha_y = 0,
output_at_each_step = 1,
statistics = 1,
concat_order = 3,
filename = %s.twi,
&end
&sdds_beam
input = "{bunchFileName}",
&end
&track
&end
&stop
&end
'''
def __init__(self, parent=None):
QtGui.QWidget.__init__(self)
self.ui = Ui_ELE()
self.ui.setupUi(self)
self.parent = parent
self.fileExtension = '.ele'
self._setup_widgets()
self.update_widget_state()
self.process = QtCore.QProcess(self)
self.process.readyReadStandardOutput.connect(self._process_stdout)
self.process.readyReadStandardError.connect(self._process_stderr)
self.process.started.connect(self._process_started)
self.process.finished.connect(self._process_finished)
# states: 'summary' or 'full'
self.status_mode = 'summary'
self.summary_html = ''
self.progressIndex = 0
self.beamlineNames = []
self.container = self
# This tab only pulls together data from other sources.
# It has no information to save. exportToFile() creates
# an empty file to signal to RadTrack to create a new
# tab when a project is opened that had this tab.
def exportToFile(self, fileName):
with open(fileName, 'w'):
pass
def importFile(self, fileName):
pass
def append_status(self, line):
"""Formats and appends the line to the status field"""
status = self.ui.simulationStatusTextEdit
line += '\n'
if self._is_error_text(line):
html = '<strong>{}</strong>'.format(cgi.escape(line))
else:
html = cgi.escape(line)
html = re.sub(r'\n', '<br>', html)
status.moveCursor(QtGui.QTextCursor.End)
status.insertHtml(html)
scroll_bar = status.verticalScrollBar()
scroll_bar.setValue(scroll_bar.maximum())
self.summary_html += html
def clear_beam_lines(self):
self.ui.beamLineComboBox.clear()
def get_tab_by_name(self, tab_name):
"""Returns an app tab widget by text value"""
tab = self.parent.tabWidget
for i in range(tab.count()):
if tab.tabText(i) == tab_name:
return getRealWidget(tab.widget(i))
return None
def get_tab_names_for_type(self, tab_type):
"""Returns a list of app tab names with the specified type"""
tab_names = []
for i in range(self.parent.tabWidget.count()):
name = self.parent.tabWidget.tabText(i)
if type(self.get_tab_by_name(name)) == tab_type:
tab_names.append(name)
return tab_names
def session_file(self, file_name=None, suffix=None):
"""Creates a file path in the sessionDirectory."""
if not file_name:
file_name = '{}.{}'.format(self.ELEGANT_BASE_NAME, suffix)
return os.path.join(self.parent.sessionDirectory, file_name)
def set_beam_lines(self, items):
"""Load the Beam Line combo with items"""
self.clear_beam_lines()
self.ui.beamLineComboBox.addItems(items)
self.ui.beamLineComboBox.setCurrentIndex(
self.ui.beamLineComboBox.count() - 1)
def show_warning_box(self, message):
QtGui.QMessageBox.warning(self, 'RadTrack', message)
def update_sources_from_tabs(self):
"""Called from RbGlobal when the tab state changes. Resync the
bunch source and beam line source combo items."""
self.bunch_source_manager.update_sources_from_tabs()
self.beam_line_source_manager.update_sources_from_tabs()
self.beam_line_source_manager._beam_line_source_changed()
def update_widget_state(self):
"""Set the enabled status on widgets based on current selections"""
has_beam_line_source = self.beam_line_source_manager.has_selection()
self.ui.beamLineLabel.setEnabled(has_beam_line_source)
self.ui.beamLineComboBox.setEnabled(has_beam_line_source)
show_momentum = self.bunch_source_manager.is_momentum_required()
self.ui.momentumLabel.setVisible(show_momentum)
self.ui.momentumLineEdit.setVisible(show_momentum)
enable_button = self.bunch_source_manager.has_selection() \
and self.beam_line_source_manager.has_selection()
if show_momentum and not self.ui.momentumLineEdit.text():
enable_button = False
self.ui.simulateButton.setEnabled(enable_button)
def validate_momentum(self):
"""Ensure the momentum value is valid"""
if self.bunch_source_manager.is_momentum_required():
try:
momentum = float(self.ui.momentumLineEdit.text())
except ValueError:
try:
momentum = convertUnitsStringToNumber(
self.ui.momentumLineEdit.text(), 'MeV')
except ValueError:
self.show_warning_box('Unable to parse momentum')
self.ui.momentumLineEdit.setFocus()
return None
else:
bunchTab = self.bunch_source_manager.get_tab_widget()
bunchTab.generateBunch() # make sure any new settings get read
try:
momentum = convertUnitsNumber(
bunchTab.myBunch.getDesignMomentumEV(), 'eV', 'MeV')
except ValueError:
self.show_warning_box(
'Invalid momentum value on Bunch Tab')
return None
return momentum
def _abort_simulation(self):
self.process.kill()
def _add_menu_actions(self, menu, file_name, tab_type, name):
"""Adds the context menu actions for the specified tab_type"""
for tab_name in self.get_tab_names_for_type(tab_type):
menu.addAction(
'Open in {} tab'.format(tab_name),
lambda: self._load_tab(tab_name, file_name))
menu.addAction(
'Open in new {} tab'.format(name),
lambda: self._new_tab(tab_type, file_name))
def _add_result_file(self, text, file_name):
"""Adds the file entry to the simulation results list"""
if not os.path.isfile(file_name):
cpr('missing result file: {}', file_name)
return
results = self.ui.simulationResultsListWidget
icon = results.style().standardIcon(QtGui.QStyle.SP_FileIcon)
item = QtGui.QListWidgetItem(icon, text)
item.setData(QtCore.Qt.UserRole, file_name)
results.addItem(item)
def _add_result_files(self):
"""Adds output files from lattice elements and elegant template"""
loader = self.beam_line_source_manager.get_lattice_element_loader()
for element in loader.elementDictionary.values():
if element.isBeamline():
continue
for output_parameter in element.outputFileParameters:
i = element.parameterNames.index(output_parameter)
if element.data[i]:
self._add_result_file(
'{}: {}'.format(element.name, type(element).__name__),
self._autocomplete_file_name(element.data[i]))
matches = re.findall(r'\%s\.\w+', self.ELEGANT_TEMPLATE)
for match in matches:
file_name = self._autocomplete_file_name(match)
if os.path.isfile(file_name):
text = self._sdds_description(file_name)
if not text:
text = os.path.basename(file_name)
self._add_result_file(text, file_name)
def _autocomplete_file_name(self, file_name):
"""Removes leading/trailing quotes and autocompletes %s.xxx files"""
file_name = re.sub(r'\'|"', '', file_name)
file_name = re.sub(r'\%s', self.ELEGANT_BASE_NAME, file_name)
if not os.path.isfile(file_name):
file_name = self.session_file(file_name=file_name)
return file_name
def _enable_parameters(self, is_enabled):
"""Enable or disable all the parameter fields"""
for w in (self.ui.bunchSourceLabel,
self.ui.bunchSourceComboBox,
self.ui.beamLineSourceLabel,
self.ui.beamLineSourceComboBox,
self.ui.beamLineLabel,
self.ui.beamLineComboBox,
self.ui.momentumLabel,
self.ui.momentumLineEdit,
self.ui.simulateButton):
w.setEnabled(is_enabled)
def _enable_status_and_results(self):
"""Enable the simulation status and results"""
self.ui.simulationStatusTextEdit.clear()
self.status_mode = 'summary'
self.summary_html = ''
self.ui.simulationResultsListWidget.clear()
for w in (self.ui.simulationStatusLabel,
self.ui.simulationStatusTextEdit,
self.ui.simulationResultsLabel,
self.ui.simulationResultsListWidget):
w.setEnabled(True)
self._enable_parameters(False)
def _is_bunch_file(self, file_name):
if re.search('\.csv$', file_name, re.IGNORECASE):
return True
if not self._is_sdds_file(file_name):
return False
"""Scans the SDDS header for the 6D field names"""
search_columns = ['x', 'xp', 'y', 'yp', 't', 'p']
for line in self._read_sdds_header(file_name):
match = re.search(r'^\&column\s.*?name\=(\w+)', line)
if match:
name = match.group(1).lower()
if name in search_columns:
search_columns.remove(name)
if len(search_columns):
return False
return True
def _is_error_text(self, text):
return re.search(r'^warn|^error|wrong units', text, re.IGNORECASE)
def _is_sdds_file(self, file_name):
"""Returns True if the file has the SDDS header"""
with open(file_name, 'r') as input_file:
if re.search(r'^SDDS', input_file.readline()):
return True
return False
def _load_tab(self, tab_name, file_name):
"""Load a parent tab with data from the specified file"""
target = self.get_tab_by_name(tab_name)
target.importFile(file_name)
self.parent.tabWidget.setCurrentWidget(target)
def _new_tab(self, tab_type, file_name):
"""Create a new parent tab and load the data from the specified file"""
self.parent.newTab(tab_type)
getRealWidget(self.parent.tabWidget.currentWidget()).importFile(file_name)
def _process_finished(self, code, status):
"""Callback when simulation process has finished"""
self._enable_parameters(True)
self.ui.abortButton.setEnabled(False)
self.ui.progressBar.setValue(self.ui.progressBar.maximum())
self.output_file.close()
self.error_file.close()
self.append_status('')
# always add the bunch source and beam line source files
# so they can be opened in the editor tabs to verify them
if not self.bunch_source_manager.is_tab_choice():
file_name = self.bunch_source_manager.get_file_name()
self._add_result_file(os.path.basename(file_name), file_name)
if not self.beam_line_source_manager.is_tab_choice():
file_name = self.beam_line_source_manager.get_file_name()
self._add_result_file(os.path.basename(file_name), file_name)
if status == 0:
if code == 0:
self.append_status('Simulation completed successfully')
self._add_result_files()
else:
self.append_status('Simulation failed with errors')
else:
self.append_status('Simulation was terminated')
def _process_started(self):
"""Callback when simulation process has started"""
self.ui.abortButton.setEnabled(True)
beamlineName=self.ui.beamLineComboBox.currentText()
self.beamlineNames = self.beam_line_source_manager.get_lattice_element_loader().elementDictionary[beamlineName].fullElementNameList()
self.ui.progressBar.setMaximum(len(self.beamlineNames))
self.ui.progressBar.reset()
self.progressIndex = 0
def _process_stderr(self):
"""Callback with simulation stderr text"""
out = str(self.process.readAllStandardError())
for line in out.split("\n"):
self.append_status(line)
self.error_file.write(out)
def _process_stdout(self):
"""Callback with simulation stdout text"""
out = str(self.process.readAllStandardOutput())
for line in out.split("\n"):
if self._is_error_text(line):
self.append_status(line)
# print_statistics lists each element traversed by the beam.
# Use this to track progress of simulation.
if line.startswith('tracking through'):
outputElementName = line.split()[-1]
if re.match('M\d+', outputElementName): # skips generated names of combined magnets
continue
elementDictionary = self.beam_line_source_manager.get_lattice_element_loader().elementDictionary
if outputElementName in elementDictionary:
while outputElementName != self.beamlineNames[self.progressIndex]:
self.progressIndex += 1
self.ui.progressBar.setValue(self.progressIndex + 1)
self.progressIndex += 1
self.output_file.write(out)
def _read_sdds_header(self, file_name):
"""Returns the header lines from a sdds file."""
# doesn't use sdds module, avoids reading in entire file
lines = []
with open(file_name, 'r') as input_file:
line = input_file.readline()
while line != '':
if re.search(r'^\&data ', line):
break
lines.append(line)
if len(lines) > 300:
break
line = input_file.readline()
return lines
def _results_context_menu(self, position):
"""Show the context menu for a result item."""
results = self.ui.simulationResultsListWidget
if results.currentItem():
file_name = results.currentItem().data(
QtCore.Qt.UserRole).toString()
menu = QtGui.QMenu()
ext = os.path.splitext(file_name)[1].strip('.')
for tabType in self.parent.availableTabTypes:
if ext in tabType.acceptsFileTypes:
self._add_menu_actions(
menu, file_name, tabType, tabType.defaultTitle)
menu.exec_(results.mapToGlobal(position))
def _run_simulation(self):
"""Generate the input files and start the Elegant process."""
momentum = self.validate_momentum()
if not momentum:
return
if not os.getenv('RPN_DEFNS', None):
os.environ['RPN_DEFNS'] = resource.filename('defns.rpn')
self.error_file = open(
self.session_file(file_name=self.ERROR_FILE_NAME), 'w')
self.output_file = open(
self.session_file(file_name=self.OUTPUT_FILE_NAME), 'w')
self._enable_status_and_results()
# The replace() command escapes backslashes
# since Elegant interprets \x as a special character.
elegant_input_file = self._write_simulation_input_files(momentum).replace('\\', '\\\\')
self.append_status('Running simulation ...\n')
self.process.start('elegant', [elegant_input_file])
def _sdds_description(self, file_name):
"""Return the sdds file description if present"""
for line in self._read_sdds_header(file_name):
match = re.search(r'^\&description .*?text="(.*?)"', line)
if match:
description = match.group(1)
description = re.sub(r'\-\-input.*', '', description)
description = re.sub(r'(\\|\/).*', '', description)
return description
return None
def _setup_widgets(self):
"""Setup initial widget state"""
self.ui.simulateButton.setIcon(
self.ui.simulateButton.style().standardIcon(
QtGui.QStyle.SP_MediaPlay))
self.ui.simulateButton.clicked.connect(self._run_simulation)
self.ui.abortButton.setIcon(
self.ui.abortButton.style().standardIcon(
QtGui.QStyle.SP_MediaStop))
self.ui.abortButton.clicked.connect(self._abort_simulation)
self.ui.momentumLineEdit.textChanged.connect(self.update_widget_state)
self.ui.abortButton.setEnabled(False)
self.ui.simulationStatusTextEdit.customContextMenuRequested.connect(
self._status_context_menu)
self.ui.simulationResultsListWidget.customContextMenuRequested.connect(
self._results_context_menu)
self.bunch_source_manager = BunchSourceManager(
self, self.ui.bunchSourceComboBox)
self.beam_line_source_manager = BeamLineSourceManager(
self, self.ui.beamLineSourceComboBox)
def _show_full_status(self):
"""Show the full process stdout and stderr text"""
self.status_mode = 'full'
status = self.ui.simulationStatusTextEdit
status.clear()
text = ''
for name in (self.OUTPUT_FILE_NAME, self.ERROR_FILE_NAME):
with open(self.session_file(file_name=name), 'r') as input_file:
text += input_file.read()
text += '\n\n'
status.moveCursor(QtGui.QTextCursor.End)
status.insertPlainText(text)
scroll_bar = status.verticalScrollBar()
scroll_bar.setValue(scroll_bar.maximum())
def _show_summary_status(self):
"""Show the summary status text"""
self.status_mode = 'summary'
self.ui.simulationStatusTextEdit.clear()
self.ui.simulationStatusTextEdit.insertHtml(self.summary_html)
def _status_context_menu(self, position):
"""Show the context menu for the status panel."""
status = self.ui.simulationStatusTextEdit
if status.isEnabled() and self.ui.simulateButton.isEnabled():
menu = QtGui.QMenu()
full = menu.addAction('Full Output', self._show_full_status)
summary = menu.addAction(
'Summary Output', self._show_summary_status)
selected = full if self.status_mode == 'full' else summary
selected.setCheckable(True)
selected.setChecked(True)
menu.exec_(status.mapToGlobal(position))
def _write_simulation_input_files(self, momentum):
"""Generates and writes simulation input files"""
self.append_status('Writing Elegant simulation file ...')
elegant_file_name = self.session_file(suffix='ele')
with open(elegant_file_name, 'w') as output_file:
output_file.write(self.ELEGANT_TEMPLATE.format(
latticeFileName=os.path.basename(self.beam_line_source_manager.get_lattice_file_name()),
beamlineName=self.ui.beamLineComboBox.currentText(),
momentum=str(momentum),
bunchFileName=os.path.basename(self.bunch_source_manager.get_bunch_file_name())
))
for fileName in [self.beam_line_source_manager.get_lattice_file_name(),
self.bunch_source_manager.get_bunch_file_name()]:
if os.path.dirname(fileName) != self.parent.sessionDirectory:
shutil.copy2(fileName, self.parent.sessionDirectory)
return elegant_file_name