Implement double peak fit

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

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