Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : import numpy as np
       5           1 : import pairinteraction.real as pi
       6           1 : import pytest
       7             : 
       8             : 
       9           1 : @pytest.fixture
      10           1 : def basis() -> pi.BasisAtom:
      11             :     """Create a test basis with a few states around Rb 60S."""
      12           1 :     ket = pi.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      13           1 :     energy_min = ket.get_energy(unit="GHz") - 100
      14           1 :     energy_max = ket.get_energy(unit="GHz") + 100
      15           1 :     return pi.BasisAtom("Rb", n=(58, 62), l=(0, 2), energy=(energy_min, energy_max), energy_unit="GHz")
      16             : 
      17             : 
      18           1 : def test_basis_creation(basis: pi.BasisAtom) -> None:
      19             :     """Test basic properties of created basis."""
      20           1 :     assert basis.species == "Rb"
      21           1 :     assert basis.number_of_kets == 80
      22           1 :     assert basis.number_of_states == basis.number_of_kets
      23           1 :     assert len(basis.kets) == basis.number_of_kets
      24           1 :     assert basis.number_of_kets < pi.BasisAtom("Rb", n=(58, 62), l=(0, 2)).number_of_kets
      25           1 :     assert all(x in str(basis) for x in ["BasisAtom", "n=(58, 62)", "l=(0, 2)"])
      26             : 
      27             : 
      28           1 : def test_coefficients(basis: pi.BasisAtom) -> None:
      29             :     """Test coefficient matrix properties."""
      30           1 :     coeffs = basis.get_coefficients()
      31           1 :     assert coeffs.shape == (basis.number_of_kets, basis.number_of_states)
      32           1 :     assert pytest.approx(coeffs.diagonal()) == 1.0  # NOSONAR
      33           1 :     assert pytest.approx(coeffs.sum()) == basis.number_of_kets  # NOSONAR
      34             : 
      35             : 
      36           1 : def test_get_amplitudes_and_overlaps(basis: pi.BasisAtom) -> None:
      37             :     """Test amplitude and overlap calculations."""
      38             :     # Test with ket
      39           1 :     test_ket = basis.kets[0]
      40           1 :     amplitudes = basis.get_amplitudes(test_ket)
      41           1 :     assert len(amplitudes) == basis.number_of_states
      42           1 :     assert pytest.approx(amplitudes[0]) == 1.0  # NOSONAR
      43           1 :     overlaps = basis.get_overlaps(test_ket)
      44           1 :     assert len(overlaps) == basis.number_of_states
      45           1 :     assert pytest.approx(overlaps[0]) == 1.0  # NOSONAR
      46             : 
      47             :     # Test with state
      48           1 :     test_state = basis.states[0]
      49           1 :     amplitudes = basis.get_amplitudes(test_state)
      50           1 :     assert len(amplitudes) == basis.number_of_states
      51           1 :     assert pytest.approx(amplitudes[0]) == 1.0  # NOSONAR
      52           1 :     overlaps = basis.get_overlaps(test_state)
      53           1 :     assert len(overlaps) == basis.number_of_states
      54           1 :     assert pytest.approx(overlaps[0]) == 1.0  # NOSONAR
      55             : 
      56             :     # Test with basis
      57           1 :     matrix_amplitudes = basis.get_amplitudes(basis)
      58           1 :     assert matrix_amplitudes.shape == (basis.number_of_kets, basis.number_of_states)
      59           1 :     assert pytest.approx(matrix_amplitudes.diagonal()) == 1.0  # NOSONAR
      60           1 :     matrix_overlaps = basis.get_overlaps(basis)
      61           1 :     assert matrix_overlaps.shape == (basis.number_of_states, basis.number_of_states)
      62           1 :     assert pytest.approx(matrix_overlaps.diagonal()) == 1.0  # NOSONAR
      63             : 
      64             : 
      65           1 : def test_get_matrix_elements(basis: pi.BasisAtom) -> None:
      66             :     """Test matrix element calculations."""
      67             :     # Test with ket
      68           1 :     test_ket = basis.kets[0]
      69           1 :     elements_dipole = basis.get_matrix_elements(test_ket, "electric_dipole", q=0, unit="e * a0")
      70           1 :     assert elements_dipole.shape == (basis.number_of_states,)
      71           1 :     assert np.count_nonzero(elements_dipole) > 0
      72           1 :     assert np.count_nonzero(elements_dipole) < basis.number_of_states
      73             : 
      74             :     # Test with state
      75           1 :     test_state = basis.states[0]
      76           1 :     elements_dipole = basis.get_matrix_elements(test_state, "electric_dipole", q=0, unit="e * a0")
      77           1 :     assert elements_dipole.shape == (basis.number_of_states,)
      78           1 :     assert np.count_nonzero(elements_dipole) > 0
      79           1 :     assert np.count_nonzero(elements_dipole) < basis.number_of_states
      80             : 
      81             :     # Test with basis
      82           1 :     matrix_elements = basis.get_matrix_elements(basis, "electric_dipole", q=0, unit="e * a0")
      83           1 :     assert matrix_elements.shape == (basis.number_of_states, basis.number_of_states)
      84           1 :     assert np.count_nonzero(matrix_elements.toarray()) > 0
      85           1 :     assert np.count_nonzero(matrix_elements.toarray()) < basis.number_of_states**2
      86             : 
      87             : 
      88           1 : def test_error_handling(basis: pi.BasisAtom) -> None:
      89             :     """Test error cases."""
      90           1 :     with pytest.raises(TypeError):
      91           1 :         basis.get_amplitudes("not a ket")  # type: ignore [call-overload]
      92             : 
      93           1 :     with pytest.raises(TypeError):
      94           1 :         basis.get_overlaps("not a ket")  # type: ignore [call-overload]
      95             : 
      96           1 :     with pytest.raises(TypeError):
      97           1 :         basis.get_matrix_elements("not a ket", "energy", 0)  # type: ignore [call-overload]