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 : from typing import TYPE_CHECKING
       7              : 
       8            1 : from PySide6.QtWidgets import QHBoxLayout, QMenu, QPushButton, QStyle
       9              : 
      10            1 : from pairinteraction_gui.calculate.calculate_two_atoms import ParametersTwoAtoms, calculate_two_atoms
      11            1 : from pairinteraction_gui.config import (
      12              :     BasisConfigTwoAtoms,
      13              :     CalculationConfig,
      14              :     KetConfigTwoAtoms,
      15              :     SystemConfigTwoAtoms,
      16              : )
      17            1 : from pairinteraction_gui.page.base_page import CalculationPage
      18              : 
      19              : if TYPE_CHECKING:
      20              :     from pairinteraction_gui.calculate.calculate_two_atoms import ResultsTwoAtoms
      21              : 
      22            1 : logger = logging.getLogger(__name__)
      23              : 
      24              : 
      25            1 : class TwoAtomsPage(CalculationPage):
      26              :     """Page for configuring and analyzing pair systems."""
      27              : 
      28            1 :     title = "Two\nAtoms"
      29            1 :     tooltip = "Configure and analyze pair systems"
      30              : 
      31            1 :     def setupWidget(self) -> None:
      32            1 :         super().setupWidget()
      33              : 
      34              :         # all attributes of instance BaseConfig will be added to the toolbox in postSetupWidget
      35            1 :         self.ket_config = KetConfigTwoAtoms(self)
      36            1 :         self.basis_config = BasisConfigTwoAtoms(self)
      37            1 :         self.system_config = SystemConfigTwoAtoms(self)
      38            1 :         self.calculation_config = CalculationConfig(self)
      39              : 
      40              :         # Set some better defaults for the two atoms page
      41            1 :         self.calculation_config.energy_range.setValues(-0.5, 0.5)
      42            1 :         self.calculation_config.energy_range.setChecked(False)
      43              : 
      44            1 :         self.ket_config.signal_species_changed.connect(self.basis_config.on_species_changed)
      45            1 :         self.ket_config.signal_species_changed.connect(self.plotwidget.clear)
      46              : 
      47              :         # Create a fit button with menu
      48            1 :         fit_button = QPushButton("Fit")
      49            1 :         fit_button.setIcon(self.style().standardIcon(QStyle.StandardPixmap.SP_DialogResetButton))
      50            1 :         fit_menu = QMenu(self)
      51            1 :         fit_menu.addAction("Fit c3", lambda: self.plotwidget.fit("c3"))
      52            1 :         fit_menu.addAction("Fit c6", lambda: self.plotwidget.fit("c6"))
      53            1 :         fit_menu.addAction("Fit c3+c6", lambda: self.plotwidget.fit("c3+c6"))
      54            1 :         fit_button.setMenu(fit_menu)
      55            1 :         fit_button.setFixedHeight(50)
      56            1 :         self.findChild(QHBoxLayout, name="bottomLayout").insertWidget(1, fit_button, stretch=1)
      57              : 
      58            1 :     def before_calculate(self) -> None:
      59            0 :         self.basis_config.clear_basis_pair_label()
      60            0 :         return super().before_calculate()
      61              : 
      62            1 :     def calculate(self) -> tuple[ParametersTwoAtoms, ResultsTwoAtoms]:
      63            1 :         parameters = ParametersTwoAtoms.from_page(self)
      64            1 :         results = calculate_two_atoms(parameters)
      65            1 :         return parameters, results
      66              : 
      67            1 :     def update_plot(self, parameters: ParametersTwoAtoms, results: ResultsTwoAtoms) -> None:  # type: ignore[override]
      68            0 :         super().update_plot(parameters, results)
      69              : 
      70            0 :         if results.basis_0_label is not None:
      71            0 :             self.basis_config.update_basis_pair_label(results.basis_0_label)
      72              : 
      73            1 :     def _get_export_replacements(self) -> dict[str, str]:
      74            1 :         parameters = ParametersTwoAtoms.from_page(self)
      75            1 :         return parameters.to_replacement_dict()
      76              : 
      77            1 :     def _get_export_notebook_template_name(self) -> str:
      78            1 :         ranges = self.system_config.get_ranges_dict()
      79            1 :         if all(v[0] == v[-1] for k, v in ranges.items() if k in ["Ex", "Ey", "Ez", "Bx", "By", "Bz"]):
      80            1 :             return "two_atoms.ipynb"
      81            0 :         return "two_atoms_variable_fields.ipynb"