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
11 1 : from pairinteraction.ket.ket_atom import KetAtom
12 1 : from pairinteraction.state.state_atom import StateAtom
13 1 : from scipy.sparse import csr_matrix
14 :
15 : if TYPE_CHECKING:
16 : from .utils import PairinteractionModule
17 :
18 :
19 1 : @pytest.fixture
20 1 : def basis(pi_module: PairinteractionModule) -> BasisAtom:
21 1 : return pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2))
22 :
23 :
24 1 : @pytest.fixture
25 1 : def basis2(pi_module: PairinteractionModule) -> BasisAtom:
26 1 : return pi_module.BasisAtom("Rb", n=(58, 62), l=(2, 3))
27 :
28 :
29 1 : def test_basis_creation(pi_module: PairinteractionModule) -> None:
30 : """Test basic properties of created basis."""
31 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
32 1 : energy_min = ket.get_energy(unit="GHz") - 100
33 1 : energy_max = ket.get_energy(unit="GHz") + 100
34 1 : basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2), energy=(energy_min, energy_max), energy_unit="GHz")
35 1 : assert basis.species == "Rb"
36 1 : assert basis.number_of_kets == 80
37 1 : assert basis.number_of_states == basis.number_of_kets
38 1 : assert len(basis.kets) == basis.number_of_kets
39 1 : assert basis.number_of_kets < pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2)).number_of_kets
40 1 : assert all(x in str(basis) for x in ["BasisAtom", "n=(58, 62)", "l=(0, 2)"])
41 :
42 :
43 1 : def test_restriction_mode(pi_module: PairinteractionModule) -> None:
44 : """Test exact, fuzzy, and numeric restrictions for expectation-value quantum numbers."""
45 1 : l_range = (1, 1)
46 1 : fuzzy_basis = pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5))
47 1 : assert fuzzy_basis.number_of_kets > 0
48 1 : assert any(ket.l > l_range[1] for ket in fuzzy_basis.kets)
49 1 : assert any(ket.l < l_range[0] for ket in fuzzy_basis.kets)
50 :
51 1 : exact_basis = pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode="exact")
52 1 : assert exact_basis.number_of_kets > 0
53 1 : assert all(l_range[0] <= ket.l <= l_range[1] for ket in exact_basis.kets)
54 1 : assert exact_basis.number_of_kets < fuzzy_basis.number_of_kets
55 :
56 1 : factor_basis = pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode=100)
57 1 : assert factor_basis.number_of_kets > fuzzy_basis.number_of_kets
58 :
59 1 : with pytest.raises(ValueError, match="mode"):
60 1 : pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode="invalid") # type: ignore[arg-type]
61 :
62 1 : with pytest.raises(ValueError, match="non-negative"):
63 1 : pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode=-1)
64 :
65 :
66 1 : def test_coefficients(basis: BasisAtom) -> None:
67 : """Test coefficient matrix properties."""
68 1 : coeffs = basis.get_coefficients()
69 1 : assert coeffs.shape == (basis.number_of_kets, basis.number_of_states)
70 1 : assert pytest.approx(coeffs.diagonal()) == 1.0 # NOSONAR
71 1 : assert pytest.approx(coeffs.sum()) == basis.number_of_kets # NOSONAR
72 :
73 :
74 1 : def test_get_corresponding_ket_and_state(basis: BasisAtom) -> None:
75 : """Test both objects and indices for a non-canonical state and a canonical ket."""
76 1 : expected_ket_index = 1
77 1 : other_ket_index = 2
78 1 : expected_ket = basis.get_ket(expected_ket_index)
79 1 : other_ket = basis.get_ket(other_ket_index)
80 1 : mixed_state = (2 * basis.get_state(expected_ket_index) + basis.get_state(other_ket_index)).normalize()
81 :
82 1 : assert mixed_state.get_corresponding_ket_index() == expected_ket_index
83 1 : assert mixed_state.get_corresponding_ket() == expected_ket
84 1 : assert basis.get_corresponding_ket_index(mixed_state) == expected_ket_index
85 1 : assert basis.get_corresponding_ket(mixed_state) == expected_ket
86 :
87 1 : assert basis.get_corresponding_state_index(other_ket) == other_ket_index
88 1 : assert basis.get_corresponding_state(other_ket).get_corresponding_ket() == other_ket
89 :
90 :
91 1 : def _get_expected_shape(other: KetAtom | StateAtom | BasisAtom, target_basis: BasisAtom) -> tuple[int, ...]:
92 1 : if isinstance(other, (KetAtom, StateAtom)):
93 1 : return (target_basis.number_of_states,)
94 1 : if isinstance(other, BasisAtom):
95 1 : return (other.number_of_states, target_basis.number_of_states)
96 0 : raise ValueError("Invalid basis_like type")
97 :
98 :
99 1 : @pytest.mark.parametrize(
100 : "other_key",
101 : [
102 : "ket_from_basis",
103 : "ket_from_basis2",
104 : "other_ket_from_basis2",
105 : "state_from_basis",
106 : "state_from_basis2",
107 : "basis",
108 : "basis2",
109 : ],
110 : )
111 1 : def test_get_methods(basis: BasisAtom, basis2: BasisAtom, other_key: str) -> None:
112 : """Test amplitude, overlap and matrix element calculations with another ket, state and basis."""
113 1 : ind = 5
114 1 : other_dict: dict[str, KetAtom | StateAtom | BasisAtom] = {
115 : "ket_from_basis": basis.get_ket(ind),
116 : "ket_from_basis2": next(ket for ket in basis2.kets if ket in basis.kets),
117 : "other_ket_from_basis2": next(ket for ket in basis2.kets if ket not in basis.kets and ket.j_ryd < 3),
118 : "state_from_basis": basis.get_state(ind),
119 : "state_from_basis2": basis2.get_state(0),
120 : "basis": basis,
121 : "basis2": basis2,
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 :
128 1 : overlaps = basis.get_overlaps(other)
129 1 : assert overlaps.shape == _get_expected_shape(other, basis)
130 :
131 1 : matrix_elements = basis.get_matrix_elements(other, "electric_dipole", q=0, unit="e * a0")
132 1 : assert matrix_elements.shape == _get_expected_shape(other, basis)
133 :
134 1 : if other_key in ["ket_from_basis", "state_from_basis"]:
135 1 : assert pytest.approx(amplitudes[ind]) == 1.0 # NOSONAR
136 1 : assert pytest.approx(overlaps[ind]) == 1.0 # NOSONAR
137 :
138 1 : if other_key == "ket_from_basis2":
139 1 : assert isinstance(amplitudes, np.ndarray)
140 1 : assert isinstance(overlaps, np.ndarray)
141 1 : assert pytest.approx(np.max(amplitudes)) == 1.0 # NOSONAR
142 1 : assert pytest.approx(np.max(overlaps)) == 1.0 # NOSONAR
143 :
144 1 : if other_key == "other_ket_from_basis2":
145 1 : assert isinstance(amplitudes, np.ndarray)
146 1 : assert isinstance(overlaps, np.ndarray)
147 1 : assert np.count_nonzero(amplitudes) == 0
148 1 : assert np.count_nonzero(overlaps) == 0
149 :
150 1 : if other_key == "basis":
151 1 : assert isinstance(amplitudes, csr_matrix)
152 1 : assert isinstance(overlaps, csr_matrix)
153 1 : assert pytest.approx(amplitudes.diagonal()) == 1.0 # NOSONAR
154 1 : assert pytest.approx(overlaps.diagonal()) == 1.0 # NOSONAR
155 :
156 1 : if other_key == "basis2":
157 1 : assert isinstance(amplitudes, csr_matrix)
158 1 : assert isinstance(overlaps, csr_matrix)
159 1 : n_matching_kets = len([ket for ket in basis2.kets if ket in basis.kets])
160 1 : assert np.count_nonzero(amplitudes.toarray()) == n_matching_kets
161 1 : assert np.count_nonzero(overlaps.toarray()) == n_matching_kets
162 :
163 1 : if other_key.startswith("basis"):
164 1 : assert isinstance(matrix_elements, csr_matrix)
165 1 : assert 0 < np.count_nonzero(matrix_elements.toarray()) < basis.number_of_states**2
166 : else:
167 1 : assert isinstance(matrix_elements, np.ndarray)
168 1 : assert 0 < np.count_nonzero(matrix_elements) < basis.number_of_states
169 :
170 :
171 1 : def test_error_handling(basis: BasisAtom) -> None:
172 : """Test error cases."""
173 1 : with pytest.raises(TypeError):
174 1 : basis.get_amplitudes("not a ket") # type: ignore [call-overload]
175 :
176 1 : with pytest.raises(TypeError):
177 1 : basis.get_overlaps("not a ket") # type: ignore [call-overload]
178 :
179 1 : with pytest.raises(TypeError):
180 1 : basis.get_matrix_elements("not a ket", "energy", 0) # type: ignore [call-overload]
181 :
182 :
183 1 : def test_from_kets(pi_module: PairinteractionModule) -> None:
184 : """Test BasisAtom.from_kets."""
185 : # single ket
186 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
187 1 : basis = pi_module.BasisAtom.from_kets(
188 : ket,
189 : delta_n=2,
190 : delta_nu=3,
191 : delta_nui=3,
192 : delta_l=2,
193 : delta_s=1,
194 : delta_j=3,
195 : delta_l_ryd=2,
196 : delta_j_ryd=3,
197 : delta_f=3,
198 : delta_m=2,
199 : delta_energy=100,
200 : delta_energy_unit="GHz",
201 : )
202 1 : assert basis.species == "Rb"
203 1 : assert all(58 <= k.n <= 62 for k in basis.kets)
204 1 : assert any(k.n == 62 for k in basis.kets)
205 1 : assert any(k.n == 58 for k in basis.kets)
206 1 : assert any(k == ket for k in basis.kets)
207 :
208 : # multiple kets
209 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
210 1 : ket2 = pi_module.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)
211 1 : basis = pi_module.BasisAtom.from_kets([ket1, ket2], delta_n=2)
212 1 : assert all(58 <= k.n <= 63 for k in basis.kets)
213 1 : assert any(k.n == 63 for k in basis.kets)
214 1 : assert any(k.n == 58 for k in basis.kets)
215 1 : assert any(k == ket1 for k in basis.kets)
216 1 : assert any(k == ket2 for k in basis.kets)
217 :
218 : # test that from_kets is consistent with direct constructor
219 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
220 1 : basis_from = pi_module.BasisAtom.from_kets(ket, delta_n=2)
221 1 : basis_direct = pi_module.BasisAtom("Rb", n=(58, 62))
222 1 : assert basis_from.number_of_kets == basis_direct.number_of_kets
223 :
224 : # test error cases
225 1 : with pytest.raises(ValueError, match="empty"):
226 1 : pi_module.BasisAtom.from_kets([])
227 :
228 1 : ket_rb = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
229 1 : ket_sr = pi_module.KetAtom("Sr88_singlet", n=60, l=1, j=1, m=0)
230 1 : with pytest.raises(ValueError, match="species"):
231 1 : pi_module.BasisAtom.from_kets([ket_rb, ket_sr])
232 :
233 :
234 1 : def test_merge(pi_module: PairinteractionModule) -> None:
235 1 : basis1 = pi_module.BasisAtom("Rb", n=(60, 60), l=(0, 1), m=(0.5, 0.5))
236 1 : basis2 = pi_module.BasisAtom("Rb", n=(60, 60), l=(1, 2), m=(0.5, 0.5))
237 :
238 1 : merged = basis1.merge(basis2)
239 1 : expected_kets = set(basis1.kets) | set(basis2.kets)
240 :
241 1 : assert set(merged.kets) == expected_kets
242 1 : assert merged.number_of_kets == len(expected_kets)
243 1 : assert merged.number_of_states == merged.number_of_kets
244 1 : np.testing.assert_allclose(merged.get_coefficients().toarray(), np.eye(merged.number_of_kets))
245 1 : assert basis1.number_of_kets < merged.number_of_kets
246 1 : assert basis2.number_of_kets < merged.number_of_kets
247 :
248 1 : merged_with_self = basis1.merge(basis1)
249 1 : assert merged_with_self.kets == basis1.kets
|
