Add fitting functions

fitting.py

+import scipy
+import numpy as np
+
+
+# Fitting
+def fit_gaussian(S, V, **kwargs):
+    gauss = lambda s, v0, vp, s0, sigma: v0+vp*np.exp(-0.5*((s-s0)/sigma)**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(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)
+
+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)
+
+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())
+    X = np.linspace(min(x), max(x), 1000)
+    return param, param_error, function, [X, function(X, *param)]
+
+def fit_exponential(S, V, **kwargs):
+    exponential = lambda s, v0, vp, s0: v0 + vp*np.exp(s/s0)
+    v0 = np.min(V)
+    vp = np.max(V)-v0
+    s0 = 1
+    return fit_function(exponential, S, V, p0=(v0, vp, s0), **kwargs)
+
+    
\ No newline at end of file