first try at PoissonGroup

Author: floriankozikowski

debug_poisson_group.py

@@ -0,0 +1,30 @@
+import numpy as np
+from skglm import GeneralizedLinearEstimator
+from skglm.datafits.group import PoissonGroup
+from skglm.penalties import WeightedGroupL2
+from skglm.solvers import GroupProxNewton
+from sklearn.metrics import mean_squared_error
+
+# Sample data and group structure
+n_samples, n_features = 20, 10
+X = np.random.randn(n_samples, n_features)
+y = np.random.poisson(np.abs(X[:, 0] + X[:, 5]))
+
+grp_ptr = np.array([0, 3, 5, 8, 10], dtype=np.int32)
+grp_indices = np.arange(n_features, dtype=np.int32)
+
+# Estimator setup
+estimator = GeneralizedLinearEstimator(
+    datafit=PoissonGroup(grp_ptr, grp_indices),
+    penalty=WeightedGroupL2(alpha=0.1, grp_ptr=grp_ptr, grp_indices=grp_indices,
+                            weights=np.ones(len(grp_ptr) - 1)),
+    solver=GroupProxNewton()
+)
+
+estimator.fit(X, y)
+print("Coefficients:", estimator.coef_)
+print("Intercept:", estimator.intercept_)
+y_pred = estimator.predict(X)
+print("First 5 predictions:", y_pred[:5])
+print("First 5 true values:", y[:5])
+print("MSE:", mean_squared_error(y, np.exp(y_pred)))

skglm/datafits/group.py

@@ -3,7 +3,7 @@
 from numba import int32, float64
 
 from skglm.datafits.base import BaseDatafit
-from skglm.datafits.single_task import Logistic
+from skglm.datafits.single_task import Logistic, Poisson
 from skglm.utils.sparse_ops import spectral_norm, sparse_columns_slice
 
 
@@ -161,3 +161,52 @@ def gradient_g(self, X, y, w, Xw, g):
             grad_g[idx] = X[:, j] @ raw_grad_val
 
         return grad_g
+
+
+class PoissonGroup(Poisson):
+    r"""Poisson datafit used with group penalties.
+
+    The datafit reads:
+
+    .. math:: 1 / n_"samples" \sum_{i=1}^{n_"samples"} (\exp((Xw)_i) - y_i (Xw)_i)
+
+    Attributes
+    ----------
+    grp_indices : array, shape (n_features,)
+        The group indices stacked contiguously
+        ``[grp1_indices, grp2_indices, ...]``.
+
+    grp_ptr : array, shape (n_groups + 1,)
+        The group pointers such that two consecutive elements delimit
+        the indices of a group in ``grp_indices``.
+    """
+
+    def __init__(self, grp_ptr, grp_indices):
+        self.grp_ptr, self.grp_indices = grp_ptr, grp_indices
+
+    def get_spec(self):
+        return (
+            ('grp_ptr', int32[:]),
+            ('grp_indices', int32[:]),
+        )
+
+    def params_to_dict(self):
+        return dict(grp_ptr=self.grp_ptr, grp_indices=self.grp_indices)
+
+    def gradient_g(self, X, y, w, Xw, g):
+        grp_ptr, grp_indices = self.grp_ptr, self.grp_indices
+        grp_g_indices = grp_indices[grp_ptr[g]: grp_ptr[g+1]]
+        raw_grad_val = self.raw_grad(y, Xw)
+        grad_g = np.zeros(len(grp_g_indices))
+        for idx, j in enumerate(grp_g_indices):
+            grad_g[idx] = X[:, j] @ raw_grad_val
+        return grad_g
+
+    def gradient_g_sparse(self, X_data, X_indptr, X_indices, y, w, Xw, g):
+        grp_ptr, grp_indices = self.grp_ptr, self.grp_indices
+        grp_g_indices = grp_indices[grp_ptr[g]: grp_ptr[g+1]]
+        grad_g = np.zeros(len(grp_g_indices))
+        for idx, j in enumerate(grp_g_indices):
+            grad_g[idx] = self.gradient_scalar_sparse(
+                X_data, X_indptr, X_indices, y, Xw, j)
+        return grad_g