Line data Source code
1 : # SPDX-FileCopyrightText: 2025 Pairinteraction Developers
2 : # SPDX-License-Identifier: LGPL-3.0-or-later
3 :
4 : """Test receiving matrix elements from the databases."""
5 :
6 1 : from collections.abc import Generator
7 1 : from pathlib import Path
8 1 : from typing import Union
9 :
10 1 : import duckdb
11 1 : import numpy as np
12 1 : import pairinteraction.real as pi
13 1 : import pytest
14 1 : from packaging.version import Version
15 1 : from sympy.physics.wigner import wigner_3j
16 :
17 1 : from tests.constants import GAUSS_IN_ATOMIC_UNITS, HARTREE_IN_GHZ, SPECIES_TO_NUCLEAR_SPIN, SUPPORTED_SPECIES
18 :
19 :
20 1 : def fetch_id(
21 : n: int, l: float, f: float, s: float, connection: duckdb.DuckDBPyConnection, table: Union[str, Path]
22 : ) -> int:
23 1 : result = connection.execute(
24 : f"SELECT id FROM '{table}' WHERE n = {n} AND f = {f} ORDER BY (exp_s - {s})^2+(exp_l - {l})^2 LIMIT 1" # noqa: S608
25 : ).fetchone()
26 1 : return result[0] if result else -1
27 :
28 :
29 1 : def fetch_wigner_element(
30 : f_initial: float,
31 : f_final: float,
32 : m_initial: float,
33 : m_final: float,
34 : kappa: int,
35 : q: int,
36 : connection: duckdb.DuckDBPyConnection,
37 : table: Union[str, Path],
38 : ) -> float:
39 1 : result = connection.execute(
40 : f"SELECT val FROM '{table}' WHERE f_initial = {f_initial} AND f_final = {f_final} AND " # noqa: S608
41 : f"m_initial = {m_initial} AND m_final = {m_final} AND kappa = {kappa} AND q = {q}"
42 : ).fetchone()
43 1 : return result[0] if result else 0
44 :
45 :
46 1 : def fetch_reduced_matrix_element(
47 : id_initial: int, id_final: int, connection: duckdb.DuckDBPyConnection, table: Union[str, Path]
48 : ) -> float:
49 1 : result = connection.execute(
50 : f"SELECT val FROM '{table}' WHERE id_initial = {id_initial} AND id_final = {id_final}" # noqa: S608
51 : ).fetchone()
52 1 : return result[0] if result else 0
53 :
54 :
55 1 : @pytest.fixture(scope="module")
56 1 : def connection() -> Generator[duckdb.DuckDBPyConnection]:
57 1 : with duckdb.connect(":memory:") as connection:
58 1 : yield connection
59 :
60 :
61 1 : @pytest.mark.parametrize("swap_states", [False, True])
62 1 : def test_database(connection: duckdb.DuckDBPyConnection, swap_states: bool) -> None: # noqa: PLR0915
63 : """Test receiving matrix elements from the databases."""
64 1 : database = pi.Database.get_global_database()
65 1 : bfield_in_gauss = 1500
66 :
67 : # Define initial and final quantum states
68 1 : n_initial, n_final = 54, 54
69 1 : l_initial, l_final = 1, 1
70 1 : f_initial, f_final = 1, 0
71 1 : m_initial, m_final = 0, 0
72 1 : s_initial, s_final = 0.6, 1.0
73 :
74 : # Swap states if required by the test parameter
75 1 : if swap_states:
76 1 : n_initial, n_final = n_final, n_initial
77 1 : l_initial, l_final = l_final, l_initial
78 1 : f_initial, f_final = f_final, f_initial
79 1 : m_initial, m_final = m_final, m_initial
80 1 : s_initial, s_final = s_final, s_initial
81 :
82 : # Get the Zeeman interaction operator from the database using pairinteraction
83 1 : ket_initial = pi.KetAtom("Yb174_mqdt", n=n_initial, l=l_initial, f=f_initial, m=m_initial, s=s_initial)
84 1 : ket_final = pi.KetAtom("Yb174_mqdt", n=n_final, l=l_final, f=f_final, m=m_final, s=s_final)
85 1 : basis = pi.BasisAtom("Yb174_mqdt", additional_kets=[ket_initial, ket_final])
86 1 : operator = (
87 : pi.SystemAtom(basis)
88 : .set_magnetic_field([0, 0, bfield_in_gauss], unit="G")
89 : .set_diamagnetism_enabled(True)
90 : .get_hamiltonian(unit="GHz")
91 : ).toarray()
92 1 : operator -= np.diag(np.sort([ket_initial.get_energy(unit="GHz"), ket_final.get_energy(unit="GHz")]))
93 1 : expected_operator = np.array([[3.58588117, 1.66420213], [1.66420213, 4.16645123]])
94 1 : assert np.allclose(operator, expected_operator)
95 :
96 : # Get the latest parquet files from the database directory
97 1 : parquet_files: dict[str, Path] = {}
98 1 : parquet_versions: dict[str, Version] = {}
99 1 : for path in list(Path(database.database_dir).rglob("*.parquet")):
100 1 : species, version_str = path.parent.name.rsplit("_v", 1)
101 1 : table = path.stem
102 1 : name = f"{species}_{table}"
103 1 : version = Version(version_str)
104 1 : if name not in parquet_files or version > parquet_versions[name]:
105 1 : parquet_files[name] = path
106 1 : parquet_versions[name] = version
107 1 : assert "misc_wigner" in parquet_files
108 1 : assert "Yb174_mqdt_states" in parquet_files
109 1 : assert "Yb174_mqdt_matrix_elements_mu" in parquet_files
110 1 : assert "Yb174_mqdt_matrix_elements_q" in parquet_files
111 1 : assert "Yb174_mqdt_matrix_elements_q0" in parquet_files
112 :
113 : # Obtain the ids of the initial and final states
114 1 : id_initial = fetch_id(n_initial, l_initial, f_initial, s_initial, connection, parquet_files["Yb174_mqdt_states"])
115 1 : assert id_initial == 362 if swap_states else 363
116 :
117 1 : id_final = fetch_id(n_final, l_final, f_final, s_final, connection, parquet_files["Yb174_mqdt_states"])
118 1 : assert id_final == 363 if swap_states else 362
119 :
120 : # Obtain a matrix element of the magnetic dipole operator (for the chosen kets, it is non-zero iff initial != final)
121 1 : kappa, q = 1, 0
122 1 : wigner_element = fetch_wigner_element(
123 : f_initial, f_final, m_initial, m_final, kappa, q, connection, parquet_files["misc_wigner"]
124 : )
125 1 : assert np.isclose(
126 : wigner_element,
127 : float((-1) ** (f_final - m_final) * wigner_3j(f_final, kappa, f_initial, -m_final, q, m_initial)),
128 : )
129 :
130 1 : me_mu = fetch_reduced_matrix_element(
131 : id_initial, id_final, connection, parquet_files["Yb174_mqdt_matrix_elements_mu"]
132 : )
133 1 : matrix_element = -wigner_element * me_mu * bfield_in_gauss * GAUSS_IN_ATOMIC_UNITS * HARTREE_IN_GHZ
134 1 : assert np.isclose(matrix_element, operator[0, 1])
135 :
136 : # Obtain a matrix element of the diamagnetic operator (for the chosen kets, it is non-zero iff initial == final)
137 1 : n_final = n_initial
138 1 : l_final = l_initial
139 1 : f_final = f_initial
140 1 : m_final = m_initial
141 1 : s_final = s_initial
142 1 : id_final = id_initial
143 :
144 1 : kappa, q = 0, 0
145 1 : wigner_element = fetch_wigner_element(
146 : f_initial, f_final, m_initial, m_final, kappa, q, connection, parquet_files["misc_wigner"]
147 : )
148 1 : assert np.isclose(
149 : wigner_element,
150 : float((-1) ** (f_final - m_final) * wigner_3j(f_final, kappa, f_initial, -m_final, q, m_initial)),
151 : )
152 :
153 1 : me_q0 = fetch_reduced_matrix_element(
154 : id_initial, id_final, connection, parquet_files["Yb174_mqdt_matrix_elements_q0"]
155 : )
156 1 : matrix_element = 1 / 12 * wigner_element * me_q0 * (bfield_in_gauss * GAUSS_IN_ATOMIC_UNITS) ** 2 * HARTREE_IN_GHZ
157 :
158 1 : kappa, q = 2, 0
159 1 : wigner_element = fetch_wigner_element(
160 : f_initial, f_final, m_initial, m_final, kappa, q, connection, parquet_files["misc_wigner"]
161 : )
162 1 : assert np.isclose(
163 : wigner_element,
164 : float((-1) ** (f_final - m_final) * wigner_3j(f_final, kappa, f_initial, -m_final, q, m_initial)),
165 : )
166 :
167 1 : me_q = fetch_reduced_matrix_element(id_initial, id_final, connection, parquet_files["Yb174_mqdt_matrix_elements_q"])
168 1 : matrix_element -= 1 / 12 * wigner_element * me_q * (bfield_in_gauss * GAUSS_IN_ATOMIC_UNITS) ** 2 * HARTREE_IN_GHZ
169 1 : assert np.isclose(matrix_element, operator[0, 0] if swap_states else operator[1, 1])
170 :
171 :
172 1 : @pytest.mark.parametrize("species", SUPPORTED_SPECIES)
173 1 : def test_obtaining_kets(species: str) -> None:
174 : """Test obtaining kets from the database."""
175 1 : is_mqdt = species.endswith("_mqdt")
176 1 : is_single_valence_electron = species in ["Rb"]
177 1 : is_triplet = species in ["Sr88_triplet"]
178 :
179 1 : quantum_number_i = SPECIES_TO_NUCLEAR_SPIN[species] if is_mqdt else 0
180 1 : quantum_number_s = 0.5 if is_single_valence_electron else (1 if is_triplet else 0)
181 1 : quantum_number_f = quantum_number_i + quantum_number_s
182 1 : quantum_number_m = quantum_number_i + quantum_number_s
183 :
184 : # Obtain a ket from the database
185 1 : ket = pi.KetAtom(species, n=60, l=0, f=quantum_number_f, m=quantum_number_m, s=quantum_number_s)
186 :
187 : # Check the result
188 1 : assert ket.species == species
189 1 : assert ket.n == 60 if not is_mqdt else abs(ket.n - 60) < 1
190 1 : assert ket.l == 0 if not is_mqdt else abs(ket.l - 0) < 1
191 1 : assert ket.f == quantum_number_f
192 1 : assert ket.m == quantum_number_m
193 1 : assert ket.s == quantum_number_s if not is_mqdt else abs(ket.s - quantum_number_s) < 1
194 1 : assert {
195 : "Rb": "Rb:60,S_1/2,1/2",
196 : "Sr88_singlet": "Sr88_singlet:60,S_0,0",
197 : "Sr88_triplet": "Sr88_triplet:60,S_1,1",
198 : "Sr88_mqdt": "Sr88:S=0.0,nu=56.7,L=0.0,J=0,0",
199 : "Sr87_mqdt": "Sr87:S=0.5,nu=56.6,L=0.0,F=9/2,9/2",
200 : "Yb174_mqdt": "Yb174:S=0.0,nu=55.7,L=0.0,J=0,0",
201 : "Yb171_mqdt": "Yb171:S=0.3,nu=55.7,L=0.0,F=1/2,1/2",
202 : "Yb173_mqdt": "Yb173:S=0.4,nu=56.4,L=0.0,F=5/2,5/2",
203 : }[species] == repr(ket)
|
