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 pathlib import Path
      16             : 
      17             :     from pairinteraction_gui.page import LifetimesPage, OneAtomPage
      18             :     from pairinteraction_gui.page.two_atoms_page import TwoAtomsPage
      19             :     from pytestqt.qtbot import QtBot
      20             : 
      21             : 
      22           1 : @pytest.fixture
      23           1 : def base_window(qtbot: QtBot, tmp_path: Path) -> MainWindow:
      24           1 :     window = MainWindow(cache_dir=tmp_path)
      25           1 :     window.show()
      26           1 :     qtbot.addWidget(window)
      27           1 :     return window
      28             : 
      29             : 
      30           1 : @pytest.fixture
      31           1 : def window_starkmap(base_window: MainWindow) -> MainWindow:
      32           1 :     one_atom_page: OneAtomPage = base_window.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
      33           1 :     one_atom_page.ket_config.species_combo_list[0].setCurrentText("Rb")
      34           1 :     ket_qn = one_atom_page.ket_config.stacked_qn_list[0].currentWidget()
      35           1 :     ket_qn.items["n"].setValue(60)
      36           1 :     ket_qn.items["l"].setValue(0)
      37           1 :     ket_qn.items["m"].setValue(0.5)
      38             : 
      39           1 :     basis_qn = one_atom_page.basis_config.stacked_basis_list[0].currentWidget()
      40           1 :     basis_qn.items["n"].setValue(2)
      41           1 :     basis_qn.items["l"].setValue(2)
      42           1 :     basis_qn.items["m"].setChecked(False)
      43             : 
      44           1 :     calculation_config = one_atom_page.calculation_config
      45           1 :     calculation_config.steps.setValue(11)
      46           1 :     system_config = one_atom_page.system_config
      47           1 :     system_config.Ez.spinboxes[1].setValue(10)
      48             : 
      49           1 :     return base_window
      50             : 
      51             : 
      52           1 : @pytest.fixture
      53           1 : def window_pair_potential(base_window: MainWindow) -> MainWindow:
      54           1 :     two_atoms_page: TwoAtomsPage = base_window.stacked_pages.getNamedWidget("TwoAtomsPage")  # type: ignore [assignment]
      55           1 :     two_atoms_page.ket_config.species_combo_list[0].setCurrentText("Rb")
      56           1 :     for ket_qn_stacked in two_atoms_page.ket_config.stacked_qn_list:
      57           1 :         ket_qn = ket_qn_stacked.currentWidget()
      58           1 :         ket_qn.items["n"].setValue(60)
      59           1 :         ket_qn.items["l"].setValue(0)
      60           1 :         ket_qn.items["m"].setValue(0.5)
      61             : 
      62           1 :     for basis_qn_stacked in two_atoms_page.basis_config.stacked_basis_list:
      63           1 :         basis_qn = basis_qn_stacked.currentWidget()
      64           1 :         basis_qn.items["n"].setValue(2)
      65           1 :         basis_qn.items["l"].setValue(2)
      66           1 :         basis_qn.items["m"].setChecked(False)
      67             : 
      68           1 :     two_atoms_page.basis_config.pair_delta_energy.setValue(3)
      69           1 :     two_atoms_page.basis_config.pair_m_range.setValues(1, 1)
      70             : 
      71           1 :     calculation_config = two_atoms_page.calculation_config
      72           1 :     calculation_config.steps.setValue(5)
      73           1 :     system_config = two_atoms_page.system_config
      74           1 :     system_config.distance.setValues(1, 5)
      75           1 :     return base_window
      76             : 
      77             : 
      78           1 : @pytest.fixture
      79           1 : def window_lifetimes(base_window: MainWindow) -> MainWindow:
      80           1 :     lifetimes_page: LifetimesPage = base_window.stacked_pages.getNamedWidget("LifetimesPage")  # type: ignore [assignment]
      81           1 :     lifetimes_page.ket_config.species_combo_list[0].setCurrentText("Rb")
      82           1 :     ket_qn = lifetimes_page.ket_config.stacked_qn_list[0].currentWidget()
      83           1 :     ket_qn.items["n"].setValue(60)
      84           1 :     ket_qn.items["l"].setValue(0)
      85           1 :     ket_qn.items["m"].setValue(0.5)
      86           1 :     lifetimes_page.ket_config.item_temperature.setValue(300)
      87           1 :     return base_window
      88             : 
      89             : 
      90           1 : def test_main_window_basic(qtbot: QtBot, window_starkmap: MainWindow) -> None:
      91             :     """Test basic main window functionality."""
      92           1 :     one_atom_page: OneAtomPage = window_starkmap.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
      93           1 :     qn_item = one_atom_page.ket_config.stacked_qn_list[0].currentWidget().items["n"]
      94           1 :     qn_item.setValue(60)
      95             : 
      96           1 :     ket_label = one_atom_page.ket_config.ket_label_list[0].text()
      97           1 :     assert all(x in ket_label for x in ["Rb", "60", "S", "1/2"])
      98           1 :     assert qn_item.label.text() == "n"
      99           1 :     assert qn_item.value() == 60
     100             : 
     101           1 :     qn_item.setValue(61)
     102           1 :     ket_label = one_atom_page.ket_config.ket_label_list[0].text()
     103           1 :     assert qn_item.value() == 61
     104           1 :     assert all(x in ket_label for x in ["Rb", "61", "S", "1/2"])
     105             : 
     106             :     # make the basis smaller for faster test
     107           1 :     basis_qn = one_atom_page.basis_config.stacked_basis_list[0].currentWidget()
     108           1 :     basis_qn.items["n"].setValue(1)
     109           1 :     basis_qn.items["l"].setValue(1)
     110           1 :     basis_qn.items["m"].setValue(0)
     111             : 
     112           1 :     one_atom_page.calculate_and_abort.getNamedWidget("Calculate").click()
     113           1 :     qtbot.waitUntil(lambda: one_atom_page._calculation_finished, timeout=30_000)  # ci macOS-15-intel is very slow
     114           1 :     qtbot.waitUntil(lambda: one_atom_page._plot_finished, timeout=5_000)
     115           1 :     window_starkmap.close()
     116             : 
     117             : 
     118           1 : def test_one_atom_page(window_starkmap: MainWindow) -> None:
     119           1 :     _test_calculate_page(window_starkmap, "OneAtomPage", "stark_map")
     120             : 
     121             : 
     122           1 : def test_two_atoms_page(window_pair_potential: MainWindow) -> None:
     123           1 :     _test_calculate_page(window_pair_potential, "TwoAtomsPage", "pair_potential")
     124             : 
     125             : 
     126           1 : def test_lifetimes_page(window_lifetimes: MainWindow) -> None:
     127           1 :     page: LifetimesPage = window_lifetimes.stacked_pages.getNamedWidget("LifetimesPage")  # type: ignore [assignment]
     128           1 :     parameters, results = page.calculate()
     129           1 :     page.plotwidget.plot(parameters, results)
     130             : 
     131           1 :     assert results.lifetime > 0
     132           1 :     assert len(results.kets_sp) == len(results.transition_rates_sp)
     133           1 :     assert len(results.kets_bbr) == len(results.transition_rates_bbr)
     134             : 
     135           1 :     python_code = page._create_python_code()
     136           1 :     python_code = python_code.replace("plt.show()", "plt.close()")  # HACK, otherwise it will block the test
     137             : 
     138           1 :     locals_globals: dict[str, Any] = {}
     139           1 :     exec(python_code, locals_globals, locals_globals)  # noqa: S102
     140             : 
     141           1 :     rates_dict = {
     142             :         key: [sum(r for _, r in page.plotwidget.sorted_rates[label][n]) for n in locals_globals["n_list"]]
     143             :         for key, label in [("SP", "Spontaneous Decay"), ("BBR", "Black Body Radiation")]
     144             :     }
     145             : 
     146           1 :     assert np.isclose(locals_globals["lifetime"], results.lifetime)
     147           1 :     for key, rates in rates_dict.items():
     148           1 :         assert np.allclose(locals_globals["rates_summed"][key], rates)
     149             : 
     150             : 
     151           1 : def _test_calculate_page(
     152             :     window: MainWindow,
     153             :     page_name: Literal["OneAtomPage", "TwoAtomsPage"],
     154             :     reference_name: str,
     155             : ) -> None:
     156           1 :     page: OneAtomPage | TwoAtomsPage = window.stacked_pages.getNamedWidget(page_name)  # type: ignore [assignment]
     157           1 :     ket_energy_0 = sum(page.ket_config.get_ket_atom(i).get_energy("GHz") for i in range(page.ket_config.n_atoms))
     158             : 
     159             :     # Test calculation with fast mode off
     160           1 :     page.calculation_config.fast_mode.setChecked(False)
     161           1 :     _parameters, results = page.calculate()
     162           1 :     energies = np.array(results.energies) + ket_energy_0
     163           1 :     compare_eigensystem_to_reference(REFERENCE_PATHS[reference_name], energies, np.array(results.ket_overlaps))
     164             : 
     165             :     # Test calculation with fast mode on and diagonalze energy_range
     166             :     # NOTE: with fast mode, the overlaps are different, so we don't compare them
     167           1 :     page.calculation_config.fast_mode.setChecked(True)
     168           1 :     page.calculation_config.energy_range.setChecked(True)
     169           1 :     page.calculation_config.energy_range.setValues(-200, 200)
     170           1 :     parameters, results = page.calculate()
     171           1 :     energies = np.array(results.energies) + ket_energy_0
     172           1 :     compare_eigensystem_to_reference(REFERENCE_PATHS[reference_name], energies)
     173           1 :     page.plotwidget.plot(parameters, results)
     174             : 
     175           1 :     if page_name == "TwoAtomsPage":
     176             :         # Test fitting of c3/c6/c3+c6 potential curves
     177             :         # We need to pass the data to the plotwidget.
     178           1 :         page.plotwidget.fit("c3")
     179             :         # We run a test fit twice, as it should iterate through the potential curves
     180           1 :         page.plotwidget.fit("c3")
     181           1 :         page.plotwidget.fit("c6")
     182           1 :         page.plotwidget.fit("c3+c6")
     183             :         # One could consider adding some return values to check if the estimated values are reasonable.
     184             :         # Currently fit does not return anything, it just plots and shows the fit parameters in UI
     185             : 
     186             :     # Test export to Python code
     187           1 :     python_code = page._create_python_code()
     188           1 :     python_code = python_code.replace("plt.show()", "plt.close()")  # HACK, otherwise it will block the test
     189             : 
     190             :     # HACK, see also https://stackoverflow.com/questions/45132645/list-comprehension-in-exec-with-empty-locals-nameerror
     191           1 :     locals_globals: dict[str, Any] = {}
     192           1 :     exec(python_code, locals_globals, locals_globals)  # noqa: S102
     193           1 :     energies = np.array(locals_globals["energies_list"]) + ket_energy_0
     194           1 :     compare_eigensystem_to_reference(REFERENCE_PATHS[reference_name], energies)
     195             : 
     196             : 
     197           1 : def test_save_and_restore_settings(qtbot: QtBot, tmp_path: Path) -> None:
     198           1 :     window = MainWindow(cache_dir=tmp_path)
     199           1 :     window.show()
     200           1 :     qtbot.addWidget(window)
     201           1 :     one_atom_page: OneAtomPage = window.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
     202           1 :     ket_qn = one_atom_page.ket_config.stacked_qn_list[0].currentWidget()
     203           1 :     ket_qn.items["n"].setValue(222)
     204           1 :     ket_qn.items["l"].setValue(999)
     205           1 :     window.close()
     206             : 
     207           1 :     window = MainWindow(cache_dir=tmp_path)
     208           1 :     window.show()
     209           1 :     qtbot.addWidget(window)
     210           1 :     one_atom_page = window.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
     211           1 :     ket_qn = one_atom_page.ket_config.stacked_qn_list[0].currentWidget()
     212             : 
     213           1 :     assert ket_qn.items["n"].value() == 222
     214           1 :     assert ket_qn.items["l"].value() == 999