Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : from typing import TYPE_CHECKING, Any, Literal, Union
       5             : 
       6           1 : import numpy as np
       7           1 : import pytest
       8           1 : from pairinteraction_gui.main_window import MainWindow
       9             : 
      10           1 : from .compare_utils import REFERENCE_PATHS, compare_eigensystem_to_reference
      11             : 
      12             : if TYPE_CHECKING:
      13             :     from pairinteraction_gui.page import OneAtomPage
      14             :     from pairinteraction_gui.page.two_atoms_page import TwoAtomsPage
      15             :     from pytestqt.qtbot import QtBot
      16             : 
      17             : 
      18           1 : @pytest.fixture
      19           1 : def base_window(qtbot: "QtBot") -> MainWindow:
      20           1 :     window = MainWindow()
      21           1 :     window.show()
      22           1 :     qtbot.addWidget(window)
      23           1 :     return window
      24             : 
      25             : 
      26           1 : @pytest.fixture
      27           1 : def window_starkmap(base_window: MainWindow) -> MainWindow:
      28           1 :     one_atom_page: OneAtomPage = base_window.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
      29           1 :     one_atom_page.ket_config.species_combo_list[0].setCurrentText("Rb")
      30           1 :     ket_qn = one_atom_page.ket_config.stacked_qn_list[0].currentWidget()
      31           1 :     ket_qn.items["n"].setValue(60)
      32           1 :     ket_qn.items["l"].setValue(0)
      33           1 :     ket_qn.items["m"].setValue(0.5)
      34             : 
      35           1 :     basis_qn = one_atom_page.basis_config.stacked_basis_list[0].currentWidget()
      36           1 :     basis_qn.items["n"].setValue(2)
      37           1 :     basis_qn.items["l"].setValue(2)
      38           1 :     basis_qn.items["m"].setChecked(False)
      39             : 
      40           1 :     calculation_config = one_atom_page.calculation_config
      41           1 :     calculation_config.steps.setValue(11)
      42           1 :     system_config = one_atom_page.system_config
      43           1 :     system_config.Ez.spinboxes[1].setValue(10)
      44             : 
      45           1 :     return base_window
      46             : 
      47             : 
      48           1 : @pytest.fixture
      49           1 : def window_pair_potential(base_window: MainWindow) -> MainWindow:
      50           1 :     two_atoms_page: TwoAtomsPage = base_window.stacked_pages.getNamedWidget("TwoAtomsPage")  # type: ignore [assignment]
      51           1 :     two_atoms_page.ket_config.species_combo_list[0].setCurrentText("Rb")
      52           1 :     for ket_qn_stacked in two_atoms_page.ket_config.stacked_qn_list:
      53           1 :         ket_qn = ket_qn_stacked.currentWidget()
      54           1 :         ket_qn.items["n"].setValue(60)
      55           1 :         ket_qn.items["l"].setValue(0)
      56           1 :         ket_qn.items["m"].setValue(0.5)
      57             : 
      58           1 :     for basis_qn_stacked in two_atoms_page.basis_config.stacked_basis_list:
      59           1 :         basis_qn = basis_qn_stacked.currentWidget()
      60           1 :         basis_qn.items["n"].setValue(2)
      61           1 :         basis_qn.items["l"].setValue(2)
      62           1 :         basis_qn.items["m"].setChecked(False)
      63             : 
      64           1 :     two_atoms_page.basis_config.pair_delta_energy.setValue(3)
      65           1 :     two_atoms_page.basis_config.pair_m_range.setValues(1, 1)
      66             : 
      67           1 :     calculation_config = two_atoms_page.calculation_config
      68           1 :     calculation_config.steps.setValue(5)
      69           1 :     system_config = two_atoms_page.system_config
      70           1 :     system_config.distance.setValues(1, 5)
      71           1 :     return base_window
      72             : 
      73             : 
      74           1 : def test_main_window_basic(qtbot: "QtBot", window_starkmap: "MainWindow") -> None:
      75             :     """Test basic main window functionality."""
      76           1 :     one_atom_page: OneAtomPage = window_starkmap.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
      77           1 :     qn_item = one_atom_page.ket_config.stacked_qn_list[0].currentWidget().items["n"]
      78           1 :     qn_item.setValue(60)
      79             : 
      80           1 :     ket_label = one_atom_page.ket_config.ket_label_list[0].text()
      81           1 :     assert all(x in ket_label for x in ["Rb", "60", "S", "1/2"])
      82           1 :     assert qn_item.label.text() == "n"
      83           1 :     assert qn_item.value() == 60
      84             : 
      85           1 :     qn_item.setValue(61)
      86           1 :     ket_label = one_atom_page.ket_config.ket_label_list[0].text()
      87           1 :     assert qn_item.value() == 61
      88           1 :     assert all(x in ket_label for x in ["Rb", "61", "S", "1/2"])
      89             : 
      90             :     # make the basis smaller for faster test
      91           1 :     basis_qn = one_atom_page.basis_config.stacked_basis_list[0].currentWidget()
      92           1 :     basis_qn.items["n"].setValue(1)
      93           1 :     basis_qn.items["l"].setValue(1)
      94           1 :     basis_qn.items["m"].setValue(0)
      95             : 
      96           1 :     one_atom_page.calculate_and_abort.getNamedWidget("Calculate").click()
      97           1 :     qtbot.waitUntil(lambda: one_atom_page._calculation_finished, timeout=30_000)  # ci macOS-13 is very slow
      98           1 :     qtbot.waitUntil(lambda: one_atom_page._plot_finished, timeout=5_000)
      99           1 :     window_starkmap.close()
     100             : 
     101             : 
     102           1 : def test_one_atom_page(window_starkmap: "MainWindow") -> None:
     103           1 :     _test_calculate_page(window_starkmap, "OneAtomPage", "stark_map")
     104             : 
     105             : 
     106           1 : def test_two_atoms_page(window_pair_potential: "MainWindow") -> None:
     107           1 :     _test_calculate_page(window_pair_potential, "TwoAtomsPage", "pair_potential")
     108             : 
     109             : 
     110           1 : def _test_calculate_page(
     111             :     window: MainWindow,
     112             :     page_name: Literal["OneAtomPage", "TwoAtomsPage"],
     113             :     reference_name: str,
     114             : ) -> None:
     115           1 :     page: Union[OneAtomPage, TwoAtomsPage] = window.stacked_pages.getNamedWidget(page_name)  # type: ignore [assignment]
     116           1 :     ket_energy_0 = sum(page.ket_config.get_ket_atom(i).get_energy("GHz") for i in range(page.ket_config.n_atoms))
     117             : 
     118             :     # Test calculation with fast mode off
     119           1 :     page.calculation_config.fast_mode.setChecked(False)
     120           1 :     _parameters, results = page.calculate()
     121           1 :     energies = np.array(results.energies) + ket_energy_0
     122           1 :     compare_eigensystem_to_reference(REFERENCE_PATHS[reference_name], energies, np.array(results.ket_overlaps))
     123             : 
     124             :     # Test calculation with fast mode on
     125             :     # NOTE: with fast mode, the overlaps are different, so we don't compare them
     126           1 :     page.calculation_config.fast_mode.setChecked(True)
     127           1 :     _parameters, results = page.calculate()
     128           1 :     energies = np.array(results.energies) + ket_energy_0
     129           1 :     compare_eigensystem_to_reference(REFERENCE_PATHS[reference_name], energies)
     130             : 
     131             :     # Test export to Python code
     132           1 :     python_code = page._create_python_code()
     133           1 :     python_code = python_code.replace("plt.show()", "")  # HACK, otherwise it will block the test
     134             : 
     135             :     # HACK, see also https://stackoverflow.com/questions/45132645/list-comprehension-in-exec-with-empty-locals-nameerror
     136           1 :     locals_globals: dict[str, Any] = {}
     137           1 :     exec(python_code, locals_globals, locals_globals)  # noqa: S102
     138           1 :     energies = np.array(locals_globals["energies_list"]) + ket_energy_0
     139           1 :     compare_eigensystem_to_reference(REFERENCE_PATHS[reference_name], energies)