Moving gauss fit for linear tune change

fitting.py

@@ -96,6 +96,39 @@ def fit_double_lorenzian(S, V, *, s0=None, dry_run=False, **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_moving_gaussian(S, T, VV, extend=0, **kwargs):
+    """Fit a gaussian function with linearly moving centroid position to the 2D data
+    
+    :param S: data x values
+    :param T: data y values (axis along which centroid varies)
+    :param VV: data z values of shape [x,y]
+    :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
+    """
+    
+    #r = np.repeat([np.linspace(0, 1, m)], n, axis=0).T
+    ST = np.tile(S, (len(T), 1)).T
+    r = np.linspace(0, 1, ST.shape[1]).reshape(1, -1)
+    moving_gauss = lambda s, v0, vp, s0, ds, sigma: v0+vp*np.exp(-0.5*((s-(s0 + ds*r))/sigma)**2)
+
+    V0 = np.average(VV[:, :VV.shape[1]//5], axis=1)  # average over first 5% of data for initial mean
+    V1 = np.average(VV[:, -VV.shape[1]//5:], axis=1)  # average over last 5% of data for final mean
+    v0 = np.min(VV)
+    vp = np.max(VV)-v0
+    s0 = S[np.argmax(V0)]
+    s1 = S[np.argmax(V1)]
+    ds = s1 - s0
+    sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * V0) / np.sum(V0)))
+    result = fit_function(lambda *a, **b: moving_gauss(*a, **b).ravel(), ST, VV.ravel(), p0=(v0, vp, s0, ds, sigma), extend=None, **kwargs) # flatten 2D arrays for fit
+    param, param_error, function, _ = result
+    param[4] = np.abs(param[4]) # return the positive sigma (could use bounds instead, but that's more likely to fail)
+    mi, ma = np.min(T), np.max(T)
+    X = np.linspace(mi - (ma-mi)*extend, ma + (ma-mi)*extend, 1000)
+    Y = param[2] + param[3]*np.linspace(0, 1, 1000)
+    return param, param_error, function, [X, Y]
     
 
 def fit_function(function, x, y, *, xerr=None, yerr=None, p0=None, odr=None, extend=0):
@@ -126,9 +159,12 @@ def fit_function(function, x, y, *, xerr=None, yerr=None, p0=None, odr=None, ext
         param, cov = scipy.optimize.curve_fit(function, x, y, p0)
         param_error = np.sqrt(np.abs(cov.diagonal()))
     
-    mi, ma = min(x), max(x)
-    X = np.linspace(mi - (ma-mi)*extend, ma + (ma-mi)*extend, 1000)
-    return param, param_error, function, [X, function(X, *param)]
+    if extend is not None:
+        mi, ma = np.min(x), np.max(x)
+        X = np.linspace(mi - (ma-mi)*extend, ma + (ma-mi)*extend, 1000)
+        return param, param_error, function, [X, function(X, *param)]
+    else:
+        return param, param_error, function, None
 
 def fit_exponential(S, V, **kwargs):
     """Fit an exponential function to the data

plotting.py

@@ -304,7 +304,7 @@ def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tu
         
 def plot_tune_spectrogram(ax, libera_data, xy, *, nperseg=2**12, noverlap=None, ninterpol=None, over_time=True, colorbar=False,
                           turn_range=None, time_range=None, tune_range=None, excitation=None, vmin=None, vmax=None,
-                          cmap='gist_heat_r', bunches=None, show_nperseg=True):
+                          cmap='gist_heat_r', bunches=None, show_nperseg=True, fit=False, fitkwarg=dict()):
     """Plot a tune spectrogram based on turn-by-turn data
     
     :param ax: Axis to plot onto
@@ -428,11 +428,27 @@ def plot_tune_spectrogram(ax, libera_data, xy, *, nperseg=2**12, noverlap=None,
             ax.plot(xdata, ex_q + ex_dq/2*np.sin(2*np.pi*ex_fc*time + ex_pc), **ex_style)
         else:
             raise NotImpementedError(f'Excitation type {ex_type} not implemented')
-            
-            
-            
+
     
     ax.set(ylim=tune_range, ylabel=f'Tune $q_{xy}$', xlabel='Time / s' if over_time else 'Turn') #  or $1-q_{xy}$')
+    
+    if fit:
+        dq_dx = None
+        try:
+            fitr = fit(tune, xdata, mag)
+            if fit in (fit_moving_gaussian, ):
+                v0, vp, q0, dq, sigma = fitr[0]
+                dx = xdata[-1]-xdata[0]
+                dq = SI.Measurement(fitr[0][3], fitr[1][3], '1/s' if over_time else '')
+                dq_dx = dq/dx
+                ax.plot(*fitr[-1], **dict(ls='--', label=f'Fit $\\partial q_{xy}/\\partial '+('t' if over_time else 'n')+f'={dq_dx:~L}$', zorder=50, color=cmap_petroff_10(9), **fitkwarg))
+                
+        except RuntimeError:
+            print('Warning: fit failed')
+        
+        return xdata, tune, mag, frev, dq_dx
+            
+    
     return xdata, tune, mag, frev