Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           0 : import logging
       5           0 : from typing import Any
       6             : 
       7           0 : from pairinteraction_gui.calculate.calculate_one_atom import ParametersOneAtom, ResultsOneAtom, calculate_one_atom
       8           0 : from pairinteraction_gui.config import (
       9             :     BasisConfigOneAtom,
      10             :     KetConfigOneAtom,
      11             :     SystemConfigOneAtom,
      12             : )
      13           0 : from pairinteraction_gui.page.base_page import CalculationPage
      14             : 
      15           0 : logger = logging.getLogger(__name__)
      16             : 
      17             : 
      18           0 : class OneAtomPage(CalculationPage):
      19             :     """Page for configuring and analyzing single atom systems."""
      20             : 
      21           0 :     title = "One Atom"
      22           0 :     tooltip = "Configure and analyze single atom systems"
      23             : 
      24           0 :     def setupWidget(self) -> None:
      25           0 :         super().setupWidget()
      26             : 
      27             :         # all attributes of instance BaseConfig will be added to the toolbox in postSetupWidget
      28           0 :         self.ket_config = KetConfigOneAtom(self)
      29           0 :         self.basis_config = BasisConfigOneAtom(self)
      30           0 :         self.system_config = SystemConfigOneAtom(self)
      31             : 
      32           0 :     def calculate(self) -> tuple[ParametersOneAtom, ResultsOneAtom]:
      33           0 :         parameters = ParametersOneAtom.from_page(self)
      34           0 :         results = calculate_one_atom(parameters)
      35           0 :         return parameters, results
      36             : 
      37           0 :     def update_plot(self, parameters: ParametersOneAtom, results: ResultsOneAtom) -> None:  # type: ignore[override]
      38           0 :         super().update_plot(parameters, results)
      39             : 
      40           0 :         self.add_short_labels(results)
      41           0 :         self.plotwidget.canvas.draw()
      42             : 
      43           0 :     def add_short_labels(
      44             :         self,
      45             :         results: ResultsOneAtom,
      46             :     ) -> None:
      47           0 :         ax = self.plotwidget.canvas.ax
      48           0 :         x_lim = ax.get_xlim()
      49           0 :         ax.set_xlim(x_lim[0] - (x_lim[1] - x_lim[0]) * 0.1, x_lim[1])
      50             : 
      51           0 :         used = set()
      52           0 :         for ket_label, energy in zip(results.state_labels_0, results.energies[0]):
      53           0 :             short_label = ket_label[1:-1]
      54           0 :             short_label = short_label.split(":", 1)[-1]
      55           0 :             components = short_label.split(",")
      56           0 :             short_label = ",".join(components[:-1])
      57           0 :             short_label = short_label.split("_", 1)[0]
      58           0 :             if short_label in used:
      59           0 :                 continue
      60           0 :             used.add(short_label)
      61           0 :             self.plotwidget.canvas.ax.text(x_lim[0], energy, short_label, va="center", ha="right")
      62             : 
      63           0 :     def _get_export_notebook_template_name(self) -> str:
      64           0 :         return "one_atom.ipynb"
      65             : 
      66           0 :     def _get_export_replacements(self) -> dict[str, Any]:
      67           0 :         parameters = ParametersOneAtom.from_page(self)
      68           0 :         return parameters.to_replacement_dict()