Plot fitted tune peak

fitting.py

@@ -9,22 +9,29 @@ def fit_gaussian(S, V, **kwargs):
     vp = np.max(V)-v0
     s0 = S[np.argmax(V)]
     sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * V) / np.sum(V)))
-    return fit_function(gauss, S, V, p0=(v0, vp, s0, sigma), bounds=((-np.inf, -np.inf, -np.inf, 0), (np.inf, np.inf, np.inf, np.inf)), **kwargs)
+    result = fit_function(gauss, S, V, p0=(v0, vp, s0, sigma), **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, lorenzian(X, *param)]
 
 def fit_lorenzian(S, V, log=False, **kwargs):
     lorenzian = lambda s, v0, vp, s0, gamma: v0 + vp/(1+((s-s0)/gamma)**2)
     v0 = np.min(V)
     vp = np.max(V)-v0
     s0 = S[np.argmax(V)]
-    sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * V) / np.sum(V)))
-    return fit_function(lorenzian, S, V, p0=(v0, vp, s0, sigma), bounds=((-np.inf, -np.inf, -np.inf, 0), (np.inf, np.inf, np.inf, np.inf)), **kwargs)
+    sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * (V-v0)) / np.sum(V-v0)))
+    result = fit_function(lorenzian, S, V, p0=(v0, vp, s0, sigma), **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, lorenzian(X, *param)]
+    
 
 def fit_function(function, x, y, p0=None, **kwargs):
     """
     :param log: if the y-data is log-scaled
     """
     param, cov = scipy.optimize.curve_fit(function, x, y, p0, **kwargs)
-    param_error = np.sqrt(cov.diagonal())
+    param_error = np.sqrt(np.abs(cov.diagonal()))
     X = np.linspace(min(x), max(x), 1000)
     return param, param_error, function, [X, function(X, *param)]
 

plotting.py

@@ -2,10 +2,19 @@ import scipy
 import numpy as np
 import matplotlib as mpl
 import matplotlib.pyplot as plt
+import pint
 
 from .data import *
+from .fitting import *
 
 
+# Setup pint
+SI = pint.UnitRegistry()
+SI.setup_matplotlib()
+SI.default_format = '~P'
+
+
+# Colors
 # https://arxiv.org/abs/2107.02270
 petroff_colors = ["#3f90da", "#ffa90e", "#bd1f01", "#94a4a2", "#832db6", "#a96b59", "#e76300", "#b9ac70", "#717581", "#92dadd"]
 cmap_petroff_10 = mpl.colors.ListedColormap(petroff_colors, 'Petroff 10')
@@ -22,6 +31,8 @@ cmap_petroff_bipolar.set_over(petroff_colors[8])
 
 
 
+
+
 def add_vline(axes, r, color='k', text=None, label=None, order=None, lw=1, text_top=True, text_vertical=True, zorder=-100, **kwargs):
     for i, a in enumerate(axes):
         if a is None: continue
@@ -86,7 +97,8 @@ def plot_tbt(ax, libera_data, over_time=True, turn_range=None, time_range=None):
     ax2.legend([lf,ls], ['Revolution frequency', 'Pickup sum signal'], loc='center right', fontsize='small')
 
 
-def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, **kwargs):
+
+def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, fit=False, **kwargs):
     """Plot a tune spectrum based on turn-by-turn data
     
     :param ax: Axis to plot onto
@@ -112,7 +124,21 @@ def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tu
     ax.plot(freq, mag, **kwargs)
     
     ax.set(xlim=tune_range, xlabel=f'Tune $q_{xy}$',
-           ylabel='Magnitude / a.u.')
+           ylabel='a.u.')
+    
+    if fit:
+        try:
+            fitr = (fit if callable(fit) else fit_lorenzian)(freq, mag)
+            q, w = fitr[0][2], fitr[0][3]
+            if q<0 or q>0.5 or w > 0.01:
+                raise RuntimeError('Fit failed')
+        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}$')
+
+
 
 
 def plot_tune_spectrogram(ax, libera_data, xy, nperseg=2**12, noverlap=None, ninterpol=4, smoothing=0, over_time=True, colorbar=False,
@@ -135,7 +161,7 @@ def plot_tune_spectrogram(ax, libera_data, xy, nperseg=2**12, noverlap=None, nin
     tbt_data = getattr(libera_data, xy)[turn_range]
     
     tune, turn, value = scipy.signal.stft(tbt_data, fs=1, nperseg=nperseg, noverlap=noverlap, window='boxcar', boundary=None, padded=False)
-    time = libera_data.t[turn.astype(int)]
+    time = libera_data.t[turn_range][turn.astype(int)]
     mag = np.abs(value)
     #mag[0,:] = 0 # supress DC