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