dedebenui/linemeasurement
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

initial commit

Browse Source
This commit is contained in:
Benoît Sierro
2022-04-06 13:41:07 +02:00
commit 4dd36f8dc2
5 changed files with 1139 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

400
deconvolution.py Normal file
View File

@@ -0,0 +1,400 @@
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()