diff --git a/fitting.py b/fitting.py
index 35fbbd6e0bf70744bb93ff9b36f8306751b1e306..0ac0636ee354745df30d40337f6abad5d022719e 100644
--- a/fitting.py
+++ b/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
diff --git a/plotting.py b/plotting.py
index d404abdc88945d1963c40147c2643d6acaa32a83..a0f3111283e3df3c59a05ea6ab524ee28365762c 100644
--- a/plotting.py
+++ b/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