QuantileRegression.py

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

class QuantileRegression(BaseEstimator, RegressorMixin):
    """Sklearn wrapper for statsmodels Quantile Regression
    """
    def __init__(self, quantile=0.5, **kwargs):
        self.quantile = quantile
        self.kwargs = kwargs
        self.model = None
        self.fitted = None
    
    def fit(self, X, y=None):
        X = sm.add_constant(X)
        self.model = sm.QuantReg(endog=y, exog=X, **self.kwargs)
        self.fitted = self.model.fit(q=self.quantile)
    
    def predict(self, X, y=None):
        X = sm.add_constant(X)
        return self.fitted.predict(X)