Line data Source code
1 : # SPDX-FileCopyrightText: 2024 PairInteraction Developers
2 : # SPDX-License-Identifier: LGPL-3.0-or-later
3 1 : from __future__ import annotations
4 :
5 1 : from typing import TYPE_CHECKING, Any, Literal, overload
6 :
7 1 : import numpy as np
8 1 : from scipy.sparse import csr_matrix
9 :
10 1 : from pairinteraction import _backend
11 1 : from pairinteraction.basis.basis_base import BasisBase
12 1 : from pairinteraction.database import Database
13 1 : from pairinteraction.enums import get_cpp_operator_type, parity_to_int
14 1 : from pairinteraction.ket import KetAtom, KetAtomReal
15 1 : from pairinteraction.state import StateAtom, StateAtomReal
16 1 : from pairinteraction.units import QuantityArray, QuantityScalar, QuantitySparse
17 :
18 : if TYPE_CHECKING:
19 : from collections.abc import Sequence
20 :
21 : from typing_extensions import Self
22 :
23 : from pairinteraction.enums import OperatorType, Parity
24 : from pairinteraction.units import NDArray, PintArray, PintFloat, PintSparse
25 :
26 :
27 1 : class BasisAtom(BasisBase[KetAtom, StateAtom]):
28 : """Basis for a single atom.
29 :
30 : Add all KetAtom objects that match the given quantum numbers to the basis.
31 : The initial coefficients matrix is a unit matrix, i.e. the first basis state is the first ket, etc.
32 : The BasisAtom coefficients matrix will always be square,
33 : i.e. the number of kets is equal to the number of states.
34 :
35 : Examples:
36 : >>> import pairinteraction as pi
37 : >>> ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
38 : >>> energy_min, energy_max = ket.get_energy(unit="GHz") - 100, ket.get_energy(unit="GHz") + 100
39 : >>> basis = pi.BasisAtom("Rb", n=(57, 63), l=(0, 3), energy=(energy_min, energy_max), energy_unit="GHz")
40 : >>> print(basis)
41 : BasisAtom('Rb', n=(57, 63), l=(0, 3), energy=(1008911.9216, 1009111.9216), energy_unit='GHz')
42 :
43 : """
44 :
45 1 : _cpp: _backend.BasisAtomComplex
46 1 : _cpp_creator = _backend.BasisAtomCreatorComplex
47 1 : _ket_class = KetAtom
48 1 : _state_class = StateAtom
49 :
50 1 : _args: dict[str, Any] | None = None
51 :
52 1 : def __init__( # noqa: C901, PLR0912
53 : self,
54 : species: str,
55 : n: tuple[int, int] | None = None,
56 : nu: tuple[float, float] | None = None,
57 : nui: tuple[float, float] | None = None,
58 : l: tuple[float, float] | None = None,
59 : s: tuple[float, float] | None = None,
60 : j: tuple[float, float] | None = None,
61 : l_ryd: tuple[float, float] | None = None,
62 : j_ryd: tuple[float, float] | None = None,
63 : f: tuple[float, float] | None = None,
64 : m: tuple[float, float] | None = None,
65 : energy: tuple[float, float] | tuple[PintFloat, PintFloat] | None = None,
66 : energy_unit: str | None = None,
67 : parity: Parity | None = None,
68 : additional_kets: Sequence[KetAtom] | None = None,
69 : *,
70 : database: Database | None = None,
71 : mode: Literal["exact", "fuzzy"] | float = "fuzzy",
72 : ) -> None:
73 : """Create a basis for a single atom.
74 :
75 : Args:
76 : species: The species of the atom.
77 : n: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
78 : nu: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
79 : nui: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
80 : l: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
81 : s: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
82 : j: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
83 : l_ryd: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
84 : j_ryd: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
85 : f: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
86 : m: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
87 : energy: tuple of (min, max) value for the energy. Default None, i.e. add all available states.
88 : energy_unit: In which unit the energy values are given, e.g. "GHz".
89 : Default None, i.e. energy is provided as pint object.
90 : parity: The parity of the states to consider. Default None, i.e. add all available states.
91 : additional_kets: List of additional kets to add to the basis. Default None.
92 : database: Which database to use. Default None, i.e. use the global database instance.
93 : mode: Specifies how restrictions on expectation-value quantum numbers are applied.
94 : ``"fuzzy"`` is equal to 2 and includes states whose expectation-value overlaps the requested range
95 : within two standard deviations.
96 : ``"exact"`` is equal to 0 and includes only states whose expectation-value itself lie in the range.
97 : A non-negative number sets the factor applied to the standard deviation explicitly.
98 : Default ``"fuzzy"``.
99 :
100 : """
101 1 : self._args = {"species": species}
102 :
103 1 : creator = self._cpp_creator()
104 1 : creator.set_species(species)
105 :
106 1 : if n is not None and not all(isinstance(x, int) or x.is_integer() for x in n):
107 0 : raise ValueError("Quantum numbers n must be integers.")
108 :
109 1 : quantum_numbers = {
110 : "n": n,
111 : "nu": nu,
112 : "nui": nui,
113 : "s": s,
114 : "l": l,
115 : "j": j,
116 : "l_ryd": l_ryd,
117 : "j_ryd": j_ryd,
118 : "f": f,
119 : "m": m,
120 : }
121 1 : for name, value in quantum_numbers.items():
122 1 : if value is not None:
123 1 : self._args[name] = value
124 1 : creator.restrict_quantum_number(name, *value)
125 :
126 1 : if parity is not None:
127 0 : self._args["parity"] = parity
128 0 : parity_int = parity_to_int(parity)
129 0 : creator.restrict_quantum_number("parity", parity_int, parity_int)
130 :
131 1 : if energy is not None:
132 1 : self._args.update({"energy": energy, "energy_unit": energy_unit})
133 1 : min_energy_au = QuantityScalar.convert_user_to_au(energy[0], energy_unit, "energy")
134 1 : max_energy_au = QuantityScalar.convert_user_to_au(energy[1], energy_unit, "energy")
135 1 : creator.restrict_energy(min_energy_au, max_energy_au)
136 :
137 1 : if database is None:
138 1 : if Database.get_global_database() is None:
139 0 : Database.initialize_global_database()
140 1 : database = Database.get_global_database()
141 :
142 1 : if additional_kets is not None:
143 1 : self._args["additional_kets"] = additional_kets
144 1 : for ket in additional_kets:
145 1 : creator.add_ket(ket._cpp)
146 :
147 1 : if mode == "fuzzy":
148 1 : quantum_number_standard_deviation_factor = 2.0
149 1 : elif mode == "exact":
150 1 : self._args["mode"] = mode
151 1 : quantum_number_standard_deviation_factor = 0.0
152 : else:
153 1 : self._args["mode"] = mode
154 1 : msg = "mode must be 'exact', 'fuzzy', or a non-negative number."
155 1 : try:
156 1 : quantum_number_standard_deviation_factor = float(mode)
157 1 : except (TypeError, ValueError) as err:
158 1 : raise ValueError(msg) from err
159 1 : if quantum_number_standard_deviation_factor < 0:
160 1 : raise ValueError(msg)
161 1 : creator.set_quantum_number_standard_deviation_factor(quantum_number_standard_deviation_factor)
162 :
163 1 : self._cpp = creator.create(database._cpp)
164 1 : self._post_init()
165 :
166 1 : @classmethod
167 1 : def from_kets(
168 : cls: type[Self],
169 : kets: KetAtom | Sequence[KetAtom],
170 : delta_n: int | None = None,
171 : delta_nu: float | None = None,
172 : delta_nui: float | None = None,
173 : delta_l: float | None = None,
174 : delta_s: float | None = None,
175 : delta_j: float | None = None,
176 : delta_l_ryd: float | None = None,
177 : delta_j_ryd: float | None = None,
178 : delta_f: int | None = None,
179 : delta_m: int | None = None,
180 : delta_energy: float | PintFloat | None = None,
181 : delta_energy_unit: str | None = None,
182 : parity: Parity | None = None,
183 : database: Database | None = None,
184 : additional_kets: Sequence[KetAtom] | None = None,
185 : *,
186 : mode: Literal["exact", "fuzzy"] | float = "fuzzy",
187 : ) -> Self:
188 : """Create a BasisAtom from one or more kets and quantum number deltas.
189 :
190 : Currently a single big basis including all kets for the quantum numbers
191 : from min_value - delta to max_value + delta is returned.
192 : In the future this might change to return a basis including all states around the given kets +/- delta,
193 : but not necessarily all states between the given kets.
194 :
195 : For each quantum number, pass the corresponding ``delta_*`` argument to include
196 : all states within ``[min_value - delta, max_value + delta]``, where
197 : ``min_value`` / ``max_value`` are the extremes across all provided kets.
198 : If no ``delta_*`` is given for a quantum number, that quantum number is left
199 : unrestricted.
200 :
201 : Args:
202 : kets: The ket(s) around which the basis should be centered.
203 : delta_n: Half-width of the n window (integer steps).
204 : Default None means no restriction on n.
205 : delta_nu: Half-width of the nu window.
206 : Default None means no restriction on nu.
207 : delta_nui: Half-width of the nui window.
208 : Default None means no restriction on nui.
209 : delta_l: Half-width of the l window.
210 : Default None means no restriction on l.
211 : delta_s: Half-width of the s window.
212 : Default None means no restriction on s.
213 : delta_j: Half-width of the j window.
214 : Default None means no restriction on j.
215 : delta_l_ryd: Half-width of the l_ryd window.
216 : Default None means no restriction on l_ryd.
217 : delta_j_ryd: Half-width of the j_ryd window.
218 : Default None means no restriction on j_ryd.
219 : delta_f: Half-width of the f window (integer steps).
220 : Default None means no restriction on f.
221 : delta_m: Half-width of the m window (integer steps).
222 : Default None means no restriction on m.
223 : delta_energy: Half-width of the energy window around the energies of the
224 : provided kets. Default None means no energy restriction.
225 : delta_energy_unit: Unit for ``delta_energy`` (e.g. ``"GHz"``).
226 : Default None means pint quantities are used.
227 : parity: Restrict to states with this parity.
228 : Default None means no parity restriction.
229 : database: Database instance to use.
230 : Default None uses the global database.
231 : additional_kets: Extra kets to force-include in the basis.
232 : Default None.
233 : mode: Passed to :class:`BasisAtom`. Default ``"fuzzy"``.
234 :
235 : Returns:
236 : A new :class:`BasisAtom` centered around the provided kets.
237 :
238 : Examples:
239 : >>> import pairinteraction as pi
240 : >>> ket1 = pi.KetAtom("Rb", n=60, l=0, m=0.5)
241 : >>> ket2 = pi.KetAtom("Rb", n=59, l=0, m=0.5)
242 : >>> basis = pi.BasisAtom.from_kets([ket1, ket2], delta_n=2, delta_l=1)
243 : >>> basis.species
244 : 'Rb'
245 : >>> all(57 <= k.n <= 62 for k in basis.kets)
246 : True
247 :
248 : """
249 1 : if isinstance(kets, KetAtom):
250 1 : kets = [kets]
251 1 : kets = list(kets)
252 1 : if len(kets) == 0:
253 1 : raise ValueError("kets must not be empty.")
254 1 : if len({ket.species for ket in kets}) > 1:
255 1 : raise ValueError(f"All kets must have the same species, but got: {sorted({ket.species for ket in kets})}.")
256 :
257 1 : def get_range(name: str, delta: float | None) -> tuple[float, float] | None:
258 1 : if delta is None:
259 1 : return None
260 1 : if name == "energy":
261 1 : values = [ket.get_energy(unit=delta_energy_unit) for ket in kets]
262 : else:
263 1 : values = [getattr(ket, name) for ket in kets]
264 1 : return (min(values) - delta, max(values) + delta)
265 :
266 1 : return cls(
267 : species=kets[0].species,
268 : n=get_range("n", delta_n), # type: ignore [arg-type]
269 : nu=get_range("nu", delta_nu),
270 : nui=get_range("nui", delta_nui),
271 : l=get_range("l", delta_l),
272 : s=get_range("s", delta_s),
273 : j=get_range("j", delta_j),
274 : l_ryd=get_range("l_ryd", delta_l_ryd),
275 : j_ryd=get_range("j_ryd", delta_j_ryd),
276 : f=get_range("f", delta_f),
277 : m=get_range("m", delta_m),
278 : energy=get_range("energy", delta_energy), # type: ignore [arg-type]
279 : energy_unit=delta_energy_unit,
280 : parity=parity,
281 : database=database,
282 : additional_kets=additional_kets,
283 : mode=mode,
284 : )
285 :
286 1 : def __repr__(self) -> str:
287 1 : if self._args is None:
288 1 : return super().__repr__()
289 :
290 1 : args_str: list[str] = []
291 1 : for k, v in self._args.items():
292 1 : if k == "species":
293 1 : args_str.append(repr(v))
294 1 : elif k == "energy":
295 1 : args_str.append(f"{k}=({v[0]:.4f}, {v[1]:.4f})")
296 : else:
297 1 : args_str.append(f"{k}={v!r}")
298 1 : return f"{type(self).__name__}({', '.join(args_str)})"
299 :
300 1 : @property
301 1 : def database(self) -> Database:
302 : """The database used for this object."""
303 0 : return self.get_ket(0).database
304 :
305 1 : @property
306 1 : def species(self) -> str:
307 : """The atomic species."""
308 1 : return self.get_ket(0).species
309 :
310 : @overload
311 : def get_amplitudes(self, other: KetAtom | StateAtom) -> NDArray: ...
312 :
313 : @overload
314 : def get_amplitudes(self, other: BasisAtom) -> csr_matrix: ...
315 :
316 1 : def get_amplitudes(self, other: KetAtom | StateAtom | BasisAtom) -> NDArray | csr_matrix:
317 1 : return self.get_matrix_elements(other, "identity", 0, unit="")
318 :
319 : @overload
320 : def get_overlaps(self, other: KetAtom | StateAtom) -> NDArray: ...
321 :
322 : @overload
323 : def get_overlaps(self, other: BasisAtom) -> csr_matrix: ...
324 :
325 1 : def get_overlaps(self, other: KetAtom | StateAtom | BasisAtom) -> NDArray | csr_matrix:
326 1 : amplitudes = self.get_amplitudes(other)
327 1 : if isinstance(amplitudes, csr_matrix):
328 1 : return amplitudes.multiply(amplitudes.conj()).real # type: ignore [no-any-return]
329 1 : return np.abs(amplitudes) ** 2
330 :
331 : @overload
332 : def get_matrix_elements(
333 : self, other: KetAtom | StateAtom, operator: OperatorType, q: int, unit: None = None
334 : ) -> PintArray: ...
335 :
336 : @overload
337 : def get_matrix_elements(self, other: KetAtom | StateAtom, operator: OperatorType, q: int, unit: str) -> NDArray: ...
338 :
339 : @overload
340 : def get_matrix_elements(
341 : self, other: BasisAtom, operator: OperatorType, q: int, unit: None = None
342 : ) -> PintSparse: ...
343 :
344 : @overload
345 : def get_matrix_elements(self, other: BasisAtom, operator: OperatorType, q: int, unit: str) -> csr_matrix: ...
346 :
347 1 : def get_matrix_elements(
348 : self, other: KetAtom | StateAtom | BasisAtom, operator: OperatorType, q: int, unit: str | None = None
349 : ) -> NDArray | PintArray | csr_matrix | PintSparse:
350 1 : cpp_op = get_cpp_operator_type(operator)
351 :
352 : matrix_elements_au: NDArray
353 1 : if isinstance(other, KetAtom):
354 1 : other = other.to_state()
355 1 : if isinstance(other, StateAtom):
356 1 : matrix_elements_au = self._cpp.get_matrix_elements(other._cpp, cpp_op, q).toarray().ravel()
357 1 : matrix_elements_au = np.real_if_close(matrix_elements_au)
358 1 : return QuantityArray.convert_au_to_user(matrix_elements_au, operator, unit)
359 1 : if isinstance(other, BasisAtom):
360 1 : matrix_elements_sparse_au = self._cpp.get_matrix_elements(other._cpp, cpp_op, q)
361 1 : matrix_elements_sparse_au.data = np.real_if_close(matrix_elements_sparse_au.data)
362 1 : return QuantitySparse.convert_au_to_user(matrix_elements_sparse_au, operator, unit)
363 1 : raise TypeError(f"Unknown type: {type(other)=}")
364 :
365 :
366 1 : class BasisAtomReal(BasisAtom):
367 1 : _cpp: _backend.BasisAtomReal # type: ignore [assignment]
368 1 : _cpp_creator = _backend.BasisAtomCreatorReal # type: ignore [assignment]
369 1 : _ket_class = KetAtomReal
370 1 : _state_class = StateAtomReal
371 :
372 :
373 1 : def get_cpp_basis_atom_from_kets(kets: Sequence[KetAtom], *, real: bool) -> _backend.BasisAtomComplex:
374 : """Create a cpp BasisAtom object containing only the given kets.
375 :
376 : Like for the _cpp attributes, the return type is annotated as the complex variant,
377 : although the real variant is returned if real=True.
378 : """
379 1 : if len(kets) == 0:
380 0 : raise ValueError("Cannot create a basis with zero kets.")
381 1 : creator = _backend.BasisAtomCreatorReal() if real else _backend.BasisAtomCreatorComplex()
382 1 : for ket in kets:
383 1 : creator.add_ket(ket._cpp)
384 1 : return creator.create(kets[0].database._cpp) # type: ignore [return-value]
|
