statgpu 0.1.0__py3-none-any.whl

statgpu/solvers/_lbfgs.py

@@ -0,0 +1,206 @@
+"""Limited-memory BFGS solver for smooth penalised objectives.
+
+Generic solver -- works with any loss that implements fused_value_and_gradient().
+Keeps parameters, gradients, and curvature history on the input backend.
+GPU-optimised path uses:
+- loss.fused_value_and_gradient to avoid redundant X@coef
+- _dot_dev / _norm2_dev to stay on device
+- _sync_scalars to batch GPU-to-CPU transfers
+- _device_leq for device-side line search
+"""
+
+from __future__ import annotations
+
+__all__ = ["lbfgs_solver"]
+
+import warnings
+import numpy as np
+
+from statgpu.backends import _resolve_backend
+from statgpu.backends._array_ops import (
+    _copy_arr,
+    _device_gt,
+    _device_leq,
+    _dot_dev,
+    _norm2_dev,
+    _sync_scalars,
+    _zeros,
+)
+
+from ._convergence import ConvergenceWarning
+from ._utils import (
+    _smooth_penalty_gradient,
+    _smooth_penalty_value_dev,
+    _validate_uniform_sample_weight,
+)
+
+
+def lbfgs_solver(
+    loss: "GLMLoss",
+    penalty: "Penalty | None",
+    X,
+    y,
+    max_iter: int = 100,
+    tol: float = 1e-4,
+    init_coef=None,
+    history_size: int = 10,
+    sample_weight=None,
+) -> tuple:
+    """Limited-memory BFGS for smooth objectives.
+
+    Works with any loss that implements ``fused_value_and_gradient(X, y, coef)``
+    returning ``(value, gradient)``.  Supports numpy / cupy / torch backends
+    via auto-detection of *X*.
+
+    Parameters
+    ----------
+    loss : object
+        Loss with ``fused_value_and_gradient(X, y, coef)`` and
+        ``preprocess(X, y)`` methods.
+    penalty : object or None
+        Smooth penalty (l2, elasticnet, none).
+    X, y : array-like
+        Design matrix and response vector.
+    max_iter : int
+        Maximum number of L-BFGS iterations.
+    tol : float
+        Convergence tolerance on gradient norm and step norm.
+    init_coef : array-like or None
+        Initial coefficient vector.  Zeros if *None*.
+    history_size : int
+        Number of past (s, y) pairs to store.
+    sample_weight : array-like or None
+        Sample weights.  Must be uniform (all equal) for this solver.
+
+    Returns
+    -------
+    params : array
+        Optimised coefficient vector.
+    n_iter : int
+        Number of iterations performed.
+    """
+    backend = _resolve_backend("auto", X)
+    X_proc, y_proc = loss.preprocess(X, y)
+    n_features = X_proc.shape[1]
+    _validate_uniform_sample_weight(sample_weight, X_proc.shape[0], "lbfgs_solver")
+
+    if init_coef is not None:
+        params = (
+            _copy_arr(init_coef)
+            if hasattr(init_coef, "copy") or hasattr(init_coef, "clone")
+            else np.array(init_coef).copy()
+        )
+    else:
+        params = _zeros(n_features, backend, ref_tensor=X)
+
+    s_hist = []
+    y_hist = []
+    rho_hist = []
+
+    # Initial gradient (fused to avoid redundant X@coef)
+    _init_val_dev, grad = loss.fused_value_and_gradient(X_proc, y_proc, params)
+    grad = grad + _smooth_penalty_gradient(penalty, params)
+
+    if backend == "torch":
+        import torch
+        tol_dev = torch.tensor(tol, dtype=torch.float64, device=params.device)
+    else:
+        tol_dev = tol
+    iteration = -1  # default if max_iter=0
+
+    for iteration in range(max_iter):
+        grad_norm_dev = _norm2_dev(grad)
+
+        # Two-loop recursion -- all dot products stay on device
+        q = _copy_arr(grad)
+        alphas = []
+        for s_vec, y_vec, rho in reversed(list(zip(s_hist, y_hist, rho_hist))):
+            alpha = rho * _dot_dev(s_vec, q)
+            alphas.append(alpha)
+            q = q - alpha * y_vec
+
+        if y_hist:
+            sy = _dot_dev(s_hist[-1], y_hist[-1])
+            yy = _dot_dev(y_hist[-1], y_hist[-1])
+            gamma = sy / yy if _device_gt(yy, 1e-30) else 1.0
+        else:
+            gamma = 1.0
+        r = gamma * q
+
+        for s_vec, y_vec, rho, alpha in zip(
+            s_hist, y_hist, rho_hist, reversed(alphas)
+        ):
+            beta = rho * _dot_dev(y_vec, r)
+            r = r + s_vec * (alpha - beta)
+
+        direction = -r
+        gdd_dev = _dot_dev(grad, direction)
+
+        # Batch sync: grad_norm + grad_dot_dir
+        gn, gdd = _sync_scalars(grad_norm_dev, gdd_dev, backend=backend)
+        if gn < tol:
+            break
+        if gdd >= 0:
+            direction = -grad
+            gdd = -gn  # -||grad||^2
+
+        # Line search -- stays on device
+        old_val_dev, _ = loss.fused_value_and_gradient(X_proc, y_proc, params)
+        old_val_dev = old_val_dev + _smooth_penalty_value_dev(penalty, params)
+
+        step = 1.0
+        params_new = params
+        _ls_accepted = False
+        for _ in range(25):
+            candidate = params + step * direction
+            cand_val_dev, _ = loss.fused_value_and_gradient(X_proc, y_proc, candidate)
+            cand_val_dev = cand_val_dev + _smooth_penalty_value_dev(penalty, candidate)
+            # Device-side comparison -- single sync for the bool
+            if _device_leq(cand_val_dev, old_val_dev + 1e-4 * step * gdd):
+                params_new = candidate
+                _ls_accepted = True
+                break
+            step *= 0.5
+        if not _ls_accepted:
+            warnings.warn(
+                "lbfgs_solver: line search failed to find a descent step "
+                f"after 25 backtracking steps (iteration {iteration}). "
+                "Solver may stagnate.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
+        # Update gradient (fused)
+        _, grad_new = loss.fused_value_and_gradient(X_proc, y_proc, params_new)
+        grad_new = grad_new + _smooth_penalty_gradient(penalty, params_new)
+
+        s_vec = params_new - params
+        y_vec = grad_new - grad
+        ys_dev = _dot_dev(y_vec, s_vec)
+        s_norm_dev = _norm2_dev(s_vec)
+
+        # Batch sync: ys + s_norm
+        ys, s_norm = _sync_scalars(ys_dev, s_norm_dev, backend=backend)
+        if ys > 1e-12:
+            s_hist.append(s_vec)
+            y_hist.append(y_vec)
+            rho_hist.append(1.0 / ys)
+            if len(s_hist) > history_size:
+                s_hist.pop(0)
+                y_hist.pop(0)
+                rho_hist.pop(0)
+
+        params = params_new
+        grad = grad_new
+        if s_norm < tol:
+            break
+
+    n_iter = iteration + 1
+    if n_iter >= max_iter:
+        warnings.warn(
+            f"lbfgs_solver did not converge within {max_iter} iterations "
+            f"(loss={getattr(loss, 'name', '?')}, penalty={getattr(penalty, 'name', '?')}).",
+            ConvergenceWarning,
+            stacklevel=2,
+        )
+    return params, n_iter

statgpu/solvers/_newton.py

@@ -0,0 +1,149 @@
+"""Newton-Raphson solver with Armijo backtracking line search.
+
+Generic solver — works with any loss that implements hessian() and gradient().
+"""
+
+from __future__ import annotations
+
+__all__ = ["newton_solver"]
+
+import warnings
+import numpy as np
+
+from statgpu.backends import _resolve_backend
+from statgpu.backends._array_ops import (
+    _copy_arr,
+    _dot_dev,
+    _norm2_dev,
+    _sync_scalars,
+    _zeros,
+    _device_leq,
+)
+from statgpu.backends._utils import _to_float_scalar
+
+from ._convergence import ConvergenceWarning
+from ._utils import (
+    _validate_uniform_sample_weight,
+    _smooth_penalty_gradient,
+    _smooth_penalty_hessian,
+    _smooth_penalty_value_dev,
+)
+
+
+def newton_solver(
+    loss: "GLMLoss",
+    penalty: "Penalty | None",
+    X,
+    y,
+    max_iter: int = 100,
+    tol: float = 1e-4,
+    init_coef=None,
+    sample_weight=None,
+) -> tuple:
+    """Newton-Raphson solver with Armijo backtracking line search.
+
+    Supports numpy / cupy / torch backends via auto-detection of X.
+
+    For losses with constant Hessian (e.g. Gamma log link), the Hessian
+    doesn't change across iterations, so the Newton step is always valid
+    and line search is skipped.
+
+    Requires: loss has hessian() and penalty is smooth.
+    """
+    backend = _resolve_backend("auto", X)
+    X_proc, y_proc = loss.preprocess(X, y)
+    n_features = X_proc.shape[1]
+
+    if init_coef is not None:
+        params = (
+            _copy_arr(init_coef)
+            if hasattr(init_coef, "copy") or hasattr(init_coef, "clone")
+            else np.array(init_coef).copy()
+        )
+    else:
+        params = _zeros(n_features, backend, ref_tensor=X_proc)
+
+    # Constant-Hessian detection via loss attribute (generic, not loss-name based)
+    _const_hessian = getattr(loss, "_has_constant_hessian", False)
+
+    _fixed_hess = None
+    if _const_hessian:
+        _fixed_hess = loss.hessian(X_proc, y_proc, params) + _smooth_penalty_hessian(
+            penalty, params
+        )
+
+    _validate_uniform_sample_weight(sample_weight, X_proc.shape[0], "newton_solver")
+    iteration = -1
+
+    for iteration in range(max_iter):
+        params_old = _copy_arr(params)
+        grad = loss.gradient(X_proc, y_proc, params) + _smooth_penalty_gradient(
+            penalty, params
+        )
+        hess = _fixed_hess if _fixed_hess is not None else (
+            loss.hessian(X_proc, y_proc, params) + _smooth_penalty_hessian(penalty, params)
+        )
+
+        try:
+            if backend == "numpy":
+                direction = np.linalg.solve(hess, grad)
+            elif backend == "cupy":
+                import cupy as cp
+
+                direction = cp.linalg.solve(hess, grad)
+            else:
+                import torch
+
+                direction = torch.linalg.solve(hess, grad.unsqueeze(1))
+                direction = direction.squeeze(1)
+        except (np.linalg.LinAlgError, ValueError, RuntimeError):
+            if backend == "numpy":
+                direction = np.linalg.lstsq(hess, grad, rcond=None)[0]
+            elif backend == "cupy":
+                import cupy as cp
+
+                direction = cp.linalg.lstsq(hess, grad)[0]
+            else:
+                import torch
+
+                direction = torch.linalg.lstsq(hess, grad.unsqueeze(1)).solution
+                direction = direction.squeeze(1)
+
+        # Armijo backtracking — use loss.fused_value_and_gradient (generic interface)
+        obj_old_dev, _ = loss.fused_value_and_gradient(X_proc, y_proc, params_old)
+        obj_old_dev = obj_old_dev + _smooth_penalty_value_dev(penalty, params_old)
+        gdd_dev = _dot_dev(grad, direction)
+        gdd = _to_float_scalar(gdd_dev)
+
+        step = 1.0
+        for _bt in range(20):
+            params_try = params_old - step * direction
+            try:
+                obj_try_dev, _ = loss.fused_value_and_gradient(X_proc, y_proc, params_try)
+                obj_try_dev = obj_try_dev + _smooth_penalty_value_dev(
+                    penalty, params_try
+                )
+                if _device_leq(obj_try_dev, obj_old_dev + 1e-4 * step * gdd):
+                    params = params_try
+                    break
+            except (ValueError, RuntimeError, FloatingPointError):
+                pass
+            step *= 0.5
+        else:
+            params = params_old - step * direction
+
+        norm_diff_dev = _norm2_dev(params - params_old)
+        (nd,) = _sync_scalars(norm_diff_dev, backend=backend)
+        if nd < tol:  # _norm2_dev returns L2 norm (not squared)
+            break
+
+    n_iter = iteration + 1
+    if n_iter >= max_iter:
+        warnings.warn(
+            f"newton_solver did not converge within {max_iter} iterations "
+            f"(loss={getattr(loss, 'name', '?')}, "
+            f"penalty={getattr(penalty, 'name', '?')}).",
+            ConvergenceWarning,
+            stacklevel=2,
+        )
+    return params, n_iter

statgpu/solvers/_utils.py

@@ -0,0 +1,277 @@
+"""Shared utility functions for solvers.
+
+Validation helpers, penalty value/gradient/hessian utilities,
+and objective function helpers. All work with generic loss/penalty interfaces.
+"""
+
+import numpy as np
+
+from statgpu.backends import _resolve_backend, _to_numpy
+from statgpu.backends._utils import _to_float_scalar, _get_xp
+from statgpu.backends._array_ops import (
+    _abs_sum,
+    _abs_sum_dev,
+    _copy_arr,
+    _dot,
+    _dot_dev,
+    _eye_like,
+    _norm2,
+    _norm2_dev,
+    _sum_sq,
+    _sum_sq_dev,
+    _sync_scalars,
+    _zeros,
+    _zeros_like,
+)
+
+
+def _validate_uniform_sample_weight(sample_weight, n_samples, solver_name):
+    if sample_weight is None:
+        return
+    _sw = _to_numpy(sample_weight)
+    if _sw.ndim != 1 or _sw.shape[0] != n_samples:
+        raise ValueError("sample_weight must be a 1D array with length n_samples")
+    if not np.all(np.isfinite(_sw)):
+        raise ValueError("sample_weight must contain only finite values")
+    if np.any(_sw < 0):
+        raise ValueError("sample_weight must be non-negative")
+    if np.sum(_sw) <= 0.0:
+        raise ValueError("sample_weight must contain at least one positive value")
+    if not np.allclose(_sw, _sw[0]):
+        raise ValueError(
+            f"{solver_name} does not support non-uniform sample_weight yet; "
+            "use solver='irls' for weighted GLM fits."
+        )
+
+
+def _validate_sample_weight(sample_weight, n_samples):
+    if sample_weight is None:
+        return
+    _sw = _to_numpy(sample_weight)
+    if _sw.ndim != 1 or _sw.shape[0] != n_samples:
+        raise ValueError("sample_weight must be 1D with length n_samples")
+    if not np.all(np.isfinite(_sw)):
+        raise ValueError("sample_weight must contain only finite values")
+    if np.any(_sw < 0):
+        raise ValueError("sample_weight must be non-negative")
+    if np.sum(_sw) <= 0:
+        raise ValueError("sample_weight must contain at least one positive value")
+
+
+def _as_backend_vector(arr, backend, ref):
+    from statgpu.backends._utils import xp_asarray
+    xp = _get_xp(backend)
+    dtype = getattr(ref, "dtype", np.float64)
+    return xp_asarray(arr, dtype=dtype, xp=xp, ref_arr=ref)
+
+
+def _call_with_weight(fn, *args, sample_weight=None, **kwargs):
+    """Call fn with sample_weight if it accepts it, without otherwise.
+
+    Avoids the repeated try/except TypeError pattern. Inspects the
+    function signature once to decide whether to pass sample_weight.
+    """
+    import inspect
+    try:
+        sig = inspect.signature(fn)
+        if 'sample_weight' in sig.parameters:
+            return fn(*args, sample_weight=sample_weight, **kwargs)
+    except (ValueError, TypeError):
+        pass
+    return fn(*args, **kwargs)
+
+
+def _nesterov_momentum(t_k, beta_cap=None):
+    """Compute Nesterov momentum parameters.
+
+    Parameters
+    ----------
+    t_k : float
+        Current momentum parameter.
+    beta_cap : float, optional
+        Maximum allowed momentum (e.g. 0.5 for CV stability).
+
+    Returns
+    -------
+    beta : float
+        Momentum coefficient.
+    t_new : float
+        Updated momentum parameter.
+    """
+    import math
+    t_new = (1.0 + math.sqrt(1.0 + 4.0 * t_k * t_k)) / 2.0
+    beta = (t_k - 1.0) / t_new
+    if beta_cap is not None:
+        beta = min(beta, beta_cap)
+    return beta, t_new
+
+
+def _nesterov_update(coef, coef_old, t_k, beta_cap=None):
+    """Nesterov momentum update: compute extrapolated point y_k and new t.
+
+    Parameters
+    ----------
+    coef : array
+        Current iterate.
+    coef_old : array
+        Previous iterate.
+    t_k : float
+        Current momentum parameter.
+    beta_cap : float, optional
+        Maximum allowed momentum (e.g. 0.5 for CV stability).
+
+    Returns
+    -------
+    y_k : array
+        Extrapolated point: coef + beta * (coef - coef_old).
+    t_new : float
+        Updated momentum parameter.
+    """
+    beta, t_new = _nesterov_momentum(t_k, beta_cap)
+    y_k = coef + beta * (coef - coef_old)
+    return y_k, t_new
+
+
+def _penalty_name(penalty):
+    return str(getattr(penalty, "name", "none")).lower()
+
+
+def _smooth_penalty_value(penalty, coef):
+    if penalty is None:
+        return 0.0
+    if hasattr(penalty, "smooth_value"):
+        return float(_to_numpy(penalty.smooth_value(coef)))
+    if _penalty_name(penalty) in ("none", "null"):
+        return 0.0
+    if _penalty_name(penalty) == "l2":
+        return 0.5 * float(getattr(penalty, "alpha", 0.0)) * _sum_sq(coef)
+    if _penalty_name(penalty) == "elasticnet":
+        alpha = float(getattr(penalty, "alpha", 0.0))
+        l1_ratio = float(getattr(penalty, "l1_ratio", 1.0))
+        return 0.5 * alpha * (1.0 - l1_ratio) * _sum_sq(coef)
+    raise ValueError(
+        f"solver requires a smooth penalty, got penalty='{_penalty_name(penalty)}'."
+    )
+
+
+def _tracking_penalty_value(penalty, coef):
+    pen_name = _penalty_name(penalty)
+    if penalty is None or pen_name in ("none", "null"):
+        return 0.0
+    n_features = getattr(penalty, "n_features", None)
+    if n_features is not None:
+        coef_eval = coef[: int(n_features)]
+        backend = _resolve_backend("auto", coef_eval)
+        if pen_name == "l1":
+            if backend in ("torch", "cupy"):
+                abs_sum, = _sync_scalars(_abs_sum_dev(coef_eval), backend=backend)
+            else:
+                abs_sum = _abs_sum(coef_eval)
+            return float(getattr(penalty, "alpha", 0.0)) * abs_sum
+        if pen_name in ("elasticnet", "en"):
+            alpha = float(getattr(penalty, "alpha", 0.0))
+            l1_ratio = float(getattr(penalty, "l1_ratio", 1.0))
+            if backend in ("torch", "cupy"):
+                abs_sum, sum_sq = _sync_scalars(
+                    _abs_sum_dev(coef_eval), _sum_sq_dev(coef_eval), backend=backend,
+                )
+            else:
+                abs_sum = _abs_sum(coef_eval)
+                sum_sq = _sum_sq(coef_eval)
+            return alpha * (l1_ratio * abs_sum + 0.5 * (1.0 - l1_ratio) * sum_sq)
+    try:
+        return float(penalty.value(coef))
+    except (ValueError, TypeError, AttributeError):
+        pass
+    try:
+        return float(penalty.value(_to_numpy(coef)))
+    except (ValueError, TypeError, AttributeError):
+        pass
+    return 0.0
+
+
+def _abs_mean_max(y, backend):
+    backend = _resolve_backend(backend, y)
+    xp = _get_xp(backend)
+    y_abs = xp.abs(y)
+    mean_val, max_val = _sync_scalars(xp.mean(y_abs), xp.max(y_abs), backend=backend)
+    return mean_val, max_val
+
+
+def _smooth_penalty_gradient(penalty, coef):
+    if penalty is None or _penalty_name(penalty) in ("none", "null"):
+        return _zeros_like(coef)
+    if hasattr(penalty, "smooth_gradient"):
+        return penalty.smooth_gradient(coef)
+    if _penalty_name(penalty) == "l2":
+        return float(getattr(penalty, "alpha", 0.0)) * coef
+    if _penalty_name(penalty) == "elasticnet":
+        alpha = float(getattr(penalty, "alpha", 0.0))
+        l1_ratio = float(getattr(penalty, "l1_ratio", 1.0))
+        return alpha * (1.0 - l1_ratio) * coef
+    raise ValueError(
+        f"solver requires a smooth penalty, got penalty='{_penalty_name(penalty)}'."
+    )
+
+
+def _smooth_penalty_hessian(penalty, coef):
+    if penalty is None or _penalty_name(penalty) in ("none", "null"):
+        return 0.0
+    n = coef.shape[0]
+    if hasattr(penalty, "smooth_hessian"):
+        return penalty.smooth_hessian(coef)
+    if _penalty_name(penalty) == "l2":
+        return float(getattr(penalty, "alpha", 0.0)) * _eye_like(n, coef)
+    if _penalty_name(penalty) in ("elasticnet", "en"):
+        # ElasticNet Hessian is the L2 component only (L1 is non-smooth)
+        alpha = float(getattr(penalty, "alpha", 0.0))
+        l1_ratio = float(getattr(penalty, "l1_ratio", 0.5))
+        return alpha * (1.0 - l1_ratio) * _eye_like(n, coef)
+    raise ValueError(
+        f"solver requires a smooth penalty, got penalty='{_penalty_name(penalty)}'."
+    )
+
+
+def _objective_value(loss, penalty, X, y, coef):
+    return float(_to_numpy(loss.value(X, y, coef))) + _smooth_penalty_value(penalty, coef)
+
+
+def _objective_gradient(loss, penalty, X, y, coef):
+    return loss.gradient(X, y, coef) + _smooth_penalty_gradient(penalty, coef)
+
+
+def _smooth_penalty_lipschitz(penalty):
+    if penalty is None:
+        return 0.0
+    _pname = _penalty_name(penalty)
+    if _pname in ("none", "null", "l1", "scad", "mcp", "adaptive_l1", "adaptive_lasso",
+                  "group_lasso", "group_mcp", "group_scad", "gl", "gmcp", "gscad"):
+        return 0.0
+    alpha = float(getattr(penalty, 'alpha', 0.0))
+    l1_ratio = float(getattr(penalty, 'l1_ratio', 0.0))
+    return alpha * (1.0 - l1_ratio)
+
+
+def _smooth_penalty_value_dev(penalty, coef):
+    if penalty is None:
+        return 0.0
+    pname = _penalty_name(penalty)
+    if pname in ("none", "null"):
+        return 0.0
+    if pname == "l2":
+        return 0.5 * float(getattr(penalty, "alpha", 0.0)) * _sum_sq_dev(coef)
+    if pname == "elasticnet":
+        alpha = float(getattr(penalty, "alpha", 0.0))
+        l1_ratio = float(getattr(penalty, "l1_ratio", 1.0))
+        return 0.5 * alpha * (1.0 - l1_ratio) * _sum_sq_dev(coef)
+    raise ValueError(
+        f"smooth_penalty_value_dev requires a smooth penalty, got '{pname}'."
+    )
+
+
+def _objective_value_dev(loss, penalty, X, y, coef):
+    val = loss.value(X, y, coef)
+    pen_val = _smooth_penalty_value_dev(penalty, coef)
+    if isinstance(pen_val, (int, float)) and pen_val == 0.0:
+        return val
+    return val + pen_val

statgpu/survival/__init__.py

@@ -0,0 +1,14 @@
+"""
+Survival analysis models.
+
+.. rubric:: Naming conventions
+
+- ``_cuda`` — CUDA RawKernel (pre-compiled CUDA C kernels).
+- ``_cupy`` — CuPy array operations (GPU via CuPy).
+- ``_triton`` — Triton kernel (GPU via OpenAI Triton).
+"""
+
+from ._cox import CoxPH
+from ._cox_cv import CoxPHCV
+
+__all__ = ['CoxPH', 'CoxPHCV']