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
       5             : 
       6           1 : import numpy as np
       7           1 : import pytest
       8           1 : from pairinteraction_gui.calculate.calculate_two_atoms import ParametersTwoAtoms, _calculate_two_atoms
       9           1 : from pairinteraction_gui.main_window import MainWindow
      10             : 
      11           1 : from .test_pair_potential import compare_pair_potential_to_reference
      12           1 : from .test_starkmap import compare_starkmap_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 window_starkmap(qtbot: "QtBot") -> MainWindow:
      22           1 :     window = MainWindow()
      23           1 :     window.show()
      24           1 :     qtbot.addWidget(window)
      25             : 
      26           1 :     one_atom_page: OneAtomPage = window.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
      27           1 :     one_atom_page.ket_config.species_combo[0].setCurrentText("Rb")
      28           1 :     ket_qn = one_atom_page.ket_config.stacked_qn[0].currentWidget()
      29           1 :     ket_qn.items["n"].setValue(60)
      30           1 :     ket_qn.items["l"].setValue(0)
      31           1 :     ket_qn.items["m"].setValue(0.5)
      32             : 
      33           1 :     basis_qn = one_atom_page.basis_config.stacked_basis[0].currentWidget()
      34           1 :     basis_qn.items["n"].setValue(2)
      35           1 :     basis_qn.items["l"].setValue(2)
      36           1 :     basis_qn.items["m"].setChecked(False)
      37             : 
      38           1 :     calculation_config = one_atom_page.calculation_config
      39           1 :     calculation_config.steps.setValue(11)
      40           1 :     system_config = one_atom_page.system_config
      41           1 :     system_config.Ez.spinboxes[1].setValue(10)
      42             : 
      43           1 :     return window
      44             : 
      45             : 
      46           1 : def test_main_window_basic(qtbot: "QtBot", window_starkmap: "MainWindow") -> None:
      47             :     """Test basic main window functionality."""
      48           1 :     one_atom_page: OneAtomPage = window_starkmap.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
      49           1 :     qn_item = one_atom_page.ket_config.stacked_qn[0].currentWidget().items["n"]
      50           1 :     qn_item.setValue(60)
      51             : 
      52           1 :     ket_label = one_atom_page.ket_config.ket_label[0].text()
      53           1 :     assert all(x in ket_label for x in ["Rb", "60", "S", "1/2"])
      54           1 :     assert qn_item.label.text() == "n"
      55           1 :     assert qn_item.value() == 60
      56             : 
      57           1 :     qn_item.setValue(61)
      58           1 :     ket_label = one_atom_page.ket_config.ket_label[0].text()
      59           1 :     assert qn_item.value() == 61
      60           1 :     assert all(x in ket_label for x in ["Rb", "61", "S", "1/2"])
      61             : 
      62             :     # make the basis smaller for faster test
      63           1 :     basis_qn = one_atom_page.basis_config.stacked_basis[0].currentWidget()
      64           1 :     basis_qn.items["n"].setValue(1)
      65           1 :     basis_qn.items["l"].setValue(1)
      66           1 :     basis_qn.items["m"].setValue(0)
      67             : 
      68           1 :     one_atom_page.calculate_and_abort.getNamedWidget("Calculate").click()
      69           1 :     qtbot.waitUntil(lambda: one_atom_page._calculation_finished, timeout=20_000)  # ci macOS-13 is very slow
      70           1 :     qtbot.waitUntil(lambda: one_atom_page._plot_finished, timeout=5_000)
      71           1 :     window_starkmap.close()
      72             : 
      73             : 
      74           1 : def test_calculate_one_atom(window_starkmap: "MainWindow") -> None:
      75           1 :     one_atom_page: OneAtomPage = window_starkmap.stacked_pages.getNamedWidget("OneAtomPage")  # type: ignore [assignment]
      76           1 :     ket = one_atom_page.ket_config.get_ket_atom(0)
      77             : 
      78           1 :     one_atom_page.calculation_config.fast_mode.setChecked(False)
      79           1 :     _parameters, results = one_atom_page.calculate()
      80           1 :     energies = np.array(results.energies) + ket.get_energy("GHz")
      81           1 :     compare_starkmap_to_reference(energies, np.array(results.ket_overlaps))
      82             : 
      83           1 :     one_atom_page.calculation_config.fast_mode.setChecked(True)
      84           1 :     _parameters, results = one_atom_page.calculate()
      85           1 :     energies = np.array(results.energies) + ket.get_energy("GHz")
      86           1 :     compare_starkmap_to_reference(energies)  # with fast mode, the overlaps are different, so we don't compare them
      87             : 
      88             : 
      89           1 : def test_calculate_two_atoms(qtbot: "QtBot") -> None:
      90           1 :     window = MainWindow()
      91           1 :     window.show()
      92           1 :     qtbot.addWidget(window)
      93             : 
      94           1 :     two_atoms_page: TwoAtomsPage = window.stacked_pages.getNamedWidget("TwoAtomsPage")  # type: ignore [assignment]
      95           1 :     two_atoms_page.ket_config.species_combo[0].setCurrentText("Rb")
      96           1 :     for ket_qn_stacked in two_atoms_page.ket_config.stacked_qn:
      97           1 :         ket_qn = ket_qn_stacked.currentWidget()
      98           1 :         ket_qn.items["n"].setValue(60)
      99           1 :         ket_qn.items["l"].setValue(0)
     100           1 :         ket_qn.items["m"].setValue(0.5)
     101             : 
     102           1 :     for basis_qn_stacked in two_atoms_page.basis_config.stacked_basis:
     103           1 :         basis_qn = basis_qn_stacked.currentWidget()
     104           1 :         basis_qn.items["n"].setValue(2)
     105           1 :         basis_qn.items["l"].setValue(2)
     106           1 :         basis_qn.items["m"].setChecked(False)
     107             : 
     108           1 :     two_atoms_page.basis_config.pair_delta_energy.setValue(3)
     109           1 :     two_atoms_page.basis_config.pair_m_range.setValues(1, 1)
     110             : 
     111           1 :     calculation_config = two_atoms_page.calculation_config
     112           1 :     calculation_config.steps.setValue(5)
     113           1 :     system_config = two_atoms_page.system_config
     114           1 :     system_config.distance.setValues(1, 5)
     115             : 
     116             :     # since no fields are applied the energy offset is simply given by the energies of the kets
     117           1 :     ket_pair_energy_0 = sum(two_atoms_page.ket_config.get_ket_atom(i).get_energy("GHz") for i in range(2))
     118             : 
     119           1 :     two_atoms_page.calculation_config.fast_mode.setChecked(False)
     120           1 :     parameters = ParametersTwoAtoms.from_page(two_atoms_page)
     121           1 :     results = _calculate_two_atoms(parameters)
     122           1 :     energies = np.array(results.energies) + ket_pair_energy_0
     123           1 :     compare_pair_potential_to_reference(energies, np.array(results.ket_overlaps))
     124             : 
     125           1 :     two_atoms_page.calculation_config.fast_mode.setChecked(True)
     126           1 :     parameters = ParametersTwoAtoms.from_page(two_atoms_page)
     127           1 :     results = _calculate_two_atoms(parameters)
     128           1 :     energies = np.array(results.energies) + ket_pair_energy_0
     129           1 :     compare_pair_potential_to_reference(
     130             :         energies
     131             :     )  # with fast mode, the overlaps are different, so we don't compare them