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 collections.abc import Iterable, Sequence
6 1 : from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, Union
7 :
8 1 : import numpy as np
9 1 : import pint
10 1 : from pint import UnitRegistry
11 1 : from pint.facets.plain import PlainQuantity
12 1 : from scipy.sparse import csr_matrix
13 :
14 : if TYPE_CHECKING:
15 : from typing import TypeAlias
16 :
17 : import numpy.typing as npt
18 : from pint.facets.plain import PlainUnit
19 : from typing_extensions import Self
20 :
21 : NDArray: TypeAlias = "npt.NDArray[Any]"
22 : ArrayLike: TypeAlias = "npt.NDArray[Any] | Sequence[float]"
23 : PintFloat: TypeAlias = "PlainQuantity[float]"
24 : PintArray: TypeAlias = "PlainQuantity[NDArray]"
25 : PintArrayLike: TypeAlias = "PintArray | Sequence[float | PintFloat]"
26 : # type ignore here and also below for PlainQuantity[ValueType] because pint has no type support for scipy.csr_matrix
27 : PintSparse: TypeAlias = "PlainQuantity[csr_matrix]"
28 : # and also for complex
29 : PintComplex: TypeAlias = "PlainQuantity[complex]"
30 :
31 1 : ureg = UnitRegistry(system="atomic")
32 :
33 1 : Dimension = Literal[
34 : "electric_field",
35 : "magnetic_field",
36 : "distance",
37 : "inverse_distance",
38 : "energy",
39 : "charge",
40 : "velocity",
41 : "temperature",
42 : "time",
43 : "transition_rate",
44 : "electric_dipole",
45 : "electric_quadrupole",
46 : "electric_quadrupole_zero",
47 : "electric_octupole",
48 : "magnetic_dipole",
49 : "c3",
50 : "c6",
51 : "green_tensor_00",
52 : "scaled_green_tensor_00",
53 : "identity",
54 : "arbitrary",
55 : "zero",
56 : ]
57 1 : DimensionLike = Dimension | Iterable[Dimension]
58 :
59 : # some abbreviations: au_time: atomic_unit_of_time; au_current: atomic_unit_of_current; m_e: electron_mass
60 1 : _CommonUnits: dict[Dimension, str] = {
61 : "electric_field": "V/cm", # 1 V/cm = 1.9446903811524456e-10 bohr * m_e / au_current / au_time ** 3
62 : "magnetic_field": "T", # 1 T = 4.254382157342044e-06 m_e / au_current / au_time ** 2
63 : "distance": "micrometer", # 1 mum = 18897.26124622279 bohr
64 : "inverse_distance": "1 / micrometer",
65 : "energy": "hartree", # 1 hartree = 1 bohr ** 2 * m_e / au_time ** 2
66 : "charge": "e", # 1 e = 1 au_current * au_time
67 : "velocity": "speed_of_light", # 1 c = 137.03599908356244 bohr / au_time
68 : "temperature": "K", # 1 K = 3.1668115634555572e-06 atomic_unit_of_temperature
69 : "time": "s", # 1 s = 4.134137333518244e+16 au_time
70 : "transition_rate": "1/s", # 1 / s = 2.4188843265856806e-17 * 1 / au_time
71 : "electric_dipole": "e * a0", # 1 e * a0 = 1 au_current * au_time * bohr
72 : "electric_quadrupole": "e * a0^2", # 1 e * a0^2 = 1 au_current * au_time * bohr ** 2
73 : "electric_quadrupole_zero": "e * a0^2", # 1 e * a0^2 = 1 au_current * au_time * bohr ** 2
74 : "electric_octupole": "e * a0^3", # 1 e * a0^3 = 1 au_current * au_time * bohr ** 3
75 : "magnetic_dipole": "hbar e / m_e", # 1 hbar e / m_e = 1 au_current * bohr ** 2
76 : "c3": "hartree * bohr^3", # 1 hartree * bohr^3 = 1 bohr ** 3 * m_e / au_time ** 2
77 : "c6": "hartree * bohr^6", # 1 hartree * bohr^6 = 1 bohr ** 6 * m_e / au_time ** 2
78 : "green_tensor_00": "meter", # unit for green tensor with kappa1 = kappa2 = 0
79 : "scaled_green_tensor_00": "hartree / e^2", # unit for scaled green tensor with kappa1 = kappa2 = 1
80 : "identity": "", # 1 dimensionless
81 : "arbitrary": "", # 1 dimensionless
82 : "zero": "", # 1 dimensionless
83 : }
84 1 : AtomicUnits: dict[Dimension, PlainUnit] = {
85 : k: ureg.Quantity(1, unit).to_base_units().units for k, unit in _CommonUnits.items()
86 : }
87 :
88 1 : Context = Literal["spectroscopy", "Gaussian"]
89 1 : BaseContexts: dict[Dimension, Context] = {
90 : "magnetic_field": "Gaussian",
91 : "energy": "spectroscopy",
92 : "c3": "spectroscopy",
93 : "c6": "spectroscopy",
94 : }
95 :
96 1 : ValueType = TypeVar("ValueType", bound=Union[float, "NDArray", "csr_matrix"])
97 1 : ValueTypeLike = TypeVar("ValueTypeLike", bound=Union[float, "ArrayLike", "csr_matrix"])
98 :
99 :
100 1 : class QuantityAbstract(Generic[ValueTypeLike, ValueType]):
101 1 : def __init__(self, pint_qty: PlainQuantity[ValueType], dimension: DimensionLike) -> None: # type: ignore [type-var]
102 1 : if not isinstance(pint_qty, ureg.Quantity):
103 0 : raise TypeError(f"pint_qty must be a ureg.Quantity, not {type(pint_qty)}")
104 1 : self._quantity = pint_qty
105 1 : self.dimension: DimensionLike = dimension
106 1 : self.check_value_type()
107 :
108 1 : def check_value_type(self) -> None:
109 0 : raise NotImplementedError("This method must be implemented in the derived classes.")
110 :
111 1 : @classmethod
112 1 : def get_atomic_unit(cls, dimension: DimensionLike) -> str:
113 1 : if isinstance(dimension, str):
114 1 : return str(AtomicUnits[dimension])
115 : # dimension isinstance Iterable[Dimension]
116 1 : return " * ".join(str(AtomicUnits[d]) for d in dimension)
117 :
118 1 : @classmethod
119 1 : def get_contexts(cls, dimension: DimensionLike) -> list[Context]:
120 1 : if isinstance(dimension, str):
121 1 : return [BaseContexts[dimension]] if dimension in BaseContexts else []
122 1 : contexts: set[Context] = {BaseContexts[d] for d in dimension if d in BaseContexts}
123 1 : return list(contexts)
124 :
125 1 : @classmethod
126 1 : def from_pint(
127 : cls: type[Self],
128 : value: PlainQuantity[ValueType], # type: ignore [type-var]
129 : dimension: DimensionLike,
130 : ) -> Self:
131 : """Initialize a Quantity from a ureg.Quantity."""
132 1 : if isinstance(value, ureg.Quantity):
133 1 : return cls(value, dimension)
134 0 : if isinstance(value, PlainQuantity):
135 0 : raise TypeError(
136 : "Only use pint quantities genereated by pairinteraction.ureg and not by a different pint.UnitRegistry."
137 : )
138 0 : raise ValueError("method from_pint: value must be a pint.Quantity")
139 :
140 1 : @classmethod
141 1 : def from_unit(
142 : cls: type[Self],
143 : value: ValueTypeLike,
144 : unit: str,
145 : dimension: DimensionLike,
146 : ) -> Self:
147 : """Initialize a Quantity from a value and a unit given as string."""
148 1 : if isinstance(value, PlainQuantity):
149 0 : raise TypeError("method from_unit: value must be a scalar or an array, not a pint.Quantity")
150 1 : return cls(ureg.Quantity(value, unit), dimension)
151 :
152 1 : @classmethod
153 1 : def from_au(
154 : cls: type[Self],
155 : value: ValueTypeLike,
156 : dimension: DimensionLike,
157 : ) -> Self:
158 : """Initialize a Quantity from a value in atomic units (a.u.) and a (list of) dimension(s)."""
159 1 : unit = cls.get_atomic_unit(dimension)
160 1 : return cls(ureg.Quantity(value, unit), dimension)
161 :
162 1 : @classmethod
163 1 : def from_pint_or_unit(
164 : cls: type[Self],
165 : value: PlainQuantity[ValueType] | ValueTypeLike, # type: ignore [type-var]
166 : unit: str | None,
167 : dimension: DimensionLike,
168 : ) -> Self:
169 1 : if unit is None:
170 1 : if isinstance(value, PlainQuantity):
171 1 : return cls.from_pint(value, dimension)
172 1 : if np.all(value == 0):
173 1 : return cls.from_au(value, dimension)
174 0 : raise ValueError("unit must be given if value is not a pint.Quantity")
175 1 : assert not isinstance(value, PlainQuantity)
176 1 : return cls.from_unit(value, unit, dimension)
177 :
178 1 : def to_pint(self) -> PlainQuantity[ValueType]: # type: ignore [type-var]
179 : """Return the pint.Quantity object."""
180 1 : contexts = self.get_contexts(self.dimension)
181 1 : atomic_unit = self.get_atomic_unit(self.dimension)
182 1 : return self._quantity.to(atomic_unit, *contexts)
183 :
184 1 : def to_unit(
185 : self,
186 : unit: str,
187 : ) -> ValueType:
188 : """Return the value of the quantity in the given unit."""
189 1 : contexts = self.get_contexts(self.dimension)
190 1 : try:
191 1 : return self._quantity.to(unit, *contexts).magnitude # type: ignore [no-any-return] # also a problem with pint with sparse matrix
192 1 : except pint.errors.DimensionalityError:
193 : # pint uses e.g. the context "spectroscopy" to convert "hartree" -> "GHz"
194 : # however, something like "hartree * bohr^3" -> "GHz * bohr^3" does not work
195 : # the following is a workaround for this kind of conversions
196 1 : if "spectroscopy" in contexts:
197 1 : q = self._quantity * ureg.Quantity(1, "GHz") / ureg.Quantity(1, "GHz").to("hartree", "spectroscopy")
198 1 : return q.to(unit, *contexts).magnitude # type: ignore [no-any-return]
199 0 : raise
200 :
201 1 : def to_au(self) -> ValueType:
202 : """Return the value of the quantity in atomic units (a.u.)."""
203 1 : value = self.to_pint().to_base_units()
204 1 : return value.magnitude
205 :
206 1 : def to_pint_or_unit(self, unit: str | None) -> ValueType | PlainQuantity[ValueType]: # type: ignore [type-var]
207 1 : if unit is None:
208 1 : return self.to_pint()
209 1 : return self.to_unit(unit)
210 :
211 1 : @classmethod
212 1 : def convert_user_to_au(
213 : cls,
214 : value: PlainQuantity[ValueType] | ValueTypeLike, # type: ignore [type-var]
215 : unit: str | None,
216 : dimension: DimensionLike,
217 : ) -> ValueType:
218 1 : return cls.from_pint_or_unit(value, unit, dimension).to_au()
219 :
220 1 : @classmethod
221 1 : def convert_user_to_pint(
222 : cls,
223 : value: PlainQuantity[ValueType] | ValueTypeLike, # type: ignore [type-var]
224 : unit: str | None,
225 : dimension: DimensionLike,
226 : ) -> PlainQuantity[ValueType]: # type: ignore [type-var]
227 1 : return cls.from_pint_or_unit(value, unit, dimension).to_pint()
228 :
229 1 : @classmethod
230 1 : def convert_au_to_user(
231 : cls,
232 : values_au: ValueTypeLike,
233 : dimension: DimensionLike,
234 : unit: str | None,
235 : ) -> ValueType | PlainQuantity[ValueType]: # type: ignore [type-var]
236 1 : return cls.from_au(values_au, dimension).to_pint_or_unit(unit)
237 :
238 1 : @classmethod
239 1 : def convert_pint_to_user(
240 : cls,
241 : value_pint: PlainQuantity[ValueType], # type: ignore [type-var]
242 : dimension: DimensionLike,
243 : unit: str | None,
244 : ) -> ValueType | PlainQuantity[ValueType]: # type: ignore [type-var]
245 1 : return cls.from_pint(value_pint, dimension).to_pint_or_unit(unit)
246 :
247 :
248 1 : class QuantityScalar(QuantityAbstract[float, float]):
249 1 : def check_value_type(self) -> None:
250 1 : magnitude = self._quantity.magnitude
251 1 : if not np.isscalar(magnitude):
252 0 : raise TypeError(f"value must be a scalar, not {type(magnitude)}")
253 :
254 :
255 1 : class QuantityArray(QuantityAbstract["ArrayLike", "NDArray"]):
256 1 : def check_value_type(self) -> None:
257 1 : magnitude = self._quantity.magnitude
258 1 : if not isinstance(magnitude, Sequence) and not isinstance(magnitude, np.ndarray):
259 0 : raise TypeError(f"value must be an np.ndarray (or a Sequence), not {type(magnitude)}")
260 :
261 :
262 1 : class QuantitySparse(QuantityAbstract["csr_matrix", "csr_matrix"]):
263 1 : def check_value_type(self) -> None:
264 1 : magnitude = self._quantity.magnitude
265 1 : if not isinstance(magnitude, csr_matrix):
266 0 : raise TypeError(f"value must be a scipy.sparse.csr_matrix, not {type(magnitude)}")
|
