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 : import logging
6 1 : from functools import cached_property
7 1 : from typing import TYPE_CHECKING, Literal, overload
8 :
9 1 : import numpy as np
10 1 : from scipy.special import exprel
11 :
12 1 : from pairinteraction import _backend
13 1 : from pairinteraction.database import Database
14 1 : from pairinteraction.enums import OperatorType, Parity, int_to_parity, parity_to_int
15 1 : from pairinteraction.ket.ket_base import KetBase
16 1 : from pairinteraction.ket.utils import format_half_integer, get_l_label
17 1 : from pairinteraction.units import QuantityArray, QuantityScalar, ureg
18 :
19 : if TYPE_CHECKING:
20 : from typing_extensions import Self
21 :
22 : from pairinteraction.enums import OperatorType, Parity
23 : from pairinteraction.state import StateAtom, StateAtomReal
24 : from pairinteraction.units import NDArray, PintArray, PintComplex, PintFloat
25 :
26 :
27 1 : logger = logging.getLogger(__name__)
28 :
29 :
30 1 : class KetAtom(KetBase):
31 : """Ket for an atomic basis state.
32 :
33 : Each KetAtom object uniquely represents a single-atom basis state
34 : (and therefore all KetAtom objects are orthogonal).
35 : When initializing a KetAtom you have to provide the species of the atom and a combination of quantum numbers,
36 : which uniquely define a single-atom basis state (this always includes providing a magnetic quantum number m).
37 :
38 : SQDT (Single Channel Quantum Defect Theory) for one valence electron (alkali atoms):
39 : The quantum numbers n (int), l (int), j (half-int) and m (half-int)
40 : should be used to define the desired atomic basis state.
41 : All other quantum numbers are trivially derived from these:
42 : s = 1/2, f = j (we neglect hyperfine interaction for SQDT),
43 : nu = n - delta, l_ryd = l, j_ryd = j.
44 :
45 : SQDT (Single Channel Quantum Defect Theory) for two valence electrons (alkaline-earth atoms):
46 : The quantum numbers n (int), l_ryd (int), j (int) and m (int)
47 : should be used to define the desired atomic basis state.
48 : The spin quantum number s is taken from the species label,
49 : which must end either with "_singlet" (s=0) or "_triplet" (s=1).
50 : Again we neglect hyperfine interaction, thus f = j. And nu = n - delta.
51 : All other quantum numbers are not necessarily eigenvalues anymore and are given as expectation values.
52 :
53 : MQDT (Multi Channel Quantum Defect Theory) for two valence electrons (alkaline-earth atoms):
54 : The quantum numbers nu (float), f (int or half-int) and m (int or half-int) are still good quantum numbers.
55 : All other quantum numbers (like l, s, j, l_ryd, j_ryd) are not necessarily eigenvalues anymore.
56 : You can still provide them to specify the atomic basis state,
57 : whose expectation value is closest to the provided value.
58 :
59 : Examples:
60 : >>> import pairinteraction as pi
61 : >>> ket_s = pi.KetAtom("Rb", n=60, l=0, m=0.5)
62 : >>> (ket_s.species, ket_s.n, ket_s.l, ket_s.j, ket_s.m, ket_s.s)
63 : ('Rb', 60, 0.0, 0.5, 0.5, 0.5)
64 : >>> print(ket_s)
65 : |Rb:60,S_1/2,1/2⟩
66 : >>> print(ket_s.to_state())
67 : 1.00 |Rb:60,S_1/2,1/2⟩
68 : >>> ket_p = pi.KetAtom("Rb", n=60, l=1, j=0.5, m=0.5)
69 : >>> print((2 * ket_p - ket_s).normalize())
70 : 0.89 |Rb:60,P_1/2,1/2⟩ - 0.45 |Rb:60,S_1/2,1/2⟩
71 : >>> ket_mqdt = pi.KetAtom("Yb174_mqdt", nu=60, l=1, f=1, m=1)
72 : >>> (ket_mqdt.species, round(ket_mqdt.nu, 3), ket_mqdt.f, ket_mqdt.m)
73 : ('Yb174_mqdt', 60.049, 1.0, 1.0)
74 : >>> print(ket_mqdt)
75 : |Yb174:S=0.0,nu=60.0,L=1.0,J=1,1⟩
76 :
77 : """
78 :
79 1 : _cpp: _backend.KetAtom
80 :
81 1 : def __init__(
82 : self,
83 : species: str,
84 : n: int | None = None,
85 : nu: float | None = None,
86 : nui: float | None = None,
87 : l: float | None = None,
88 : s: float | None = None,
89 : j: float | None = None,
90 : l_ryd: float | None = None,
91 : j_ryd: float | None = None,
92 : f: float | None = None,
93 : m: float | None = None,
94 : energy: float | PintFloat | None = None,
95 : energy_unit: str | None = None,
96 : parity: Parity | None = None,
97 : database: Database | None = None,
98 : ) -> None:
99 : """Create a single-atom canonical basis state, which is defined by its species and quantum numbers.
100 :
101 : Args:
102 : species: See attribute.
103 : n: See attribute. Default None, i.e. load from the database.
104 : nu: See attribute. Default None, i.e. load from the database.
105 : nui: See attribute. Default None, i.e. load from the database.
106 : l: See attribute. Default None, i.e. load from the database.
107 : s: See attribute. Default None, i.e. load from the database.
108 : j: See attribute. Default None, i.e. load from the database.
109 : l_ryd: See attribute. Default None, i.e. load from the database.
110 : j_ryd: See attribute. Default None, i.e. load from the database.
111 : f: See attribute. Default None, i.e. load from the database.
112 : m: See attribute. This should always be provided.
113 : energy: See attribute. Default None, i.e. load from the database.
114 : energy_unit: In which unit the energy is given, e.g. "GHz".
115 : Default None, i.e. energy is provided as pint object.
116 : parity: See attribute. Default None, i.e. load from the database.
117 : database: Which database to use. Default None, i.e. use the global database instance.
118 :
119 : """
120 1 : creator = _backend.KetAtomCreator()
121 1 : creator.set_species(species)
122 1 : if energy is not None:
123 0 : energy_au = QuantityScalar.convert_user_to_au(energy, energy_unit, "energy")
124 0 : creator.set_energy(energy_au)
125 1 : if n is not None and not (isinstance(n, int) or n.is_integer()):
126 0 : raise ValueError("Quantum number n must be an integer.")
127 1 : quantum_numbers = {
128 : "f": f,
129 : "m": m,
130 : "n": n,
131 : "nu": nu,
132 : "nui": nui,
133 : "l": l,
134 : "s": s,
135 : "j": j,
136 : "l_ryd": l_ryd,
137 : "j_ryd": j_ryd,
138 : "parity": parity_to_int(parity) if parity is not None else None,
139 : }
140 1 : for name, value in quantum_numbers.items():
141 1 : if value is not None:
142 1 : creator.set_quantum_number(name, value)
143 1 : if database is None:
144 1 : if Database.get_global_database() is None:
145 0 : Database.initialize_global_database()
146 1 : database = Database.get_global_database()
147 1 : try:
148 1 : self._cpp = creator.create(database._cpp)
149 1 : except _backend.KetNotUniqueError as err:
150 0 : candidates = [type(self)._from_cpp_object(ket) for ket in err.kets] # type: ignore [attr-defined]
151 0 : labels = "\n".join(ket.get_label("ket") for ket in candidates)
152 0 : raise ValueError(f"The ket is not uniquely specified. Possible kets are:\n{labels}") from None
153 :
154 1 : def _get_raw_label(self) -> str:
155 1 : s, l, f, m = self.s, self.l, self.f, self.m
156 :
157 1 : label = self.species.split("_", 1)[0]
158 1 : label = label[0].upper() + label[1:]
159 :
160 1 : if not self.species.endswith("_mqdt"):
161 1 : if s == 0:
162 0 : label += "_singlet"
163 1 : elif s == 1:
164 0 : label += "_triplet"
165 1 : elif s != 0.5:
166 0 : logger.error("Unexpected spin quantum number s=%f for species %s.", s, self.species)
167 :
168 1 : label += ":"
169 :
170 1 : if self.species.endswith("_mqdt"):
171 1 : label += f"S={s:.1f},nu={self.nu:.1f},L={l:.1f},"
172 1 : label += "J=" if self.is_j_total_momentum else "F="
173 : else:
174 1 : label += f"{self.n:d},"
175 1 : label += get_l_label(l)
176 1 : label += "_"
177 :
178 1 : label += format_half_integer(f)
179 1 : label += "," + format_half_integer(m)
180 :
181 1 : return label
182 :
183 1 : @cached_property
184 1 : def database(self) -> Database:
185 : """The database from which the KetAtom was loaded."""
186 1 : database_cpp = self._cpp.get_database()
187 1 : return Database._from_cpp_object(database_cpp)
188 :
189 1 : @property
190 1 : def m(self) -> float:
191 : """The magnetic quantum number m (int or half-int)."""
192 1 : return self._cpp.get_quantum_number("m")
193 :
194 1 : @property
195 1 : def f(self) -> float:
196 : """The total momentum quantum number f (int or half-int)."""
197 1 : return self._cpp.get_quantum_number("f")
198 :
199 1 : @property
200 1 : def parity(self) -> Parity:
201 : """The parity of the ket."""
202 1 : return int_to_parity(int(self._cpp.get_quantum_number("parity")))
203 :
204 1 : @property
205 1 : def species(self) -> str:
206 : """The atomic species."""
207 1 : return self._cpp.get_species()
208 :
209 1 : @property
210 1 : def n(self) -> int:
211 : """The principal quantum number n."""
212 1 : return int(self._cpp.get_quantum_number("n"))
213 :
214 1 : @property
215 1 : def nu(self) -> float:
216 : """The effective principal quantum number nu."""
217 1 : return self._cpp.get_quantum_number("nu")
218 :
219 1 : @property
220 1 : def nui(self) -> float:
221 : """The expectation value of the effective principal quantum numbers nu_i of the channels."""
222 1 : return self._cpp.get_quantum_number("nui")
223 :
224 1 : @property
225 1 : def l(self) -> float: # noqa: E743
226 : """The expectation value of the orbital quantum number l of all valence electrons."""
227 1 : return self._cpp.get_quantum_number("l")
228 :
229 1 : @property
230 1 : def s(self) -> float:
231 : """The expectation value of the total spin quantum number s of all valence electrons."""
232 1 : return self._cpp.get_quantum_number("s")
233 :
234 1 : @property
235 1 : def j(self) -> float:
236 : """The expectation value of the total angular quantum number j of all valence electrons."""
237 1 : return self._cpp.get_quantum_number("j")
238 :
239 1 : @property
240 1 : def l_ryd(self) -> float:
241 : """The expectation value of the orbital quantum number l_{Ryd} of the Rydberg electron."""
242 1 : return self._cpp.get_quantum_number("l_ryd")
243 :
244 1 : @property
245 1 : def j_ryd(self) -> float:
246 : """The expectation value of the total angular quantum number j_{Ryd} of the Rydberg electron."""
247 1 : return self._cpp.get_quantum_number("j_ryd")
248 :
249 1 : @property
250 1 : def nui_std(self) -> float:
251 : """The standard deviation of the effective principal quantum numbers nu_i of the channels."""
252 0 : return self._cpp.get_quantum_number_std("nui")
253 :
254 1 : @property
255 1 : def l_std(self) -> float:
256 : """The standard deviation of the orbital quantum number l of all valence electrons."""
257 0 : return self._cpp.get_quantum_number_std("l")
258 :
259 1 : @property
260 1 : def s_std(self) -> float:
261 : """The standard deviation of the total spin quantum number s of all valence electrons."""
262 0 : return self._cpp.get_quantum_number_std("s")
263 :
264 1 : @property
265 1 : def j_std(self) -> float:
266 : """The standard deviation of the total angular quantum number j of all valence electrons."""
267 0 : return self._cpp.get_quantum_number_std("j")
268 :
269 1 : @property
270 1 : def l_ryd_std(self) -> float:
271 : """The standard deviation of the orbital quantum number l_{Ryd} of the Rydberg electron."""
272 0 : return self._cpp.get_quantum_number_std("l_ryd")
273 :
274 1 : @property
275 1 : def j_ryd_std(self) -> float:
276 : """The standard deviation of the total angular quantum number j_{Ryd} of the Rydberg electron."""
277 0 : return self._cpp.get_quantum_number_std("j_ryd")
278 :
279 1 : @property
280 1 : def is_j_total_momentum(self) -> bool:
281 : """Whether j is the total momentum quantum number, otherwise f is the total momentum quantum number."""
282 1 : return bool(self._cpp.get_quantum_number("is_j_total_momentum"))
283 :
284 1 : @property
285 1 : def is_calculated_with_mqdt(self) -> bool:
286 : """Whether the state was calculated with multi-channel quantum defect theory."""
287 1 : return bool(self._cpp.get_quantum_number("is_calculated_with_mqdt"))
288 :
289 1 : @property
290 1 : def underspecified_channel_contribution(self) -> float:
291 : """The contribution of channels whose quantum numbers are not exactly known."""
292 0 : return self._cpp.get_quantum_number("underspecified_channel_contribution")
293 :
294 1 : def to_state(self) -> StateAtom:
295 : """Create a canonical state representing the single ket.
296 :
297 : The returned state has a minimal basis consisting only of this ket and a single coefficient equal to one.
298 :
299 : Returns:
300 : A state object representing the ket.
301 :
302 : """
303 1 : from pairinteraction.state import StateAtom
304 :
305 1 : return StateAtom([1], [self])
306 :
307 1 : def __add__(self, other: KetAtom | StateAtom) -> StateAtom:
308 : """Build the superposition of this ket and another ket or state.
309 :
310 : The ket is converted to a state via :meth:`to_state`, the resulting superposition is in general not normalized.
311 : """
312 1 : return self.to_state() + other
313 :
314 1 : def __sub__(self, other: KetAtom | StateAtom) -> StateAtom:
315 : """Build the superposition of this ket and the negative of another ket or state.
316 :
317 : The ket is converted to a state via :meth:`to_state`, the resulting superposition is in general not normalized.
318 : """
319 1 : return self.to_state() - other
320 :
321 1 : def __mul__(self, factor: complex) -> StateAtom:
322 : """Scale the ket by a complex amplitude, e.g. to build superpositions like ``2 * ket_s - 1j * ket_p``."""
323 1 : return self.to_state() * factor
324 :
325 1 : def __truediv__(self, factor: complex) -> StateAtom:
326 : """Scale the ket by the inverse of a complex amplitude."""
327 1 : return self.to_state() / factor
328 :
329 1 : def __neg__(self) -> StateAtom:
330 : """Flip the sign of the amplitude of the ket."""
331 1 : return -self.to_state()
332 :
333 1 : __rmul__ = __mul__ # for reverse multiplication, i.e. scalar * ket will use ket.__rmul__
334 :
335 : @overload
336 : def get_matrix_element(
337 : self, ket: Self | StateAtom, operator: OperatorType, q: int, unit: None = None
338 : ) -> PintFloat | PintComplex: ...
339 :
340 : @overload
341 : def get_matrix_element(
342 : self, ket: Self | StateAtom, operator: OperatorType, q: int, unit: str
343 : ) -> float | complex: ...
344 :
345 1 : def get_matrix_element(
346 : self, ket: Self | StateAtom, operator: OperatorType, q: int, unit: str | None = None
347 : ) -> PintFloat | PintComplex | float | complex:
348 : """Get the matrix element between two atomic basis states from the database.
349 :
350 : Args:
351 : ket: The second atomic basis state to calculate the matrix element with.
352 : operator: The operator, for which to calculate the matrix element.
353 : q: The index for the matrix element.
354 : unit: The unit to return the matrix element in. Default None will return a `pint.Quantity`.
355 :
356 : Returns:
357 : The matrix element between the two states in the given unit or as a `pint.Quantity`.
358 :
359 : """
360 1 : return self.to_state().get_matrix_element(ket, operator, q, unit=unit)
361 :
362 : @overload
363 : def get_spontaneous_transition_rates(self, unit: None = None) -> tuple[list[KetAtom], PintArray]: ...
364 :
365 : @overload
366 : def get_spontaneous_transition_rates(self, unit: str) -> tuple[list[KetAtom], NDArray]: ...
367 :
368 1 : def get_spontaneous_transition_rates(self, unit: str | None = None) -> tuple[list[KetAtom], NDArray | PintArray]:
369 : """Calculate the spontaneous transition rates for the KetAtom.
370 :
371 : The spontaneous transition rates are given by the Einstein A coefficients.
372 :
373 : Args:
374 : unit: The unit to which to convert the result.
375 : Default None will return a `pint.Quantity`.
376 :
377 : Returns:
378 : The relevant states and the transition rates.
379 :
380 : """
381 1 : relevant_kets, transition_rates_au = self._get_transition_rates("spontaneous")
382 1 : transition_rates = QuantityArray.convert_au_to_user(transition_rates_au, "transition_rate", unit)
383 1 : return relevant_kets, transition_rates
384 :
385 : @overload
386 : def get_black_body_transition_rates(
387 : self, temperature: float | PintFloat, temperature_unit: str | None = None, unit: None = None
388 : ) -> tuple[list[KetAtom], PintArray]: ...
389 :
390 : @overload
391 : def get_black_body_transition_rates(
392 : self, temperature: PintFloat, *, unit: str
393 : ) -> tuple[list[KetAtom], NDArray]: ...
394 :
395 : @overload
396 : def get_black_body_transition_rates(
397 : self, temperature: float, temperature_unit: str, unit: str
398 : ) -> tuple[list[KetAtom], NDArray]: ...
399 :
400 1 : def get_black_body_transition_rates(
401 : self, temperature: float | PintFloat, temperature_unit: str | None = None, unit: str | None = None
402 : ) -> tuple[list[KetAtom], NDArray | PintArray]:
403 : """Calculate the black body transition rates of the KetAtom.
404 :
405 : The black body transition rates are given by the Einstein B coefficients,
406 : with a weight factor given by Planck's law.
407 :
408 : Args:
409 : temperature: The temperature, for which to calculate the black body transition rates.
410 : temperature_unit: The unit of the temperature.
411 : Default None will assume the temperature is given as `pint.Quantity`.
412 : unit: The unit to which to convert the result.
413 : Default None will return a `pint.Quantity`.
414 :
415 : Returns:
416 : The relevant states and the transition rates.
417 :
418 : """
419 1 : temperature_au = QuantityScalar.convert_user_to_au(temperature, temperature_unit, "temperature")
420 1 : relevant_kets, transition_rates_au = self._get_transition_rates("black_body", temperature_au)
421 1 : transition_rates = QuantityArray.convert_au_to_user(transition_rates_au, "transition_rate", unit)
422 1 : return relevant_kets, transition_rates
423 :
424 : @overload
425 : def get_lifetime(
426 : self,
427 : temperature: float | PintFloat | None = None,
428 : temperature_unit: str | None = None,
429 : unit: None = None,
430 : ) -> PintFloat: ...
431 :
432 : @overload
433 : def get_lifetime(self, *, unit: str) -> float: ...
434 :
435 : @overload
436 : def get_lifetime(self, temperature: PintFloat, *, unit: str) -> float: ...
437 :
438 : @overload
439 : def get_lifetime(self, temperature: float, temperature_unit: str, unit: str) -> float: ...
440 :
441 1 : def get_lifetime(
442 : self,
443 : temperature: float | PintFloat | None = None,
444 : temperature_unit: str | None = None,
445 : unit: str | None = None,
446 : ) -> float | PintFloat:
447 : """Calculate the lifetime of the KetAtom.
448 :
449 : The lifetime is the inverse of the sum of all transition rates.
450 :
451 : Args:
452 : temperature: The temperature, for which to calculate the black body transition rates.
453 : Default None will not include black body transitions.
454 : temperature_unit: The unit of the temperature.
455 : Default None will assume the temperature is given as `pint.Quantity`.
456 : unit: The unit to which to convert the result.
457 : Default None will return a `pint.Quantity`.
458 :
459 : Returns:
460 : The lifetime of the state.
461 :
462 : """
463 1 : _, transition_rates = self.get_spontaneous_transition_rates()
464 1 : transition_rates_au = transition_rates.to_base_units().magnitude
465 1 : if temperature is not None:
466 1 : _, black_body_transition_rates = self.get_black_body_transition_rates(temperature, temperature_unit)
467 1 : transition_rates_au = np.append(transition_rates_au, black_body_transition_rates.to_base_units().magnitude)
468 :
469 1 : lifetime_au = 1 / np.sum(transition_rates_au)
470 :
471 1 : return QuantityScalar.convert_au_to_user(lifetime_au, "time", unit)
472 :
473 1 : def _get_transition_rates(
474 : self, which_transitions: Literal["spontaneous", "black_body"], temperature_au: float | None = None
475 : ) -> tuple[list[KetAtom], NDArray]:
476 1 : if not isinstance(self, KetAtomReal):
477 1 : from pairinteraction.basis import BasisAtom
478 :
479 1 : basis_atom_class = BasisAtom
480 : else:
481 1 : from pairinteraction.basis import BasisAtomReal
482 :
483 1 : basis_atom_class = BasisAtomReal
484 :
485 1 : assert which_transitions in ["spontaneous", "black_body"]
486 1 : is_spontaneous = which_transitions == "spontaneous"
487 1 : n_max = self.n + 30
488 :
489 1 : energy_range = None
490 1 : if is_spontaneous:
491 1 : energy_range = (-1, self.get_energy("hartree"))
492 :
493 1 : basis = basis_atom_class(
494 : self.species,
495 : n=(1, n_max),
496 : l=(self.l - 1, self.l + 1),
497 : m=(self.m - 1, self.m + 1),
498 : energy=energy_range,
499 : energy_unit="hartree",
500 : additional_kets=[self], # needed to make get_matrix_elements(self, ...) work
501 : database=self.database,
502 : )
503 :
504 1 : relevant_kets = basis.kets
505 1 : energy_differences_au = np.abs(
506 : self.get_energy("hartree") - np.array([ket_cpp.get_energy() for ket_cpp in basis._cpp.get_kets()])
507 : )
508 1 : electric_dipole_moments_au = np.zeros(len(basis.kets), dtype=complex)
509 1 : for q in [-1, 0, 1]:
510 : # the different entries are only at most once nonzero -> we can just add the arrays
511 1 : el_di_m = basis.get_matrix_elements(self, "electric_dipole", q)
512 1 : electric_dipole_moments_au += el_di_m.to_base_units().magnitude
513 :
514 1 : transition_rates_au = (
515 : (4 / 3)
516 : * np.abs(electric_dipole_moments_au) ** 2
517 : * energy_differences_au**2
518 : / ureg.Quantity(1, "speed_of_light").to_base_units().magnitude ** 3
519 : )
520 :
521 1 : if is_spontaneous:
522 1 : transition_rates_au *= energy_differences_au
523 : else:
524 1 : assert temperature_au is not None, "Temperature must be given for black body transitions."
525 1 : if temperature_au == 0:
526 0 : transition_rates_au *= 0
527 : else: # for numerical stability we use 1 / exprel(x) = x / (exp(x) - 1)
528 1 : transition_rates_au *= temperature_au / exprel(energy_differences_au / temperature_au)
529 :
530 1 : mask = transition_rates_au != 0
531 1 : relevant_kets = [ket for ket, is_relevant in zip(relevant_kets, mask, strict=True) if is_relevant]
532 1 : transition_rates_au = transition_rates_au[mask]
533 1 : return relevant_kets, transition_rates_au
534 :
535 :
536 1 : class KetAtomReal(KetAtom):
537 1 : def to_state(self) -> StateAtomReal:
538 1 : from pairinteraction.state import StateAtomReal
539 :
540 1 : return StateAtomReal([1], [self])
|
