Line data Source code
1 : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
2 : # SPDX-License-Identifier: LGPL-3.0-or-later
3 1 : from __future__ import annotations
4 :
5 1 : import logging
6 1 : import time
7 1 : from pathlib import Path
8 1 : from typing import TYPE_CHECKING, Any
9 :
10 1 : import nbformat
11 1 : from nbconvert import PythonExporter
12 1 : from PySide6.QtCore import Qt
13 1 : from PySide6.QtGui import QIcon, QPixmap
14 1 : from PySide6.QtWidgets import (
15 : QFileDialog,
16 : QHBoxLayout,
17 : QMenu,
18 : QPushButton,
19 : QStyle,
20 : QToolBox,
21 : )
22 :
23 1 : import pairinteraction
24 1 : from pairinteraction_gui.config import BaseConfig
25 1 : from pairinteraction_gui.plotwidget.plotwidget import PlotEnergies
26 1 : from pairinteraction_gui.qobjects import NamedStackedWidget, WidgetV, show_status_tip
27 1 : from pairinteraction_gui.worker import MultiThreadWorker
28 :
29 : if TYPE_CHECKING:
30 : from collections.abc import Callable
31 :
32 : from PySide6.QtGui import QHideEvent, QShowEvent
33 :
34 : from pairinteraction_gui.calculate.calculate_base import Parameters, Results
35 : from pairinteraction_gui.config.calculation_config import CalculationConfig
36 : from pairinteraction_gui.config.ket_config import KetConfig
37 : from pairinteraction_gui.plotwidget.plotwidget import PlotWidget
38 :
39 1 : logger = logging.getLogger(__name__)
40 :
41 :
42 1 : class BasePage(WidgetV):
43 : """Base class for all pages in this application."""
44 :
45 1 : margin = (20, 20, 20, 20)
46 1 : spacing = 15
47 :
48 1 : title: str
49 1 : tooltip: str
50 1 : icon_path: Path | None = None
51 :
52 1 : def showEvent(self, event: QShowEvent) -> None:
53 : """Show event."""
54 1 : super().showEvent(event)
55 1 : self.window().setWindowTitle(
56 : f"PairInteraction v{pairinteraction.__version__} - " + self.title.replace("\n", " ")
57 : )
58 :
59 :
60 1 : class SimulationPage(BasePage):
61 : """Base class for all simulation pages in this application."""
62 :
63 1 : ket_config: KetConfig
64 :
65 1 : plotwidget: PlotWidget
66 :
67 1 : def setupWidget(self) -> None:
68 1 : self.toolbox = QToolBox()
69 :
70 : # Create a dummy icon to allow adjusting the height of the toolbox tabs,
71 : # see https://stackoverflow.com/questions/48503645/customizing-qtoolbox-tab-height
72 1 : px = QPixmap(1, 1)
73 1 : px.fill(Qt.GlobalColor.transparent)
74 1 : self._toolbox_dummy_icon = QIcon(px)
75 :
76 1 : def postSetupWidget(self) -> None:
77 1 : for attr in self.__dict__.values():
78 1 : if isinstance(attr, BaseConfig):
79 1 : self.toolbox.addItem(attr, self._toolbox_dummy_icon, attr.title)
80 :
81 1 : for i, species_combo in enumerate(self.ket_config.species_combo_list):
82 1 : self.ket_config.signal_species_changed.emit(i, species_combo.currentText())
83 :
84 1 : def showEvent(self, event: QShowEvent) -> None:
85 1 : super().showEvent(event)
86 1 : self.window().dockwidget.setWidget(self.toolbox)
87 1 : self.window().dockwidget.setVisible(True)
88 1 : self.toolbox.show()
89 :
90 1 : def hideEvent(self, event: QHideEvent) -> None:
91 1 : super().hideEvent(event)
92 1 : self.window().dockwidget.setVisible(False)
93 :
94 :
95 1 : class CalculationPage(SimulationPage):
96 : """Base class for all pages with a calculation button."""
97 :
98 1 : plotwidget: PlotEnergies
99 1 : _calculation_finished = False
100 1 : _plot_finished = False
101 :
102 1 : def setupWidget(self) -> None:
103 1 : super().setupWidget()
104 :
105 : # Plot Panel
106 1 : self.plotwidget = self._create_plot_widget()
107 1 : self.layout().addWidget(self.plotwidget)
108 :
109 : # Control panel below the plot
110 1 : bottom_layout = QHBoxLayout()
111 1 : bottom_layout.setObjectName("bottomLayout")
112 :
113 : # Calculate/Abort stacked buttons
114 1 : self.calculate_and_abort = NamedStackedWidget[QPushButton](self)
115 :
116 1 : calculate_button = QPushButton("Calculate")
117 1 : calculate_button.setIcon(self.style().standardIcon(QStyle.StandardPixmap.SP_BrowserReload))
118 1 : calculate_button.clicked.connect(self.calculate_clicked)
119 1 : self.calculate_and_abort.addNamedWidget(calculate_button, "Calculate")
120 :
121 1 : abort_button = QPushButton("Abort")
122 1 : abort_button.setIcon(self.style().standardIcon(QStyle.StandardPixmap.SP_BrowserStop))
123 1 : abort_button.clicked.connect(self.abort_clicked)
124 1 : self.calculate_and_abort.addNamedWidget(abort_button, "Abort")
125 :
126 1 : self.calculate_and_abort.setFixedHeight(50)
127 1 : bottom_layout.addWidget(self.calculate_and_abort, stretch=2)
128 :
129 : # Create export button with menu
130 1 : export_button = QPushButton("Export")
131 1 : export_button.setObjectName("Export")
132 1 : export_button.setIcon(self.style().standardIcon(QStyle.StandardPixmap.SP_DialogSaveButton))
133 1 : export_menu = QMenu(self)
134 1 : for label, handler in self._get_export_actions():
135 1 : export_menu.addAction(label, handler)
136 1 : export_button.setMenu(export_menu)
137 1 : export_button.setFixedHeight(50)
138 1 : bottom_layout.addWidget(export_button, stretch=1)
139 :
140 1 : self.layout().addLayout(bottom_layout)
141 :
142 1 : def calculate_clicked(self) -> None:
143 1 : self._calculation_finished = False
144 1 : self._plot_finished = False
145 1 : self.before_calculate()
146 :
147 1 : def update_plot(
148 : parameters_and_results: tuple[Parameters[Any], Results],
149 : ) -> None:
150 1 : worker_plot = MultiThreadWorker(self.update_plot, *parameters_and_results)
151 1 : worker_plot.signals.progress.connect(lambda message: show_status_tip(self, message))
152 1 : worker_plot.signals.finished.connect(lambda _: setattr(self, "_plot_finished", True))
153 1 : worker_plot.start()
154 :
155 1 : worker = MultiThreadWorker(self.calculate)
156 1 : if hasattr(self, "calculation_config"):
157 1 : calculation_config: CalculationConfig = self.calculation_config
158 1 : number_of_steps = calculation_config.steps.value()
159 1 : worker.enable_busy_indicator(self.plotwidget, add_progress_label=True, number_of_steps=number_of_steps)
160 : else:
161 0 : worker.enable_busy_indicator(self.plotwidget)
162 1 : worker.signals.progress.connect(lambda message: show_status_tip(self, message))
163 1 : worker.signals.result.connect(update_plot)
164 1 : worker.signals.finished.connect(self.after_calculate)
165 1 : worker.signals.finished.connect(lambda _: setattr(self, "_calculation_finished", True))
166 1 : worker.start()
167 :
168 1 : def before_calculate(self) -> None:
169 1 : show_status_tip(self, "Calculating... Please wait.", logger=logger)
170 1 : self.calculate_and_abort.setCurrentNamedWidget("Abort")
171 1 : self.plotwidget.clear()
172 :
173 1 : self._start_time = time.perf_counter()
174 :
175 1 : def after_calculate(self, status: str) -> None:
176 1 : time_needed = time.perf_counter() - self._start_time
177 1 : show_status_tip(self, f"{status} after {time_needed:.2f} seconds.", logger=logger)
178 1 : self.calculate_and_abort.setCurrentNamedWidget("Calculate")
179 :
180 1 : def calculate(self) -> tuple[Parameters[Any], Results]:
181 0 : raise NotImplementedError("Subclasses must implement this method")
182 :
183 1 : def update_plot(self, parameters: Parameters[Any], results: Results) -> None:
184 1 : self.plotwidget.canvas.draw() # draw once before, to avoid displaying artifacts during plotting
185 1 : self.plotwidget.plot(parameters, results)
186 1 : self.plotwidget.setup_annotations(parameters, results)
187 1 : self.plotwidget.canvas.draw()
188 1 : self.plotwidget.navigation_toolbar.reset_home_view()
189 1 : show_status_tip(self, "Finished updating plot. Tip: Click on the plot to see state information.", logger=logger)
190 :
191 1 : def export_png(self) -> None:
192 : """Export the current plot as a PNG file."""
193 0 : logger.debug("Exporting results as PNG...")
194 :
195 0 : filename, _ = QFileDialog.getSaveFileName(self, "Save Plot", "", "PNG Files (*.png)")
196 :
197 0 : if filename:
198 0 : filename = filename.removesuffix(".png") + ".png"
199 0 : self.plotwidget.canvas.fig.savefig(filename, dpi=300, facecolor="white", edgecolor="none")
200 0 : logger.info("Plot saved as %s", filename)
201 :
202 1 : def _create_python_code(self) -> str:
203 1 : template_path = Path(__file__).parent.parent / "export_templates" / self._get_export_notebook_template_name()
204 1 : with Path(template_path).open() as f:
205 1 : notebook = nbformat.read(f, as_version=4)
206 :
207 1 : exporter = PythonExporter(exclude_output_prompt=True, exclude_input_prompt=True)
208 1 : content, _ = exporter.from_notebook_node(notebook)
209 :
210 1 : replacements = self._get_export_replacements()
211 1 : for key, value in replacements.items():
212 1 : content = content.replace(key, str(value))
213 :
214 1 : return content
215 :
216 1 : def export_python(self) -> None:
217 : """Export the current calculation as a Python script."""
218 0 : logger.debug("Exporting results as Python script...")
219 0 : filename, _ = QFileDialog.getSaveFileName(self, "Save Python Script", "", "Python Files (*.py)")
220 0 : if filename:
221 0 : filename = filename.removesuffix(".py") + ".py"
222 0 : content = self._create_python_code()
223 0 : with Path(filename).open("w") as f:
224 0 : f.write(content)
225 0 : logger.info("Python script saved as %s", filename)
226 :
227 1 : def export_notebook(self) -> None:
228 : """Export the current calculation as a Jupyter notebook."""
229 0 : logger.debug("Exporting results as Jupyter notebook...")
230 :
231 0 : filename, _ = QFileDialog.getSaveFileName(self, "Save Jupyter Notebook", "", "Jupyter Notebooks (*.ipynb)")
232 :
233 0 : if filename:
234 0 : filename = filename.removesuffix(".ipynb") + ".ipynb"
235 :
236 0 : template_path = (
237 : Path(__file__).parent.parent / "export_templates" / self._get_export_notebook_template_name()
238 : )
239 0 : with Path(template_path).open() as f:
240 0 : notebook = nbformat.read(f, as_version=4)
241 :
242 0 : replacements = self._get_export_replacements()
243 0 : for cell in notebook.cells:
244 0 : if cell.cell_type == "code":
245 0 : source = cell.source
246 0 : for key, value in replacements.items():
247 0 : source = source.replace(key, str(value))
248 0 : cell.source = source
249 :
250 0 : nbformat.write(notebook, filename)
251 :
252 0 : logger.info("Jupyter notebook saved as %s", filename)
253 :
254 1 : def _get_export_notebook_template_name(self) -> str:
255 0 : raise NotImplementedError("Subclasses must implement this method")
256 :
257 1 : def _get_export_replacements(self) -> dict[str, str]:
258 : # Override this method in subclasses to provide specific replacements for the export
259 0 : return {}
260 :
261 1 : def abort_clicked(self) -> None:
262 : """Handle abort button click."""
263 0 : logger.debug("Aborting calculation.")
264 0 : MultiThreadWorker.terminate_all()
265 0 : self.after_calculate("Calculation aborted.")
266 :
267 1 : def _create_plot_widget(self) -> PlotEnergies:
268 1 : return PlotEnergies(self)
269 :
270 1 : def _get_export_actions(self) -> list[tuple[str, Callable[[], None]]]:
271 1 : return [
272 : ("Export as PNG", self.export_png),
273 : ("Export as Python script", self.export_python),
274 : ("Export as Jupyter notebook", self.export_notebook),
275 : ]
|
