Add linear fit

fitting.py

@@ -3,6 +3,11 @@ import numpy as np
 
 
 # Fitting
+def fit_linear(S, V, **kwargs):
+    result = fit_function(lambda s, v0, m: v0+m*s, S, V, p0=(np.mean(V), np.mean(np.diff(V)/np.diff(S))), **kwargs)
+    param, param_error, function, [X, _] = result
+    return param, param_error, function, [X, function(X, *param)]
+
 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)
@@ -12,7 +17,7 @@ def fit_gaussian(S, V, **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)]
+    return param, param_error, function, [X, function(X, *param)]
 
 def fit_lorenzian(S, V, log=False, **kwargs):
     lorenzian = lambda s, v0, vp, s0, gamma: v0 + vp/(1+((s-s0)/gamma)**2)
@@ -23,16 +28,17 @@ def fit_lorenzian(S, V, log=False, **kwargs):
     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)]
+    return param, param_error, function, [X, function(X, *param)]
     
 
-def fit_function(function, x, y, p0=None, **kwargs):
+def fit_function(function, x, y, p0=None, extend=0, **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(np.abs(cov.diagonal()))
-    X = np.linspace(min(x), max(x), 1000)
+    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)]
 
 def fit_exponential(S, V, **kwargs):
@@ -42,4 +48,3 @@ def fit_exponential(S, V, **kwargs):
     s0 = 1
     return fit_function(exponential, S, V, p0=(v0, vp, s0), **kwargs)
 
-    
\ No newline at end of file