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)
|