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 collections.abc import Sequence
7 1 : from typing import TYPE_CHECKING, Any, TypeGuard, cast, overload
8 :
9 1 : import numpy as np
10 1 : from scipy.sparse import csr_matrix
11 1 : from typing_extensions import Self, TypeAliasType, deprecated
12 :
13 1 : from pairinteraction import _backend
14 1 : from pairinteraction.basis.basis_atom import BasisAtom, get_cpp_basis_atom_from_kets
15 1 : from pairinteraction.basis.basis_base import BasisBase
16 1 : from pairinteraction.enums import OperatorType, Parity, get_cpp_operator_type, get_cpp_parity
17 1 : from pairinteraction.ket import KetPair, KetPairReal, is_ket_atom_tuple
18 1 : from pairinteraction.ket.ket_pair import get_ketpairlike_energy, get_ketpairlike_m, is_ket_pair_like
19 1 : from pairinteraction.state import StatePair, StatePairReal
20 1 : from pairinteraction.state.state_pair import is_state_atom_tuple, is_state_pair_like
21 1 : from pairinteraction.units import QuantityArray, QuantityScalar, QuantitySparse
22 :
23 : if TYPE_CHECKING:
24 : from pairinteraction.basis.basis_base import UnionCPPBasis
25 : from pairinteraction.enums import OperatorType, Parity
26 : from pairinteraction.ket import KetAtomTuple, KetPairLike
27 : from pairinteraction.state.state_pair import StatePairLike
28 : from pairinteraction.system import SystemAtom
29 : from pairinteraction.units import NDArray, PintArray, PintFloat, PintSparse
30 :
31 1 : logger = logging.getLogger(__name__)
32 :
33 :
34 1 : def is_basis_pair_like(obj: Any) -> TypeGuard[BasisPairLike]:
35 1 : return isinstance(obj, BasisPair) or is_basis_atom_tuple(obj)
36 :
37 :
38 1 : def is_basis_atom_tuple(obj: Any) -> TypeGuard[tuple[BasisAtom, BasisAtom]]:
39 1 : return hasattr(obj, "__len__") and len(obj) == 2 and all(isinstance(x, BasisAtom) for x in obj)
40 :
41 :
42 1 : class BasisPair(BasisBase[KetPair, StatePair]):
43 : """Basis for a pair of atoms.
44 :
45 : Add all product states of the eigenstates of two given SystemAtom objects to the basis,
46 : which pair energy is within the given energy range.
47 : You can also specify which total magnetic quantum number m the pair should have (if it is conserved)
48 : and which parities under inversion and permutation should be used for symmetrization.
49 : Due to the possible restrictions of the basis states, the BasisPair coefficients matrix will in general
50 : not be square but (n x d),
51 : where n is the number of all involved kets (typically basis1.number_of_kets * basis2.number_of_kets)
52 : and d is the number of basis states (after applying the restrictions).
53 :
54 : Examples:
55 : >>> import pairinteraction as pi
56 : >>> ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
57 : >>> basis = pi.BasisAtom("Rb", n=(58, 63), l=(0, 3))
58 : >>> system = pi.SystemAtom(basis).set_magnetic_field([0, 0, 1], unit="G").diagonalize()
59 : >>> pair_energy = 2 * system.get_corresponding_energy(ket, unit="GHz")
60 : >>> pair_basis = pi.BasisPair(
61 : ... [system, system],
62 : ... energy=(pair_energy - 3, pair_energy + 3),
63 : ... energy_unit="GHz",
64 : ... )
65 : >>> print(pair_basis)
66 : BasisPair(|Rb:59,S_1/2,-1/2; Rb:61,S_1/2,-1/2⟩ ... |Rb:58,F_7/2,7/2; Rb:59,S_1/2,1/2⟩)
67 :
68 : """
69 :
70 1 : _cpp: _backend.BasisPairComplex
71 1 : _cpp_creator = _backend.BasisPairCreatorComplex
72 1 : _ket_class = KetPair
73 1 : _state_class = StatePair
74 :
75 1 : system_atoms: tuple[SystemAtom, SystemAtom]
76 1 : """The two SystemAtom objects, from which the BasisPair is build."""
77 :
78 1 : def __init__(
79 : self,
80 : system_atoms: Sequence[SystemAtom],
81 : m: tuple[float, float] | None = None,
82 : parity_under_inversion: Parity | None = None,
83 : parity_under_permutation: Parity | None = None,
84 : energy: tuple[float, float] | tuple[PintFloat, PintFloat] | None = None,
85 : energy_unit: str | None = None,
86 : ) -> None:
87 : """Create a basis for a pair of atoms.
88 :
89 : Args:
90 : system_atoms: tuple of two SystemAtom objects, which define the two atoms, from which the BasisPair is build
91 : Both system_atoms have to be diagonalized before creating the BasisPair.
92 : m: tuple of (min, max) values for the total magnetic quantum number m of the pair state.
93 : Default None, i.e. no restriction.
94 : parity_under_inversion: Restrict to pair states with this parity under inversion.
95 : Default None, i.e. do not apply inversion symmetrization.
96 : Requires the same SystemAtom to be passed for both atoms, since symmetrization is
97 : only defined for two identical atoms.
98 : parity_under_permutation: Restrict to pair states with this parity under permutation.
99 : Default None, i.e. do not apply permutation symmetrization.
100 : Requires the same SystemAtom to be passed for both atoms, since symmetrization is
101 : only defined for two identical atoms.
102 : energy: tuple of (min, max) value for the pair energy. Default None, i.e. add all available states.
103 : energy_unit: In which unit the energy values are given, e.g. "GHz".
104 : Default None, i.e. energy is provided as pint object.
105 :
106 : """
107 1 : assert len(system_atoms) == 2, "BasisPair requires exactly two SystemAtom objects."
108 1 : creator = self._cpp_creator()
109 1 : for system in system_atoms:
110 1 : creator.add(system._cpp)
111 1 : if m is not None:
112 1 : creator.restrict_quantum_number_m(*m)
113 1 : if parity_under_inversion is not None:
114 1 : creator.restrict_parity_under_inversion(get_cpp_parity(parity_under_inversion))
115 1 : if parity_under_permutation is not None:
116 1 : creator.restrict_parity_under_permutation(get_cpp_parity(parity_under_permutation))
117 1 : if energy is not None:
118 1 : min_energy_au = QuantityScalar.convert_user_to_au(energy[0], energy_unit, "energy")
119 1 : max_energy_au = QuantityScalar.convert_user_to_au(energy[1], energy_unit, "energy")
120 : # in atomic units all energies should be on the order of -0.5 * Z^2 to 0
121 : # so choosing some very large values for the limits should be fine
122 : # (we cant use np.inf here, since this is passed to cpp code)
123 1 : min_energy_au = np.clip(min_energy_au, -1e10, 1e10) # FIXME
124 1 : max_energy_au = np.clip(max_energy_au, -1e10, 1e10) # FIXME
125 1 : creator.restrict_energy(min_energy_au, max_energy_au)
126 1 : self._cpp = creator.create()
127 :
128 1 : self.system_atoms = tuple(system_atoms) # type: ignore [assignment]
129 :
130 1 : self._post_init()
131 :
132 1 : @classmethod
133 1 : def from_kets( # noqa: C901
134 : cls: type[Self],
135 : kets: KetPairLike | Sequence[KetPairLike],
136 : system_atoms: Sequence[SystemAtom],
137 : delta_m: float | None = None,
138 : parity_under_inversion: Parity | None = None,
139 : parity_under_permutation: Parity | None = None,
140 : delta_energy: float | PintFloat | None = None,
141 : delta_energy_unit: str | None = None,
142 : number_of_kets: int | None = None,
143 : *,
144 : warn_number_of_kets: bool = True,
145 : ) -> Self:
146 : """Create a BasisPair from one or more pairs of kets with optional energy/m windows.
147 :
148 : Currently a single big basis including all kets for the quantum numbers
149 : from min_value - delta to max_value + delta is returned.
150 : In the future this might change to return a basis including all states around the given kets +/- delta,
151 : but not necessarily all states between the given kets.
152 :
153 : You can either give a ``delta_energy`` to restrict the basis size in energy,
154 : or give an (approximate) number_of_kets, which is used to construct a basis centered around
155 : the provided ket pairs with approximately that many kets. If there are multiple states with the same energy,
156 : the actual number of kets may be a bit higher than the specified number.
157 :
158 : Args:
159 : kets: A single ket pair (given either as tuple ``(ket1, ket2)`` of
160 : :class:`~pairinteraction.KetAtom` objects or as a :class:`~pairinteraction.KetPair` object)
161 : or a list of ket pairs.
162 : Must not be empty.
163 : system_atoms: A collection of exactly two diagonalized
164 : :class:`~pairinteraction.SystemAtom` objects, one per atom.
165 : delta_m: Half-width of the total magnetic quantum number window
166 : ``m = m1 + m2``. Default None means no m restriction.
167 : parity_under_inversion: Restrict to pair states with this parity under inversion.
168 : Default None means no inversion symmetrization.
169 : Requires the same SystemAtom to be passed for both atoms, since symmetrization is
170 : only defined for two identical atoms.
171 : parity_under_permutation: Restrict to pair states with this parity under permutation.
172 : Default None means no permutation symmetrization.
173 : Requires the same SystemAtom to be passed for both atoms, since symmetrization is
174 : only defined for two identical atoms.
175 : delta_energy: Half-width of the energy window. Mutually exclusive with
176 : ``number_of_kets``. Default None means no energy restriction.
177 : delta_energy_unit: Unit for ``delta_energy`` and the pair energies
178 : (e.g. ``"GHz"``). Default None means pint quantities are used.
179 : number_of_kets: Target number of pair kets to include. The method
180 : keeps the ``number_of_kets`` states closest in energy to the
181 : reference ket pairs. Mutually exclusive with ``delta_energy``.
182 : Default None means no count restriction.
183 : warn_number_of_kets: Don't warn about the possible issues with using number_of_kets.
184 :
185 : Returns:
186 : A new :class:`BasisPair` whose energy (and optionally m) range is
187 : determined by the provided ket pairs and the chosen restrictions.
188 :
189 : Examples:
190 : >>> import pairinteraction as pi
191 : >>> ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
192 : >>> basis_atom = pi.BasisAtom("Rb", n=(58, 62), l=(0, 2))
193 : >>> system = pi.SystemAtom(basis_atom).diagonalize()
194 : >>> pair_basis = pi.BasisPair.from_kets(
195 : ... (ket, ket), [system, system], delta_energy=3, delta_energy_unit="GHz"
196 : ... )
197 :
198 : """
199 1 : assert len(system_atoms) == 2, "BasisPair requires exactly two SystemAtom objects."
200 :
201 1 : if number_of_kets is not None:
202 1 : if delta_energy is not None:
203 1 : raise ValueError("Cannot specify both number_of_kets and delta_energy. Please choose one of them.")
204 1 : if number_of_kets <= 0:
205 0 : raise ValueError("number_of_kets must be positive.")
206 1 : if warn_number_of_kets:
207 1 : logger.warning(
208 : "Use number_of_kets with caution and only if you know what you are doing! "
209 : "It might lead to unexpected results. "
210 : "E.g. loosening other restrictions while keeping number_of_kets fixed can lead to worse results!"
211 : )
212 :
213 1 : if is_ket_atom_tuple(kets) or isinstance(kets, KetPair):
214 1 : kets = [kets]
215 1 : if not all(is_ket_pair_like(t) for t in kets):
216 0 : raise ValueError("kets must be a KetPairLike or a list of KetPairLike.")
217 1 : if len(kets) == 0:
218 1 : raise ValueError("kets must not be empty.")
219 1 : kets = cast("Sequence[KetAtomTuple | KetPair]", kets)
220 :
221 1 : energy_range = None
222 1 : if delta_energy is not None:
223 1 : pair_energies = [get_ketpairlike_energy(ket, system_atoms, delta_energy_unit) for ket in kets]
224 1 : energy_range = (min(pair_energies) - delta_energy, max(pair_energies) + delta_energy)
225 :
226 1 : m_range = None
227 1 : if delta_m is not None:
228 1 : pair_ms = [get_ketpairlike_m(ket) for ket in kets]
229 1 : m_range = (min(pair_ms) - delta_m, max(pair_ms) + delta_m)
230 :
231 1 : basis_pair = cls(
232 : system_atoms,
233 : m=m_range,
234 : parity_under_inversion=parity_under_inversion,
235 : parity_under_permutation=parity_under_permutation,
236 : energy=energy_range,
237 : energy_unit=delta_energy_unit,
238 : )
239 1 : if number_of_kets is None or number_of_kets >= basis_pair.number_of_kets:
240 1 : return basis_pair
241 :
242 1 : pair_energies_au = [get_ketpairlike_energy(ket, system_atoms, "hartree") for ket in kets]
243 1 : min_energy_au, max_energy_au = min(pair_energies_au), max(pair_energies_au)
244 :
245 1 : cpp_kets = basis_pair._cpp.get_kets()
246 1 : all_energies_au = np.array([ket_cpp.get_energy() for ket_cpp in cpp_kets])
247 1 : deltas = np.maximum(np.maximum(min_energy_au - all_energies_au, all_energies_au - max_energy_au), 0)
248 1 : delta_energy = float(np.partition(deltas, number_of_kets - 1)[number_of_kets - 1]) + 1e-10
249 :
250 1 : return cls(
251 : system_atoms,
252 : m=m_range,
253 : parity_under_inversion=parity_under_inversion,
254 : parity_under_permutation=parity_under_permutation,
255 : energy=(min_energy_au - delta_energy, max_energy_au + delta_energy),
256 : energy_unit="hartree",
257 : )
258 :
259 1 : @classmethod
260 1 : @deprecated("Use `BasisPair.from_kets` instead.")
261 1 : def from_ket_atoms(
262 : cls: type[Self],
263 : ket_atom_tuples: KetAtomTuple | Sequence[KetAtomTuple],
264 : system_atoms: Sequence[SystemAtom],
265 : delta_m: float | None = None,
266 : parity_under_inversion: Parity | None = None,
267 : parity_under_permutation: Parity | None = None,
268 : delta_energy: float | PintFloat | None = None,
269 : delta_energy_unit: str | None = None,
270 : number_of_kets: int | None = None,
271 : ) -> Self:
272 0 : return cls.from_kets(
273 : ket_atom_tuples,
274 : system_atoms,
275 : delta_m=delta_m,
276 : parity_under_inversion=parity_under_inversion,
277 : parity_under_permutation=parity_under_permutation,
278 : delta_energy=delta_energy,
279 : delta_energy_unit=delta_energy_unit,
280 : number_of_kets=number_of_kets,
281 : )
282 :
283 1 : @classmethod
284 1 : def _from_cpp_object(cls: type[Self], cpp_obj: UnionCPPBasis, system_atoms: tuple[SystemAtom, SystemAtom]) -> Self:
285 1 : obj = super()._from_cpp_object(cpp_obj)
286 1 : obj.system_atoms = system_atoms
287 1 : return obj
288 :
289 1 : def _from_cpp_object_additional_args(self) -> tuple[Any, ...]:
290 1 : return (self.system_atoms,)
291 :
292 1 : def get_corresponding_state(self, ket: KetPairLike) -> StatePair: # type: ignore [override]
293 : # override the accepted ket type
294 1 : return super().get_corresponding_state(ket) # type: ignore [arg-type]
295 :
296 1 : def get_corresponding_state_index(self, ket: KetPairLike) -> int: # type: ignore [override]
297 : # override the accepted ket type
298 1 : return super().get_corresponding_state_index(ket) # type: ignore [arg-type]
299 :
300 1 : def get_corresponding_ket(self, state: StatePairLike) -> KetPair:
301 : # override the accepted state type
302 1 : return super().get_corresponding_ket(state) # type: ignore [arg-type]
303 :
304 1 : def get_corresponding_ket_index(self, state: StatePairLike) -> int:
305 : # override the accepted state type
306 1 : return super().get_corresponding_ket_index(state) # type: ignore [arg-type]
307 :
308 : @overload
309 : def get_amplitudes(self, other: KetPairLike | StatePairLike) -> NDArray: ...
310 :
311 : @overload
312 : def get_amplitudes(self, other: BasisPairLike) -> csr_matrix: ...
313 :
314 1 : def get_amplitudes(self, other: KetPairLike | StatePairLike | BasisPairLike) -> NDArray | csr_matrix:
315 1 : return self.get_matrix_elements(other, ("identity", "identity"), (0, 0), unit="")
316 :
317 : @overload
318 : def get_overlaps(self, other: KetPairLike | StatePairLike) -> NDArray: ...
319 :
320 : @overload
321 : def get_overlaps(self, other: BasisPairLike) -> csr_matrix: ...
322 :
323 1 : def get_overlaps(self, other: KetPairLike | StatePairLike | BasisPairLike) -> NDArray | csr_matrix:
324 1 : amplitudes = self.get_amplitudes(other)
325 1 : if isinstance(amplitudes, csr_matrix):
326 1 : return amplitudes.multiply(amplitudes.conj()).real # type: ignore [no-any-return]
327 1 : return np.abs(amplitudes) ** 2
328 :
329 : @overload
330 : def get_matrix_elements(
331 : self,
332 : other: KetPairLike | StatePairLike,
333 : operators: tuple[OperatorType, OperatorType],
334 : qs: tuple[int, int],
335 : unit: None = None,
336 : ) -> PintArray: ...
337 :
338 : @overload
339 : def get_matrix_elements(
340 : self,
341 : other: KetPairLike | StatePairLike,
342 : operators: tuple[OperatorType, OperatorType],
343 : qs: tuple[int, int],
344 : unit: str,
345 : ) -> NDArray: ...
346 :
347 : @overload
348 : def get_matrix_elements(
349 : self,
350 : other: BasisPairLike,
351 : operators: tuple[OperatorType, OperatorType],
352 : qs: tuple[int, int],
353 : unit: None = None,
354 : ) -> PintSparse: ...
355 :
356 : @overload
357 : def get_matrix_elements(
358 : self,
359 : other: BasisPairLike,
360 : operators: tuple[OperatorType, OperatorType],
361 : qs: tuple[int, int],
362 : unit: str,
363 : ) -> csr_matrix: ...
364 :
365 1 : def get_matrix_elements(
366 : self,
367 : other: KetPairLike | StatePairLike | BasisPairLike,
368 : operators: tuple[OperatorType, OperatorType],
369 : qs: tuple[int, int],
370 : unit: str | None = None,
371 : ) -> NDArray | PintArray | csr_matrix | PintSparse:
372 1 : operators_cpp = (get_cpp_operator_type(operators[0]), get_cpp_operator_type(operators[1]))
373 1 : is_real = isinstance(self._cpp, _backend.BasisPairReal)
374 :
375 1 : if is_ket_pair_like(other) or is_state_pair_like(other):
376 : # KetPair like
377 1 : if isinstance(other, KetPair):
378 1 : other_cpp = get_cpp_basis_pair_from_atom_bases(other._cpp.get_atomic_states(), real=is_real)
379 1 : elif is_ket_atom_tuple(other):
380 1 : other_cpp = get_cpp_basis_pair_from_atom_bases(
381 : [get_cpp_basis_atom_from_kets([ket], real=is_real) for ket in other], real=is_real
382 : )
383 : # StatePair like
384 1 : elif isinstance(other, StatePair):
385 1 : other_cpp = other._cpp
386 1 : elif is_state_atom_tuple(other):
387 1 : other_cpp = get_cpp_basis_pair_from_atom_bases([state._cpp for state in other], real=is_real)
388 : else:
389 0 : raise TypeError(f"Unknown type: {type(other)=}")
390 :
391 1 : matrix_elements_au = self._cpp.get_matrix_elements(other_cpp, *operators_cpp, *qs).toarray().ravel()
392 1 : matrix_elements_au = np.real_if_close(matrix_elements_au)
393 1 : return QuantityArray.convert_au_to_user(matrix_elements_au, operators, unit)
394 :
395 : # BasisPair like
396 1 : if is_basis_pair_like(other):
397 1 : if isinstance(other, BasisPair):
398 1 : other_cpp = other._cpp
399 1 : elif is_basis_atom_tuple(other):
400 1 : other_cpp = get_cpp_basis_pair_from_atom_bases([basis._cpp for basis in other], real=is_real)
401 : else:
402 0 : raise TypeError(f"Unknown type: {type(other)=}")
403 :
404 1 : matrix_elements_sparse_au = self._cpp.get_matrix_elements(other_cpp, *operators_cpp, *qs)
405 1 : matrix_elements_sparse_au.data = np.real_if_close(matrix_elements_sparse_au.data)
406 1 : return QuantitySparse.convert_au_to_user(matrix_elements_sparse_au, operators, unit)
407 :
408 1 : raise TypeError(f"Unknown type: {type(other)=}")
409 :
410 :
411 1 : class BasisPairReal(BasisPair):
412 1 : _cpp: _backend.BasisPairReal # type: ignore [assignment]
413 1 : _cpp_creator = _backend.BasisPairCreatorReal # type: ignore [assignment]
414 1 : _ket_class = KetPairReal
415 1 : _state_class = StatePairReal
416 :
417 :
418 1 : BasisPairLike = TypeAliasType("BasisPairLike", BasisPair | tuple[BasisAtom, BasisAtom] | Sequence[BasisAtom])
419 :
420 :
421 1 : def get_cpp_basis_pair_from_atom_bases(
422 : basis_atoms_cpp: Sequence[_backend.BasisAtomComplex], *, real: bool
423 : ) -> _backend.BasisPairComplex:
424 : """Create a cpp BasisPair object from the product of two cpp BasisAtom objects.
425 :
426 : Like for the _cpp attributes, the cpp objects are annotated as the complex variants,
427 : although for BasisPairReal the real variants are used.
428 : """
429 1 : if not real:
430 1 : cpp_system_atom_class = _backend.SystemAtomComplex
431 1 : creator = _backend.BasisPairCreatorComplex()
432 : else:
433 1 : cpp_system_atom_class = _backend.SystemAtomReal # type: ignore [assignment]
434 1 : creator = _backend.BasisPairCreatorReal() # type: ignore [assignment]
435 :
436 1 : systems = [cpp_system_atom_class(basis_cpp) for basis_cpp in basis_atoms_cpp]
437 1 : for system in systems:
438 1 : creator.add(system)
439 1 : return creator.create()
|
