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()