Line data    Source code

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