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, Literal, overload
6 :
7 1 : import numpy as np
8 :
9 1 : from pairinteraction import _backend
10 1 : from pairinteraction.green_tensor.green_tensor_base import GreenTensorBase
11 1 : from pairinteraction.units import QuantityArray, QuantityScalar
12 :
13 : if TYPE_CHECKING:
14 : from collections.abc import Collection
15 :
16 : from typing_extensions import Self
17 :
18 : from pairinteraction.units import NDArray, PintArray, PintFloat
19 :
20 1 : Coordinates = Literal["cartesian", "spherical"]
21 :
22 :
23 1 : class GreenTensorInterpolator:
24 : """Green tensor interpolator for the multipole pair interactions.
25 :
26 : This class is mainly used internally, a user usually only has to access the GreenTensor classes.
27 : It allows to define constant or frequency-dependent Green tensor interpolators,
28 : which can then be used for the interaction of a :class:`SystemPair`.
29 :
30 : Examples:
31 : >>> from pairinteraction.green_tensor.green_tensor_interpolator import GreenTensorInterpolator
32 : >>> gt = GreenTensorInterpolator()
33 : >>> distance_mum = 5
34 : >>> transition_energy = 1 # planck_constant * GHz
35 : >>> tensor = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]]) / (distance_mum**3)
36 : >>> tensor_unit = "hartree / (e * micrometer)^2"
37 : >>> gt.set_constant(1, 1, tensor, tensor_unit)
38 : GreenTensorInterpolator(...)
39 : >>> tensor_sph = gt.get(1, 1, transition_energy, "planck_constant * GHz", unit=tensor_unit, scaled=True)
40 : >>> print(tensor_sph.diagonal())
41 : [ 0.008 -0.016 0.008]
42 :
43 : """
44 :
45 1 : _cpp: _backend.GreenTensorInterpolatorComplex
46 1 : _cpp_type = _backend.GreenTensorInterpolatorComplex
47 :
48 1 : def __init__(self) -> None:
49 : """Initialize a new Green tensor interpolator object.
50 :
51 : The actual tensor can be set afterwards via the
52 : :meth:`set_constant` or :meth:`set_list` method.
53 : """
54 1 : self._cpp = self._cpp_type()
55 :
56 1 : def __repr__(self) -> str:
57 1 : return f"{type(self).__name__}(...)"
58 :
59 1 : def __str__(self) -> str:
60 0 : return self.__repr__()
61 :
62 1 : def set_constant(
63 : self,
64 : kappa1: int,
65 : kappa2: int,
66 : tensor: NDArray | PintArray,
67 : tensor_unit: str | None = None,
68 : *,
69 : coordinates: Coordinates = "cartesian",
70 : ) -> Self:
71 : r"""Set the scaled Green tensor to a constant entry.
72 :
73 : Constant means, that :math:`\omega^2 G(\omega)` (which is the quantity that enters the interaction)
74 : is constant and independent of omega.
75 :
76 : Args:
77 : kappa1: The rank of the first multipole operator.
78 : kappa2: The rank of the second multipole operator.
79 : tensor: The scaled green tensor including the prefactor for the interaction strength
80 : (see :meth:`GreenTensorBase._get_prefactor_au`).
81 : tensor_unit: The unit of the tensor.
82 : Default None, which means that the tensor must be given as pint object.
83 : coordinates: The coordinate system in which the tensor is given.
84 : Default "cartesian".
85 :
86 : """
87 1 : if coordinates != "cartesian":
88 0 : raise NotImplementedError("Only cartesian coordinates are currently implemented for set_constant.")
89 1 : if tensor.shape != (3**kappa1, 3**kappa2) or tensor.ndim != 2: # type: ignore [union-attr]
90 0 : raise ValueError("The tensor must be a 2D array of shape (3**kappa1, 3**kappa2).")
91 :
92 1 : dimension = GreenTensorBase._get_dimension(kappa1, kappa2, scaled=True)
93 1 : scaled_tensor_au = QuantityArray.convert_user_to_au(tensor, tensor_unit, dimension)
94 1 : self._cpp.create_entries_from_cartesian(kappa1, kappa2, scaled_tensor_au)
95 1 : return self
96 :
97 1 : def set_list(
98 : self,
99 : kappa1: int,
100 : kappa2: int,
101 : tensors: Collection[PintArray] | Collection[NDArray],
102 : transition_energies: Collection[PintFloat] | Collection[float],
103 : tensors_unit: str | None = None,
104 : transition_energies_unit: str | None = None,
105 : *,
106 : coordinates: Coordinates = "cartesian",
107 : scaled: bool = False,
108 : ) -> Self:
109 : """Set the entries of the Green tensor for specified transition energies.
110 :
111 : Args:
112 : kappa1: The rank of the first multipole operator.
113 : kappa2: The rank of the second multipole operator.
114 : tensors: A list of frequency-dependent green tensors in cartesian coordinates.
115 : transition_energies: A list of transition energies at which the green tensors are defined.
116 : tensors_unit: The unit of the tensor.
117 : Default None, which means that the tensor must be given as pint object.
118 : transition_energies_unit: The unit of the transition energies.
119 : Default None, which means that the transition energies must be given as pint object.
120 : coordinates: The coordinate system in which the tensor is given.
121 : Default "cartesian".
122 : scaled: Whether the prefactor for the interaction strength
123 : (see :meth:`GreenTensorBase._get_prefactor_au`) is already included in the given tensor.
124 : The unit has to be adjusted accordingly. Default False.
125 :
126 : """
127 1 : if coordinates != "cartesian":
128 0 : raise NotImplementedError("Only cartesian coordinates are currently implemented for set_list.")
129 1 : if not all(t.ndim == 2 for t in tensors):
130 0 : raise ValueError("The tensor must be a list of 2D arrays.")
131 1 : if not all(t.shape == (3**kappa1, 3**kappa2) for t in tensors): # type: ignore [union-attr]
132 0 : raise ValueError("The tensors must be of shape (3**kappa1, 3**kappa2).")
133 :
134 1 : omegas_au = [
135 : QuantityScalar.convert_user_to_au(omega, transition_energies_unit, "energy")
136 : for omega in transition_energies
137 : ]
138 1 : prefactors = [1.0] * len(tensors)
139 1 : if not scaled:
140 1 : prefactors = [GreenTensorBase._get_prefactor_au(kappa1, kappa2, omega) for omega in omegas_au]
141 1 : dimension = GreenTensorBase._get_dimension(kappa1, kappa2, scaled=scaled)
142 1 : tensors_au = [QuantityArray.convert_user_to_au(t, tensors_unit, dimension) for t in tensors]
143 1 : scaled_tensors_au = [prefactor * t for prefactor, t in zip(prefactors, tensors_au, strict=True)]
144 1 : self._cpp.create_entries_from_cartesian(kappa1, kappa2, scaled_tensors_au, omegas_au) # type: ignore [arg-type]
145 1 : return self
146 :
147 : @overload
148 : def get(
149 : self,
150 : kappa1: int,
151 : kappa2: int,
152 : transition_energy: float | PintFloat,
153 : transition_energy_unit: str | None = None,
154 : unit: None = None,
155 : *,
156 : scaled: bool = False,
157 : coordinates: Coordinates = "spherical",
158 : ) -> PintArray: ...
159 :
160 : @overload
161 : def get(
162 : self,
163 : kappa1: int,
164 : kappa2: int,
165 : transition_energy: float,
166 : transition_energy_unit: str,
167 : unit: str,
168 : *,
169 : scaled: bool = False,
170 : coordinates: Coordinates = "spherical",
171 : ) -> NDArray: ...
172 :
173 : @overload
174 : def get(
175 : self,
176 : kappa1: int,
177 : kappa2: int,
178 : transition_energy: PintFloat,
179 : *,
180 : unit: str,
181 : scaled: bool = False,
182 : coordinates: Coordinates = "spherical",
183 : ) -> NDArray: ...
184 :
185 1 : def get(
186 : self,
187 : kappa1: int,
188 : kappa2: int,
189 : transition_energy: float | PintFloat,
190 : transition_energy_unit: str | None = None,
191 : unit: str | None = None,
192 : *,
193 : scaled: bool = False,
194 : coordinates: Coordinates = "spherical",
195 : ) -> PintArray | NDArray:
196 : """Get the Green tensor in the given coordinates for the given ranks kappa1, kappa2 and transition energy.
197 :
198 : kappa = 1 corresponds to dipole operator with the basis
199 : - spherical: [p_{1,-1}, p_{1,0}, p_{1,1}]
200 : kappa = 2 corresponds to quadrupole operator with the basis
201 : - spherical: [p_{2,-2}, p_{2,-1}, p_{2,0}, p_{2,1}, p_{2,2}, p_{0,0}]
202 :
203 : Args:
204 : kappa1: The rank of the first multipole operator.
205 : kappa2: The rank of the second multipole operator.
206 : transition_energy: The transition energy at which to evaluate the Green tensor.
207 : Use transition_energy=0 for the static limit.
208 : transition_energy_unit: The unit of the transition energy.
209 : Default None, which means that the transition energy must be given as pint object (or is 0).
210 : unit: The unit to which to convert the result.
211 : Default None, which means that the result is returned as pint object.
212 : scaled: If True, the Green tensor is returned with the prefactor for the interaction
213 : already included (the unit has to be adopted accordingly).
214 : Default False returns the bare Green tensor.
215 : coordinates: The coordinate system in which to return the tensor.
216 : Default "spherical".
217 :
218 : Returns:
219 : The Green tensor as a 2D array.
220 :
221 : """
222 1 : if coordinates != "spherical":
223 0 : raise NotImplementedError("Only spherical coordinates are currently implemented for get.")
224 :
225 1 : omega_au = QuantityScalar.convert_user_to_au(transition_energy, transition_energy_unit, "energy")
226 :
227 1 : entries_cpp = self._cpp.get_spherical_entries(kappa1, kappa2)
228 1 : kappa_to_dim = {1: 3, 2: 6}
229 1 : dim1, dim2 = kappa_to_dim[kappa1], kappa_to_dim[kappa2]
230 1 : tensor_au = np.zeros((dim1, dim2), dtype=complex)
231 1 : for entry_cpp in entries_cpp:
232 1 : if isinstance(entry_cpp, (_backend.ConstantEntryReal, _backend.ConstantEntryComplex)):
233 1 : val = entry_cpp.val()
234 : else:
235 1 : val = entry_cpp.val(omega_au)
236 1 : tensor_au[entry_cpp.row(), entry_cpp.col()] = val
237 1 : tensor_au = np.real_if_close(tensor_au)
238 :
239 1 : prefactor = 1 if scaled else GreenTensorBase._get_prefactor_au(kappa1, kappa2, omega_au)
240 1 : dimension = GreenTensorBase._get_dimension(kappa1, kappa2, scaled=scaled)
241 1 : return QuantityArray.convert_au_to_user(tensor_au / prefactor, dimension, unit)
242 :
243 :
244 1 : class GreenTensorInterpolatorReal(GreenTensorInterpolator):
245 1 : _cpp: _backend.GreenTensorInterpolatorReal # type: ignore [assignment]
246 1 : _cpp_type = _backend.GreenTensorInterpolatorReal # type: ignore [assignment]
|
