LCOV - code coverage report
Current view: top level - tests - test_perturbative.py (source / functions) Hit Total Coverage
Test: coverage.info Lines: 89 90 98.9 %
Date: 2025-06-06 09:09:03 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.perturbative 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 :     hamiltonian_eff, eig_perturb = _calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, order=0)
      86           1 :     assert np.any(hamiltonian_eff == np.array([[0, 0], [0, 1]]))
      87           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.eye(2, 5, k=0)))
      88             : 
      89             :     # Order 1
      90           1 :     hamiltonian_eff, eig_perturb = _calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, order=1)
      91           1 :     assert np.any(hamiltonian_eff == np.array([[0, 1], [1, 1]]))
      92           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix([[1, 0, 0, 0, 0], [0, 1, 0, 0, 0]]))
      93             : 
      94             :     # Order 2
      95           1 :     hamiltonian_eff, eig_perturb = _calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, order=2)
      96           1 :     a_00 = 0 + 1 * 1 / (0 - 10) + 2 * 2 / (0 - 12)
      97           1 :     a_01 = 1 + 2 * 3 / (0 - 12)
      98           1 :     a_10 = 1 + 3 * 2 / (1 - 12)
      99           1 :     a_11 = 1 + 1 * 1 / (1 - 11) + 3 * 3 / (1 - 12)
     100           1 :     hamiltonian_new = np.array([[a_00, (a_01 + a_10) / 2], [(a_01 + a_10) / 2, a_11]])
     101           1 :     assert np.any(hamiltonian_eff == hamiltonian_new)
     102             : 
     103           1 :     v0 = (
     104             :         sparse.eye(1, 5, k=0, format="csr")
     105             :         + 1 / (0 - 10) * sparse.eye(1, 5, k=2, format="csr")
     106             :         + 2 / (0 - 12) * sparse.eye(1, 5, k=4, format="csr")
     107             :     )
     108           1 :     v1 = (
     109             :         sparse.eye(1, 5, k=1, format="csr")
     110             :         + 1 / (1 - 11) * sparse.eye(1, 5, k=3, format="csr")
     111             :         + 3 / (1 - 12) * sparse.eye(1, 5, k=4, format="csr")
     112             :     )
     113           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.vstack([v0, v1])))
     114             : 
     115             :     # Order 3
     116           1 :     with caplog.at_level(logging.ERROR):
     117           1 :         hamiltonian_eff, eig_perturb = _calculate_perturbative_hamiltonian(hamiltonian, model_space_indices, order=3)
     118           1 :     a_00 -= 2 * 3 * 1 / ((0 - 12) * (1 - 12))
     119           1 :     a_01 += (
     120             :         1 * 1 * 1 / ((1 - 10) * (1 - 11))
     121             :         + 2 * 1 * 1 / ((1 - 11) * (1 - 12))
     122             :         - 1 * 1 * 1 / ((0 - 10) * (1 - 10))
     123             :         - 2 * 2 * 1 / ((1 - 12) * (0 - 12))
     124             :     )
     125           1 :     a_10 += (
     126             :         1 * 1 * 1 / ((0 - 10) * (0 - 11))
     127             :         + 2 * 1 * 1 / ((0 - 11) * (0 - 12))
     128             :         - 1 * 1 * 1 / ((0 - 11) * (1 - 11))
     129             :         - 3 * 3 * 1 / ((0 - 12) * (1 - 12))
     130             :     )
     131           1 :     a_11 += 1 * 1 * 3 / ((1 - 11) * (1 - 12)) + 3 * 1 * 1 / ((1 - 11) * (1 - 12)) - 3 * 2 * 1 / ((1 - 12) * (0 - 12))
     132           1 :     hamiltonian_new = np.array([[a_00, (a_01 + a_10) / 2], [(a_01 + a_10) / 2, a_11]])
     133           1 :     assert np.any(hamiltonian_eff == hamiltonian_new)
     134             : 
     135           1 :     v0 += (
     136             :         -0.5 * (1 * 1 / (0 - 10) ** 2 + 2 * 2 / (0 - 12) ** 2) * sparse.eye(1, 5, k=0, format="csr")
     137             :         + (1 * 1 / ((0 - 10) * (0 - 11)) + 2 * 1 / ((0 - 11) * (0 - 12))) * sparse.eye(1, 5, k=3, format="csr")
     138             :         + 1 * 1 / ((0 - 11) * (0 - 1)) * sparse.eye(1, 5, k=3, format="csr")
     139             :         + 3 * 1 / ((0 - 1) * (0 - 12)) * sparse.eye(1, 5, k=4, format="csr")
     140             :         + 3 * 2 / ((0 - 1) * (0 - 12)) * sparse.eye(1, 5, k=1, format="csr")
     141             :     )
     142           1 :     v1 += (
     143             :         -0.5 * (1 * 1 / (1 - 11) ** 2 + 3 * 3 / (1 - 12) ** 2) * sparse.eye(1, 5, k=1, format="csr")
     144             :         + 1 * 1 / ((1 - 10) * (1 - 11)) * sparse.eye(1, 5, k=2, format="csr")
     145             :         + 1 * 3 / ((1 - 11) * (1 - 12)) * sparse.eye(1, 5, k=3, format="csr")
     146             :         + 1 * 1 / ((1 - 11) * (1 - 12)) * sparse.eye(1, 5, k=4, format="csr")
     147             :         + 1 * 1 / ((1 - 0) * (1 - 10)) * sparse.eye(1, 5, k=2, format="csr")
     148             :         + 2 * 1 / ((1 - 0) * (1 - 12)) * sparse.eye(1, 5, k=4, format="csr")
     149             :         + 3 * 2 / ((1 - 0) * (1 - 12)) * sparse.eye(1, 5, k=0, format="csr")
     150             :     )
     151           1 :     assert _check_sparse_matrices_equal(eig_perturb, sparse.csr_matrix(sparse.vstack([v0, v1])))
     152             : 
     153             : 
     154           1 : def test_c3_with_system() -> None:
     155             :     """Test whether the C3 coefficient with a given system is calculated correctly."""
     156           1 :     system_pair = _create_system_pair_sample()
     157             : 
     158           1 :     ket_tuple_list = [
     159             :         (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)),
     160             :         (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)),
     161             :     ]
     162           1 :     c3 = perturbative.get_c3_from_system(
     163             :         system_pair=system_pair, ket_tuple_list=ket_tuple_list, unit="planck_constant * gigahertz * micrometer^3"
     164             :     )
     165           1 :     assert np.isclose(-0.5 * c3, 3.1515)
     166             : 
     167             : 
     168           1 : def test_c3_create_system() -> None:
     169             :     """Test whether the C3 coefficient with automatically constructed system is calculated correctly."""
     170           1 :     ket_tuple_list = [
     171             :         (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)),
     172             :         (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)),
     173             :     ]
     174           1 :     magnetic_field = ureg.Quantity([0, 0, 10], "gauss")
     175           1 :     electric_field = ureg.Quantity([0, 0, 0], "volt/cm")
     176           1 :     distance_vector = ureg.Quantity([0, 0, 500], "micrometer")
     177             : 
     178           1 :     system = perturbative.create_system_for_perturbative(
     179             :         ket_tuple_list, electric_field, magnetic_field, distance_vector
     180             :     )
     181             : 
     182           1 :     c3 = perturbative.get_c3_from_system(
     183             :         system_pair=system, ket_tuple_list=ket_tuple_list, unit="planck_constant * gigahertz * micrometer^3"
     184             :     )
     185           1 :     assert np.isclose(-0.5 * c3, 3.2188)
     186             : 
     187             : 
     188           1 : def test_c6_with_system() -> None:
     189             :     """Test whether the C6 coefficient with a given system is calculated correctly."""
     190           1 :     system_pair = _create_system_pair_sample()
     191           1 :     ket_atom = pi.KetAtom(species="Rb", n=61, l=0, j=0.5, m=0.5)
     192           1 :     c6 = perturbative.get_c6_from_system(
     193             :         ket_tuple=(ket_atom, ket_atom), system_pair=system_pair, unit="planck_constant * gigahertz * micrometer^6"
     194             :     )
     195           1 :     assert np.isclose(c6, 167.92)
     196             : 
     197             : 
     198           1 : def test_c6_create_system() -> None:
     199             :     """Test whether the C6 coefficient with automatically constructed system is calculated correctly."""
     200           1 :     magnetic_field = ureg.Quantity([0, 0, 10], "gauss")
     201           1 :     electric_field = ureg.Quantity([0, 0, 0], "volt/cm")
     202           1 :     distance_vector = ureg.Quantity([0, 0, 500], "micrometer")
     203           1 :     ket_atom = pi.KetAtom(species="Rb", n=61, l=0, j=0.5, m=0.5)
     204             : 
     205           1 :     system = perturbative.create_system_for_perturbative(
     206             :         [(ket_atom, ket_atom)], electric_field, magnetic_field, distance_vector
     207             :     )
     208             : 
     209           1 :     c6 = perturbative.get_c6_from_system(
     210             :         ket_tuple=(ket_atom, ket_atom), system_pair=system, unit="planck_constant * gigahertz * micrometer^6"
     211             :     )
     212           1 :     assert np.isclose(c6, 169.189)
     213             : 
     214             : 
     215           1 : def test_resonance_detection(caplog: pytest.LogCaptureFixture) -> None:
     216             :     """Test whether resonance is correctly detected."""
     217           1 :     system_pair = _create_system_pair_sample()
     218           1 :     ket_tuple_list = [
     219             :         (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))
     220             :     ]
     221           1 :     with (
     222             :         pytest.raises(ValueError, match="Perturbative Calculation not possible due to resonances."),
     223             :         caplog.at_level(logging.CRITICAL),
     224             :     ):
     225           1 :         perturbative.get_effective_hamiltonian_from_system(
     226             :             ket_tuple_list=ket_tuple_list, system_pair=system_pair, order=2, required_overlap=0.8
     227             :         )

Generated by: LCOV version 1.16