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 : @pytest.fixture
      19           1 : def state(basis: pi.BasisAtom) -> pi.StateAtom:
      20             :     """Create a test state."""
      21           1 :     ket = pi.KetAtom("Rb", n=60, l=1, j=1.5, m=-0.5)
      22           1 :     return basis.get_corresponding_state(ket)
      23             : 
      24             : 
      25           1 : def test_state_creation(state: pi.StateAtom) -> None:
      26             :     """Test basic properties of created state."""
      27           1 :     assert state.species == "Rb"
      28           1 :     assert state.number_of_kets == 80
      29           1 :     assert len(state.kets) == state.number_of_kets
      30           1 :     assert all(x in str(state) for x in ["StateAtom", "60", "S", "3/2", "-1/2"])
      31           1 :     assert state.is_canonical
      32             : 
      33             : 
      34           1 : def test_coefficients(state: pi.StateAtom) -> None:
      35             :     """Test coefficient matrix properties."""
      36           1 :     coeffs = state.get_coefficients()
      37           1 :     assert coeffs.shape == (state.number_of_kets,)
      38           1 :     assert np.count_nonzero(coeffs) == 1
      39           1 :     assert pytest.approx(coeffs.sum()) == 1.0  # NOSONAR
      40             : 
      41             : 
      42           1 : def test_get_amplitude_and_overlap(state: pi.StateAtom) -> None:
      43             :     """Test amplitude and overlap calculations."""
      44             :     # Test with ket
      45           1 :     test_ket = state.get_corresponding_ket()
      46           1 :     amplitude = state.get_amplitude(test_ket)
      47           1 :     assert np.isscalar(amplitude)
      48           1 :     assert pytest.approx(amplitude) == 1.0  # NOSONAR
      49           1 :     overlap = state.get_overlap(test_ket)
      50           1 :     assert np.isscalar(overlap)
      51           1 :     assert pytest.approx(overlap) == 1.0  # NOSONAR
      52             : 
      53             :     # Test with state
      54           1 :     amplitude = state.get_amplitude(state)
      55           1 :     assert np.isscalar(amplitude)
      56           1 :     assert pytest.approx(amplitude) == 1.0  # NOSONAR
      57           1 :     overlap = state.get_overlap(state)
      58           1 :     assert np.isscalar(overlap)
      59           1 :     assert pytest.approx(overlap) == 1.0  # NOSONAR
      60             : 
      61             : 
      62           1 : def test_get_matrix_element(basis: pi.BasisAtom) -> None:
      63             :     """Test matrix element calculations."""
      64           1 :     ket1 = pi.KetAtom("Rb", n=60, l=1, j=1.5, m=-0.5)
      65           1 :     ket2 = pi.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      66           1 :     state1 = basis.get_corresponding_state(ket1)
      67           1 :     state2 = basis.get_corresponding_state(ket2)
      68             : 
      69             :     # Test with ket
      70           1 :     element_dipole_ket = state1.get_matrix_element(ket2, "electric_dipole", q=1, unit="e * a0")
      71           1 :     assert np.isscalar(element_dipole_ket)
      72           1 :     assert element_dipole_ket != 0
      73             : 
      74             :     # Test with state
      75           1 :     element_dipole_state = state1.get_matrix_element(state2, "electric_dipole", q=1, unit="e * a0")
      76           1 :     assert np.isscalar(element_dipole_state)
      77           1 :     assert pytest.approx(element_dipole_ket) == element_dipole_state  # NOSONAR
      78           1 :     assert state1.get_matrix_element(state1, "electric_dipole", q=1, unit="e * a0") == 0
      79           1 :     assert state1.get_matrix_element(state2, "electric_dipole", q=0, unit="e * a0") == 0
      80             : 
      81             : 
      82           1 : def test_error_handling(state: pi.StateAtom) -> None:
      83             :     """Test error cases."""
      84           1 :     with pytest.raises(TypeError):
      85           1 :         state.get_amplitude("not a ket")  # type: ignore [arg-type]
      86             : 
      87           1 :     with pytest.raises(TypeError):
      88           1 :         state.get_overlap("not a ket")  # type: ignore [arg-type]
      89             : 
      90           1 :     with pytest.raises(TypeError):
      91           1 :         state.get_matrix_element("not a ket", "energy", 0)  # type: ignore [call-overload]