Line data    Source code

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