LCOV - code coverage report
Current view: top level - tests - test_perturbative.py (source / functions) Hit Total Coverage
Test: coverage.info Lines: 116 117 99.1 %
Date: 2025-08-29 20:47:05 Functions: 9 18 50.0 %

          Line data    Source code
       1             : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : import logging
       5           1 : from typing import TYPE_CHECKING
       6             : 
       7           1 : import numpy as np
       8           1 : import pairinteraction.real as pi
       9           1 : import pytest
      10           1 : from pairinteraction import perturbative
      11           1 : from pairinteraction.perturbative.perturbation_theory import calculate_perturbative_hamiltonian
      12           1 : from scipy import sparse
      13             : 
      14           1 : from .compare_utils import no_log_propagation
      15             : 
      16             : if TYPE_CHECKING:
      17             :     from scipy.sparse import csr_matrix
      18             : 
      19             : 
      20           1 : def _check_sparse_matrices_equal(matrix_a: "csr_matrix", matrix_b: "csr_matrix") -> bool:
      21             :     """Check for equality of sparse matrices efficiently.
      22             : 
      23             :     This functions compares two sparse matrices in compressed sparse row format on their equality.
      24             : 
      25             :     Args:
      26             :         matrix_a: A sparse matrix in csr format.
      27             :         matrix_b: A sparse matrix in csr format.
      28             : 
      29             :     Returns:
      30             :         bool: True if matrices are equal, False if not.
      31             : 
      32             :     """
      33           1 :     matrix_a.sort_indices()
      34           1 :     matrix_b.sort_indices()
      35           1 :     if not (
      36             :         matrix_a.format == "csr"
      37             :         and matrix_b.format == "csr"
      38             :         and len(matrix_a.indices) == len(matrix_b.indices)
      39             :         and len(matrix_a.indptr) == len(matrix_b.indptr)
      40             :         and len(matrix_a.data) == len(matrix_b.data)
      41             :     ):
      42           0 :         return False
      43             : 
      44           1 :     return bool(
      45             :         np.all(matrix_a.indices == matrix_b.indices)
      46             :         and np.all(matrix_a.indptr == matrix_b.indptr)
      47             :         and np.allclose(matrix_a.data, matrix_b.data, rtol=0, atol=1e-14)
      48             :     )
      49             : 
      50             : 
      51           1 : @pytest.fixture
      52           1 : def system_pair_sample() -> pi.SystemPair:
      53           1 :     basis = pi.BasisAtom(
      54             :         species="Rb",
      55             :         n=(59, 63),
      56             :         l=(0, 1),
      57             :         m=(-1.5, 1.5),
      58             :     )
      59           1 :     system = pi.SystemAtom(basis=basis)
      60           1 :     system.set_diamagnetism_enabled(False)
      61           1 :     system.set_magnetic_field([0, 0, 1e-3], "gauss")
      62           1 :     if not system.is_diagonal:
      63           1 :         pi.diagonalize([system], diagonalizer="eigen", sort_by_energy=False)
      64           1 :     basis_pair = pi.BasisPair([system, system])
      65           1 :     system_pair = pi.SystemPair(basis_pair)
      66           1 :     theta = 0
      67           1 :     r = 12
      68           1 :     system_pair.set_distance_vector(r * np.array([np.sin(theta), 0, np.cos(theta)]), "micrometer")
      69           1 :     system_pair.set_interaction_order(3)
      70           1 :     return system_pair
      71             : 
      72             : 
      73           1 : def test_perturbative_calculation1(caplog: pytest.LogCaptureFixture) -> None:
      74             :     """Test of mathematic functionality."""
      75           1 :     hamiltonian = sparse.csr_matrix(
      76             :         [[0, 1, 1, 0, 2], [1, 1, 0, 1, 3], [1, 0, 10, 1, 0], [0, 1, 1, 11, 1], [2, 3, 0, 1, 12]]
      77             :     )
      78           1 :     model_space_indices = [0, 1]
      79             : 
      80             :     # Order 0
      81           1 :     h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, perturbation_order=0)
      82           1 :     h_eff = sum(h for h in h_eff_dict.values())
      83           1 :     assert np.any(h_eff == np.array([[0, 0], [0, 1]]))
      84           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.eye(2, 5, k=0)))
      85             : 
      86             :     # Order 1
      87           1 :     h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, perturbation_order=1)
      88           1 :     h_eff = sum(h for h in h_eff_dict.values())
      89           1 :     assert np.any(h_eff == np.array([[0, 1], [1, 1]]))
      90           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix([[1, 0, 0, 0, 0], [0, 1, 0, 0, 0]]))
      91             : 
      92             :     # Order 2
      93           1 :     h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, perturbation_order=2)
      94           1 :     h_eff = sum(h for h in h_eff_dict.values())
      95           1 :     a_00 = 0 + 1 * 1 / (0 - 10) + 2 * 2 / (0 - 12)
      96           1 :     a_01 = 1 + 2 * 3 / (0 - 12)
      97           1 :     a_10 = 1 + 3 * 2 / (1 - 12)
      98           1 :     a_11 = 1 + 1 * 1 / (1 - 11) + 3 * 3 / (1 - 12)
      99           1 :     hamiltonian_new = np.array([[a_00, (a_01 + a_10) / 2], [(a_01 + a_10) / 2, a_11]])
     100           1 :     assert np.any(h_eff == hamiltonian_new)
     101             : 
     102           1 :     v0 = sparse.eye(1, 5, k=0) + 1 / (0 - 10) * sparse.eye(1, 5, k=2) + 2 / (0 - 12) * sparse.eye(1, 5, k=4)
     103           1 :     v1 = sparse.eye(1, 5, k=1) + 1 / (1 - 11) * sparse.eye(1, 5, k=3) + 3 / (1 - 12) * sparse.eye(1, 5, k=4)
     104           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.vstack([v0, v1])))
     105             : 
     106             :     # Order 3
     107           1 :     with caplog.at_level(logging.ERROR):
     108           1 :         h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(
     109             :             hamiltonian, model_space_indices, perturbation_order=3
     110             :         )
     111           1 :     h_eff = sum(h for h in h_eff_dict.values())
     112           1 :     a_00 -= 2 * 3 * 1 / ((0 - 12) * (1 - 12))
     113           1 :     a_01 += (
     114             :         1 * 1 * 1 / ((1 - 10) * (1 - 11))
     115             :         + 2 * 1 * 1 / ((1 - 11) * (1 - 12))
     116             :         - 1 * 1 * 1 / ((0 - 10) * (1 - 10))
     117             :         - 2 * 2 * 1 / ((1 - 12) * (0 - 12))
     118             :     )
     119           1 :     a_10 += (
     120             :         1 * 1 * 1 / ((0 - 10) * (0 - 11))
     121             :         + 2 * 1 * 1 / ((0 - 11) * (0 - 12))
     122             :         - 1 * 1 * 1 / ((0 - 11) * (1 - 11))
     123             :         - 3 * 3 * 1 / ((0 - 12) * (1 - 12))
     124             :     )
     125           1 :     a_11 += 1 * 1 * 3 / ((1 - 11) * (1 - 12)) + 3 * 1 * 1 / ((1 - 11) * (1 - 12)) - 3 * 2 * 1 / ((1 - 12) * (0 - 12))
     126           1 :     hamiltonian_new = np.array([[a_00, (a_01 + a_10) / 2], [(a_01 + a_10) / 2, a_11]])
     127           1 :     assert np.any(h_eff == hamiltonian_new)
     128             : 
     129           1 :     v0 = v0 + (
     130             :         -0.5 * (1 * 1 / (0 - 10) ** 2 + 2 * 2 / (0 - 12) ** 2) * sparse.eye(1, 5, k=0)
     131             :         + (1 * 1 / ((0 - 10) * (0 - 11)) + 2 * 1 / ((0 - 11) * (0 - 12))) * sparse.eye(1, 5, k=3)
     132             :         + 1 * 1 / ((0 - 11) * (0 - 1)) * sparse.eye(1, 5, k=3)
     133             :         + 3 * 1 / ((0 - 1) * (0 - 12)) * sparse.eye(1, 5, k=4)
     134             :         + 3 * 2 / ((0 - 1) * (0 - 12)) * sparse.eye(1, 5, k=1)
     135             :     )
     136           1 :     v1 = v1 + (
     137             :         -0.5 * (1 * 1 / (1 - 11) ** 2 + 3 * 3 / (1 - 12) ** 2) * sparse.eye(1, 5, k=1)
     138             :         + 1 * 1 / ((1 - 10) * (1 - 11)) * sparse.eye(1, 5, k=2)
     139             :         + 1 * 3 / ((1 - 11) * (1 - 12)) * sparse.eye(1, 5, k=3)
     140             :         + 1 * 1 / ((1 - 11) * (1 - 12)) * sparse.eye(1, 5, k=4)
     141             :         + 1 * 1 / ((1 - 0) * (1 - 10)) * sparse.eye(1, 5, k=2)
     142             :         + 2 * 1 / ((1 - 0) * (1 - 12)) * sparse.eye(1, 5, k=4)
     143             :         + 3 * 2 / ((1 - 0) * (1 - 12)) * sparse.eye(1, 5, k=0)
     144             :     )
     145           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.vstack([v0, v1])))
     146             : 
     147             : 
     148           1 : def test_c3_with_sample_system(system_pair_sample: pi.SystemPair) -> None:
     149             :     """Test whether the C3 coefficient with a given system is calculated correctly."""
     150           1 :     ket1 = pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)
     151           1 :     ket2 = pi.KetAtom("Rb", n=61, l=1, j=1.5, m=0.5)
     152           1 :     c3_obj = perturbative.C3(ket1, ket2)
     153           1 :     c3_obj._distance_vector = None  # avoid warning due when setting system pair
     154           1 :     c3_obj.system_pair = system_pair_sample
     155             : 
     156           1 :     c3 = c3_obj.get(unit="planck_constant * gigahertz * micrometer^3")
     157           1 :     assert np.isclose(-0.5 * c3, 3.1515)
     158             : 
     159             : 
     160           1 : def test_c3() -> None:
     161             :     """Test whether the C3 coefficient with automatically constructed system is calculated correctly."""
     162           1 :     ket1 = pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)
     163           1 :     ket2 = pi.KetAtom("Rb", n=61, l=1, j=1.5, m=0.5)
     164           1 :     c3_obj = perturbative.C3(ket1, ket2)
     165             : 
     166           1 :     c3_obj.set_electric_field([0, 0, 0], "volt/cm")
     167           1 :     c3_obj.set_magnetic_field([0, 0, 10], "gauss")
     168             : 
     169           1 :     c3 = c3_obj.get(unit="planck_constant * gigahertz * micrometer^3")
     170           1 :     assert np.isclose(-0.5 * c3, 3.2188)
     171             : 
     172             : 
     173           1 : def test_c6_with_sample_system(system_pair_sample: pi.SystemPair) -> None:
     174             :     """Test whether the C6 coefficient with a given system is calculated correctly."""
     175           1 :     ket = pi.KetAtom(species="Rb", n=61, l=0, j=0.5, m=0.5)
     176           1 :     c6_obj = perturbative.C6(ket, ket)
     177           1 :     c6_obj._distance_vector = None  # avoid warning due when setting system pair
     178           1 :     c6_obj.system_pair = system_pair_sample
     179             : 
     180           1 :     c6 = c6_obj.get(unit="planck_constant * gigahertz * micrometer^6")
     181           1 :     assert np.isclose(c6, 167.880)
     182             : 
     183             : 
     184           1 : def test_c6() -> None:
     185             :     """Test whether the C6 coefficient with automatically constructed system is calculated correctly."""
     186           1 :     ket = pi.KetAtom(species="Rb", n=61, l=0, j=0.5, m=0.5)
     187           1 :     c6_obj = perturbative.C6(ket, ket)
     188             : 
     189           1 :     c6_obj.set_electric_field([0, 0, 0], "volt/cm")
     190           1 :     c6_obj.set_magnetic_field([0, 0, 10], "gauss")
     191             : 
     192           1 :     c6 = c6_obj.get(unit="planck_constant * gigahertz * micrometer^6")
     193           1 :     assert np.isclose(c6, 169.149)
     194             : 
     195             : 
     196           1 : def test_exact_resonance_detection(system_pair_sample: pi.SystemPair, capsys: pytest.CaptureFixture[str]) -> None:
     197             :     """Test whether resonance with infinite admixture is correctly detected."""
     198           1 :     ket1 = pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)
     199           1 :     ket2 = pi.KetAtom("Rb", n=61, l=1, j=1.5, m=0.5)
     200           1 :     eff_system = perturbative.EffectiveSystemPair([(ket1, ket2)])
     201           1 :     eff_system.system_pair = system_pair_sample
     202             : 
     203             :     # workaround to test for errors, without showing them in the std output
     204           1 :     with no_log_propagation("pairinteraction"):
     205           1 :         eff_system.get_effective_hamiltonian()
     206           1 :     captured = capsys.readouterr()
     207           1 :     assert "Detected 'inf' entries" in captured.err
     208           1 :     assert "|Rb:61,P_3/2,1/2; Rb:61,S_1/2,1/2⟩ has infinite admixture" in captured.err
     209             : 
     210             : 
     211           1 : def test_near_resonance_detection(capsys: pytest.CaptureFixture[str]) -> None:
     212             :     """Test whether a near resonance is correctly detected."""
     213           1 :     ket1 = pi.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
     214           1 :     ket2 = pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)
     215           1 :     eff_system = perturbative.EffectiveSystemPair([(ket1, ket2), (ket2, ket1)])
     216           1 :     eff_system.set_magnetic_field([0, 0, 245], "gauss")
     217           1 :     eff_system.set_minimum_number_of_ket_pairs(100)
     218           1 :     eff_system.set_distance(10, 35.1, "micrometer")
     219           1 :     eff_system.get_effective_hamiltonian()
     220             : 
     221             :     # workaround to test for errors, without showing them in the std output
     222           1 :     with no_log_propagation("pairinteraction"):
     223           1 :         eff_system.check_for_resonances(0.99)
     224           1 :     captured = capsys.readouterr()
     225           1 :     assert "The most perturbing states are" in captured.err
     226           1 :     assert "Rb:60,P_3/2,1/2; Rb:60,P_3/2,3/2" in captured.err

Generated by: LCOV version 1.16