fitting.py

import scipy
import scipy.odr
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)
    vp = np.max(V)-v0
    s0 = S[np.argmax(V)]
    sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * V) / np.sum(V)))
    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, 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)
    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-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, function(X, *param)]
    

def fit_function(function, x, y, *, xerr=None, yerr=None, p0=None, odr=None, extend=0, **kwargs):
    """
    :param log: if the y-data is log-scaled
    """
    if odr is None:
        odr = xerr is not None or yerr is not None
    
    if odr:
        # orthogonal distance regression
        data = scipy.odr.RealData(x, y, sx=xerr, sy=yerr)
        model = scipy.odr.Model(lambda beta, x: function(x, *beta))
        odr = scipy.odr.ODR(data, model, p0)
        # if not odr: odr.set_job(fit_type=2) # ordinary least-squares
        output = odr.run()
        param, param_error = output.beta, output.sd_beta
    else:
        # non-linear least squares
        param, cov = scipy.optimize.curve_fit(function, x, y, p0, **kwargs)
        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)]

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)