mirror of
https://github.com/gwm17/spspy.git
synced 2024-09-20 22:47:27 -04:00
304 lines
13 KiB
Python
304 lines
13 KiB
Python
from .SPSPlot import SPSPlot
|
|
from .SPSReaction import RxnParameters
|
|
from .ui.MPLCanvas import MPLCanvas
|
|
from .ui.ReactionDialog import ReactionDialog
|
|
from .ui.TargetDialog import TargetDialog
|
|
from .ui.ExcitationDialog import ExcitationDialog
|
|
|
|
from PySide6.QtWidgets import QApplication, QWidget, QMainWindow
|
|
from PySide6.QtWidgets import QLabel, QTabWidget, QTableWidget, QTableWidgetItem
|
|
from PySide6.QtWidgets import QHBoxLayout, QVBoxLayout, QGroupBox
|
|
from PySide6.QtWidgets import QPushButton, QRadioButton
|
|
from PySide6.QtWidgets import QDoubleSpinBox
|
|
from PySide6.QtWidgets import QFileDialog
|
|
from PySide6.QtGui import QAction
|
|
|
|
from qdarktheme import load_stylesheet
|
|
from enum import Enum, auto
|
|
import matplotlib as mpl
|
|
import sys
|
|
import pickle
|
|
|
|
DEFAULT_RHO_MIN: float = 69.0
|
|
DEFAULT_RHO_MAX: float = 87.0
|
|
|
|
class PlotType(Enum):
|
|
PLOT_EX = auto()
|
|
PLOT_KE = auto()
|
|
PLOT_Z = auto()
|
|
|
|
class SPSPlotGUI(QMainWindow):
|
|
def __init__(self, parent=None) :
|
|
super().__init__(parent)
|
|
self.setWindowTitle("SPSPlot")
|
|
self.sps = SPSPlot()
|
|
|
|
self.plotLayout = QVBoxLayout()
|
|
self.targetLayout = QVBoxLayout()
|
|
self.generalLayout = QVBoxLayout()
|
|
self.centralWidget = QTabWidget(self)
|
|
self.setCentralWidget(self.centralWidget)
|
|
self.centralWidget.setLayout(self.generalLayout)
|
|
|
|
self.targetTab = QWidget(self.centralWidget)
|
|
self.targetTab.setLayout(self.targetLayout)
|
|
self.plotTab = QWidget(self.centralWidget)
|
|
self.plotTab.setLayout(self.plotLayout)
|
|
self.centralWidget.addTab(self.plotTab, "Plot")
|
|
self.centralWidget.addTab(self.targetTab, "Targets")
|
|
|
|
self.plotType = PlotType.PLOT_EX
|
|
self.create_menus()
|
|
self.create_canvas()
|
|
self.create_inputs()
|
|
self.create_target_table()
|
|
self.update_plot()
|
|
self.show()
|
|
|
|
def create_canvas(self) -> None:
|
|
self.canvas = MPLCanvas(self.plotTab, width=14, height=5, dpi=100)
|
|
self.plotLayout.addWidget(self.canvas, 4)
|
|
|
|
def create_menus(self) -> None:
|
|
self.fileMenu = self.menuBar().addMenu("&File")
|
|
saveAction = QAction("&Save...",self)
|
|
openAction = QAction("&Open...",self)
|
|
self.fileMenu.addAction(saveAction)
|
|
self.fileMenu.addAction(openAction)
|
|
self.fileMenu.addAction("&Exit", self.close)
|
|
saveAction.triggered.connect(self.handle_save)
|
|
openAction.triggered.connect(self.handle_open)
|
|
|
|
self.addMenu = self.menuBar().addMenu("&New")
|
|
newTargetAction = QAction("New target...", self)
|
|
newReactionAction = QAction("New reaction...", self)
|
|
newStateAction = QAction("New state...", self)
|
|
self.addMenu.addAction(newTargetAction)
|
|
self.addMenu.addAction(newReactionAction)
|
|
self.addMenu.addAction(newStateAction)
|
|
newStateAction.triggered.connect(self.handle_new_state)
|
|
newReactionAction.triggered.connect(self.handle_new_reaction)
|
|
newTargetAction.triggered.connect(self.handle_new_target)
|
|
|
|
self.exportMenu = self.menuBar().addMenu("&Export")
|
|
exportLevels = QAction("Export levels to csv...", self)
|
|
self.exportMenu.addAction(exportLevels)
|
|
exportLevels.triggered.connect(self.handle_export_levels)
|
|
|
|
def create_inputs(self) -> None:
|
|
inputLayout = QVBoxLayout()
|
|
self.inputGroupBox = QGroupBox("Adjustable Inputs", self.plotTab)
|
|
|
|
self.spsGroupBox = QGroupBox("SPS Parameters", self.inputGroupBox)
|
|
spsGroupLayout = QHBoxLayout()
|
|
rhoMinLabel = QLabel("<p>ρ<sub>Min<\sub><\p>", self.spsGroupBox)
|
|
self.rhoMinInput = QDoubleSpinBox(self.spsGroupBox)
|
|
self.rhoMinInput.setRange(0.0, 150.0)
|
|
self.rhoMinInput.setValue(DEFAULT_RHO_MIN)
|
|
self.sps.rhoMin = DEFAULT_RHO_MIN
|
|
self.rhoMinInput.setSuffix(" cm")
|
|
rhoMaxLabel = QLabel("<p>ρ<sub>Max<\sub><\p>", self.spsGroupBox)
|
|
self.rhoMaxInput = QDoubleSpinBox(self.spsGroupBox)
|
|
self.rhoMaxInput.setRange(0.0,150.0)
|
|
self.rhoMaxInput.setValue(DEFAULT_RHO_MAX)
|
|
self.sps.rhoMax = DEFAULT_RHO_MAX
|
|
self.rhoMaxInput.setSuffix(" cm")
|
|
bkeLabel = QLabel("<p>E<sub>beam<\sub><\p>", self.spsGroupBox)
|
|
self.bkeInput = QDoubleSpinBox(self.spsGroupBox)
|
|
self.bkeInput.setRange(0.0, 500.0)
|
|
self.bkeInput.setSuffix(" MeV")
|
|
bfieldLabel = QLabel("B", self.spsGroupBox)
|
|
self.bfieldInput = QDoubleSpinBox(self.spsGroupBox)
|
|
self.bfieldInput.setRange(0.0, 17.0)
|
|
self.bfieldInput.setSuffix(" kG")
|
|
angleLabel = QLabel("<p>θ<sub>SPS<\sub><\p>", self.spsGroupBox)
|
|
self.angleInput = QDoubleSpinBox(self.spsGroupBox)
|
|
self.angleInput.setRange(0.0, 180.0)
|
|
self.angleInput.setSuffix(" deg")
|
|
self.runButton = QPushButton("Set", self.spsGroupBox)
|
|
self.runButton.clicked.connect(self.handle_run)
|
|
spsGroupLayout.addWidget(rhoMinLabel, 1)
|
|
spsGroupLayout.addWidget(self.rhoMinInput, 2)
|
|
spsGroupLayout.addWidget(rhoMaxLabel,1 )
|
|
spsGroupLayout.addWidget(self.rhoMaxInput, 2)
|
|
spsGroupLayout.addWidget(bkeLabel, 1)
|
|
spsGroupLayout.addWidget(self.bkeInput, 2)
|
|
spsGroupLayout.addWidget(bfieldLabel, 1)
|
|
spsGroupLayout.addWidget(self.bfieldInput, 2)
|
|
spsGroupLayout.addWidget(angleLabel, 1)
|
|
spsGroupLayout.addWidget(self.angleInput, 2)
|
|
spsGroupLayout.addWidget(self.runButton, 1)
|
|
self.spsGroupBox.setLayout(spsGroupLayout)
|
|
|
|
self.energyButtonGroup = QGroupBox("Labels",self.plotTab)
|
|
buttonLayout = QHBoxLayout()
|
|
self.exButton = QRadioButton("Excitation Energy(MeV)", self.energyButtonGroup)
|
|
self.exButton.toggled.connect(self.handle_ex_switch)
|
|
self.exButton.toggle()
|
|
self.keButton = QRadioButton("Ejectile KE(MeV)", self.energyButtonGroup)
|
|
self.keButton.toggled.connect(self.handle_ke_switch)
|
|
self.zButton = QRadioButton("FocalPlane Z-Shift(cm)", self.energyButtonGroup)
|
|
self.zButton.toggled.connect(self.handle_z_switch)
|
|
buttonLayout.addWidget(self.exButton)
|
|
buttonLayout.addWidget(self.keButton)
|
|
buttonLayout.addWidget(self.zButton)
|
|
self.energyButtonGroup.setLayout(buttonLayout)
|
|
|
|
inputLayout.addWidget(self.spsGroupBox)
|
|
inputLayout.addWidget(self.energyButtonGroup)
|
|
self.inputGroupBox.setLayout(inputLayout)
|
|
|
|
self.plotLayout.addWidget(self.inputGroupBox, 1)
|
|
|
|
def create_target_table(self) -> None:
|
|
self.targetGroup = QGroupBox("Targets", self.targetTab)
|
|
tableLayout = QVBoxLayout()
|
|
self.targetTable = QTableWidget(self.targetGroup)
|
|
self.targetTable.setColumnCount(6)
|
|
self.targetTable.setHorizontalHeaderLabels(["L1 Thickness(ug/cm^2)", "L1 Compound",
|
|
"L2 Thickness(ug/cm^2)", "L2 Compound",
|
|
"L3 Thickness(ug/cm^2)", "Layer3 Compound"])
|
|
tableLayout.addWidget(self.targetTable)
|
|
self.targetGroup.setLayout(tableLayout)
|
|
self.targetLayout.addWidget(self.targetGroup)
|
|
self.targetTable.resizeColumnsToContents()
|
|
self.targetTable.cellDoubleClicked.connect(self.handle_update_target)
|
|
|
|
def handle_save(self) -> None:
|
|
fileName = QFileDialog.getSaveFileName(self, "Save Input","./","SPSPlot Files (*.sps)")
|
|
if fileName[0]:
|
|
with open(fileName[0], "wb") as file:
|
|
pickle.dump(self.sps, file, pickle.HIGHEST_PROTOCOL)
|
|
|
|
def handle_open(self) -> None:
|
|
fileName = QFileDialog.getOpenFileName(self, "Open Input","./","SPSPlot Files (*.sps)")
|
|
if fileName[0]:
|
|
with open(fileName[0], "rb") as file:
|
|
self.sps = pickle.load(file)
|
|
self.update_inputs()
|
|
self.update_plot()
|
|
self.update_target_table()
|
|
|
|
def handle_new_state(self) -> None:
|
|
stDlg = ExcitationDialog(self, self.sps.data.keys())
|
|
stDlg.new_level.connect(self.sps.add_excitation)
|
|
if stDlg.exec():
|
|
self.update_plot()
|
|
|
|
def handle_new_reaction(self) -> None:
|
|
rxnDlg = ReactionDialog(parent=self, targets=self.sps.targets.keys())
|
|
rxnDlg.new_reaction.connect(self.add_reaction)
|
|
rxnDlg.exec()
|
|
|
|
def handle_new_target(self) -> None:
|
|
targDlg = TargetDialog(self)
|
|
targDlg.new_target.connect(self.sps.add_target)
|
|
if targDlg.exec():
|
|
self.update_target_table()
|
|
|
|
def handle_update_target(self, row: int, col: int) -> None:
|
|
targName = self.targetTable.verticalHeaderItem(row).text()
|
|
targDia = TargetDialog(self, target=self.sps.targets[targName])
|
|
targDia.new_target.connect(self.sps.add_target)
|
|
if targDia.exec():
|
|
self.update_target_table()
|
|
self.sps.update_reactions()
|
|
self.update_plot() #in case a reaction is using the target
|
|
|
|
def handle_run(self) -> None:
|
|
self.sps.beamEnergy = self.bkeInput.value()
|
|
self.sps.spsAngle = self.angleInput.value()
|
|
self.sps.magneticField = self.bfieldInput.value()
|
|
self.sps.update_reactions()
|
|
self.sps.rhoMin = self.rhoMinInput.value()
|
|
self.sps.rhoMax = self.rhoMaxInput.value()
|
|
self.update_plot()
|
|
|
|
def handle_ex_switch(self) -> None:
|
|
if self.exButton.isChecked() and self.plotType != PlotType.PLOT_EX:
|
|
self.plotType = PlotType.PLOT_EX
|
|
self.update_plot()
|
|
|
|
def handle_ke_switch(self) -> None:
|
|
if self.keButton.isChecked() and self.plotType != PlotType.PLOT_KE:
|
|
self.plotType = PlotType.PLOT_KE
|
|
self.update_plot()
|
|
|
|
def handle_z_switch(self) -> None:
|
|
if self.zButton.isChecked() and self.plotType != PlotType.PLOT_Z:
|
|
self.plotType = PlotType.PLOT_Z
|
|
self.update_plot()
|
|
|
|
def handle_export_levels(self) -> None:
|
|
fileName = QFileDialog.getSaveFileName(self, "Export Levels to CSV","./","Comma-Separated Values File (*.csv)")
|
|
if fileName[0]:
|
|
self.sps.export_reaction_data(fileName[0])
|
|
|
|
def add_reaction(self, rxnParams: RxnParameters, targName: str) -> None:
|
|
rxnParams.beamEnergy = self.bkeInput.value()
|
|
rxnParams.spsAngle = self.angleInput.value()
|
|
rxnParams.magneticField = self.bfieldInput.value()
|
|
self.sps.add_reaction(rxnParams, targName)
|
|
self.update_plot()
|
|
|
|
def update_plot(self) -> None:
|
|
rhos = []
|
|
exs = []
|
|
kes = []
|
|
zs = []
|
|
rxns = []
|
|
for rxnNumber, rxn in enumerate(self.sps.data.values()):
|
|
for point in rxn.excitations:
|
|
rxns.append(rxnNumber+1)
|
|
rhos.append(point.rho)
|
|
exs.append(point.excitation)
|
|
kes.append(point.kineticEnergy)
|
|
zs.append(point.fpZ)
|
|
self.canvas.axes.cla()
|
|
self.canvas.axes.plot(rhos, rxns, marker="o", linestyle="None")
|
|
for i in range(len(rxns)):
|
|
y = rxns[i]
|
|
x = rhos[i]
|
|
label = ''
|
|
if self.plotType == PlotType.PLOT_EX:
|
|
label = "{:.2f}".format(exs[i])
|
|
elif self.plotType == PlotType.PLOT_KE:
|
|
label = "{:.2f}".format(kes[i])
|
|
elif self.plotType == PlotType.PLOT_Z:
|
|
label = "{:.2f}".format(zs[i])
|
|
self.canvas.axes.annotate(label, (x,y), textcoords="offset points",xytext=(0,10),ha="center",rotation="vertical")
|
|
ylabels = [r.rxn.get_latex_rep() for r in self.sps.data.values()]
|
|
ylabels.append("Reactions")
|
|
self.canvas.axes.set_xlim(self.sps.rhoMin, self.sps.rhoMax)
|
|
self.canvas.axes.set_yticks(range(1,len(self.sps.data)+2))
|
|
self.canvas.axes.set_yticklabels(ylabels)
|
|
self.canvas.axes.set_xlabel(r"$\rho$ (cm)")
|
|
self.canvas.draw()
|
|
|
|
def update_inputs(self):
|
|
self.rhoMinInput.setValue(self.sps.rhoMin)
|
|
self.rhoMaxInput.setValue(self.sps.rhoMax)
|
|
self.bfieldInput.setValue(self.sps.magneticField)
|
|
self.bkeInput.setValue(self.sps.beamEnergy)
|
|
self.angleInput.setValue(self.sps.spsAngle)
|
|
|
|
def update_target_table(self):
|
|
self.targetTable.setRowCount(len(self.sps.targets))
|
|
self.targetTable.setVerticalHeaderLabels(self.sps.targets.keys())
|
|
for row, key in enumerate(self.sps.targets):
|
|
for col, layer in enumerate(self.sps.targets[key].layer_details) :
|
|
self.targetTable.setItem(row, col*2, QTableWidgetItem(str(layer.thickness)))
|
|
self.targetTable.setCellWidget(row, 1+col*2, QLabel(str(layer)))
|
|
self.targetTable.resizeColumnsToContents()
|
|
self.targetTable.resizeRowsToContents()
|
|
|
|
|
|
|
|
def run_spsplot_ui():
|
|
mpl.use("Qt5Agg")
|
|
app = QApplication.instance()
|
|
if not app:
|
|
app = QApplication(sys.argv)
|
|
app.setStyleSheet(load_stylesheet())
|
|
window = SPSPlotGUI()
|
|
sys.exit(app.exec_()) |