LCOV - code coverage report
Current view: top level - tests - test_perturbative.py (source / functions) Hit Total Coverage
Test: coverage.info Lines: 93 94 98.9 %
Date: 2025-06-17 15:25:52 Functions: 8 16 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 pairinteraction.units import ureg
      13           1 : from scipy import sparse
      14             : 
      15             : if TYPE_CHECKING:
      16             :     from scipy.sparse import csr_matrix
      17             : 
      18             : 
      19             : # ---------------------------------------------------------------------------------------
      20             : # Helper functions
      21             : # ---------------------------------------------------------------------------------------
      22           1 : def _check_sparse_matrices_equal(matrix_a: "csr_matrix", matrix_b: "csr_matrix") -> bool:
      23             :     """Check for equality of sparse matrices efficiently.
      24             : 
      25             :     This functions compares two sparse matrices in compressed sparse row format on their equality.
      26             : 
      27             :     Args:
      28             :         matrix_a: A sparse matrix in csr format.
      29             :         matrix_b: A sparse matrix in csr format.
      30             : 
      31             :     Returns:
      32             :         bool: True if matrices are equal, False if not.
      33             : 
      34             :     """
      35           1 :     matrix_a.sort_indices()
      36           1 :     matrix_b.sort_indices()
      37           1 :     if not (
      38             :         matrix_a.format == "csr"
      39             :         and matrix_b.format == "csr"
      40             :         and len(matrix_a.indices) == len(matrix_b.indices)
      41             :         and len(matrix_a.indptr) == len(matrix_b.indptr)
      42             :         and len(matrix_a.data) == len(matrix_b.data)
      43             :     ):
      44           0 :         return False
      45             : 
      46           1 :     return bool(
      47             :         np.all(matrix_a.indices == matrix_b.indices)
      48             :         and np.all(matrix_a.indptr == matrix_b.indptr)
      49             :         and np.allclose(matrix_a.data, matrix_b.data, rtol=0, atol=1e-14)
      50             :     )
      51             : 
      52             : 
      53           1 : def _create_system_pair_sample() -> pi.SystemPair:
      54           1 :     basis = pi.BasisAtom(
      55             :         species="Rb",
      56             :         n=(59, 63),
      57             :         l=(0, 1),
      58             :         m=(-1.5, 1.5),
      59             :     )
      60           1 :     system = pi.SystemAtom(basis=basis)
      61           1 :     system.set_diamagnetism_enabled(False)
      62           1 :     system.set_magnetic_field([0, 0, 1e-3], "gauss")
      63           1 :     if not system.is_diagonal:
      64           1 :         pi.diagonalize([system], diagonalizer="eigen", sort_by_energy=False)
      65           1 :     basis_pair = pi.BasisPair([system, system])
      66           1 :     system_pair = pi.SystemPair(basis_pair)
      67           1 :     theta = 0
      68           1 :     r = 12
      69           1 :     system_pair.set_distance_vector(r * np.array([np.sin(theta), 0, np.cos(theta)]), "micrometer")
      70           1 :     system_pair.set_interaction_order(3)
      71           1 :     return system_pair
      72             : 
      73             : 
      74             : # ---------------------------------------------------------------------------------------
      75             : # Tests
      76             : # ---------------------------------------------------------------------------------------
      77           1 : def test_perturbative_calculation1(caplog: pytest.LogCaptureFixture) -> None:
      78             :     """Test of mathematic functionality."""
      79           1 :     hamiltonian = sparse.csr_matrix(
      80             :         [[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]]
      81             :     )
      82           1 :     model_space_indices = [0, 1]
      83             : 
      84             :     # Order 0
      85           1 :     h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, perturbation_order=0)
      86           1 :     h_eff = sum(h for h in h_eff_dict.values())
      87           1 :     assert np.any(h_eff == np.array([[0, 0], [0, 1]]))
      88           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.eye(2, 5, k=0)))
      89             : 
      90             :     # Order 1
      91           1 :     h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, perturbation_order=1)
      92           1 :     h_eff = sum(h for h in h_eff_dict.values())
      93           1 :     assert np.any(h_eff == np.array([[0, 1], [1, 1]]))
      94           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix([[1, 0, 0, 0, 0], [0, 1, 0, 0, 0]]))
      95             : 
      96             :     # Order 2
      97           1 :     h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, perturbation_order=2)
      98           1 :     h_eff = sum(h for h in h_eff_dict.values())
      99           1 :     a_00 = 0 + 1 * 1 / (0 - 10) + 2 * 2 / (0 - 12)
     100           1 :     a_01 = 1 + 2 * 3 / (0 - 12)
     101           1 :     a_10 = 1 + 3 * 2 / (1 - 12)
     102           1 :     a_11 = 1 + 1 * 1 / (1 - 11) + 3 * 3 / (1 - 12)
     103           1 :     hamiltonian_new = np.array([[a_00, (a_01 + a_10) / 2], [(a_01 + a_10) / 2, a_11]])
     104           1 :     assert np.any(h_eff == hamiltonian_new)
     105             : 
     106           1 :     v0 = (
     107             :         sparse.eye(1, 5, k=0, format="csr")
     108             :         + 1 / (0 - 10) * sparse.eye(1, 5, k=2, format="csr")
     109             :         + 2 / (0 - 12) * sparse.eye(1, 5, k=4, format="csr")
     110             :     )
     111           1 :     v1 = (
     112             :         sparse.eye(1, 5, k=1, format="csr")
     113             :         + 1 / (1 - 11) * sparse.eye(1, 5, k=3, format="csr")
     114             :         + 3 / (1 - 12) * sparse.eye(1, 5, k=4, format="csr")
     115             :     )
     116           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.vstack([v0, v1])))
     117             : 
     118             :     # Order 3
     119           1 :     with caplog.at_level(logging.ERROR):
     120           1 :         h_eff_dict, eig_perturb = calculate_perturbative_hamiltonian(
     121             :             hamiltonian, model_space_indices, perturbation_order=3
     122             :         )
     123           1 :     h_eff = sum(h for h in h_eff_dict.values())
     124           1 :     a_00 -= 2 * 3 * 1 / ((0 - 12) * (1 - 12))
     125           1 :     a_01 += (
     126             :         1 * 1 * 1 / ((1 - 10) * (1 - 11))
     127             :         + 2 * 1 * 1 / ((1 - 11) * (1 - 12))
     128             :         - 1 * 1 * 1 / ((0 - 10) * (1 - 10))
     129             :         - 2 * 2 * 1 / ((1 - 12) * (0 - 12))
     130             :     )
     131           1 :     a_10 += (
     132             :         1 * 1 * 1 / ((0 - 10) * (0 - 11))
     133             :         + 2 * 1 * 1 / ((0 - 11) * (0 - 12))
     134             :         - 1 * 1 * 1 / ((0 - 11) * (1 - 11))
     135             :         - 3 * 3 * 1 / ((0 - 12) * (1 - 12))
     136             :     )
     137           1 :     a_11 += 1 * 1 * 3 / ((1 - 11) * (1 - 12)) + 3 * 1 * 1 / ((1 - 11) * (1 - 12)) - 3 * 2 * 1 / ((1 - 12) * (0 - 12))
     138           1 :     hamiltonian_new = np.array([[a_00, (a_01 + a_10) / 2], [(a_01 + a_10) / 2, a_11]])
     139           1 :     assert np.any(h_eff == hamiltonian_new)
     140             : 
     141           1 :     v0 += (
     142             :         -0.5 * (1 * 1 / (0 - 10) ** 2 + 2 * 2 / (0 - 12) ** 2) * sparse.eye(1, 5, k=0, format="csr")
     143             :         + (1 * 1 / ((0 - 10) * (0 - 11)) + 2 * 1 / ((0 - 11) * (0 - 12))) * sparse.eye(1, 5, k=3, format="csr")
     144             :         + 1 * 1 / ((0 - 11) * (0 - 1)) * sparse.eye(1, 5, k=3, format="csr")
     145             :         + 3 * 1 / ((0 - 1) * (0 - 12)) * sparse.eye(1, 5, k=4, format="csr")
     146             :         + 3 * 2 / ((0 - 1) * (0 - 12)) * sparse.eye(1, 5, k=1, format="csr")
     147             :     )
     148           1 :     v1 += (
     149             :         -0.5 * (1 * 1 / (1 - 11) ** 2 + 3 * 3 / (1 - 12) ** 2) * sparse.eye(1, 5, k=1, format="csr")
     150             :         + 1 * 1 / ((1 - 10) * (1 - 11)) * sparse.eye(1, 5, k=2, format="csr")
     151             :         + 1 * 3 / ((1 - 11) * (1 - 12)) * sparse.eye(1, 5, k=3, format="csr")
     152             :         + 1 * 1 / ((1 - 11) * (1 - 12)) * sparse.eye(1, 5, k=4, format="csr")
     153             :         + 1 * 1 / ((1 - 0) * (1 - 10)) * sparse.eye(1, 5, k=2, format="csr")
     154             :         + 2 * 1 / ((1 - 0) * (1 - 12)) * sparse.eye(1, 5, k=4, format="csr")
     155             :         + 3 * 2 / ((1 - 0) * (1 - 12)) * sparse.eye(1, 5, k=0, format="csr")
     156             :     )
     157           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.vstack([v0, v1])))
     158             : 
     159             : 
     160           1 : def test_c3_with_system() -> None:
     161             :     """Test whether the C3 coefficient with a given system is calculated correctly."""
     162           1 :     system_pair = _create_system_pair_sample()
     163             : 
     164           1 :     ket_tuple_list = [
     165             :         (pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5), pi.KetAtom("Rb", n=61, l=1, j=1.5, m=0.5)),
     166             :         (pi.KetAtom("Rb", n=61, l=1, j=1.5, m=0.5), pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)),
     167             :     ]
     168           1 :     c3 = perturbative.get_c3_from_system(
     169             :         system_pair=system_pair, ket_tuple_list=ket_tuple_list, unit="planck_constant * gigahertz * micrometer^3"
     170             :     )
     171           1 :     assert np.isclose(-0.5 * c3, 3.1515)
     172             : 
     173             : 
     174           1 : def test_c3_create_system() -> None:
     175             :     """Test whether the C3 coefficient with automatically constructed system is calculated correctly."""
     176           1 :     ket_tuple_list = [
     177             :         (pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5), pi.KetAtom("Rb", n=61, l=1, j=1.5, m=0.5)),
     178             :         (pi.KetAtom("Rb", n=61, l=1, j=1.5, m=0.5), pi.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)),
     179             :     ]
     180           1 :     magnetic_field = ureg.Quantity([0, 0, 10], "gauss")
     181           1 :     electric_field = ureg.Quantity([0, 0, 0], "volt/cm")
     182           1 :     distance_vector = ureg.Quantity([0, 0, 500], "micrometer")
     183             : 
     184           1 :     system = perturbative.create_system_for_perturbative(
     185             :         ket_tuple_list, electric_field, magnetic_field, distance_vector
     186             :     )
     187             : 
     188           1 :     c3 = perturbative.get_c3_from_system(
     189             :         system_pair=system, ket_tuple_list=ket_tuple_list, unit="planck_constant * gigahertz * micrometer^3"
     190             :     )
     191           1 :     assert np.isclose(-0.5 * c3, 3.2188)
     192             : 
     193             : 
     194           1 : def test_c6_with_system() -> None:
     195             :     """Test whether the C6 coefficient with a given system is calculated correctly."""
     196           1 :     system_pair = _create_system_pair_sample()
     197           1 :     ket_atom = pi.KetAtom(species="Rb", n=61, l=0, j=0.5, m=0.5)
     198           1 :     c6 = perturbative.get_c6_from_system(
     199             :         ket_tuple=(ket_atom, ket_atom), system_pair=system_pair, unit="planck_constant * gigahertz * micrometer^6"
     200             :     )
     201           1 :     assert np.isclose(c6, 167.92)
     202             : 
     203             : 
     204           1 : def test_c6_create_system() -> None:
     205             :     """Test whether the C6 coefficient with automatically constructed system is calculated correctly."""
     206           1 :     magnetic_field = ureg.Quantity([0, 0, 10], "gauss")
     207           1 :     electric_field = ureg.Quantity([0, 0, 0], "volt/cm")
     208           1 :     distance_vector = ureg.Quantity([0, 0, 500], "micrometer")
     209           1 :     ket_atom = pi.KetAtom(species="Rb", n=61, l=0, j=0.5, m=0.5)
     210             : 
     211           1 :     system = perturbative.create_system_for_perturbative(
     212             :         [(ket_atom, ket_atom)], electric_field, magnetic_field, distance_vector
     213             :     )
     214             : 
     215           1 :     c6 = perturbative.get_c6_from_system(
     216             :         ket_tuple=(ket_atom, ket_atom), system_pair=system, unit="planck_constant * gigahertz * micrometer^6"
     217             :     )
     218           1 :     assert np.isclose(c6, 169.189)
     219             : 
     220             : 
     221           1 : def test_resonance_detection(caplog: pytest.LogCaptureFixture) -> None:
     222             :     """Test whether resonance is correctly detected."""
     223           1 :     system_pair = _create_system_pair_sample()
     224           1 :     ket_tuple_list = [
     225             :         (pi.KetAtom(species="Rb", n=61, l=0, j=0.5, m=0.5), pi.KetAtom(species="Rb", n=61, l=1, j=0.5, m=0.5))
     226             :     ]
     227           1 :     with (
     228             :         pytest.raises(ValueError, match="Perturbative Calculation not possible due to resonances."),
     229             :         caplog.at_level(logging.CRITICAL),
     230             :     ):
     231           1 :         perturbative.get_effective_hamiltonian_from_system(ket_tuple_list, system_pair, order=2, required_overlap=0.8)

Generated by: LCOV version 1.16