Line data Source code
1 : # SPDX-FileCopyrightText: 2024 PairInteraction Developers
2 : # SPDX-License-Identifier: LGPL-3.0-or-later
3 :
4 1 : from __future__ import annotations
5 :
6 1 : from typing import TYPE_CHECKING
7 :
8 1 : import numpy as np
9 1 : import pytest
10 1 : from pairinteraction import BasisAtom, BasisPair
11 1 : from pairinteraction.ket.ket_pair import is_ket_pair_like
12 1 : from pairinteraction.state.state_pair import is_state_pair_like
13 :
14 : if TYPE_CHECKING:
15 : from pairinteraction import KetAtom, SystemAtom
16 : from pairinteraction.basis.basis_pair import BasisPairLike
17 : from pairinteraction.ket.ket_pair import KetPairLike
18 : from pairinteraction.state.state_pair import StatePairLike
19 :
20 : from .utils import PairinteractionModule
21 :
22 1 : from .utils import no_log_propagation
23 :
24 :
25 1 : @pytest.fixture
26 1 : def basis(pi_module: PairinteractionModule, system_atom: SystemAtom, system_atom2: SystemAtom) -> BasisPair:
27 1 : system_atoms = [system_atom, system_atom2]
28 1 : ket_atoms = tuple(s.basis.get_ket(30) for s in system_atoms)
29 1 : with no_log_propagation("pairinteraction.basis.basis_pair"): # suppress number_of_kets warning
30 1 : return pi_module.BasisPair.from_kets(ket_atoms, system_atoms, number_of_kets=80)
31 :
32 :
33 1 : @pytest.fixture
34 1 : def basis2(pi_module: PairinteractionModule, system_atom: SystemAtom) -> BasisPair:
35 1 : system_atoms = [system_atom, system_atom]
36 1 : ket_atoms = tuple(s.basis.get_ket(30) for s in system_atoms)
37 1 : with no_log_propagation("pairinteraction.basis.basis_pair"): # suppress number_of_kets warning
38 1 : return pi_module.BasisPair.from_kets(ket_atoms, system_atoms, number_of_kets=80)
39 :
40 :
41 1 : @pytest.fixture
42 1 : def basis_atom(pi_module: PairinteractionModule) -> BasisAtom:
43 1 : return pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2))
44 :
45 :
46 1 : @pytest.fixture
47 1 : def basis_atom2(pi_module: PairinteractionModule) -> BasisAtom:
48 1 : return pi_module.BasisAtom("Rb", n=(58, 62), l=(2, 3))
49 :
50 :
51 1 : @pytest.fixture
52 1 : def system_atom(pi_module: PairinteractionModule, basis_atom: BasisAtom) -> SystemAtom:
53 1 : return pi_module.SystemAtom(basis_atom)
54 :
55 :
56 1 : @pytest.fixture
57 1 : def system_atom2(pi_module: PairinteractionModule, basis_atom2: BasisAtom) -> SystemAtom:
58 1 : return pi_module.SystemAtom(basis_atom2)
59 :
60 :
61 1 : def test_basis_creation(pi_module: PairinteractionModule, system_atom: SystemAtom, system_atom2: SystemAtom) -> None:
62 : """Test basic properties of created basis."""
63 1 : basis = pi_module.BasisPair([system_atom, system_atom2])
64 1 : assert basis.number_of_kets == system_atom.basis.number_of_kets * system_atom2.basis.number_of_kets
65 1 : assert basis.number_of_states == basis.number_of_kets
66 1 : assert len(basis.kets) == basis.number_of_kets
67 1 : assert all(x in str(basis) for x in ["BasisPair", "Rb:58,S_1/2,", "...", "Rb:62,D_5/2,"])
68 :
69 :
70 1 : def test_coefficients(basis: BasisPair) -> None:
71 : """Test coefficient matrix properties."""
72 1 : coeffs = basis.get_coefficients()
73 1 : assert coeffs.shape == (basis.number_of_kets, basis.number_of_states)
74 1 : assert pytest.approx(coeffs.diagonal()) == 1.0 # NOSONAR
75 1 : assert pytest.approx(coeffs.sum()) == basis.number_of_kets # NOSONAR
76 :
77 :
78 1 : def _get_expected_shape(other: KetPairLike | StatePairLike | BasisPairLike, target_basis: BasisPair) -> tuple[int, ...]:
79 1 : if is_ket_pair_like(other) or is_state_pair_like(other):
80 1 : return (target_basis.number_of_states,)
81 1 : if isinstance(other, BasisPair):
82 1 : return (other.number_of_states, target_basis.number_of_states)
83 1 : if isinstance(other, list) and all(isinstance(b, BasisAtom) for b in other):
84 1 : return (other[0].number_of_states * other[1].number_of_states, target_basis.number_of_states)
85 0 : raise ValueError("Invalid basis_like type")
86 :
87 :
88 1 : @pytest.mark.parametrize(
89 : "other_key",
90 : [
91 : "ket_from_basis",
92 : "ket_from_basis2",
93 : "ket_atom_tuple_from_basis",
94 : "ket_atom_tuple_from_basis2",
95 : "state_from_basis",
96 : "state_from_basis2",
97 : "state_atom_tuple_from_basis",
98 : "state_atom_tuple_from_basis2",
99 : "state_atom_tuple_from_ket",
100 : "basis",
101 : "basis2",
102 : "basis_atom_tuple_from_basis",
103 : "basis_atom_tuple_from_basis2",
104 : ],
105 : )
106 1 : def test_get_methods(basis: BasisPair, basis2: BasisPair, other_key: str) -> None:
107 : """Test amplitude, overlap and matrix element calculations with another ket, state and basis."""
108 1 : other_dict: dict[str, KetPairLike | StatePairLike | BasisPairLike] = {
109 : "ket_from_basis": basis.get_ket(0),
110 : "ket_from_basis2": basis2.get_ket(0),
111 : "ket_atom_tuple_from_basis": [s.basis.get_ket(0) for s in basis.system_atoms],
112 : "ket_atom_tuple_from_basis2": [s.basis.get_ket(0) for s in basis2.system_atoms],
113 : "state_from_basis": basis.get_state(0),
114 : "state_from_basis2": basis2.get_state(0),
115 : "state_atom_tuple_from_basis": [s.basis.get_state(0) for s in basis.system_atoms],
116 : "state_atom_tuple_from_basis2": [s.basis.get_state(0) for s in basis2.system_atoms],
117 : "state_atom_tuple_from_ket": [s.basis.get_ket(0).to_state() for s in basis.system_atoms],
118 : "basis": basis,
119 : "basis2": basis2,
120 : "basis_atom_tuple_from_basis": [s.basis for s in basis.system_atoms],
121 : "basis_atom_tuple_from_basis2": [s.basis for s in basis2.system_atoms],
122 : }
123 1 : other = other_dict[other_key]
124 :
125 1 : amplitudes = basis.get_amplitudes(other)
126 1 : assert amplitudes.shape == _get_expected_shape(other, basis)
127 1 : if other_key.startswith("basis_atom_tuple"):
128 : # The tuple of one-atom bases represents product states, even for identical one-atom bases
129 1 : np.testing.assert_allclose(np.abs(amplitudes.data), 1)
130 :
131 1 : overlaps = basis.get_overlaps(other)
132 1 : assert overlaps.shape == _get_expected_shape(other, basis)
133 :
134 1 : elements_dipole = basis.get_matrix_elements(
135 : other, ("electric_dipole", "electric_dipole"), qs=(0, 0), unit="e^2 * a0^2"
136 : )
137 1 : assert elements_dipole.shape == _get_expected_shape(other, basis)
138 :
139 :
140 1 : def test_get_overlaps_pair_explicitly(pi_module: PairinteractionModule, system_atom: SystemAtom) -> None:
141 : # Test with a different basis (non-identical kets) - cross-basis overlaps
142 1 : ket = system_atom.basis.get_ket(20)
143 1 : energy_min = 2 * ket.get_energy("GHz") - 1
144 1 : energy_max = 2 * ket.get_energy("GHz") + 1
145 1 : basis_pair1 = pi_module.BasisPair([system_atom, system_atom])
146 1 : basis_pair2 = pi_module.BasisPair([system_atom, system_atom], energy=(energy_min, energy_max), energy_unit="GHz")
147 1 : assert basis_pair1.number_of_states > basis_pair2.number_of_states
148 1 : assert basis_pair1.number_of_states > 1
149 :
150 1 : matrix_overlaps = basis_pair1.get_overlaps(basis_pair2)
151 1 : assert matrix_overlaps.shape == (basis_pair2.number_of_states, basis_pair1.number_of_states)
152 1 : assert (matrix_overlaps.data >= 0).all()
153 1 : assert (matrix_overlaps.data <= 1 + 1e-10).all()
154 :
155 1 : col_sums = np.array(matrix_overlaps.sum(axis=0)).flatten()
156 1 : assert np.all(col_sums <= 1.0 + 1e-10)
157 1 : row_sums = np.array(matrix_overlaps.sum(axis=1)).flatten()
158 1 : assert np.all(row_sums <= 1.0 + 1e-10)
159 :
160 1 : idx0 = basis_pair1.get_corresponding_state_index((ket, ket))
161 1 : idx2 = basis_pair2.get_corresponding_state_index((ket, ket))
162 1 : assert pytest.approx(matrix_overlaps[idx2, idx0]) == 1.0 # NOSONAR
163 :
164 :
165 1 : def test_get_corresponding_ket_and_state(basis: BasisPair) -> None:
166 : """Test pair objects, pair-like tuples, and canonicalization retaining the atom systems."""
167 1 : expected_index = 1
168 1 : expected_ket = basis.get_ket(expected_index)
169 1 : expected_state = basis.get_state(expected_index)
170 1 : ket_atoms = tuple(state.get_corresponding_ket() for state in expected_ket.state_atoms)
171 1 : state_atoms = expected_ket.state_atoms
172 :
173 1 : assert basis.get_corresponding_state_index(expected_ket) == expected_index
174 1 : assert basis.get_corresponding_state(expected_ket).get_corresponding_ket() == expected_ket
175 1 : assert basis.get_corresponding_state_index(ket_atoms) == expected_index
176 1 : assert basis.get_corresponding_state(ket_atoms).get_corresponding_ket() == expected_ket
177 :
178 1 : assert basis.get_corresponding_ket_index(expected_state) == expected_index
179 1 : assert basis.get_corresponding_ket(expected_state) == expected_ket
180 1 : assert basis.get_corresponding_ket_index(state_atoms) == expected_index
181 1 : assert basis.get_corresponding_ket(state_atoms) == expected_ket
182 :
183 :
184 1 : def test_error_handling(basis: BasisPair) -> None:
185 : """Test error cases."""
186 1 : with pytest.raises(TypeError):
187 1 : basis.get_amplitudes("not a ket") # type: ignore [arg-type]
188 :
189 1 : with pytest.raises(TypeError):
190 1 : basis.get_overlaps("not a ket") # type: ignore [arg-type]
191 :
192 1 : with pytest.raises(TypeError):
193 1 : basis.get_matrix_elements("not a ket", ("energy", "energy"), (0, 0)) # type: ignore [arg-type]
194 :
195 :
196 1 : @pytest.mark.parametrize("restriction", ["product_of_parities", "parity_under_inversion", "parity_under_permutation"])
197 1 : def test_symmetrization_is_complete(
198 : pi_module: PairinteractionModule, system_atom: SystemAtom, restriction: str
199 : ) -> None:
200 : """The parity sectors partition the canonical basis.
201 :
202 : Summing the number of states over the even and odd sector reproduces the number of states of
203 : the unsymmetrized basis, i.e. symmetrization neither loses nor duplicates states.
204 : """
205 1 : canonical = pi_module.BasisPair([system_atom, system_atom], symmetrize=False)
206 1 : sectors = [pi_module.BasisPair([system_atom, system_atom], **{restriction: parity}) for parity in ["even", "odd"]] # type: ignore [arg-type]
207 :
208 1 : assert sum(sector.number_of_states for sector in sectors) == canonical.number_of_states
209 1 : assert all(0 < sector.number_of_states < canonical.number_of_states for sector in sectors)
210 :
211 :
212 1 : @pytest.mark.parametrize("restriction", ["parity_under_inversion", "parity_under_permutation"])
213 1 : def test_symmetrization_requires_same_system(
214 : pi_module: PairinteractionModule, system_atom: SystemAtom, system_atom2: SystemAtom, restriction: str
215 : ) -> None:
216 : """A parity restriction requires the symmetrization, i.e. the identical SystemAtom object for both atoms.
217 :
218 : Two distinct SystemAtom objects represent different atoms, so symmetrization is rejected
219 : by BasisPairCreator::create.
220 : """
221 1 : with pytest.raises(ValueError, match="not well-defined"):
222 1 : pi_module.BasisPair([system_atom, system_atom2], **{restriction: "odd"}) # type: ignore [arg-type]
223 1 : with pytest.raises(ValueError, match="not well-defined"):
224 1 : pi_module.BasisPair([system_atom, system_atom], **{restriction: "odd"}, symmetrize=False) # type: ignore [arg-type]
225 :
226 : # Without a parity restriction, two different systems remain allowed.
227 1 : basis = pi_module.BasisPair([system_atom, system_atom2])
228 1 : assert basis.number_of_states > 0
229 :
230 :
231 1 : def test_product_of_parities_does_not_require_same_system(
232 : pi_module: PairinteractionModule, system_atom: SystemAtom, system_atom2: SystemAtom
233 : ) -> None:
234 : """The product of parities is well-defined without symmetrization."""
235 1 : basis = pi_module.BasisPair([system_atom, system_atom2], product_of_parities="odd")
236 1 : assert basis.number_of_states > 0
237 :
238 :
239 1 : def test_from_kets(pi_module: PairinteractionModule, system_atom: SystemAtom) -> None:
240 : """Test BasisPair.from_kets."""
241 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
242 1 : ket1 = pi_module.KetAtom("Rb", n=59, l=1, j=0.5, m=0.5)
243 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=1, j=0.5, m=0.5)
244 1 : ket_atoms_dict: dict[str, tuple[KetAtom, KetAtom] | list[tuple[KetAtom, KetAtom]]] = {
245 : "single_ket": (ket, ket),
246 : "multiple_kets": [(ket, ket), (ket1, ket2), (ket2, ket1)],
247 : }
248 :
249 1 : for ket_atoms in ket_atoms_dict.values():
250 : # delta_energy restriction
251 1 : pair_basis = pi_module.BasisPair.from_kets(
252 : ket_atoms, [system_atom, system_atom], delta_energy=3, delta_energy_unit="GHz", delta_m=1
253 : )
254 1 : assert pair_basis.number_of_kets > 0
255 :
256 : # number_of_kets
257 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
258 1 : for target in [50, 100, 1000]:
259 1 : with no_log_propagation("pairinteraction.basis.basis_pair"): # suppress number_of_kets warning
260 1 : pair_basis = pi_module.BasisPair.from_kets(
261 : ket_atoms, [system_atom, system_atom], number_of_kets=target, delta_m=1
262 : )
263 1 : assert pair_basis.number_of_kets >= target
264 1 : assert pair_basis.number_of_kets < target + 20 # allow some extra due to degeneracies
265 :
266 : # test error cases
267 1 : with pytest.raises(ValueError, match="empty"):
268 1 : pi_module.BasisPair.from_kets([], [system_atom, system_atom])
269 :
270 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
271 1 : with pytest.raises(ValueError, match="number_of_kets"):
272 1 : pi_module.BasisPair.from_kets(
273 : (ket, ket),
274 : [system_atom, system_atom],
275 : delta_energy=3,
276 : delta_energy_unit="GHz",
277 : number_of_kets=10,
278 : )
279 :
280 :
281 1 : def test_merge_pair_basis(pi_module: PairinteractionModule, system_atom: SystemAtom, system_atom2: SystemAtom) -> None:
282 1 : system_atoms = [system_atom, system_atom2]
283 1 : ket_atoms1 = tuple(s.basis.get_ket(70) for s in system_atoms)
284 1 : ket_atoms2 = tuple(s.basis.get_ket(30) for s in system_atoms)
285 1 : with no_log_propagation("pairinteraction.basis.basis_pair"): # suppress number_of_kets warning
286 1 : basis1 = pi_module.BasisPair.from_kets(ket_atoms1, system_atoms, number_of_kets=20)
287 1 : basis2 = pi_module.BasisPair.from_kets(ket_atoms2, system_atoms, number_of_kets=10)
288 :
289 1 : merged = basis1.merge(basis2)
290 :
291 1 : assert set(merged.kets) == set(basis1.kets) | set(basis2.kets)
292 1 : assert merged.number_of_states > max(basis1.number_of_states, basis2.number_of_states)
293 1 : np.testing.assert_allclose(merged.get_coefficients().toarray(), np.eye(merged.number_of_kets))
294 :
295 1 : merged_with_self = basis1.merge(basis1)
296 1 : assert merged_with_self.kets == basis1.kets
297 :
298 : # merging is symmetric, i.e. the ket order does not depend on the argument order
299 1 : assert basis2.merge(basis1).kets == merged.kets
|
