From 87b10485baa5d66a65378c785cae273dab9ab534 Mon Sep 17 00:00:00 2001
From: Philipp Niedermayer <p.niedermayer@gsi.de>
Date: Wed, 25 May 2022 11:53:59 +0200
Subject: [PATCH] Implement double peak fit
---
fitting.py | 36 ++++++++++++++++++++++++++++++++++++
plotting.py | 22 ++++++++++++++++------
2 files changed, 52 insertions(+), 6 deletions(-)
diff --git a/fitting.py b/fitting.py
index 6406cab..35fbbd6 100644
--- a/fitting.py
+++ b/fitting.py
@@ -60,6 +60,42 @@ def fit_lorenzian(S, V, **kwargs):
param, param_error, function, [X, _] = result
param[3] = np.abs(param[3]) # return the positive sigma (could use bounds instead, but that's more likely to fail)
return param, param_error, function, [X, function(X, *param)]
+
+
+def fit_double_lorenzian(S, V, *, s0=None, dry_run=False, **kwargs):
+ """Fit a lorenzian function to the data
+
+ :param S: data x values
+ :param V: data y values
+ :param xerr: data x uncertainties
+ :param yerr: data y uncertainties
+ :param s0: initial peak positions
+ :param p0: initial values of parameters
+ :param odr: Fit method to use: True means ortogonal distance regression (ODR), False means ordinary least square fit, None (default) uses ODR only if either xerr or yerr is not None
+ :param extend: fraction by which the returned fit trace extends beyond the first/last data point
+ """
+ lorenzian = lambda s, v0, vp, s0, gamma, vp2, s02, gamma2: v0 + vp/(1+((s-s0)/gamma)**2) + vp2/(1+((s-s02)/gamma2)**2)
+
+ # ! for double peak fit, initial conditions are very important
+ v0 = np.min(V)
+ vp = np.max(V)-v0
+ if s0 is None:
+ s0 = S[np.argmax(V)]
+ sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * (V-v0)) / np.sum(V-v0)))/10
+ s0 = [s0-3*sigma, s0+3*sigma]
+ sigma0 = np.sqrt(np.abs(np.sum((S - s0[0]) ** 2 * (V-v0)) / np.sum(V-v0)))/10
+ sigma1 = np.sqrt(np.abs(np.sum((S - s0[1]) ** 2 * (V-v0)) / np.sum(V-v0)))/10
+ p0=(v0, vp, s0[0], sigma0, vp, s0[1], sigma1)
+
+ if dry_run:
+ mi, ma = min(S), max(S)
+ X = np.linspace(mi, ma, 1000)
+ return p0, (0,0,0,0), 0, [X, lorenzian(X, *p0)]
+
+ result = fit_function(lorenzian, S, V, p0=p0, **kwargs)
+ param, param_error, function, [X, _] = result
+ param[3] = np.abs(param[3]) # return the positive sigma (could use bounds instead, but that's more likely to fail)
+ return param, param_error, function, [X, function(X, *param)]
def fit_function(function, x, y, *, xerr=None, yerr=None, p0=None, odr=None, extend=0):
diff --git a/plotting.py b/plotting.py
index 4f55bfd..e1dd308 100644
--- a/plotting.py
+++ b/plotting.py
@@ -236,14 +236,24 @@ def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tu
q = None
try:
fitr = (fit if callable(fit) else fit_lorenzian)(freq, mag)
- q, w = fitr[0][2], fitr[0][3]
- if q<np.min(freq) or q>np.max(freq) or w > 0.1:
- raise RuntimeError('Fit failed')
+ if fit in (fit_lorenzian, fit_gaussian):
+ q, w = fitr[0][2], fitr[0][3]
+ if q<np.min(freq) or q>np.max(freq) or w > 0.1:
+ raise RuntimeError('Fit failed')
+ q = SI.Measurement(fitr[0][2], (fitr[1][2]**2+fitr[0][3]**2+fitr[1][3]**2)**0.5, '') # conservative estimate of error including width of peak
+ ax.plot(*fitr[-1], '--', label=f'Fit $q_{xy}={q:~L}$', zorder=50)
+
+ elif fit in (fit_double_lorenzian, ):
+ q1, w1, q2, w2 = fitr[0][2], fitr[0][3], fitr[0][5], fitr[0][6]
+ if q1<np.min(freq) or q1>np.max(freq) or w1 > 0.1 or q2<np.min(freq) or q2>np.max(freq) or w2 > 0.1:
+ raise RuntimeError('Fit failed')
+ q1 = SI.Measurement(fitr[0][2], (fitr[1][2]**2+fitr[0][3]**2+fitr[1][3]**2)**0.5, '') # conservative estimate of error including width of peak
+ q2 = SI.Measurement(fitr[0][5], (fitr[1][5]**2+fitr[0][6]**2+fitr[1][6]**2)**0.5, '') # conservative estimate of error including width of peak
+ ax.plot(*fitr[-1], '--', label=f'Fit $q_{{{xy},1}}={q1:~L}$\nFit $q_{{{xy},2}}={q2:~L}$', zorder=50)
+ q = (q1, q2)
+
except RuntimeError:
print('Warning: fit failed')
- else:
- q = SI.Measurement(fitr[0][2], (fitr[1][2]**2+fitr[0][3]**2+fitr[1][3]**2)**0.5, '') # conservative estimate of error including width of peak
- ax.plot(*fitr[-1], '--', label=f'Fit $q_{xy}={q:~L}$', zorder=50)
if return_spectrum:
return freq, mag, phase, q
--
GitLab