-
Notifications
You must be signed in to change notification settings - Fork 4
/
general.py
369 lines (300 loc) · 9.37 KB
/
general.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
"""Some general functions and utilities used in fitting.
Many functions are `numba`-compiled for speed.
"""
__all__ = [
"TINY",
"bcs_gap",
"do_convolve",
"do_convolve_2d",
"dynes",
"fermi_dirac",
"fermi_dirac_linbkg",
"fermi_dirac_linbkg_broad",
"gaussian",
"gaussian_wh",
"lorentzian",
"lorentzian_wh",
"step_broad",
"step_linbkg_broad",
]
from collections.abc import Callable
import numba
import numpy as np
import numpy.typing as npt
import scipy.special
from erlab.constants import kb_eV
#: From :mod:`lmfit.lineshapes`, equal to `numpy.finfo(numpy.float64).resolution`
TINY: float = 1.0e-15
S2PI = np.sqrt(2 * np.pi)
@numba.njit(cache=True)
def _infer_meshgrid_shape(arr: np.ndarray) -> tuple[tuple[int, int], int, np.ndarray]:
if arr.ndim != 1:
raise ValueError("Array must be 1-dimensional")
axis = 1
# Find the index at which the array value decreases
change_index = np.where(np.diff(arr) <= 0)[0]
if change_index.size == 0:
raise ValueError("Array does not appear to be a flattened meshgrid")
if change_index[0] == 0:
change_index = np.where(np.diff(arr) > 0)[0]
axis = 0
# The shape of the original meshgrid
shape = len(arr) // (change_index[0] + 1), change_index[0] + 1
if axis == 0:
coord = arr.reshape(shape)[:, 0]
else:
coord = arr.reshape(shape)[0, :]
return shape, axis, coord
@numba.njit(cache=True)
def _gen_kernel(
x: npt.NDArray[np.float64], resolution: float, pad: int = 5
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
r"""Generate a Gaussian kernel for convolution.
Parameters
----------
x
The input array of x values.
resolution
The resolution of the kernel given as FWHM.
pad
Multiples of the standard deviation :math:`\sigma` to truncate the kernel at.
Returns
-------
extended
The domain of the kernel.
gauss
The gaussian kernel defined on `extended`.
"""
delta_x = x[1] - x[0]
sigma = abs(resolution) / np.sqrt(8 * np.log(2)) # resolution given in FWHM
n_pad = int(sigma * pad / delta_x + 0.5)
x_pad = n_pad * delta_x
extended = np.linspace(x[0] - x_pad, x[-1] + x_pad, 2 * n_pad + len(x))
gauss = (
delta_x
* np.exp(
-(np.linspace(-x_pad, x_pad, 2 * n_pad + 1) ** 2) / max(TINY, 2 * sigma**2)
)
/ max(TINY, S2PI * sigma)
)
return extended, gauss
def do_convolve(
x: npt.NDArray[np.float64],
func: Callable,
resolution: float,
pad: int = 5,
**kwargs,
) -> npt.NDArray[np.float64]:
r"""Convolves `func` with gaussian of FWHM `resolution` in `x`.
Parameters
----------
x
A evenly spaced array specifing where to evaluate the convolution.
func
Function to convolve.
resolution
FWHM of the gaussian kernel.
pad
Multiples of the standard deviation :math:`\sigma` to pad with.
**kwargs
Additional keyword arguments to `func`.
"""
xn, g = _gen_kernel(
np.asarray(x, dtype=np.float64), float(resolution), pad=int(pad)
)
return np.convolve(func(xn, **kwargs), g, mode="valid")
def do_convolve_2d(
x: npt.NDArray[np.float64],
y: npt.NDArray[np.float64] | float,
func: Callable,
resolution: float,
pad: int = 5,
**kwargs,
) -> npt.NDArray[np.float64]:
idx_x = None
if not np.iterable(y):
y = np.asarray([y])
try:
# check if x is a meshgrid
shape_x, idx_x, x = _infer_meshgrid_shape(np.ascontiguousarray(x))
shape_y, _, y = _infer_meshgrid_shape(np.ascontiguousarray(y))
except ValueError:
pass
else:
if shape_x != shape_y:
raise ValueError("x and y do not have matching shape")
xn, g = _gen_kernel(
np.asarray(np.squeeze(x), dtype=np.float64), resolution, pad=pad
)
convolved = np.vstack(
[
np.convolve(func(xn, yi, **kwargs), g, mode="valid")
for yi in np.asarray(y).flat
]
)
if idx_x is None:
return convolved.T
elif idx_x == 0:
return convolved.T.ravel()
else:
return convolved.ravel()
@numba.njit(cache=True)
def gaussian_wh(
x: npt.NDArray[np.float64], center: float, width: float, height: float
) -> npt.NDArray[np.float64]:
r"""Gaussian parametrized with FWHM and peak height.
Note
----
:math:`\sigma=\frac{w}{2\sqrt{2\log{2}}}`
"""
return height * np.exp(
-16 * np.log(2) * (1.0 * x - center) ** 2 / max(TINY, width**2)
)
@numba.njit(cache=True)
def gaussian(
x: npt.NDArray[np.float64], center: float, sigma: float, amplitude: float
) -> npt.NDArray[np.float64]:
return (amplitude / (max(TINY, S2PI * sigma))) * np.exp(
-((1.0 * x - center) ** 2) / max(TINY, (2 * sigma**2))
)
@numba.njit(cache=True)
def lorentzian_wh(
x: npt.NDArray[np.float64], center: float, width: float, height: float
) -> npt.NDArray[np.float64]:
r"""Lorentzian parametrized with FWHM and peak height.
Note
----
:math:`\sigma=w/2`
"""
return height / (1 + 4 * ((1.0 * x - center) / max(TINY, width)) ** 2)
@numba.njit(cache=True)
def lorentzian(
x: npt.NDArray[np.float64], center: float, sigma: float, amplitude: float
) -> npt.NDArray[np.float64]:
return (amplitude / (1 + ((1.0 * x - center) / max(TINY, sigma)) ** 2)) / max(
TINY, (np.pi * sigma)
)
@numba.njit(cache=True)
def fermi_dirac(
x: npt.NDArray[np.float64], center: float, temp: float
) -> npt.NDArray[np.float64]:
"""Fermi-dirac edge in terms of temperature."""
return 1 / (1 + np.exp((1.0 * x - center) / max(TINY, temp * kb_eV)))
# adapted and improved from KWAN Igor procedures
@numba.njit(cache=True)
def fermi_dirac_linbkg(
x: npt.NDArray[np.float64],
center: float,
temp: float,
back0: float,
back1: float,
dos0: float,
dos1: float,
) -> npt.NDArray[np.float64]:
"""Fermi-dirac edge with linear backgrounds above and below the fermi level.
Note
----
`back0` and `back1` corresponds to the linear background above and below EF (due to
non-homogeneous detector efficiency or residual intensity on the phosphor screen
during sweep mode), while `dos0` and `dos1` corresponds to the linear density of
states below EF including the linear background.
"""
return (back0 + back1 * x) + (dos0 - back0 + (dos1 - back1) * x) / (
1 + np.exp((1.0 * x - center) / max(TINY, temp * kb_eV))
)
def fermi_dirac_linbkg_broad(
x: npt.NDArray[np.float64],
center: float,
temp: float,
resolution: float,
back0: float,
back1: float,
dos0: float,
dos1: float,
) -> npt.NDArray[np.float64]:
"""Resolution-broadened Fermi edge with linear backgrounds above and below EF."""
return do_convolve(
x,
fermi_dirac_linbkg,
resolution=resolution,
center=center,
temp=temp,
back0=back0,
back1=back1,
dos0=dos0,
dos1=dos1,
)
def step_broad(
x: npt.NDArray[np.float64],
center: float = 0.0,
sigma: float = 1.0,
amplitude: float = 1.0,
):
"""Step function convolved with a Gaussian."""
return (
amplitude
* 0.5
* scipy.special.erfc((1.0 * x - center) / max(TINY, np.sqrt(2) * sigma))
)
def step_linbkg_broad(
x: npt.NDArray[np.float64],
center: float,
sigma: float,
back0: float,
back1: float,
dos0: float,
dos1: float,
):
"""Resolution broadened step function with linear backgrounds."""
return (back0 + back1 * x) + (dos0 - back0 + (dos1 - back1) * x) * (
step_broad(x, center, sigma, 1.0)
)
@numba.njit()
def bcs_gap(
x, a: float = 1.76, b: float = 1.74, tc: float = 100.0
) -> npt.NDArray[np.float64]:
r"""Interpolation formula for a temperature dependent BCS-like gap.
.. math::
\Delta(T) \simeq a \cdot k_B T_c \cdot \tanh\left(b \sqrt{\frac{T_c}{T} -
1}\right)
Parameters
----------
x : array-like
The temperature values in kelvins at which to calculate the BCS gap.
a
Proportionality constant. Default is 1.76.
b
Proportionality constant. Default is 1.74.
tc
The critical temperature in Kelvins. Default is 100.0.
"""
out = np.empty_like(x, dtype=np.float64)
for i in range(len(x)):
if x[i] < tc:
out[i] = a * kb_eV * tc * np.tanh(b * np.sqrt(tc / x[i] - 1))
else:
out[i] = 0.0
return out
def dynes(x, n0=1.0, gamma=0.003, delta=0.01):
r"""Dynes formula for superconducting density of states.
The formula is given by :cite:p:`dynes1978dynes`:
.. math::
f(x) = N_0 \text{Re}\left[\frac{|x| + i \Gamma}{\sqrt{(|x| + i \Gamma)^2 -
\Delta^2}}\right]
where :math:`x` is the binding energy, :math:`N_0` is the normal-state density of
states at the Fermi level, :math:`\Gamma` is the broadening term, and :math:`\Delta`
is the superconducting energy gap.
Parameters
----------
x : array-like
The input array of energy in eV.
n0
:math:`N_0`, by default 1.0.
gamma
:math:`\Gamma`, by default 0.003.
delta
The superconducting energy gap :math:`\Delta`, by default 0.01.
"""
return n0 * np.real(
(np.abs(x) + 1j * gamma) / (np.sqrt((np.abs(x) + 1j * gamma) ** 2 - delta**2))
)