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