ENH Add Poisson datafit (#78)

Co-authored-by: mathurinm <[email protected]>

Author: PABannier

skglm/datafits/__init__.py

@@ -1,12 +1,12 @@
 from .base import BaseDatafit, BaseMultitaskDatafit
-from .single_task import Quadratic, QuadraticSVC, Logistic, Huber
+from .single_task import Quadratic, QuadraticSVC, Logistic, Huber, Poisson
 from .multi_task import QuadraticMultiTask
 from .group import QuadraticGroup
 
 
 __all__ = [
     BaseDatafit, BaseMultitaskDatafit,
-    Quadratic, QuadraticSVC, Logistic, Huber,
+    Quadratic, QuadraticSVC, Logistic, Huber, Poisson,
     QuadraticMultiTask,
     QuadraticGroup
 ]

skglm/datafits/single_task.py

@@ -360,3 +360,66 @@ def intercept_update_step(self, y, Xw):
             else:
                 update -= np.sign(residual) * self.delta
         return update / n_samples
+
+
+class Poisson(BaseDatafit):
+    r"""Poisson datafit.
+
+    The datafit reads::
+
+    (1 / n_samples) * \sum_i (exp(Xw_i) - y_i * Xw_i)
+
+    Note:
+    ----
+    The class is jit compiled at fit time using Numba compiler.
+    This allows for faster computations.
+    """
+
+    def __init__(self):
+        pass
+
+    def get_spec(self):
+        pass
+
+    def params_to_dict(self):
+        return dict()
+
+    def initialize(self, X, y):
+        pass
+
+    def initialize_sparse(self, X_data, X_indptr, X_indices, y):
+        pass
+
+    def raw_grad(self, y, Xw):
+        """Compute gradient of datafit w.r.t ``Xw``."""
+        return (np.exp(Xw) - y) / len(y)
+
+    def raw_hessian(self, y, Xw):
+        """Compute Hessian of datafit w.r.t ``Xw``."""
+        return np.exp(Xw) / len(y)
+
+    def value(self, y, w, Xw):
+        return np.sum(np.exp(Xw) - y * Xw) / len(y)
+
+    def gradient_scalar(self, X, y, w, Xw, j):
+        return (X[:, j] @ (np.exp(Xw) - y)) / len(y)
+
+    def full_grad_sparse(self, X_data, X_indptr, X_indices, y, Xw):
+        n_features = X_indptr.shape[0] - 1
+        grad = np.zeros(n_features, dtype=X_data.dtype)
+        for j in range(n_features):
+            grad[j] = 0.
+            for i in range(X_indptr[j], X_indptr[j + 1]):
+                grad[j] += X_data[i] * (
+                    np.exp(Xw[X_indices[i]] - y[X_indices[i]])) / len(y)
+        return grad
+
+    def gradient_scalar_sparse(self, X_data, X_indptr, X_indices, y, Xw, j):
+        grad = 0.
+        for i in range(X_indptr[j], X_indptr[j + 1]):
+            idx_i = X_indices[i]
+            grad += X_data[i] * (np.exp(Xw[idx_i]) - y[idx_i])
+        return grad / len(y)
+
+    def intercept_update_self(self, y, Xw):
+        pass

skglm/tests/test_datafits.py

@@ -4,9 +4,9 @@
 from sklearn.linear_model import HuberRegressor
 from numpy.testing import assert_allclose, assert_array_less
 
-from skglm.datafits import Huber, Logistic
-from skglm.penalties import WeightedL1
-from skglm.solvers import AndersonCD
+from skglm.datafits import Huber, Logistic, Poisson
+from skglm.penalties import L1, WeightedL1
+from skglm.solvers import AndersonCD, ProxNewton
 from skglm import GeneralizedLinearEstimator
 from skglm.utils import make_correlated_data
 
@@ -56,5 +56,33 @@ def test_log_datafit():
     np.testing.assert_almost_equal(-grad * (y + n_samples * grad), hess)
 
 
+def test_poisson():
+    try:
+        from statsmodels.discrete.discrete_model import Poisson as PoissonRegressor  # noqa
+    except ImportError:
+        pytest.xfail("`statsmodels` not found. `Poisson` datafit can't be tested.")
+
+    n_samples, n_features = 10, 22
+    tol = 1e-14
+    X, y, _ = make_correlated_data(n_samples, n_features, random_state=0)
+    y = np.abs(y)
+
+    alpha_max = np.linalg.norm(X.T @ (np.ones(n_samples) - y), ord=np.inf) / n_samples
+    alpha = alpha_max * 0.1
+
+    df = Poisson()
+    pen = L1(alpha)
+
+    solver = ProxNewton(tol=tol, fit_intercept=False)
+    model = GeneralizedLinearEstimator(df, pen, solver).fit(X, y)
+
+    poisson_regressor = PoissonRegressor(y, X, offset=None)
+    res = poisson_regressor.fit_regularized(
+        method="l1", size_trim_tol=tol, alpha=alpha * n_samples, trim_mode="size")
+    w_statsmodels = res.params
+
+    assert_allclose(model.coef_, w_statsmodels, rtol=1e-4)
+
+
 if __name__ == '__main__':
     pass