Line data Source code
1 : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
2 : # SPDX-License-Identifier: LGPL-3.0-or-later
3 1 : from __future__ import annotations
4 :
5 1 : import contextlib
6 1 : import copy
7 1 : import logging
8 1 : from functools import cached_property, lru_cache
9 1 : from typing import TYPE_CHECKING, Literal, overload
10 :
11 1 : import numpy as np
12 1 : from scipy import sparse
13 1 : from typing_extensions import deprecated
14 :
15 1 : from pairinteraction.basis import BasisAtom, BasisAtomReal, BasisPair, BasisPairReal
16 1 : from pairinteraction.diagonalization import diagonalize
17 1 : from pairinteraction.perturbative.perturbation_theory import calculate_perturbative_hamiltonian
18 1 : from pairinteraction.system import SystemAtom, SystemAtomReal, SystemPair, SystemPairReal
19 1 : from pairinteraction.units import QuantityArray
20 :
21 : if TYPE_CHECKING:
22 : from collections.abc import Sequence
23 :
24 : from scipy.sparse import csr_matrix
25 : from typing_extensions import Self
26 :
27 : from pairinteraction.ket import KetAtom, KetAtomTuple # noqa: F401
28 : from pairinteraction.units import ArrayLike, NDArray, PintArray, PintFloat
29 :
30 :
31 1 : logger = logging.getLogger(__name__)
32 :
33 1 : BasisSystemLiteral = Literal["basis_atoms", "system_atoms", "basis_pair", "system_pair"]
34 :
35 :
36 1 : class EffectiveSystemPair:
37 : """Class for creating an effective SystemPair object and calculating the effective Hamiltonian.
38 :
39 : Given a subspace spanned by tuples of `KetAtom` objects (ket_tuples),
40 : this class automatically generates appropriate `BasisAtom`, `SystemAtom` objects as well as a `BasisPair` and
41 : `SystemPair` object to calculate the effective Hamiltonian in the subspace via perturbation theory.
42 :
43 : This class also allows to set magnetic and electric fields similar to the `SystemAtom` class,
44 : as well as the angle and distance between the two atoms like in the `SystemPair` class.
45 :
46 : Examples:
47 : >>> import pairinteraction as pi
48 : >>> ket_atoms = {
49 : ... "+": pi.KetAtom("Rb", n=59, l=0, j=0.5, m=0.5),
50 : ... "0": pi.KetAtom("Rb", n=58, l=1, j=1.5, m=1.5),
51 : ... "-": pi.KetAtom("Rb", n=58, l=0, j=0.5, m=0.5),
52 : ... }
53 : >>> ket_tuples = [
54 : ... (ket_atoms["+"], ket_atoms["-"]),
55 : ... (ket_atoms["0"], ket_atoms["0"]),
56 : ... (ket_atoms["-"], ket_atoms["+"]),
57 : ... ]
58 : >>> eff_system = pi.EffectiveSystemPair(ket_tuples)
59 : >>> eff_system = eff_system.set_distance(10, angle_degree=45, unit="micrometer")
60 : >>> eff_h = eff_system.get_effective_hamiltonian(unit="MHz")
61 : >>> eff_h -= np.eye(3) * eff_system.get_pair_energies("MHz")[1]
62 : >>> print(np.round(eff_h, 0), "MHz")
63 : [[292. 3. 0.]
64 : [ 3. -0. 3.]
65 : [ 0. 3. 292.]] MHz
66 :
67 : """
68 :
69 1 : _basis_atom_class = BasisAtom
70 1 : _basis_pair_class = BasisPair
71 1 : _system_atom_class = SystemAtom
72 1 : _system_pair_class = SystemPair
73 :
74 1 : def __init__(self, ket_tuples: Sequence[KetAtomTuple]) -> None:
75 1 : if not all(len(ket_tuple) == 2 for ket_tuple in ket_tuples):
76 0 : raise ValueError("All ket tuples must contain exactly two kets")
77 1 : for i in range(2):
78 1 : if not all(ket_tuple[i].species == ket_tuples[0][i].species for ket_tuple in ket_tuples):
79 0 : raise ValueError(f"All kets for atom={i} must have the same species")
80 :
81 : # Perturbation attributes
82 1 : self._ket_tuples = [tuple(kets) for kets in ket_tuples]
83 1 : self._perturbation_order = 2
84 :
85 : # BasisAtom and SystemAtom attributes
86 1 : self._delta_n: int | None = None
87 1 : self._delta_l: int | None = None
88 1 : self._delta_m: int | None = None
89 1 : self._electric_field: PintArray | None = None
90 1 : self._magnetic_field: PintArray | None = None
91 1 : self._diamagnetism_enabled: bool | None = None
92 :
93 : # BasisPair and SystemPair attributes
94 1 : self._interaction_order: int | None = None
95 1 : self._distance_vector: PintArray | None = None
96 :
97 : # misc
98 1 : self._eff_h_dict_au: dict[int, NDArray] | None = None
99 1 : self._eff_vecs: csr_matrix | None = None
100 :
101 : # misc user set stuff
102 1 : self._user_set_parts: set[BasisSystemLiteral] = set()
103 :
104 1 : def copy(self: Self) -> Self:
105 : """Create a copy of the EffectiveSystemPair object (before it has been created)."""
106 0 : if self._is_created("basis_atoms"):
107 0 : raise RuntimeError(
108 : "Cannot copy the EffectiveSystemPair object after it has been created. "
109 : "Please create a new object instead."
110 : )
111 0 : return copy.copy(self)
112 :
113 1 : def _is_created(self: Self, what: BasisSystemLiteral = "basis_atoms") -> bool:
114 : """Check if some part of the effective Hamiltonian has already been created."""
115 1 : return hasattr(self, "_" + what)
116 :
117 1 : def _ensure_not_created(self: Self, what: BasisSystemLiteral = "basis_atoms") -> None:
118 : """Ensure that some part of the effective Hamiltonian has not been created yet."""
119 1 : if self._is_created(what):
120 0 : raise RuntimeError(
121 : f"Cannot change parameters for {what} after it has already been created. "
122 : f"Please set all parameters before {what} before accessing it (or creating the effective Hamiltonian)."
123 : )
124 :
125 1 : def _delete_created(self: Self, what: BasisSystemLiteral = "basis_atoms") -> None:
126 : """Delete the created part of the effective Hamiltonian.
127 :
128 : Args:
129 : what: The part of the effective Hamiltonian to delete.
130 : Default is "basis_atoms", which means delete all parts that have been created.
131 :
132 : """
133 1 : self._eff_h_dict_au = None
134 1 : self._eff_vecs = None
135 1 : self._eff_basis = None
136 1 : with contextlib.suppress(AttributeError):
137 1 : del self.model_inds
138 :
139 1 : parts_order: list[BasisSystemLiteral] = ["system_pair", "basis_pair", "system_atoms", "basis_atoms"]
140 1 : for part in parts_order:
141 1 : if part in self._user_set_parts:
142 0 : raise RuntimeError(
143 : f"Cannot delete {part} because it has been set by the user. "
144 : "Please create a new EffectiveSystemPair object instead."
145 : )
146 1 : with contextlib.suppress(AttributeError):
147 1 : delattr(self, "_" + part)
148 1 : if part == what:
149 1 : break
150 :
151 : # # # Perturbation methods and attributes # # #
152 1 : @property
153 1 : def ket_tuples(self) -> list[KetAtomTuple]:
154 : """The tuples of kets, which form the model space for the effective Hamiltonian."""
155 1 : return self._ket_tuples # type: ignore [return-value]
156 :
157 1 : @property
158 1 : def perturbation_order(self) -> int:
159 : """The perturbation order for the effective Hamiltonian."""
160 1 : return self._perturbation_order
161 :
162 1 : def set_perturbation_order(self: Self, order: int) -> Self:
163 : """Set the perturbation order for the effective Hamiltonian."""
164 1 : self._delete_created()
165 1 : self._perturbation_order = order
166 1 : return self
167 :
168 : # # # BasisAtom methods and attributes # # #
169 1 : @property
170 1 : def basis_atoms(self) -> tuple[BasisAtom, BasisAtom]:
171 : """The basis objects for the single-atom systems."""
172 1 : if not self._is_created("basis_atoms"):
173 1 : self._create_basis_atoms()
174 1 : return self._basis_atoms # type: ignore [return-value]
175 :
176 1 : @basis_atoms.setter
177 1 : def basis_atoms(self, basis_atoms: tuple[BasisAtom, BasisAtom]) -> None:
178 1 : self._ensure_not_created()
179 1 : if self._delta_n is not None or self._delta_l is not None or self._delta_m is not None:
180 0 : logger.warning("Setting basis_atoms will overwrite parameters defined for basis_atoms.")
181 1 : self._user_set_parts.add("basis_atoms")
182 1 : self._basis_atoms = tuple(basis_atoms)
183 :
184 1 : def set_delta_n(self: Self, delta_n: int) -> Self:
185 : """Set the delta_n value for single-atom basis."""
186 0 : self._delete_created()
187 0 : self._delta_n = delta_n
188 0 : return self
189 :
190 1 : def set_delta_l(self: Self, delta_l: int) -> Self:
191 : """Set the delta_l value for single-atom basis."""
192 0 : self._delete_created()
193 0 : self._delta_l = delta_l
194 0 : return self
195 :
196 1 : def set_delta_m(self: Self, delta_m: int) -> Self:
197 : """Set the delta_m value for single-atom basis."""
198 0 : self._delete_created()
199 0 : self._delta_m = delta_m
200 0 : return self
201 :
202 1 : def _create_basis_atoms(self) -> None:
203 1 : delta_n = self._delta_n if self._delta_n is not None else 7
204 1 : delta_l = self._delta_l
205 1 : if delta_l is None:
206 1 : delta_l = self.perturbation_order * (self.interaction_order - 2)
207 1 : delta_m = self._delta_m
208 1 : if delta_m is None and self._delta_l is None and self._are_fields_along_z:
209 1 : delta_m = self.perturbation_order * (self.interaction_order - 2)
210 :
211 1 : basis_atoms: list[BasisAtom] = []
212 1 : use_real = isinstance(self, EffectiveSystemPairReal)
213 1 : for i in range(2):
214 1 : kets = [ket_tuple[i] for ket_tuple in self.ket_tuples]
215 1 : nlfm = np.transpose([[ket.n, ket.l, ket.f, ket.m] for ket in kets])
216 1 : n_range = (int(np.min(nlfm[0])) - delta_n, int(np.max(nlfm[0])) + delta_n)
217 1 : l_range = (np.min(nlfm[1]) - delta_l, np.max(nlfm[1]) + delta_l)
218 1 : if any(ket.is_calculated_with_mqdt for ket in kets) and self._delta_l is None:
219 : # for mqdt we increase the default delta_l by 1 to take into account the variance ...
220 0 : l_range = (np.min(nlfm[1]) - delta_l - 1, np.max(nlfm[1]) + delta_l + 1)
221 1 : m_range = (np.min(nlfm[3]) - delta_m, np.max(nlfm[3]) + delta_m) if delta_m is not None else None
222 1 : basis = get_basis_atom_with_cache(kets[0].species, n_range, l_range, m_range, use_real=use_real)
223 1 : basis_atoms.append(basis)
224 :
225 1 : self._basis_atoms = tuple(basis_atoms)
226 :
227 : # # # SystemAtom methods and attributes # # #
228 1 : @property
229 1 : def system_atoms(self) -> tuple[SystemAtom, SystemAtom]:
230 : """The system objects for the single-atom systems."""
231 1 : if not self._is_created("system_atoms"):
232 1 : self._create_system_atoms()
233 1 : return self._system_atoms
234 :
235 1 : @system_atoms.setter
236 1 : def system_atoms(self, system_atoms: tuple[SystemAtom, SystemAtom]) -> None:
237 1 : self._ensure_not_created()
238 1 : if (
239 : self._electric_field is not None
240 : or self._magnetic_field is not None
241 : or self._diamagnetism_enabled is not None
242 : ):
243 0 : logger.warning("Setting system_atoms will overwrite parameters defined for system_atoms.")
244 1 : self._user_set_parts.add("system_atoms")
245 1 : self._system_atoms: tuple[SystemAtom, SystemAtom] = tuple(system_atoms) # type: ignore [assignment]
246 1 : self.basis_atoms = tuple(system.basis for system in system_atoms) # type: ignore [assignment]
247 :
248 1 : @property
249 1 : def electric_field(self) -> PintArray:
250 : """The electric field for the single-atom systems."""
251 1 : if self._electric_field is None:
252 1 : self.set_electric_field([0, 0, 0], "V/cm")
253 1 : assert self._electric_field is not None
254 1 : return self._electric_field
255 :
256 1 : def set_electric_field(
257 : self: Self,
258 : electric_field: PintArray | ArrayLike,
259 : unit: str | None = None,
260 : ) -> Self:
261 : """Set the electric field for the single-atom systems.
262 :
263 : Args:
264 : electric_field: The electric field to set for the systems.
265 : unit: The unit of the electric field, e.g. "V/cm".
266 : Default None expects a `pint.Quantity`.
267 :
268 : """
269 1 : self._delete_created()
270 1 : self._electric_field = QuantityArray.convert_user_to_pint(electric_field, unit, "electric_field")
271 1 : return self
272 :
273 1 : @property
274 1 : def magnetic_field(self) -> PintArray:
275 : """The magnetic field for the single-atom systems."""
276 1 : if self._magnetic_field is None:
277 1 : self.set_magnetic_field([0, 0, 0], "gauss")
278 1 : assert self._magnetic_field is not None
279 1 : return self._magnetic_field
280 :
281 1 : def set_magnetic_field(
282 : self: Self,
283 : magnetic_field: PintArray | ArrayLike,
284 : unit: str | None = None,
285 : ) -> Self:
286 : """Set the magnetic field for the single-atom systems.
287 :
288 : Args:
289 : magnetic_field: The magnetic field to set for the systems.
290 : unit: The unit of the magnetic field, e.g. "gauss".
291 : Default None expects a `pint.Quantity`.
292 :
293 : """
294 1 : self._delete_created()
295 1 : self._magnetic_field = QuantityArray.convert_user_to_pint(magnetic_field, unit, "magnetic_field")
296 1 : return self
297 :
298 1 : @property
299 1 : def _are_fields_along_z(self) -> bool:
300 1 : return all(x == 0 for x in [*self.magnetic_field[:2], *self.electric_field[:2]]) # type: ignore [index]
301 :
302 1 : @property
303 1 : def diamagnetism_enabled(self) -> bool:
304 : """Whether diamagnetism is enabled for the single-atom systems."""
305 1 : if self._diamagnetism_enabled is None:
306 1 : self.set_diamagnetism_enabled(False)
307 1 : assert self._diamagnetism_enabled is not None
308 1 : return self._diamagnetism_enabled
309 :
310 1 : def set_diamagnetism_enabled(self: Self, enable: bool = True) -> Self:
311 : """Enable or disable diamagnetism for the system.
312 :
313 : Args:
314 : enable: Whether to enable or disable diamagnetism.
315 :
316 : """
317 1 : self._delete_created("system_atoms")
318 1 : self._diamagnetism_enabled = enable
319 1 : return self
320 :
321 1 : def _create_system_atoms(self) -> None:
322 1 : system_atoms: list[SystemAtom] = []
323 1 : for basis_atom in self.basis_atoms:
324 1 : system = self._system_atom_class(basis_atom)
325 1 : system.set_diamagnetism_enabled(self.diamagnetism_enabled)
326 1 : system.set_electric_field(self.electric_field)
327 1 : system.set_magnetic_field(self.magnetic_field)
328 1 : system_atoms.append(system)
329 1 : diagonalize(system_atoms)
330 :
331 1 : self._system_atoms = tuple(system_atoms) # type: ignore [assignment]
332 :
333 : @overload
334 : def get_pair_energies(self, unit: None = None) -> list[PintFloat]: ...
335 :
336 : @overload
337 : def get_pair_energies(self, unit: str) -> list[float]: ...
338 :
339 1 : def get_pair_energies(self, unit: str | None = None) -> list[float] | list[PintFloat]:
340 : """Get the pair energies of the ket tuples for infinite distance (i.e. no interaction).
341 :
342 : Args:
343 : unit: The unit to which to convert the energies to.
344 : Default None will return a list of `pint.Quantity`.
345 :
346 : Returns:
347 : The energies as list of float if a unit was given, otherwise as list of `pint.Quantity`.
348 :
349 : """
350 1 : return [ # type: ignore [return-value]
351 : sum(
352 : system.get_corresponding_energy(ket, unit=unit)
353 : for system, ket in zip(self.system_atoms, ket_tuple, strict=True)
354 : )
355 : for ket_tuple in self.ket_tuples
356 : ]
357 :
358 : # # # BasisPair methods and attributes # # #
359 1 : @property
360 1 : def basis_pair(self) -> BasisPair:
361 : """The basis pair object for the pair system."""
362 1 : if not self._is_created("basis_pair"):
363 1 : self.create_basis_pair()
364 1 : return self._basis_pair
365 :
366 1 : @basis_pair.setter
367 1 : def basis_pair(self, basis_pair: BasisPair) -> None:
368 1 : self._ensure_not_created()
369 1 : self._user_set_parts.add("basis_pair")
370 1 : self._basis_pair = basis_pair
371 1 : self.system_atoms = basis_pair.system_atoms
372 :
373 1 : @deprecated("set_minimum_number_of_ket_pairs is deprecated, use create_basis_pair(...) instead.")
374 1 : def set_minimum_number_of_ket_pairs(self: Self, number_of_kets: int) -> Self: # noqa: ARG002
375 0 : raise DeprecationWarning("set_minimum_number_of_ket_pairs is deprecated, use create_basis_pair(...) instead.")
376 :
377 1 : @deprecated("set_maximum_number_of_ket_pairs is deprecated, use create_basis_pair(...) instead.")
378 1 : def set_maximum_number_of_ket_pairs(self: Self, number_of_kets: int) -> Self: # noqa: ARG002
379 0 : raise DeprecationWarning("set_maximum_number_of_ket_pairs is deprecated, use create_basis_pair(...) instead.")
380 :
381 1 : def create_basis_pair(
382 : self,
383 : delta_energy: float | PintFloat | None = None,
384 : delta_energy_unit: str | None = None,
385 : number_of_kets: int | None = None,
386 : *,
387 : allow_large_basis: bool = False,
388 : ) -> None:
389 1 : if self._is_created("basis_pair"):
390 0 : raise RuntimeError("The basis_pair has already been created. Cannot create it again.")
391 :
392 1 : if delta_energy is not None or number_of_kets is not None:
393 1 : self._basis_pair = self._basis_pair_class.from_kets(
394 : self.ket_tuples,
395 : system_atoms=self.system_atoms,
396 : delta_energy=delta_energy,
397 : delta_energy_unit=delta_energy_unit,
398 : number_of_kets=number_of_kets,
399 : )
400 1 : return
401 :
402 1 : min_nu = min(ket.nu for ket_tuple in self.ket_tuples for ket in ket_tuple)
403 : # for nu = 40 use delta_energy = 8GHz and scale with nu^3 (i.e. for nu=80 use 1GHz)
404 1 : delta_energy_ghz = 8 * (40 / min_nu) ** 3
405 1 : basis_pair = self._basis_pair_class.from_kets(
406 : self.ket_tuples,
407 : system_atoms=self.system_atoms,
408 : delta_energy=delta_energy_ghz,
409 : delta_energy_unit="GHz",
410 : )
411 :
412 1 : if basis_pair.number_of_kets > 25_000:
413 0 : msg = (
414 : f"The automatically generated basis_pair contains {basis_pair.number_of_kets} kets. "
415 : "This might lead to long calculation times for the effective Hamiltonian. "
416 : )
417 0 : if not allow_large_basis:
418 0 : raise RuntimeError(
419 : msg
420 : + "If this is on purpose, consider calling `create_basis_pair(allow_large_basis=True)`. "
421 : + "If not, consider calling `create_basis_pair(delta_energy=..., delta_energy_unit=...)` "
422 : + "or `create_basis_pair(number_of_kets=...)` "
423 : + "with custom parameters to control the basis size. "
424 : )
425 0 : logger.warning(msg)
426 :
427 1 : self._basis_pair = basis_pair
428 1 : logger.debug("The pair basis for the perturbative calculations consists of %d kets.", basis_pair.number_of_kets)
429 :
430 : # # # SystemPair methods and attributes # # #
431 1 : @property
432 1 : def system_pair(self) -> SystemPair:
433 : """The system pair object for the pair system."""
434 1 : if not self._is_created("system_pair"):
435 1 : self._create_system_pair()
436 1 : return self._system_pair
437 :
438 1 : @system_pair.setter
439 1 : def system_pair(self, system_pair: SystemPair) -> None:
440 1 : self._ensure_not_created()
441 1 : if self._interaction_order is not None or self._distance_vector is not None:
442 0 : logger.warning("Setting system_pair will overwrite parameters defined for system_pair.")
443 1 : self._user_set_parts.add("system_pair")
444 1 : self._system_pair = system_pair
445 1 : self.basis_pair = system_pair.basis
446 :
447 1 : @property
448 1 : def interaction_order(self) -> int:
449 : """The interaction order for the pair system."""
450 1 : if self._interaction_order is None:
451 1 : self.set_interaction_order(3)
452 1 : return self._interaction_order # type: ignore [return-value]
453 :
454 1 : def set_interaction_order(self: Self, order: int) -> Self:
455 : """Set the interaction order of the pair system.
456 :
457 : Args:
458 : order: The interaction order to set for the pair system.
459 : The order must be 3, 4, or 5.
460 :
461 : """
462 1 : self._delete_created()
463 1 : self._interaction_order = order
464 1 : return self
465 :
466 1 : @property
467 1 : def distance_vector(self) -> PintArray:
468 : """The distance vector between the atoms in the pair system."""
469 1 : if self._distance_vector is None:
470 0 : self.set_distance_vector([0, 0, np.inf], "micrometer")
471 1 : return self._distance_vector # type: ignore [return-value]
472 :
473 1 : def set_distance(
474 : self: Self,
475 : distance: float | PintFloat,
476 : angle_degree: float = 0,
477 : unit: str | None = None,
478 : ) -> Self:
479 : """Set the distance between the atoms using the specified distance and angle.
480 :
481 : Args:
482 : distance: The distance to set between the atoms in the given unit.
483 : angle_degree: The angle between the distance vector and the z-axis in degrees.
484 : 90 degrees corresponds to the x-axis.
485 : Defaults to 0, which corresponds to the z-axis.
486 : unit: The unit of the distance, e.g. "micrometer".
487 : Default None expects a `pint.Quantity`.
488 :
489 : """
490 1 : distance_vector = [np.sin(np.deg2rad(angle_degree)) * distance, 0, np.cos(np.deg2rad(angle_degree)) * distance]
491 1 : return self.set_distance_vector(distance_vector, unit)
492 :
493 1 : def set_distance_vector(
494 : self: Self,
495 : distance: ArrayLike | PintArray,
496 : unit: str | None = None,
497 : ) -> Self:
498 : """Set the distance vector between the atoms.
499 :
500 : Args:
501 : distance: The distance vector to set between the atoms in the given unit.
502 : unit: The unit of the distance, e.g. "micrometer".
503 : Default None expects a `pint.Quantity`.
504 :
505 : """
506 1 : self._delete_created("system_pair")
507 1 : self._distance_vector = QuantityArray.convert_user_to_pint(distance, unit, "distance")
508 1 : return self
509 :
510 1 : def set_angle(
511 : self: Self,
512 : angle: float = 0,
513 : unit: Literal["degree", "radian"] = "degree",
514 : ) -> Self:
515 : """Set the angle between the atoms in degrees.
516 :
517 : Args:
518 : angle: The angle between the distance vector and the z-axis (by default in degrees).
519 : 90 degrees corresponds to the x-axis.
520 : Defaults to 0, which corresponds to the z-axis.
521 : unit: The unit of the angle, either "degree" or "radian", by default "degree".
522 :
523 : """
524 0 : assert unit in ("radian", "degree"), f"Unit {unit} is not supported for angle."
525 0 : if unit == "radian":
526 0 : angle = np.rad2deg(angle)
527 0 : distance_mum: float = np.linalg.norm(self.distance_vector.to("micrometer").magnitude) # type: ignore [assignment]
528 0 : return self.set_distance(distance_mum, angle, "micrometer")
529 :
530 1 : def _create_system_pair(self) -> None:
531 1 : system_pair = self._system_pair_class(self.basis_pair)
532 1 : system_pair.set_distance_vector(self.distance_vector)
533 1 : system_pair.set_interaction_order(self.interaction_order)
534 1 : self._system_pair = system_pair
535 :
536 : # # # Effective Hamiltonian methods and attributes # # #
537 : @overload
538 : def get_effective_hamiltonian(self, return_order: int | None = None, unit: None = None) -> PintArray: ...
539 :
540 : @overload
541 : def get_effective_hamiltonian(self, return_order: int | None = None, *, unit: str) -> NDArray: ...
542 :
543 1 : def get_effective_hamiltonian(
544 : self, return_order: int | None = None, unit: str | None = None
545 : ) -> NDArray | PintArray:
546 : """Get the effective Hamiltonian of the pair system.
547 :
548 : Args:
549 : return_order: The order of the perturbation to return.
550 : Default None, returns the sum up to the perturbation order set in the class.
551 : unit: The unit in which to return the effective Hamiltonian.
552 : If None, returns a pint array.
553 :
554 : Returns:
555 : The effective Hamiltonian of the pair system in the given unit.
556 : If unit is None, returns a pint array, otherwise returns a numpy array.
557 :
558 : """
559 1 : if self._eff_h_dict_au is None:
560 1 : self._create_effective_hamiltonian()
561 1 : assert self._eff_h_dict_au is not None
562 1 : if return_order is None:
563 1 : h_eff_au: NDArray = sum(self._eff_h_dict_au.values()) # type: ignore [assignment]
564 1 : elif return_order in self._eff_h_dict_au:
565 1 : h_eff_au = self._eff_h_dict_au[return_order]
566 : else:
567 0 : raise ValueError(
568 : f"The perturbation order {return_order} is not available in the effective Hamiltonian "
569 : f"with the specified perturbation_order {self.perturbation_order}."
570 : )
571 1 : return QuantityArray.convert_au_to_user(np.real_if_close(h_eff_au), "energy", unit)
572 :
573 1 : def get_effective_basisvectors(self) -> csr_matrix:
574 : """Get the eigenvectors of the perturbative Hamiltonian."""
575 0 : if len(self.model_inds) > 1 and self.perturbation_order > 2:
576 0 : logger.warning("For more than one state and perturbation_order > 2 the effective basis might be wrong.")
577 0 : if self._eff_vecs is None:
578 0 : self._create_effective_hamiltonian()
579 0 : assert self._eff_vecs is not None
580 0 : return self._eff_vecs
581 :
582 1 : def get_effective_basis(self) -> BasisPair:
583 : """Get the effective basis of the pair system."""
584 0 : raise NotImplementedError("The get effective basis method is not implemented yet.")
585 :
586 1 : def _create_effective_hamiltonian(self) -> None:
587 : """Calculate the perturbative Hamiltonian up to the given perturbation order."""
588 1 : hamiltonian_au = self.system_pair.get_hamiltonian(unit="hartree")
589 1 : eff_h_dict_au, eff_vecs = calculate_perturbative_hamiltonian(
590 : hamiltonian_au, self.model_inds, self.perturbation_order
591 : )
592 1 : self._eff_h_dict_au = eff_h_dict_au
593 1 : self._eff_vecs = eff_vecs
594 :
595 1 : self.check_for_resonances()
596 :
597 : # # # Other stuff # # #
598 1 : @cached_property
599 1 : def model_inds(self) -> list[int]:
600 : """The indices of the corresponding KetPairs of the given ket_tuples in the basis_pair."""
601 1 : model_inds = []
602 1 : for kets in self.ket_tuples:
603 1 : overlap = self.basis_pair.get_overlaps(kets)
604 1 : inds = np.argsort(overlap)[::-1]
605 1 : model_inds.append(int(inds[0]))
606 1 : self._warn_model_inds_overlap(overlap, inds, kets)
607 1 : return model_inds
608 :
609 1 : def _warn_model_inds_overlap(self, overlap: NDArray, inds: NDArray, kets: KetAtomTuple) -> None:
610 1 : if overlap[inds[0]] > 0.8:
611 1 : return
612 :
613 0 : if overlap[inds[0]] == 0:
614 0 : raise ValueError(f"The pairstate {kets} is not part of the basis_pair.")
615 :
616 0 : msg = ""
617 0 : accumulated = overlap[inds[0]]
618 0 : for i in inds[1:5]:
619 0 : msg += f"\n - {self.basis_pair.get_state(i)} with overlap {overlap[i]:.3e}"
620 0 : accumulated += overlap[i]
621 0 : if accumulated > 0.8:
622 0 : break
623 :
624 0 : logger.warning(
625 : "The pairstate %s has only an overlap of %.3f with its corresponding state in the basis_pair.\n"
626 : "Note that the effective hamiltonian is calculated with respect to the corresponding state %s.\n"
627 : "The most perturbing other states in the basis_pair are:\n%s",
628 : *(kets, overlap[inds[0]], self.basis_pair.get_state(inds[0]), msg),
629 : )
630 :
631 1 : def check_for_resonances(self, max_perturber_weight: float = 0.05) -> None:
632 : r"""Check if states of the model space have strong resonances with states outside the model space."""
633 : # Get the effective eigenvectors without potential warning
634 1 : if self._eff_vecs is None:
635 0 : self._create_effective_hamiltonian()
636 0 : assert self._eff_vecs is not None
637 1 : eff_vecs = self._eff_vecs
638 :
639 1 : overlaps = (eff_vecs.multiply(eff_vecs.conj())).real # elementwise multiplication
640 :
641 1 : for i, m_ind in enumerate(self.model_inds):
642 1 : overlaps_i = overlaps[i, :]
643 1 : other_weight = np.sum(overlaps_i.data) - 1
644 1 : if other_weight < max_perturber_weight:
645 1 : continue
646 :
647 1 : msg = ""
648 1 : indices = [
649 : int(index) for index in sparse.find(overlaps_i >= 0.1 * max_perturber_weight)[1] if index != m_ind
650 : ]
651 1 : indices = sorted(indices, key=lambda index, ov=overlaps_i: ov[0, index], reverse=True) # type: ignore [misc]
652 1 : overlap = 0
653 1 : for index in indices[:5]:
654 1 : admixture = overlaps_i[0, index]
655 1 : msg += f"\n - {self.basis_pair.get_state(index)} has admixture {overlaps_i[0, index]:.3e}"
656 1 : overlap += admixture
657 1 : if overlap > 0.8 * other_weight:
658 1 : break
659 :
660 1 : logger.warning(
661 : "The state (from the model space) %s gets a large dressing (%.3f overlap) "
662 : "in perturbation theory from other states from the basis_pair.\n"
663 : "Thus, treating these states perturbatively might not be accurate. "
664 : "Consider adding these states to the model space.\n"
665 : "The most perturbing states are:\n%s",
666 : *(self.basis_pair.get_state(m_ind), other_weight, msg),
667 : )
668 :
669 :
670 1 : class EffectiveSystemPairReal(EffectiveSystemPair):
671 1 : _basis_atom_class = BasisAtomReal
672 1 : _basis_pair_class = BasisPairReal
673 1 : _system_atom_class = SystemAtomReal
674 1 : _system_pair_class = SystemPairReal
675 :
676 :
677 1 : @lru_cache(maxsize=20)
678 1 : def get_basis_atom_with_cache(
679 : species: str, n: tuple[int, int], l: tuple[int, int], m: tuple[int, int], *, use_real: bool
680 : ) -> BasisAtom:
681 : """Get a BasisAtom object potentially by using a cache to avoid recomputing it."""
682 1 : if use_real:
683 1 : return BasisAtomReal(species, n=n, l=l, m=m)
684 1 : return BasisAtom(species, n=n, l=l, m=m)
|
