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 :
50 1 : exact_basis = pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode="exact")
51 1 : assert exact_basis.number_of_kets > 0
52 : # The range is widened by a small absolute epsilon to tolerate floating-point noise in the
53 : # expectation values stored in the database, see test_restriction_at_expectation_value_boundary
54 1 : assert all(l_range[0] - 1e-9 <= ket.l <= l_range[1] + 1e-9 for ket in exact_basis.kets)
55 1 : assert exact_basis.number_of_kets < fuzzy_basis.number_of_kets
56 :
57 1 : factor_basis = pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode=100)
58 1 : assert factor_basis.number_of_kets > fuzzy_basis.number_of_kets
59 :
60 1 : with pytest.raises(ValueError, match="mode"):
61 1 : pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode="invalid") # type: ignore[arg-type]
62 :
63 1 : with pytest.raises(ValueError, match="non-negative"):
64 1 : pi_module.BasisAtom("Yb171_mqdt", nu=(58, 62), l=l_range, m=(0.5, 0.5), mode=-1)
65 :
66 :
67 1 : @pytest.mark.parametrize("j_max", [1, 2, 3])
68 1 : def test_restriction_at_expectation_value_boundary(pi_module: PairinteractionModule, j_max: int) -> None:
69 : """Test that states are not dropped because of floating-point noise in the expectation values.
70 :
71 : The expectation values stored in the database are channel-weighted sums, so a value that is
72 : mathematically an integer or half-integer can be off by a few ulp, e.g. exp_j is
73 : 2.0000000000000004 for some states with j = 2. Since these states have a vanishing standard
74 : deviation, a range that ends exactly at such a quantum number must not select fewer states than
75 : a range widened by a physically irrelevant amount.
76 : """
77 1 : basis = pi_module.BasisAtom("Yb174_mqdt", nu=(49, 55), j=(0, j_max), mode="exact")
78 1 : widened_basis = pi_module.BasisAtom("Yb174_mqdt", nu=(49, 55), j=(0, j_max + 1e-9), mode="exact")
79 1 : assert basis.number_of_kets > 0
80 1 : assert basis.number_of_kets == widened_basis.number_of_kets
81 :
82 :
83 1 : def test_coefficients(basis: BasisAtom) -> None:
84 : """Test coefficient matrix properties."""
85 1 : coeffs = basis.get_coefficients()
86 1 : assert coeffs.shape == (basis.number_of_kets, basis.number_of_states)
87 1 : assert pytest.approx(coeffs.diagonal()) == 1.0 # NOSONAR
88 1 : assert pytest.approx(coeffs.sum()) == basis.number_of_kets # NOSONAR
89 :
90 :
91 1 : def test_get_corresponding_ket_and_state(basis: BasisAtom) -> None:
92 : """Test both objects and indices for a non-canonical state and a canonical ket."""
93 1 : expected_ket_index = 1
94 1 : other_ket_index = 2
95 1 : expected_ket = basis.get_ket(expected_ket_index)
96 1 : other_ket = basis.get_ket(other_ket_index)
97 1 : mixed_state = (2 * basis.get_state(expected_ket_index) + basis.get_state(other_ket_index)).normalize()
98 :
99 1 : assert mixed_state.get_corresponding_ket_index() == expected_ket_index
100 1 : assert mixed_state.get_corresponding_ket() == expected_ket
101 1 : assert basis.get_corresponding_ket_index(mixed_state) == expected_ket_index
102 1 : assert basis.get_corresponding_ket(mixed_state) == expected_ket
103 :
104 1 : assert basis.get_corresponding_state_index(other_ket) == other_ket_index
105 1 : assert basis.get_corresponding_state(other_ket).get_corresponding_ket() == other_ket
106 :
107 :
108 1 : def _get_expected_shape(other: KetAtom | StateAtom | BasisAtom, target_basis: BasisAtom) -> tuple[int, ...]:
109 1 : if isinstance(other, (KetAtom, StateAtom)):
110 1 : return (target_basis.number_of_states,)
111 1 : if isinstance(other, BasisAtom):
112 1 : return (other.number_of_states, target_basis.number_of_states)
113 0 : raise ValueError("Invalid basis_like type")
114 :
115 :
116 1 : @pytest.mark.parametrize(
117 : "other_key",
118 : [
119 : "ket_from_basis",
120 : "ket_from_basis2",
121 : "other_ket_from_basis2",
122 : "state_from_basis",
123 : "state_from_basis2",
124 : "basis",
125 : "basis2",
126 : ],
127 : )
128 1 : def test_get_methods(basis: BasisAtom, basis2: BasisAtom, other_key: str) -> None:
129 : """Test amplitude, overlap and matrix element calculations with another ket, state and basis."""
130 1 : ind = 5
131 1 : other_dict: dict[str, KetAtom | StateAtom | BasisAtom] = {
132 : "ket_from_basis": basis.get_ket(ind),
133 : "ket_from_basis2": next(ket for ket in basis2.kets if ket in basis.kets),
134 : "other_ket_from_basis2": next(ket for ket in basis2.kets if ket not in basis.kets and ket.j_ryd < 3),
135 : "state_from_basis": basis.get_state(ind),
136 : "state_from_basis2": basis2.get_state(0),
137 : "basis": basis,
138 : "basis2": basis2,
139 : }
140 1 : other = other_dict[other_key]
141 :
142 1 : amplitudes = basis.get_amplitudes(other)
143 1 : assert amplitudes.shape == _get_expected_shape(other, basis)
144 :
145 1 : overlaps = basis.get_overlaps(other)
146 1 : assert overlaps.shape == _get_expected_shape(other, basis)
147 :
148 1 : matrix_elements = basis.get_matrix_elements(other, "electric_dipole", q=0, unit="e * a0")
149 1 : assert matrix_elements.shape == _get_expected_shape(other, basis)
150 :
151 1 : if other_key in ["ket_from_basis", "state_from_basis"]:
152 1 : assert pytest.approx(amplitudes[ind]) == 1.0 # NOSONAR
153 1 : assert pytest.approx(overlaps[ind]) == 1.0 # NOSONAR
154 :
155 1 : if other_key == "ket_from_basis2":
156 1 : assert isinstance(amplitudes, np.ndarray)
157 1 : assert isinstance(overlaps, np.ndarray)
158 1 : assert pytest.approx(np.max(amplitudes)) == 1.0 # NOSONAR
159 1 : assert pytest.approx(np.max(overlaps)) == 1.0 # NOSONAR
160 :
161 1 : if other_key == "other_ket_from_basis2":
162 1 : assert isinstance(amplitudes, np.ndarray)
163 1 : assert isinstance(overlaps, np.ndarray)
164 1 : assert np.count_nonzero(amplitudes) == 0
165 1 : assert np.count_nonzero(overlaps) == 0
166 :
167 1 : if other_key == "basis":
168 1 : assert isinstance(amplitudes, csr_matrix)
169 1 : assert isinstance(overlaps, csr_matrix)
170 1 : assert pytest.approx(amplitudes.diagonal()) == 1.0 # NOSONAR
171 1 : assert pytest.approx(overlaps.diagonal()) == 1.0 # NOSONAR
172 :
173 1 : if other_key == "basis2":
174 1 : assert isinstance(amplitudes, csr_matrix)
175 1 : assert isinstance(overlaps, csr_matrix)
176 1 : n_matching_kets = len([ket for ket in basis2.kets if ket in basis.kets])
177 1 : assert np.count_nonzero(amplitudes.toarray()) == n_matching_kets
178 1 : assert np.count_nonzero(overlaps.toarray()) == n_matching_kets
179 :
180 1 : if other_key.startswith("basis"):
181 1 : assert isinstance(matrix_elements, csr_matrix)
182 1 : assert 0 < np.count_nonzero(matrix_elements.toarray()) < basis.number_of_states**2
183 : else:
184 1 : assert isinstance(matrix_elements, np.ndarray)
185 1 : assert 0 < np.count_nonzero(matrix_elements) < basis.number_of_states
186 :
187 :
188 1 : def test_error_handling(basis: BasisAtom) -> None:
189 : """Test error cases."""
190 1 : with pytest.raises(TypeError):
191 1 : basis.get_amplitudes("not a ket") # type: ignore [call-overload]
192 :
193 1 : with pytest.raises(TypeError):
194 1 : basis.get_overlaps("not a ket") # type: ignore [call-overload]
195 :
196 1 : with pytest.raises(TypeError):
197 1 : basis.get_matrix_elements("not a ket", "energy", 0) # type: ignore [call-overload]
198 :
199 :
200 1 : def test_from_kets(pi_module: PairinteractionModule) -> None:
201 : """Test BasisAtom.from_kets."""
202 : # single ket
203 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
204 1 : basis = pi_module.BasisAtom.from_kets(
205 : ket,
206 : delta_n=2,
207 : delta_nu=3,
208 : delta_nui=3,
209 : delta_l=2,
210 : delta_s=1,
211 : delta_j=3,
212 : delta_l_ryd=2,
213 : delta_j_ryd=3,
214 : delta_f=3,
215 : delta_m=2,
216 : delta_energy=100,
217 : delta_energy_unit="GHz",
218 : )
219 1 : assert basis.species == "Rb"
220 1 : assert all(58 <= k.n <= 62 for k in basis.kets)
221 1 : assert any(k.n == 62 for k in basis.kets)
222 1 : assert any(k.n == 58 for k in basis.kets)
223 1 : assert any(k == ket for k in basis.kets)
224 :
225 : # multiple kets
226 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
227 1 : ket2 = pi_module.KetAtom("Rb", n=61, l=0, j=0.5, m=0.5)
228 1 : basis = pi_module.BasisAtom.from_kets([ket1, ket2], delta_n=2)
229 1 : assert all(58 <= k.n <= 63 for k in basis.kets)
230 1 : assert any(k.n == 63 for k in basis.kets)
231 1 : assert any(k.n == 58 for k in basis.kets)
232 1 : assert any(k == ket1 for k in basis.kets)
233 1 : assert any(k == ket2 for k in basis.kets)
234 :
235 : # test that from_kets is consistent with direct constructor
236 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
237 1 : basis_from = pi_module.BasisAtom.from_kets(ket, delta_n=2)
238 1 : basis_direct = pi_module.BasisAtom("Rb", n=(58, 62))
239 1 : assert basis_from.number_of_kets == basis_direct.number_of_kets
240 :
241 : # test error cases
242 1 : with pytest.raises(ValueError, match="empty"):
243 1 : pi_module.BasisAtom.from_kets([])
244 :
245 1 : ket_rb = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
246 1 : ket_sr = pi_module.KetAtom("Sr88_sqdt", n=60, l=1, s=0, j=1, m=0)
247 1 : with pytest.raises(ValueError, match="species"):
248 1 : pi_module.BasisAtom.from_kets([ket_rb, ket_sr])
249 :
250 :
251 1 : def test_merge(pi_module: PairinteractionModule) -> None:
252 1 : basis1 = pi_module.BasisAtom("Rb", n=(60, 60), l=(0, 1), m=(0.5, 0.5))
253 1 : basis2 = pi_module.BasisAtom("Rb", n=(60, 60), l=(1, 2), m=(0.5, 0.5))
254 :
255 1 : merged = basis1.merge(basis2)
256 1 : expected_kets = set(basis1.kets) | set(basis2.kets)
257 :
258 1 : assert set(merged.kets) == expected_kets
259 1 : assert merged.number_of_kets == len(expected_kets)
260 1 : assert merged.number_of_states == merged.number_of_kets
261 1 : np.testing.assert_allclose(merged.get_coefficients().toarray(), np.eye(merged.number_of_kets))
262 1 : assert basis1.number_of_kets < merged.number_of_kets
263 1 : assert basis2.number_of_kets < merged.number_of_kets
264 :
265 1 : merged_with_self = basis1.merge(basis1)
266 1 : assert merged_with_self.kets == basis1.kets
|
