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 : import warnings
7 1 : from typing import TYPE_CHECKING, cast, overload
8 :
9 1 : import numpy as np
10 1 : from scipy.sparse import csr_matrix
11 1 : from typing_extensions import deprecated
12 :
13 1 : from pairinteraction.enums import get_cpp_operator_type
14 1 : from pairinteraction.ket import KetAtom, KetAtomReal
15 1 : from pairinteraction.state.state_base import StateBase
16 1 : from pairinteraction.units import QuantityScalar
17 :
18 : if TYPE_CHECKING:
19 : from collections.abc import Sequence
20 :
21 : from typing_extensions import Self
22 :
23 : from pairinteraction import _backend
24 : from pairinteraction.basis import BasisAtom
25 : from pairinteraction.database import Database
26 : from pairinteraction.enums import OperatorType
27 : from pairinteraction.units import PintComplex, PintFloat
28 :
29 1 : logger = logging.getLogger(__name__)
30 :
31 :
32 1 : class StateAtom(StateBase[KetAtom]):
33 : """State of a single atom.
34 :
35 : A coefficient vector and a list of kets are used to represent an arbitrary single-atom state.
36 :
37 : Examples:
38 : >>> import pairinteraction as pi
39 : >>> ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
40 : >>> state = ket.to_state()
41 : >>> print(state)
42 : 1.00 |Rb:60,S_1/2,1/2⟩
43 : >>> ket2 = pi.KetAtom("Rb", n=60, l=1, j=0.5, m=0.5)
44 : >>> state2 = pi.StateAtom([1], [ket2])
45 : >>> print(state2)
46 : 1.00 |Rb:60,P_1/2,1/2⟩
47 : >>> print((2 * state2 - state).normalize())
48 : 0.89 |Rb:60,P_1/2,1/2⟩ - 0.45 |Rb:60,S_1/2,1/2⟩
49 : >>> print(pi.StateAtom([2, 1], [ket, ket2]).normalize())
50 : 0.89 |Rb:60,S_1/2,1/2⟩ + 0.45 |Rb:60,P_1/2,1/2⟩
51 :
52 : """
53 :
54 1 : _cpp: _backend.BasisAtomComplex
55 1 : _ket_class = KetAtom
56 :
57 : @overload
58 : def __init__(
59 : self, coefficients: Sequence[complex], kets: Sequence[KetAtom], *, basis: BasisAtom | None = None
60 : ) -> None: ...
61 :
62 : @overload
63 : @deprecated("Use ket.to_state() instead of StateAtom(ket, basis).")
64 : def __init__(self, ket: KetAtom, basis: BasisAtom) -> None: ...
65 :
66 1 : def __init__( # type: ignore [misc]
67 : self,
68 : coefficients: Sequence[complex] | KetAtom,
69 : kets: Sequence[KetAtom] | BasisAtom | None = None,
70 : *,
71 : basis: BasisAtom | None = None,
72 : ) -> None:
73 : """Initialize a state object from a coefficient vector and the corresponding kets.
74 :
75 : Args:
76 : coefficients: The coefficient of each of the given kets.
77 : kets: The kets the state is composed of.
78 : basis: The basis in which the state should be expressed.
79 : If None (default), a minimal basis consisting only of the given kets is constructed.
80 : Providing a basis is only relevant if you want the state to already live in a larger Hilbert space;
81 : when adding states, their bases are merged automatically.
82 : All given kets must be part of this basis.
83 : Since the coefficients are always defined with respect to the kets,
84 : only the kets of the given basis are used and the coefficients of the basis are ignored,
85 : i.e. the basis is canonicalized first.
86 :
87 : """
88 1 : super().__init__()
89 :
90 1 : if isinstance(coefficients, KetAtom): # deprecated interface StateAtom(ket, basis)
91 0 : coefficients, kets, basis = self._unpack_deprecated_args(coefficients, kets, basis)
92 1 : if kets is None:
93 0 : raise TypeError("StateAtom.__init__() missing 1 required positional argument: 'kets'")
94 1 : kets = cast("Sequence[KetAtom]", kets)
95 :
96 1 : is_real = isinstance(self, StateAtomReal)
97 1 : coeffs = np.array(coefficients, dtype=float if is_real else complex).ravel()
98 1 : if len(coeffs) != len(kets):
99 0 : raise ValueError(f"Got {len(coeffs)} coefficients for {len(kets)} kets, these must match.")
100 1 : if len(kets) != len(set(kets)):
101 1 : raise ValueError("The given kets must be unique.")
102 :
103 1 : if basis is None:
104 1 : from pairinteraction.basis.basis_atom import get_cpp_basis_atom_from_kets
105 :
106 1 : cpp_basis = get_cpp_basis_atom_from_kets(kets, real=is_real)
107 : else:
108 1 : cpp_basis = basis._cpp.canonicalized()
109 :
110 1 : cpp_ket_to_index = {cpp_ket: i for i, cpp_ket in enumerate(cpp_basis.get_kets())}
111 1 : basis_coeffs = np.zeros((len(cpp_ket_to_index), 1), dtype=coeffs.dtype)
112 1 : for coeff, ket in zip(coeffs, kets, strict=True):
113 1 : if ket._cpp not in cpp_ket_to_index:
114 1 : raise ValueError(f"The ket {ket} is not part of the given basis.")
115 1 : basis_coeffs[cpp_ket_to_index[ket._cpp], 0] = coeff
116 :
117 1 : state_cpp = cpp_basis.get_state(0) # single-state basis, i.e. shape (n_kets, 1)
118 1 : self._cpp = state_cpp.copy_with_coefficients(csr_matrix(basis_coeffs))
119 :
120 1 : @staticmethod
121 1 : def _unpack_deprecated_args(
122 : ket: KetAtom, kets: Sequence[KetAtom] | BasisAtom | None, basis: BasisAtom | None
123 : ) -> tuple[Sequence[complex], Sequence[KetAtom], BasisAtom]:
124 : """Translate the arguments of the deprecated StateAtom(ket, basis) interface into the new interface."""
125 0 : from pairinteraction.basis import BasisAtom # imported here to avoid a circular import
126 :
127 0 : warnings.warn(
128 : "Calling StateAtom(ket, basis) is deprecated use ket.to_state() instead.",
129 : DeprecationWarning,
130 : stacklevel=3,
131 : )
132 0 : if kets is not None: # in the deprecated interface the second argument was the basis
133 0 : if basis is not None:
134 0 : raise TypeError("The basis must not be given both as positional and as keyword argument.")
135 0 : basis = cast("BasisAtom", kets)
136 0 : if not isinstance(basis, BasisAtom):
137 0 : raise TypeError("The basis must be given as a BasisAtom object when creating a StateAtom from a ket.")
138 0 : return [1], [ket], basis
139 :
140 1 : def __add__(self, other: Self | KetAtom) -> Self:
141 : """Build the superposition of this state and another state or ket.
142 :
143 : The bases are merged into a common basis, in which the coefficients are re-expressed before adding them.
144 : The resulting superposition is in general not normalized, use :meth:`normalize` to normalize it.
145 : """
146 1 : if isinstance(other, KetAtom):
147 1 : other = cast("Self", other.to_state())
148 1 : if type(self) is not type(other):
149 0 : raise TypeError(f"Cannot add/subtract {type(self)} and {type(other)}.")
150 :
151 : # merge the (canonical) bases and re-express the coefficients in the merged basis;
152 1 : merged_cpp = self._cpp.canonicalized().merge(other._cpp.canonicalized())
153 1 : cpp_op = get_cpp_operator_type("identity")
154 1 : coeffs1 = self._cpp.get_matrix_elements(merged_cpp, cpp_op, 0)
155 1 : coeffs2 = other._cpp.get_matrix_elements(merged_cpp, cpp_op, 0)
156 1 : coeffs = coeffs1 + coeffs2
157 1 : new_cpp = merged_cpp.get_state(0) # single-state basis, i.e. shape (n_kets, 1)
158 :
159 1 : new_cpp = new_cpp.copy_with_coefficients(coeffs)
160 1 : return type(self)._from_cpp_object(new_cpp)
161 :
162 1 : def __sub__(self, other: Self | KetAtom) -> Self:
163 : """Build the superposition of this state and the negative of another state or ket.
164 :
165 : As for :meth:`__add__` the bases are merged first, the resulting superposition is in general not normalized.
166 : """
167 1 : if isinstance(other, KetAtom):
168 1 : other = cast("Self", other.to_state())
169 1 : return self.__add__(-1 * other)
170 :
171 1 : def __mul__(self, factor: complex) -> Self:
172 : """Scale the state by a complex amplitude, e.g. to build superpositions like ``2 * state_s - 1j * state_p``."""
173 1 : if not isinstance(factor, (int, float, complex)):
174 0 : raise TypeError(f"Cannot multiply {type(self)} with {type(factor)}.")
175 1 : coeffs = factor * self._cpp.get_coefficients() # type: ignore [operator]
176 1 : new_cpp = self._cpp.copy_with_coefficients(coeffs)
177 1 : return type(self)._from_cpp_object(new_cpp)
178 :
179 1 : def __truediv__(self, factor: complex) -> Self:
180 : """Scale the state by the inverse of a complex amplitude."""
181 1 : return self.__mul__(1 / factor)
182 :
183 1 : def __neg__(self) -> Self:
184 : """Flip the sign of all amplitudes of the state."""
185 1 : return self.__mul__(-1)
186 :
187 1 : __rmul__ = __mul__ # for reverse multiplication, i.e. scalar * state will use state.__rmul__
188 :
189 1 : def normalize(self) -> Self:
190 : """Normalize the coefficients of the state."""
191 1 : coeffs = self._cpp.get_coefficients()
192 1 : self._cpp = self._cpp.copy_with_coefficients(coeffs / self.norm)
193 1 : return self
194 :
195 1 : def is_normalized(self, tol: float = 1e-10) -> bool:
196 : """Check if the state is normalized within a given tolerance.
197 :
198 : Args:
199 : tol: The tolerance for the normalization check. Default is 1e-10.
200 :
201 : Returns:
202 : True if the state is normalized within the given tolerance, False otherwise.
203 :
204 : """
205 1 : return abs(self.norm - 1) < tol # type: ignore [return-value] # numpy
206 :
207 1 : @property
208 1 : def database(self) -> Database:
209 : """The database used for this object."""
210 0 : return self.get_ket(0).database
211 :
212 1 : @property
213 1 : def species(self) -> str:
214 : """The atomic species."""
215 1 : return self.get_ket(0).species
216 :
217 1 : @property
218 1 : def is_canonical(self) -> bool:
219 1 : return np.count_nonzero(self.get_coefficients()) == 1 # type: ignore [no-any-return]
220 :
221 1 : def get_amplitude(self, other: Self | KetAtom) -> float | complex:
222 : """Calculate the amplitude of the state with respect to another state or ket.
223 :
224 : This means the inner product <self|other>.
225 :
226 : Args:
227 : other: Either a state or a ket for which the amplitude should be calculated.
228 :
229 : Returns:
230 : The amplitude between self and other.
231 :
232 : """
233 1 : return self.get_matrix_element(other, "identity", 0, unit="")
234 :
235 1 : def get_overlap(self, other: Self | KetAtom) -> float:
236 : r"""Calculate the overlap of the state with respect to another state or ket.
237 :
238 : This means calculate :math:`|\langle \mathrm{self} | \mathrm{other} \rangle|^2`.
239 :
240 : Args:
241 : other: Either a state or a ket for which the overlap should be calculated.
242 :
243 : Returns:
244 : The overlap between self and other.
245 :
246 : """
247 1 : return abs(self.get_amplitude(other)) ** 2
248 :
249 : @overload
250 : def get_matrix_element(
251 : self, other: KetAtom | Self, operator: OperatorType, q: int, unit: None = None
252 : ) -> PintFloat | PintComplex: ...
253 :
254 : @overload
255 : def get_matrix_element(
256 : self, other: KetAtom | Self, operator: OperatorType, q: int, unit: str
257 : ) -> float | complex: ...
258 :
259 1 : def get_matrix_element(
260 : self, other: KetAtom | Self, operator: OperatorType, q: int, unit: str | None = None
261 : ) -> PintFloat | PintComplex | float | complex:
262 : """Calculate the matrix element of the operator with respect to the state and another state or ket.
263 :
264 : This means the inner product <self|operator|other>.
265 :
266 : Args:
267 : other: Either a state or a ket for which the matrix element should be calculated.
268 : operator: The operator for which the matrix element should be calculated.
269 : q: The projection quantum number of the operator.
270 : unit: The unit in which the result should be returned.
271 : Default None will return a `pint.Quantity`.
272 :
273 : Returns:
274 : The matrix element between self and other.
275 :
276 : """
277 1 : if not self.is_normalized() or (isinstance(other, StateAtom) and not other.is_normalized()):
278 0 : logger.warning("get_matrix_element/get_overlap/get_amplitude is called with a non-normalized state.")
279 :
280 1 : cpp_op = get_cpp_operator_type(operator)
281 :
282 1 : if isinstance(other, KetAtom):
283 1 : other = cast("Self", other.to_state())
284 1 : if isinstance(other, StateAtom):
285 1 : matrix_elements_au = self._cpp.get_matrix_elements(other._cpp, cpp_op, q).toarray().ravel()
286 1 : matrix_element_au = np.real_if_close(matrix_elements_au)[0]
287 1 : return QuantityScalar.convert_au_to_user(matrix_element_au, operator, unit)
288 0 : raise TypeError(f"Unknown type: {type(other)=}")
289 :
290 :
291 1 : class StateAtomReal(StateAtom):
292 1 : _cpp: _backend.BasisAtomReal # type: ignore [assignment]
293 1 : _ket_class = KetAtomReal
|
