ExponentialDecayRegressor.py

from scipy.optimize import curve_fit
from sklearn.base import BaseEstimator, RegressorMixin
import statsmodels.api as sm

class ExponentialDecayRegressor(BaseEstimator, RegressorMixin):
    """Fits an exponential decay curve
    """
    
    def __init__(self, starting_values=[1.,1.e-5,1.], **kwargs,):
        self.starting_values = starting_values
        self.kwargs = kwargs
        self.params = None
        
    def fit(self, X, y=None):
        self.params, _ = curve_fit(self.exp_decay_f, X, y, p0=self.starting_values)
        
    def predict(self, X, y=None):
        return self.exp_decay_f(X, *self.params)
    
    @staticmethod
    def exp_decay_f(X, a, k, b):
        return a * np.exp(-k*X) + b