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 pairinteraction_gui.calculate.calculate_one_atom import ParametersOneAtom, calculate_one_atom
       9           1 : from pairinteraction_gui.config import (
      10             :     BasisConfigOneAtom,
      11             :     CalculationConfig,
      12             :     KetConfigOneAtom,
      13             :     SystemConfigOneAtom,
      14             : )
      15           1 : from pairinteraction_gui.page.base_page import CalculationPage
      16             : 
      17             : if TYPE_CHECKING:
      18             :     from pairinteraction.state import StateBase
      19             :     from pairinteraction_gui.calculate.calculate_one_atom import ResultsOneAtom
      20             : 
      21           1 : logger = logging.getLogger(__name__)
      22             : 
      23             : 
      24           1 : class OneAtomPage(CalculationPage):
      25             :     """Page for configuring and analyzing single-atom systems."""
      26             : 
      27           1 :     title = "One\nAtom"
      28           1 :     tooltip = "Configure and analyze single-atom systems"
      29             : 
      30           1 :     def setupWidget(self) -> None:
      31           1 :         super().setupWidget()
      32             : 
      33             :         # all attributes of instance BaseConfig will be added to the toolbox in postSetupWidget
      34           1 :         self.ket_config = KetConfigOneAtom(self)
      35           1 :         self.basis_config = BasisConfigOneAtom(self)
      36           1 :         self.system_config = SystemConfigOneAtom(self)
      37           1 :         self.calculation_config = CalculationConfig(self)
      38             : 
      39           1 :         self.ket_config.signal_species_changed.connect(self.basis_config.on_species_changed)
      40           1 :         self.ket_config.signal_species_changed.connect(self.plotwidget.clear)
      41             : 
      42           1 :     def calculate(self) -> tuple[ParametersOneAtom, ResultsOneAtom]:
      43           1 :         parameters = ParametersOneAtom.from_page(self)
      44           1 :         results = calculate_one_atom(parameters)
      45           1 :         return parameters, results
      46             : 
      47           1 :     def _plot_function(self, parameters: ParametersOneAtom, results: ResultsOneAtom) -> None:  # type: ignore[override]
      48           1 :         self.add_short_labels(results)
      49             : 
      50           1 :     def add_short_labels(
      51             :         self,
      52             :         results: ResultsOneAtom,
      53             :     ) -> None:
      54           1 :         if not results.systems:
      55           0 :             return
      56             : 
      57           1 :         ax = self.plotwidget.canvas.ax
      58           1 :         x_lim = ax.get_xlim()
      59           1 :         ax.set_xlim(x_lim[0] - (x_lim[1] - x_lim[0]) * 0.1, x_lim[1])
      60             : 
      61           1 :         used = set()
      62           1 :         states0: list[StateBase[Any]] = results.systems[0].get_eigenbasis().states
      63           1 :         for state, energy in zip(states0, results.energies[0], strict=True):
      64           1 :             ket_label = state.get_label()
      65           1 :             short_label = ket_label[1:-1]
      66           1 :             short_label = short_label.split(":", 1)[-1]
      67           1 :             components = short_label.split(",")
      68           1 :             short_label = ",".join(components[:-1])
      69           1 :             short_label = short_label.split("_", 1)[0]
      70           1 :             if short_label in used:
      71           1 :                 continue
      72           1 :             used.add(short_label)
      73           1 :             self.plotwidget.canvas.ax.text(x_lim[0], energy, short_label, va="center", ha="right")
      74             : 
      75           1 :     def _get_export_notebook_template_name(self) -> str:
      76           1 :         return "one_atom.ipynb"
      77             : 
      78           1 :     def _get_export_replacements(self) -> dict[str, str]:
      79           1 :         parameters = ParametersOneAtom.from_page(self)
      80           1 :         return parameters.to_replacement_dict()