401 lines
14 KiB
Python
401 lines
14 KiB
Python
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
import click
|
|
import labmodule as lab
|
|
import numpy as np
|
|
from customfunc.app import PlotApp
|
|
from numpy.fft import fft, fftshift, ifft
|
|
from scipy.optimize import curve_fit
|
|
|
|
LN2 = np.log(2)
|
|
|
|
|
|
@click.group()
|
|
def main():
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class FitResult:
|
|
curve: np.ndarray = field(repr=False)
|
|
covariance: np.ndarray
|
|
amp: float
|
|
width: float
|
|
x0: float = 0
|
|
params: tuple[float, ...] = field(default_factory=tuple)
|
|
|
|
|
|
@dataclass
|
|
class VoigtFitResult:
|
|
name: str
|
|
curve: np.ndarray = field(repr=False)
|
|
covariance: np.ndarray
|
|
gaussian_width: float
|
|
lorentzian_width: float
|
|
|
|
@property
|
|
def descr(self) -> str:
|
|
return (
|
|
f"{self.name} - Gaussian width : {self.gaussian_width:.2f}, "
|
|
f"Lorentzian width : {self.lorentzian_width:.2f}"
|
|
)
|
|
|
|
|
|
def setup_axes(n: int, lims=(-10, 10)) -> tuple[np.ndarray, float, np.ndarray, float]:
|
|
"""Creates a time-like and a frequency-like array and returns their respective spacing"""
|
|
x, dx = np.linspace(*lims, n, retstep=True)
|
|
f = np.fft.fftfreq(n, dx)
|
|
return x, dx, f, f[1] - f[0]
|
|
|
|
|
|
def freq_axis(x: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
|
|
"""
|
|
given an array with equally spaced values, returns a corresponding frequency array,
|
|
an array of indices to automatically sort ffts and the x spacing
|
|
"""
|
|
dx = np.diff(x).mean()
|
|
f = np.fft.fftfreq(len(x), dx)
|
|
ind = np.argsort(f)
|
|
return f[ind], ind, dx
|
|
|
|
|
|
def lorentzian(x: np.ndarray, amp: float, width: float, x0: float) -> np.ndarray:
|
|
"""Lorentzian function of max `amp`, FWHM `width` and position `x0`"""
|
|
gamma = (width * 0.5) ** 2
|
|
return amp * gamma / ((x - x0) ** 2 + gamma)
|
|
|
|
|
|
def unit_lorentzian(x: np.ndarray, width: float, x0: float) -> np.ndarray:
|
|
"""Normalized Lorentzian function of FWHM `width` and position `x0`"""
|
|
return lorentzian(x, 1 / (np.pi * width * 0.5), width, x0)
|
|
|
|
|
|
def peak_func(x: np.ndarray, amp: float, width: float, x0: float) -> np.ndarray:
|
|
"""2-sided decaying exp of max `amp`, FWHM `width` and position `x0`"""
|
|
a = 2 * LN2 / width
|
|
return amp * np.exp(-np.abs(x - x0) * a)
|
|
|
|
|
|
def fft_lorentzian(f: np.ndarray, dx: float, amp: float, width: float) -> np.ndarray:
|
|
"""
|
|
Same as `peak_func` but parametrized as the fourier transform
|
|
of a Lorentzian function of max `amp` and FWHM `width`
|
|
"""
|
|
gamma_sqrt = width / 2
|
|
return amp * np.pi * gamma_sqrt * np.exp(-2 * np.pi * gamma_sqrt * np.abs(f)) / dx
|
|
|
|
|
|
def gaussian(x: np.ndarray, amp: float, width: float, x0: float) -> np.ndarray:
|
|
"""Gaussian function of max `amp`, FWHM `width` and position `x0`"""
|
|
sig2 = width ** 2 / (8 * LN2)
|
|
return amp * np.exp(-((x - x0) ** 2) / (2 * sig2))
|
|
|
|
|
|
def elevated_gaussian(x: np.ndarray, amp: float, width: float, x0: float, y0: float) -> np.ndarray:
|
|
"""Gaussian function of max `amp`, FWHM `width`, position `x0` and baseline `y0`"""
|
|
sig2 = width ** 2 / (8 * LN2)
|
|
return amp * np.exp(-((x - x0) ** 2) / (2 * sig2)) + y0
|
|
|
|
|
|
def unit_gaussian(x: np.ndarray, width: float, x0: float) -> np.ndarray:
|
|
"""Normalized Gaussian function of FWHM `width` and position `x0`"""
|
|
return gaussian(x, np.sqrt(4 * LN2 / np.pi) / width, width, x0)
|
|
|
|
|
|
def fft_gaussian(f: np.ndarray, dx: float, amp: float, width: float) -> np.ndarray:
|
|
"""
|
|
Same as `gaussian` but parametrized as the fourier transform
|
|
of another Gaussian function of max `amp` and FWHM `width`
|
|
""" # exp_param =
|
|
freq_amp = amp * width / 2 * np.sqrt(np.pi / LN2) / dx
|
|
freq_width = 4 * LN2 / (width * np.pi)
|
|
return gaussian(f, freq_amp, freq_width, 0)
|
|
|
|
|
|
def voigt(x: np.ndarray, gaussian_width: float, lorentzian_width: float) -> np.ndarray:
|
|
"""Voigt profile of max 1 with specified Gaussian and Lorentzian FWHM"""
|
|
conv = convolve(unit_gaussian(x, gaussian_width, 0), unit_lorentzian(x, lorentzian_width, 0))
|
|
if conv.max() > 0:
|
|
return conv / conv.max()
|
|
return conv
|
|
|
|
|
|
def log_voigt(x: np.ndarray, gaussian_width: float, lorentzian_width: float) -> np.ndarray:
|
|
"""log of `voigt`"""
|
|
return np.log(voigt(x, gaussian_width, lorentzian_width))
|
|
|
|
|
|
def fft_voigt(
|
|
f: np.ndarray, dx: float, gaussian_width: float, lorentzian_width: float
|
|
) -> np.ndarray:
|
|
"""
|
|
returns the product of a `peak_func` and a `gaussian`
|
|
parameterized as the Fourier transform of a Voigt profile
|
|
given by `gaussian_width` and `lorentzian_width`
|
|
"""
|
|
profile = fft_gaussian(f, dx, 1, gaussian_width) * fft_lorentzian(f, dx, 1, lorentzian_width)
|
|
return profile / profile.max()
|
|
|
|
|
|
def single_func_fit(
|
|
x: np.ndarray,
|
|
y: np.ndarray,
|
|
fct: Callable[[np.ndarray, float, float, float], np.ndarray],
|
|
force_center=False,
|
|
use_weights=False,
|
|
extra_params: int = 0,
|
|
) -> FitResult:
|
|
"""fits a single bell-like function (Gaussian or Lorentian, not Voigt)
|
|
|
|
Parameters
|
|
----------
|
|
x : np.ndarray
|
|
x input
|
|
y : np.ndarray
|
|
corresponding y data
|
|
fct : Callable[[np.ndarray, float, float, float], np.ndarray]
|
|
function to fit
|
|
force_center : bool, optional
|
|
assume the bell is centered at x=0, by default False
|
|
use_weights : bool, optional
|
|
give more importance to higher y values, by default False
|
|
|
|
Returns
|
|
-------
|
|
FitResult
|
|
resuts of the fit
|
|
"""
|
|
if force_center:
|
|
|
|
def _fct(_x: np.ndarray, amp: float, width: float):
|
|
return fct(_x, amp, width, 0)
|
|
|
|
fct = _fct
|
|
am = y.argmax()
|
|
params, cov = curve_fit(
|
|
fct, x, y, sigma=1 / y if use_weights else None, p0=[y[am], 1, x[am]] + [0] * extra_params
|
|
)
|
|
curve = fct(x, *params)
|
|
return FitResult(curve, cov, params[0], params[1], params[2], tuple(params))
|
|
|
|
|
|
def linear_direct_voigt_fit(x: np.ndarray, y: np.ndarray, use_weights=False) -> VoigtFitResult:
|
|
params, cov = curve_fit(voigt, x, y, sigma=1 / y if use_weights else None)
|
|
curve = voigt(x, *params)
|
|
return VoigtFitResult("Linear direct fit", curve, cov, *params)
|
|
|
|
|
|
def log_direct_voigt_fit(x: np.ndarray, y: np.ndarray) -> VoigtFitResult:
|
|
ind = y > 0
|
|
y_fit = np.log(y[ind])
|
|
x_fit = x[ind]
|
|
params, cov = curve_fit(log_voigt, x_fit, y_fit, bounds=[1e-12, np.inf])
|
|
curve = np.exp(log_voigt(x, *params))
|
|
return VoigtFitResult("Log direct fit", curve, cov, *params)
|
|
|
|
|
|
def linear_fft_voigt_fit(f: np.ndarray, A: np.ndarray, x: np.ndarray, dx: float) -> VoigtFitResult:
|
|
params, cov = curve_fit(
|
|
lambda _f, _g, _l: fft_voigt(_f, dx, _g, _l), f, A, bounds=[1e-12, np.inf]
|
|
)
|
|
curve = voigt(x, *params)
|
|
return VoigtFitResult("Linear FFT fit", curve, cov, *params)
|
|
|
|
|
|
def convolve(gau: np.ndarray, lor: np.ndarray) -> np.ndarray:
|
|
return np.convolve(gau, lor, mode="same")
|
|
|
|
|
|
def add_noise(arr: np.ndarray, abs_fac: float, rel_fac: float):
|
|
return (
|
|
arr + abs_fac * np.random.rand(len(arr)) + rel_fac * arr * (np.random.rand(len(arr)) - 0.5)
|
|
)
|
|
|
|
|
|
@main.command(help="fit a lorentzian with a gaussian")
|
|
def demo1():
|
|
x = np.linspace(-10, 10, 501)
|
|
with PlotApp(
|
|
FWHM=np.geomspace(0.1, 10),
|
|
relative_noise=np.linspace(0, 2),
|
|
absolute_noise=[0, *np.geomspace(0.001, 1)],
|
|
) as app:
|
|
app.set_antialiasing(True)
|
|
app.params["FWHM"].value = 3.7
|
|
|
|
@app.update
|
|
def draw(FWHM, relative_noise, absolute_noise):
|
|
lo = lorentzian(x, 1, FWHM, 0)
|
|
lo = add_noise(lo, absolute_noise, relative_noise * 0.04)
|
|
fit = single_func_fit(x, lo, gaussian)
|
|
app[0].set_line_data("Lorentzian", x, lo, width=2)
|
|
app[0].set_line_data(
|
|
"gaus fit", x, fit.curve, width=2, label=f"Gaussian fit : FWHM = {fit.width}"
|
|
)
|
|
|
|
|
|
@main.command(name="default", help="illustrate the influence of noises and fit a voigt profile")
|
|
def demo2():
|
|
with PlotApp(
|
|
n=np.arange(256, 2046),
|
|
lorentz_fwhm=np.geomspace(0.1, 1),
|
|
gauss_fwhm=np.linspace(1, 10),
|
|
relative_noise=np.linspace(0, 2),
|
|
absolute_noise=[0, *np.geomspace(1e-6, 1, 201)],
|
|
) as app:
|
|
app.set_antialiasing(True)
|
|
ax = app["Main Plot"]
|
|
|
|
app.params["gauss_fwhm"].value = 5
|
|
app.params["n"].value = 1024
|
|
app.params_layout.setSpacing(0)
|
|
|
|
@app.update
|
|
def draw(n, lorentz_fwhm, gauss_fwhm, absolute_noise, relative_noise):
|
|
x, dx, f, df = setup_axes(n, (-20, 20))
|
|
gaus = gaussian(x, 1, gauss_fwhm, 0)
|
|
lore = lorentzian(x, 1, lorentz_fwhm, 0)
|
|
ax.set_line_data("Gaussian", x, gaus / gaus.max(), width=2)
|
|
ax.set_line_data("Lorentz", x, lore / lore.max(), width=2)
|
|
|
|
conv = convolve(gaus, lore)
|
|
conv /= conv.max()
|
|
conv_noise = add_noise(conv, absolute_noise, relative_noise * 0.1)
|
|
gaus_noise = add_noise(gaus, absolute_noise, relative_noise * 0.1)
|
|
ax.set_line_data("Voigt", x, conv, width=2)
|
|
ax.set_line_data("Noisy Voigt", x, conv_noise, width=2)
|
|
ax.set_line_data("Noisy Gaussian", x, gaus_noise, width=2)
|
|
|
|
lin_fit = linear_direct_voigt_fit(x, conv_noise)
|
|
log_fit = log_direct_voigt_fit(x, conv_noise)
|
|
transfo = np.fft.fft(np.fft.fftshift(conv_noise)).real
|
|
fft_fit = linear_fft_voigt_fit(f, transfo / transfo.max(), x, dx)
|
|
for fit in lin_fit, log_fit, fft_fit:
|
|
display_fit(x, fit)
|
|
|
|
def display_fit(x: np.ndarray, fit: VoigtFitResult):
|
|
ax.set_line_data(fit.name, x, fit.curve)
|
|
ax.set_line_data(
|
|
fit.name + "gauss",
|
|
x,
|
|
gaussian(x, 1, fit.gaussian_width, 0),
|
|
label=f"{fit.name} - Gaussian {fit.gaussian_width:.2f}",
|
|
width=1.5,
|
|
)
|
|
ax.set_line_data(
|
|
fit.name + "loren",
|
|
x,
|
|
lorentzian(x, 1, fit.lorentzian_width, 0),
|
|
label=f"{fit.name} - Lorentzian {fit.lorentzian_width:.2f}",
|
|
width=1.5,
|
|
)
|
|
|
|
|
|
@main.command(help="Explore fft fitting")
|
|
def demo3():
|
|
x, dx, f, df = setup_axes(1001, (-20, 20))
|
|
with PlotApp(
|
|
lorentz_fwhm=np.geomspace(0.1, 1),
|
|
gauss_fwhm=np.linspace(1, 10),
|
|
relative_noise=np.linspace(0, 2),
|
|
absolute_noise=[0, *np.geomspace(1e-6, 1, 201)],
|
|
) as app:
|
|
app.set_antialiasing(True)
|
|
|
|
ax = app[0]
|
|
ind = np.argsort(f)
|
|
|
|
@app.update
|
|
def draw(lorentz_fwhm, gauss_fwhm, absolute_noise, relative_noise):
|
|
gaus_clean = gaussian(x, 1, gauss_fwhm, 0)
|
|
lore_clean = lorentzian(x, 1, lorentz_fwhm, 0)
|
|
voig_clean = convolve(gaus_clean, lore_clean)
|
|
|
|
gaus = add_noise(gaus_clean, absolute_noise, relative_noise * 0.1)
|
|
lore = add_noise(lore_clean, absolute_noise, relative_noise * 0.1)
|
|
voig = add_noise(voig_clean, absolute_noise, relative_noise * 0.1)
|
|
|
|
gaus_f = np.fft.fft(np.fft.fftshift(gaus)).real
|
|
lore_f = np.fft.fft(np.fft.fftshift(lore)).real
|
|
voig_f = np.fft.fft(np.fft.fftshift(voig)).real
|
|
|
|
ax.set_line_data("Transformed Gaussian", f[ind], gaus_f[ind])
|
|
ax.set_line_data("Transformed Lorentz", f[ind], lore_f[ind])
|
|
ax.set_line_data("Transformed Voigt", f[ind], voig_f[ind])
|
|
|
|
gaus_ff = fft_gaussian(f, dx, 1, gauss_fwhm)
|
|
lore_ff = fft_lorentzian(f, dx, 1, lorentz_fwhm)
|
|
voig_ff = fft_voigt(f, dx, gauss_fwhm, lorentz_fwhm) * voig_clean.sum()
|
|
|
|
ax.set_line_data("Expected Gaussian", f[ind], gaus_ff[ind], width=1)
|
|
ax.set_line_data("Expected Lorentz", f[ind], lore_ff[ind], width=1)
|
|
ax.set_line_data("Expected Voigt", f[ind], voig_ff[ind], width=1)
|
|
|
|
|
|
@click.argument("file", type=click.Path(True, dir_okay=False))
|
|
@main.command(name="fit")
|
|
def fit_measurement(file):
|
|
wl, intens = lab.osa.load_osa_spectrum(file)
|
|
wl *= 1e9
|
|
gauss_fit = single_func_fit(wl, intens, elevated_gaussian, extra_params=1)
|
|
wl0 = gauss_fit.x0
|
|
baseline = gauss_fit.params[-1]
|
|
intens -= baseline
|
|
gauss_fit.curve -= baseline
|
|
intens[intens < 0] = 0
|
|
wl -= wl0
|
|
with PlotApp(n=np.arange(len(wl) // 2)[1:]) as app:
|
|
app.set_antialiasing(True)
|
|
|
|
ax1 = app["Measurement"]
|
|
ax2 = app["Fourier domain"]
|
|
ax1.set_line_data("Measured spectrum", wl, intens, width=2, color="red")
|
|
ax1.set_line_data("Initial Guassian fit", wl, gauss_fit.curve, width=2, color="blue")
|
|
ax1.plot_widget.plotItem.setLogMode(y=True)
|
|
|
|
@app.update
|
|
def draw(n):
|
|
s = slice(n, -n) if n else slice(None)
|
|
|
|
wl_fit = wl[s]
|
|
intens_fit = intens[s] - intens[s].min()
|
|
f, ind, dwl = freq_axis(wl_fit)
|
|
|
|
ax1.set_line_data("Fitted spectrum", wl_fit, intens_fit)
|
|
|
|
trans = np.abs(np.fft.fft(intens_fit))[ind]
|
|
m = trans.max()
|
|
ax2.set_line_data("Transformed", f, trans)
|
|
|
|
fft_fit = linear_fft_voigt_fit(f, trans / m, wl_fit, dwl)
|
|
lin_fit = linear_direct_voigt_fit(wl_fit, intens_fit)
|
|
log_fit = log_direct_voigt_fit(wl_fit, intens_fit)
|
|
for fit in fft_fit, lin_fit, log_fit:
|
|
display_fit(wl_fit, dwl, f, fit, m)
|
|
|
|
def display_fit(x: np.ndarray, dx: float, f: np.ndarray, fit: VoigtFitResult, amp: float):
|
|
ax2.set_line_data(
|
|
fit.name, f, amp * fft_voigt(f, dx, fit.gaussian_width, fit.lorentzian_width)
|
|
)
|
|
ax1.set_line_data(fit.name, x, fit.curve)
|
|
ax1.set_line_data(
|
|
fit.name + "gauss",
|
|
x,
|
|
gaussian(x, 1, fit.gaussian_width, 0),
|
|
label=f"{fit.name} - Gaussian {fit.gaussian_width:.2f}",
|
|
width=1.5,
|
|
)
|
|
ax1.set_line_data(
|
|
fit.name + "lorentz",
|
|
x,
|
|
lorentzian(x, 1, fit.lorentzian_width, 0),
|
|
label=f"{fit.name} - Lorentzian {fit.lorentzian_width:.2f}",
|
|
width=1.5,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|