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, Any
       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 Atoms"
      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.number_state_labels.setValue(5)
      42           1 :         self.calculation_config.number_state_labels.setChecked(False)
      43           1 :         self.calculation_config.energy_range.setValues(-0.5, 0.5)
      44           1 :         self.calculation_config.energy_range.setChecked(False)
      45             : 
      46           1 :         self.ket_config.signal_species_changed.connect(self.basis_config.on_species_changed)
      47           1 :         self.ket_config.signal_species_changed.connect(self.plotwidget.clear)
      48             : 
      49             :         # Create a fit button with menu
      50           1 :         fit_button = QPushButton("Fit")
      51           1 :         fit_button.setIcon(self.style().standardIcon(QStyle.StandardPixmap.SP_DialogResetButton))
      52           1 :         fit_menu = QMenu(self)
      53           1 :         fit_menu.addAction("Fit c3", lambda: self.plotwidget.fit("c3"))
      54           1 :         fit_menu.addAction("Fit c6", lambda: self.plotwidget.fit("c6"))
      55           1 :         fit_menu.addAction("Fit c3+c6", lambda: self.plotwidget.fit("c3+c6"))
      56           1 :         fit_button.setMenu(fit_menu)
      57           1 :         fit_button.setFixedHeight(50)
      58           1 :         self.findChild(QHBoxLayout, name="bottomLayout").insertWidget(1, fit_button, stretch=1)
      59             : 
      60           1 :     def before_calculate(self) -> None:
      61           0 :         self.basis_config.clear_basis_pair_label()
      62           0 :         return super().before_calculate()
      63             : 
      64           1 :     def calculate(self) -> tuple[ParametersTwoAtoms, ResultsTwoAtoms]:
      65           1 :         parameters = ParametersTwoAtoms.from_page(self)
      66           1 :         results = calculate_two_atoms(parameters)
      67           1 :         return parameters, results
      68             : 
      69           1 :     def update_plot(self, parameters: ParametersTwoAtoms, results: ResultsTwoAtoms) -> None:  # type: ignore[override]
      70           0 :         super().update_plot(parameters, results)
      71             : 
      72           0 :         if results.basis_0_label is not None:
      73           0 :             self.basis_config.update_basis_pair_label(results.basis_0_label)
      74             : 
      75           1 :     def _get_export_replacements(self) -> dict[str, Any]:
      76           1 :         parameters = ParametersTwoAtoms.from_page(self)
      77           1 :         return parameters.to_replacement_dict()
      78             : 
      79           1 :     def _get_export_notebook_template_name(self) -> str:
      80           1 :         ranges = self.system_config.get_ranges_dict()
      81           1 :         if all(v[0] == v[-1] for k, v in ranges.items() if k in ["Ex", "Ey", "Ez", "Bx", "By", "Bz"]):
      82           1 :             return "two_atoms.ipynb"
      83           0 :         return "two_atoms_variable_fields.ipynb"