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