statgpu 0.1.0__py3-none-any.whl

statgpu/feature_selection/_stepwise.py

@@ -0,0 +1,300 @@
+"""
+Stepwise model selection for regression models.
+Supports forward, backward, and bidirectional selection.
+"""
+
+from typing import Optional, Union, List, Literal
+import numpy as np
+from copy import deepcopy
+from joblib import Parallel, delayed
+
+from statgpu.linear_model import LinearRegression, Ridge, Lasso, LogisticRegression
+
+
+class StepwiseSelector:
+    """
+    Stepwise model selection using AIC or BIC criterion.
+    
+    Supports forward selection, backward elimination, and bidirectional search.
+    
+    Parameters
+    ----------
+    model_class : class
+        Model class to use (LinearRegression, Ridge, Lasso, LogisticRegression).
+    criterion : str, default='aic'
+        Criterion for model selection: 'aic' or 'bic'.
+    direction : str, default='both'
+        Direction of search: 'forward', 'backward', or 'both'.
+    max_features : int, optional
+        Maximum number of features to select.
+    **model_kwargs
+        Additional arguments passed to the model.
+    
+    Attributes
+    ----------
+    selected_features_ : list
+        Indices of selected features.
+    best_model_ : object
+        Fitted model with selected features.
+    aic_history_ : list
+        AIC values at each step.
+    """
+    
+    def __init__(
+        self,
+        model_class,
+        criterion: str = 'aic',
+        direction: Literal['forward', 'backward', 'both'] = 'both',
+        max_features: Optional[int] = None,
+        n_jobs: Optional[int] = None,
+        **model_kwargs
+    ):
+        self.model_class = model_class
+        self.criterion = criterion.lower()
+        self.direction = direction
+        self.max_features = max_features
+        self.n_jobs = n_jobs
+        self.model_kwargs = model_kwargs
+        
+        if self.criterion not in ('aic', 'bic'):
+            raise ValueError("criterion must be 'aic' or 'bic'")
+        
+        self.selected_features_ = None
+        self.best_model_ = None
+        self.aic_history_ = []
+        self.bic_history_ = []
+        self._score_cache = {}
+    
+    def fit(self, X, y):
+        """
+        Fit stepwise model selection.
+        
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Training data.
+        y : array-like of shape (n_samples,)
+            Target values.
+        
+        Returns
+        -------
+        self : object
+        """
+        X = np.asarray(X)
+        y = np.asarray(y)
+        n_samples, n_features = X.shape
+        
+        if self.max_features is None:
+            self.max_features = n_features
+        
+        # Initialize
+        self._score_cache = {}
+        if self.direction == 'forward':
+            selected = []
+            remaining = list(range(n_features))
+        elif self.direction == 'backward':
+            selected = list(range(n_features))
+            remaining = []
+        else:  # both
+            selected = []
+            remaining = list(range(n_features))
+        
+        # Fit initial model
+        best_score = self._fit_and_score(X, y, selected)
+        self.aic_history_.append(best_score['aic'])
+        self.bic_history_.append(best_score['bic'])
+        
+        improved = True
+        iteration = 0
+        
+        while improved and len(selected) < self.max_features:
+            improved = False
+            iteration += 1
+            
+            if self.direction in ('forward', 'both'):
+                # Try adding each remaining feature
+                candidates = [(feature, selected + [feature]) for feature in remaining[:]]
+                scores = self._evaluate_candidates(X, y, candidates)
+                for feature, score in scores:
+                    current_score = score[self.criterion]
+                    if current_score < best_score[self.criterion]:
+                        best_score = score
+                        best_feature = feature
+                        best_action = 'add'
+                        improved = True
+            
+            if self.direction in ('backward', 'both') and len(selected) > 0:
+                # Try removing each selected feature
+                candidates = [(feature, [f for f in selected if f != feature]) for feature in selected[:]]
+                scores = self._evaluate_candidates(X, y, candidates)
+                for feature, score in scores:
+                    current_score = score[self.criterion]
+                    if current_score < best_score[self.criterion]:
+                        best_score = score
+                        best_feature = feature
+                        best_action = 'remove'
+                        improved = True
+            
+            if improved:
+                if best_action == 'add':
+                    selected.append(best_feature)
+                    remaining.remove(best_feature)
+                else:
+                    selected.remove(best_feature)
+                    remaining.append(best_feature)
+                
+                self.aic_history_.append(best_score['aic'])
+                self.bic_history_.append(best_score['bic'])
+                
+                print(f"Step {iteration}: {best_action} feature {best_feature}, "
+                      f"{self.criterion.upper()}={best_score[self.criterion]:.2f}")
+        
+        # Fit final model
+        self.selected_features_ = sorted(selected)
+        if len(selected) > 0:
+            self.best_model_ = self.model_class(**self.model_kwargs)
+            self.best_model_.fit(X[:, selected], y)
+        
+        return self
+    
+    def _evaluate_candidates(self, X, y, candidates):
+        """Evaluate feature candidates in parallel with memoized scores."""
+        feature_to_cache_key = {
+            feature: tuple(sorted(feature_indices)) for feature, feature_indices in candidates
+        }
+
+        def _score_for_indices(feature_indices):
+            key = tuple(sorted(feature_indices))
+            if key in self._score_cache:
+                return key, self._score_cache[key]
+
+            score = self._fit_and_score(X, y, feature_indices)
+            return key, score
+
+        def eval_one(feature, feature_indices):
+            key, score = _score_for_indices(feature_indices)
+            self._score_cache[key] = score
+            return feature, score
+
+        def eval_one_parallel(feature, feature_indices):
+            key, score = _score_for_indices(feature_indices)
+            return feature, key, score
+
+        if self.n_jobs == 1 or self.n_jobs is None or len(candidates) <= 1:
+            return [eval_one(feature, feature_indices) for feature, feature_indices in candidates]
+
+        out = Parallel(n_jobs=self.n_jobs)(
+            delayed(eval_one_parallel)(feature, feature_indices) for feature, feature_indices in candidates
+        )
+        for feature, key, score in out:
+            expected_key = feature_to_cache_key.get(feature)
+            if expected_key is not None and key not in self._score_cache:
+                self._score_cache[key] = score
+        return [(feature, score) for feature, _, score in out]
+    
+    def _fit_and_score(self, X, y, feature_indices):
+        """Fit model and return AIC/BIC scores."""
+        key = tuple(sorted(feature_indices))
+        if key in self._score_cache:
+            return self._score_cache[key]
+
+        if len(feature_indices) == 0:
+            # Null model
+            score = {'aic': np.inf, 'bic': np.inf}
+            self._score_cache[key] = score
+            return score
+        
+        model = self.model_class(**self.model_kwargs)
+        try:
+            model.fit(X[:, feature_indices], y)
+            
+            if hasattr(model, 'aic') and model.aic is not None:
+                score = {'aic': model.aic, 'bic': model.bic}
+                self._score_cache[key] = score
+                return score
+            else:
+                # Fallback: use R²-based approximation
+                n = len(y)
+                k = len(feature_indices) + 1  # +1 for intercept
+                if hasattr(model, 'rsquared'):
+                    r2 = model.rsquared
+                    # Approximate AIC
+                    aic = n * np.log(1 - r2 + 1e-10) + 2 * k
+                    bic = n * np.log(1 - r2 + 1e-10) + k * np.log(n)
+                    score = {'aic': aic, 'bic': bic}
+                    self._score_cache[key] = score
+                    return score
+                else:
+                    score = {'aic': np.inf, 'bic': np.inf}
+                    self._score_cache[key] = score
+                    return score
+        except Exception:
+            score = {'aic': np.inf, 'bic': np.inf}
+            self._score_cache[key] = score
+            return score
+    
+    def predict(self, X):
+        """Predict using the best model."""
+        if self.best_model_ is None:
+            raise RuntimeError("Model has not been fitted yet.")
+        X = np.asarray(X)
+        return self.best_model_.predict(X[:, self.selected_features_])
+    
+    def score(self, X, y):
+        """Return R² score of the best model."""
+        if self.best_model_ is None:
+            raise RuntimeError("Model has not been fitted yet.")
+        X = np.asarray(X)
+        return self.best_model_.score(X[:, self.selected_features_], y)
+    
+    def summary(self):
+        """Print summary of stepwise selection."""
+        print("=" * 60)
+        print("Stepwise Model Selection Summary")
+        print("=" * 60)
+        print(f"Criterion: {self.criterion.upper()}")
+        print(f"Direction: {self.direction}")
+        print(f"Selected features: {self.selected_features_}")
+        print(f"Number of features: {len(self.selected_features_)}")
+        if self.aic_history_:
+            print(f"Final AIC: {self.aic_history_[-1]:.2f}")
+            print(f"Final BIC: {self.bic_history_[-1]:.2f}")
+        print("=" * 60)
+
+
+def stepwise_selection(
+    X, y,
+    model_class=LinearRegression,
+    criterion: str = 'aic',
+    direction: str = 'both',
+    **model_kwargs
+):
+    """
+    Convenience function for stepwise selection.
+    
+    Parameters
+    ----------
+    X, y : array-like
+        Training data.
+    model_class : class
+        Model class to use.
+    criterion : str, default='aic'
+        Selection criterion.
+    direction : str, default='both'
+        Search direction.
+    **model_kwargs
+        Model parameters.
+    
+    Returns
+    -------
+    selector : StepwiseSelector
+        Fitted selector.
+    """
+    selector = StepwiseSelector(
+        model_class=model_class,
+        criterion=criterion,
+        direction=direction,
+        **model_kwargs
+    )
+    selector.fit(X, y)
+    return selector

statgpu/glm_core/__init__.py

@@ -0,0 +1,81 @@
+"""
+GLM core utilities for statgpu.
+
+Usage:
+    from statgpu.glm_core import get_glm_loss, register_glm_loss
+
+    # Built-in
+    loss = get_glm_loss('squared_error')
+
+    # Custom
+    @register_glm_loss('huber')
+    class HuberLoss(GLMLoss):
+        ...
+"""
+
+from ._base import (
+    GLMLoss,
+    get_glm_loss,
+    register_glm_loss,
+    list_glm_losses,
+)
+from ._squared import SquaredErrorLoss
+from ._logistic import LogisticLoss
+from ._poisson import PoissonLoss
+from ._gamma import GammaLoss
+from ._inverse_gaussian import InverseGaussianLoss
+from ._negative_binomial import NegativeBinomialLoss
+from ._tweedie import TweedieLoss
+from ._family import (
+    GLMFamily,
+    Link,
+    Gaussian,
+    Binomial,
+    Poisson,
+    Gamma,
+    InverseGaussian,
+    NegativeBinomial,
+    Tweedie,
+)
+from ._irls import IRLSSolver
+
+# Solvers: re-export from solvers/ (generic)
+from statgpu.solvers import (
+    fista_solver,
+    fista_bb_solver,
+    fista_lla_path,
+    newton_solver,
+    lbfgs_solver,
+    admm_solver,
+    ConvergenceWarning,
+)
+
+__all__ = [
+    "GLMLoss",
+    "SquaredErrorLoss",
+    "LogisticLoss",
+    "PoissonLoss",
+    "GammaLoss",
+    "InverseGaussianLoss",
+    "NegativeBinomialLoss",
+    "TweedieLoss",
+    "GLMFamily",
+    "Link",
+    "Gaussian",
+    "Binomial",
+    "Poisson",
+    "Gamma",
+    "InverseGaussian",
+    "NegativeBinomial",
+    "Tweedie",
+    "IRLSSolver",
+    "fista_solver",
+    "fista_bb_solver",
+    "admm_solver",
+    "newton_solver",
+    "lbfgs_solver",
+    "ConvergenceWarning",
+    "get_glm_loss",
+    "register_glm_loss",
+    "list_glm_losses",
+]

statgpu/glm_core/_base.py

@@ -0,0 +1,202 @@
+"""
+Base class for GLM loss functions in statgpu.
+
+The GLM core loss framework supports 7 families:
+- Squared error (linear regression)
+- Logistic loss (binary classification)
+- Poisson loss (count data)
+- Gamma loss (positive continuous)
+- Inverse Gaussian loss (positive continuous)
+- Negative Binomial loss (overdispersed count data)
+- Tweedie loss (generalized GLM family)
+
+Structured models such as Cox, panel, and time-series models should use a
+future objective layer rather than this GLM-specific interface.
+"""
+
+__all__ = ["GLMLoss", "get_glm_loss", "register_glm_loss", "list_glm_losses"]
+
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import numpy as np
+from statgpu.backends._array_ops import _xp as _get_xp_mod
+from statgpu.backends._utils import _to_float_scalar
+
+
+class GLMLoss(ABC):
+    """GLM loss function base class.
+
+    Objective: minimize: loss(X, y, w) + penalty(w)
+
+    Subclasses implement per-sample formulas as the single source of truth.
+    The base class derives ``value()``, ``gradient()``, and
+    ``fused_value_and_gradient()`` from them automatically.
+
+    Subclass API (implement these):
+        - ``per_sample_value(eta, y)`` — per-sample loss ℓ(η, y)
+        - ``per_sample_gradient(eta, y)`` — per-sample gradient ∂ℓ/∂η
+        - ``_mu_from_eta(eta)`` — link inverse μ = g⁻¹(η), with clipping
+    """
+
+    name: str = "base"
+    y_type: str = "continuous"
+    smooth_gradient: bool = True
+    has_hessian: bool = False
+
+    # ── Optimization hints (solvers read these, subclasses can override) ──
+    _lipschitz_safety: float = 1.0       # Lipschitz safety factor
+    _lipschitz_safety_cv: float = 1.0    # Extra safety factor in CV mode
+    _lipschitz_uses_y: bool = False      # Whether Lipschitz needs y-scaling
+    _momentum_beta_cap: Optional[float] = None  # Nesterov momentum cap (None=unlimited)
+    _skip_momentum: bool = False         # Disable momentum entirely
+    _has_constant_hessian: bool = False  # Hessian is constant (Newton fast path)
+    _prefer_fista_over_bb: bool = False  # Prefer FISTA over FISTA-BB for smooth penalties
+    _is_quadratic: bool = False          # True for squared_error (XtX constant, no y-scaling)
+    _supports_cholesky: bool = False     # True for squared_error (ADMM can use Cholesky)
+    _gpu_loop_excluded: bool = False     # True for logistic (async GPU loop not suitable)
+    _conservative_momentum_with_nonsmooth: bool = False  # Cap momentum when penalty is non-smooth
+    _inverse_gaussian: bool = False      # True for inverse Gaussian (special BB handling)
+    _tweedie: bool = False               # True for Tweedie (special BB handling)
+    _poisson_like: bool = False          # True for Poisson (conservative momentum burn-in)
+    _gamma_like: bool = False            # True for Gamma (adjusted BB/momentum params)
+
+    # ── Per-sample formulas (single source of truth) ──────────────────
+
+    def per_sample_value(self, eta, y):
+        """Per-sample loss: ℓ(η, y). Returns array of shape (n,)."""
+        raise NotImplementedError(f"{self.name} does not implement per_sample_value")
+
+    def per_sample_gradient(self, eta, y):
+        """Per-sample gradient: ∂ℓ/∂η. Returns array of shape (n,)."""
+        raise NotImplementedError(f"{self.name} does not implement per_sample_gradient")
+
+    def _mu_from_eta(self, eta):
+        """Link inverse: μ = g⁻¹(η). Override for clipping."""
+        return eta  # default: identity link
+
+    # ── Derived methods (implemented once in base class) ──────────────
+
+    def value(self, X, y, coef, sample_weight=None) -> float:
+        """Loss value: (1/n) Σ ℓ(ηᵢ, yᵢ)."""
+        xp = _get_xp_mod(X)
+        eta = X @ coef
+        ps = self.per_sample_value(eta, y)
+        if sample_weight is not None:
+            return float(xp.dot(sample_weight, ps)) / float(sample_weight.sum())
+        return float(xp.sum(ps)) / X.shape[0]
+
+    def gradient(self, X, y, coef, sample_weight=None) -> np.ndarray:
+        """Gradient: X' ∂ℓ/∂η / n."""
+        xp = _get_xp_mod(X)
+        eta = X @ coef
+        resid = self.per_sample_gradient(eta, y)
+        if sample_weight is not None:
+            return X.T @ (sample_weight * resid) / float(sample_weight.sum())
+        return X.T @ resid / X.shape[0]
+
+    def fused_value_and_gradient(self, X, y, coef, sample_weight=None):
+        """Compute value and gradient in one pass (avoids redundant X @ coef).
+
+        Returns (value, gradient) tuple.
+        """
+        xp = _get_xp_mod(X)
+        eta = X @ coef
+        ps = self.per_sample_value(eta, y)
+        resid = self.per_sample_gradient(eta, y)
+        if sample_weight is not None:
+            sw_sum = float(sample_weight.sum())
+            val = float(xp.dot(sample_weight, ps)) / sw_sum
+            grad = X.T @ (sample_weight * resid) / sw_sum
+        else:
+            n = X.shape[0]
+            val = float(xp.sum(ps)) / n
+            grad = X.T @ resid / n
+        return val, grad
+
+    def hessian(self, X, y, coef) -> np.ndarray:
+        """Hessian matrix (for IRLS/Newton).
+
+        Raises NotImplementedError when auto solver falls back to FISTA.
+        """
+        raise NotImplementedError(
+            f"{self.name} does not support Hessian."
+        )
+
+    def lipschitz(self, X, coef, y=None) -> float:
+        """Lipschitz constant (for FISTA step size step=1/L)."""
+        from statgpu.backends._array_ops import _max_eigval_power
+        XtX = X.T @ X
+        return _max_eigval_power(XtX) / X.shape[0]
+
+    def preprocess(self, X, y):
+        """Preprocess y. Default returns as-is."""
+        return X, y
+
+    def predict(self, X, coef):
+        """Map from X @ coef to prediction. Default X @ coef."""
+        return X @ coef
+
+
+# ─── Registry ──────────────────────────────────────────────────────────────
+
+_GLM_LOSS_REGISTRY: dict = {}
+
+
+def get_glm_loss(name: str, **kwargs) -> GLMLoss:
+    """
+    Get a GLM loss by name from the registry.
+
+    Parameters
+    ----------
+    name : str
+        GLM loss name: 'squared_error', 'logistic', 'poisson', etc.
+    **kwargs
+        Arguments passed to the loss constructor.
+
+    Returns
+    -------
+    GLMLoss
+        Instantiated GLM loss object.
+
+    Raises
+    ------
+    ValueError
+        If loss name is not in the registry.
+    """
+    if name not in _GLM_LOSS_REGISTRY:
+        available = list(_GLM_LOSS_REGISTRY.keys())
+        raise ValueError(
+            f"Unknown GLM loss: {name}. Available GLM losses: {available}"
+        )
+    return _GLM_LOSS_REGISTRY[name](**kwargs)
+
+
+def register_glm_loss(name: str):
+    """
+    Decorator to register a custom GLM loss class.
+
+    Parameters
+    ----------
+    name : str
+        Name to register the GLM loss under.
+
+    Returns
+    -------
+    callable
+        Decorator function that registers the loss class.
+    """
+    def decorator(cls):
+        if not issubclass(cls, GLMLoss):
+            raise TypeError(
+                f"GLM loss class must inherit from GLMLoss, got {cls.__bases__}"
+            )
+        _GLM_LOSS_REGISTRY[name] = cls
+        return cls
+    return decorator
+
+
+def list_glm_losses() -> list:
+    """List all registered GLM loss names."""
+    return list(_GLM_LOSS_REGISTRY.keys())