statgpu 0.1.0__py3-none-any.whl

statgpu/backends/_torch_safe.py

@@ -0,0 +1,47 @@
+"""Safe torch import wrapper for Torch 2.8+ compatibility.
+
+Torch 2.8.0+ may raise RuntimeError('Only a single TORCH_LIBRARY can be
+used to register the namespace prims') when imported in environments where
+torch has already been loaded (e.g., Jupyter kernels, other processes).
+
+This module provides a safe import that catches this error and marks torch
+as unavailable. All torch imports in statgpu should go through this module
+via: from statgpu.backends._torch_safe import get_torch
+"""
+
+_torch = None
+_torch_available = None  # None = not checked, True/False = checked
+
+
+def get_torch():
+    """Return the torch module, or None if not available.
+
+    Catches RuntimeError from TORCH_LIBRARY registration conflicts
+    that occur on Torch 2.8+ in environments with pre-existing torch state.
+    """
+    global _torch, _torch_available
+
+    if _torch_available is True:
+        return _torch
+    if _torch_available is False:
+        return None
+
+    try:
+        import torch
+        _torch = torch
+        _torch_available = True
+        return _torch
+    except (ImportError, RuntimeError) as e:
+        # RuntimeError: TORCH_LIBRARY conflict on Torch 2.8+
+        # ImportError: torch not installed
+        _torch = None
+        _torch_available = False
+        return None
+
+
+def torch_available():
+    """Check if torch is available without importing it."""
+    global _torch_available
+    if _torch_available is None:
+        get_torch()
+    return _torch_available

statgpu/backends/_utils.py

@@ -0,0 +1,423 @@
+"""General-purpose backend utility functions.
+
+These helpers are used across statgpu submodules to avoid duplicating
+array-library detection, module resolution, and scalar conversion logic.
+"""
+
+from __future__ import annotations
+
+__all__ = ["xp_zeros", "xp_eye", "xp_full", "xp_astype", "xp_asarray", "xp_empty", "torch_compile_supported"]
+
+from typing import Any, Optional
+
+import numpy as np
+
+# Exception types raised by linalg operations on singular/ill-conditioned matrices.
+# numpy raises LinAlgError; torch raises RuntimeError for linalg failures.
+# NOTE: torch RuntimeError is overly broad (also catches OOM, autograd errors).
+# Callers should re-raise if the error message doesn't match linalg patterns.
+# NOTE: We do NOT import torch at module level to avoid TORCH_LIBRARY
+# registration conflicts on Torch 2.8+. Torch is imported lazily via _safe_import_torch().
+_LINALG_ERRORS: tuple = (np.linalg.LinAlgError,)
+
+
+def _safe_import_torch():
+    """Import torch safely, catching TORCH_LIBRARY registration conflicts.
+
+    Torch 2.8+ may raise RuntimeError when imported in environments where
+    torch has already been loaded (Jupyter kernels, other processes).
+    Returns the torch module, or None if import fails.
+    """
+    try:
+        import torch
+        return torch
+    except (ImportError, RuntimeError):
+        return None
+
+
+# Module-level check: is torch available?
+_TORCH_AVAILABLE = None  # None = not checked yet
+_TORCH_MODULE = None
+
+
+def _ensure_torch():
+    """Ensure torch is imported and available. Returns torch module or None."""
+    global _TORCH_AVAILABLE, _TORCH_MODULE, _LINALG_ERRORS
+    if _TORCH_AVAILABLE is None:
+        _TORCH_MODULE = _safe_import_torch()
+        _TORCH_AVAILABLE = _TORCH_MODULE is not None
+        if _TORCH_AVAILABLE:
+            _LINALG_ERRORS = (np.linalg.LinAlgError, RuntimeError)
+    return _TORCH_MODULE
+
+
+def _require_torch():
+    """Import torch or raise ImportError with clear message."""
+    torch = _ensure_torch()
+    if torch is None:
+        raise ImportError(
+            "Torch is not available. This may be due to a TORCH_LIBRARY "
+            "registration conflict (Torch 2.8+) or missing installation."
+        )
+    return torch
+
+
+def _get_xp(backend_name: str):
+    """Return the array module (numpy / cupy / torch) for *backend_name*.
+
+    Parameters
+    ----------
+    backend_name : str
+        One of ``'numpy'``, ``'cupy'``, or ``'torch'``.
+
+    Returns
+    -------
+    module
+        The array module (``numpy``, ``cupy``, or ``torch``).
+
+    Raises
+    ------
+    ValueError
+        If *backend_name* is not recognised.
+    ImportError
+        If the requested library is not installed.
+    """
+    if backend_name == "numpy":
+        return np
+    if backend_name == "cupy":
+        try:
+            import cupy as cp
+
+            return cp
+        except ImportError as exc:
+            raise ImportError(
+                "backend='cupy' requires CuPy, but CuPy is not installed"
+            ) from exc
+    if backend_name == "torch":
+        try:
+            torch = _require_torch()
+
+            return torch
+        except ImportError as exc:
+            raise ImportError(
+                "backend='torch' requires PyTorch, but PyTorch is not installed"
+            ) from exc
+    raise ValueError(f"Unsupported backend: {backend_name}")
+
+
+def _to_numpy(x):
+    """Convert *x* to a ``numpy.ndarray``.
+
+    Handles CuPy arrays (``.get()``) and PyTorch tensors (``.cpu().numpy()``).
+    """
+    if hasattr(x, "get"):
+        return x.get()
+    if hasattr(x, "cpu") and hasattr(x, "numpy"):
+        return x.detach().cpu().numpy() if hasattr(x, 'detach') else x.cpu().numpy()
+    return np.asarray(x)
+
+
+def _to_float_scalar(x: Any) -> float:
+    """Extract a Python ``float`` from a backend array scalar."""
+    if hasattr(x, "item"):
+        return float(x.item())
+    return float(x)
+
+
+def _get_torch_device_str() -> str:
+    """Return ``'cuda'`` if PyTorch CUDA is available, else ``'cpu'``."""
+    try:
+        torch = _require_torch()
+
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    except ImportError:
+        return "cpu"
+    except Exception as e:
+        import warnings
+        warnings.warn(f"torch.cuda.is_available() failed, falling back to CPU: {e}")
+        return "cpu"
+
+
+def _torch_on_target_device(tensor, device: Optional[str]) -> bool:
+    """Return True when a torch tensor is already on the requested device."""
+    if device is None:
+        return True
+    device = str(device)
+    tensor_device = str(getattr(tensor, "device", ""))
+    # "cuda" means any CUDA device; "cuda:0", "cuda:1" etc. require exact match
+    if device == "cuda":
+        return getattr(tensor, "device", None).type == "cuda"
+    return tensor_device == device
+
+
+def _move_torch_tensor(tensor, device: Optional[str] = None, dtype=None, pin_memory: bool = False):
+    """Move/cast a torch tensor, using pinned non-blocking H2D when useful."""
+    torch = _require_torch()
+
+    if dtype is not None and not isinstance(dtype, torch.dtype):
+        try:
+            dtype = getattr(torch, np.dtype(dtype).name)
+        except Exception:
+            pass
+
+    target = device or _get_torch_device_str()
+    needs_move = not _torch_on_target_device(tensor, target)
+    needs_dtype = dtype is not None and tensor.dtype != dtype
+    if not needs_move and not needs_dtype:
+        return tensor
+
+    if pin_memory and str(target).startswith("cuda") and tensor.device.type == "cpu":
+        try:
+            pinned = tensor.pin_memory() if not tensor.is_pinned() else tensor
+            kwargs = {"device": target, "non_blocking": True}
+            if dtype is not None:
+                kwargs["dtype"] = dtype
+            return pinned.to(**kwargs)
+        except Exception:
+            pass
+
+    kwargs = {}
+    if device is not None:
+        kwargs["device"] = target
+    if dtype is not None:
+        kwargs["dtype"] = dtype
+    return tensor.to(**kwargs) if kwargs else tensor
+
+
+def _numpy_to_torch_tensor(x, device: Optional[str] = None, dtype=None, pin_memory: bool = False):
+    """Convert NumPy-like input to torch, preserving contiguous fast paths."""
+    torch = _require_torch()
+
+    arr = np.asarray(x)
+    if not arr.flags["C_CONTIGUOUS"]:
+        arr = np.ascontiguousarray(arr)
+    tensor = torch.from_numpy(arr)
+    return _move_torch_tensor(tensor, device=device, dtype=dtype, pin_memory=pin_memory)
+
+
+def _cupy_to_torch_dlpack(x, device: Optional[str] = None):
+    """Convert a CuPy array to torch through DLPack, returning None if unsupported."""
+    try:
+        import cupy as cp
+        torch = _require_torch()
+
+        if not isinstance(x, cp.ndarray):
+            return None
+        try:
+            tensor = torch.utils.dlpack.from_dlpack(x)
+        except TypeError:
+            tensor = torch.utils.dlpack.from_dlpack(x.toDlpack())
+        return _move_torch_tensor(tensor, device=device)
+    except Exception:
+        return None
+
+
+def _torch_to_cupy_dlpack(x):
+    """Convert a CUDA torch tensor to CuPy through DLPack, returning None if unsupported."""
+    try:
+        import cupy as cp
+        torch = _require_torch()
+
+        if not isinstance(x, torch.Tensor) or not x.is_cuda:
+            return None
+        tensor = x.detach()
+        if not tensor.is_contiguous():
+            tensor = tensor.contiguous()
+        try:
+            return cp.from_dlpack(tensor)
+        except Exception:
+            return cp.fromDlpack(torch.utils.dlpack.to_dlpack(tensor))
+    except Exception:
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Device-aware array creation helpers
+# ---------------------------------------------------------------------------
+
+def _torch_dev(arr):
+    """Extract device from a torch tensor, or ``None`` for non-torch arrays."""
+    try:
+        torch = _require_torch()
+        if isinstance(arr, torch.Tensor):
+            return arr.device
+    except (ImportError, AttributeError):
+        pass
+    return None
+
+
+def xp_zeros(shape, dtype, xp, ref_arr=None):
+    """Device-aware ``xp.zeros``.  *ref_arr* provides the target device."""
+    dev = _torch_dev(ref_arr) if ref_arr is not None else None
+    if dev is not None:
+        return xp.zeros(shape, dtype=dtype, device=dev)
+    return xp.zeros(shape, dtype=dtype)
+
+
+def xp_eye(n, dtype, xp, ref_arr=None):
+    """Device-aware ``xp.eye``.  *ref_arr* provides the target device."""
+    dev = _torch_dev(ref_arr) if ref_arr is not None else None
+    if dev is not None:
+        return xp.eye(n, dtype=dtype, device=dev)
+    return xp.eye(n, dtype=dtype)
+
+
+def xp_full(shape, fill_value, dtype, xp, ref_arr=None):
+    """Device-aware ``xp.full`` with int→tuple normalisation."""
+    if isinstance(shape, int):
+        shape = (shape,)
+    dev = _torch_dev(ref_arr) if ref_arr is not None else None
+    if dev is not None:
+        return xp.full(shape, fill_value, dtype=dtype, device=dev)
+    return xp.full(shape, fill_value, dtype=dtype)
+
+
+def _np_dtype_to_torch(dtype):
+    """Convert a numpy dtype to the equivalent torch dtype."""
+    torch = _require_torch()
+    _MAP = {
+        'float32': torch.float32,
+        'float64': torch.float64,
+        'float16': torch.float16,
+        'int32': torch.int32,
+        'int64': torch.int64,
+        'int16': torch.int16,
+        'int8': torch.int8,
+        'uint8': torch.uint8,
+        'bool': torch.bool,
+    }
+    key = str(np.dtype(dtype)).split('.')[-1]
+    result = _MAP.get(key)
+    if result is None:
+        import warnings
+        warnings.warn(f"Unknown numpy dtype '{dtype}' for torch conversion, falling back to float64", stacklevel=2)
+        return torch.float64
+    return result
+
+
+def _torch_dtype_to_np(dtype):
+    """Convert a torch dtype to the equivalent numpy dtype."""
+    torch = _require_torch()
+    _MAP = {
+        torch.float32: np.dtype('float32'),
+        torch.float64: np.dtype('float64'),
+        torch.float16: np.dtype('float16'),
+        torch.int32: np.dtype('int32'),
+        torch.int64: np.dtype('int64'),
+        torch.int16: np.dtype('int16'),
+        torch.int8: np.dtype('int8'),
+        torch.uint8: np.dtype('uint8'),
+        torch.bool: np.dtype('bool'),
+    }
+    return _MAP.get(dtype, np.dtype('float64'))
+
+
+def xp_astype(arr, dtype, xp=None):
+    """Backend-safe type cast (``.to()`` for torch, ``.astype()`` otherwise).
+
+    Note: ``xp`` parameter is unused — backend is detected from ``arr`` directly.
+    Kept for backward compatibility with existing callers.
+    """
+    if _torch_dev(arr) is not None:
+        torch = _require_torch()
+        if not isinstance(dtype, torch.dtype):
+            dtype = _np_dtype_to_torch(dtype)
+        return arr.to(dtype)
+    return arr.astype(dtype)
+
+
+def xp_asarray(data, dtype=None, xp=None, ref_arr=None):
+    """Device-aware ``xp.asarray``.  *ref_arr* provides the target device."""
+    dev = _torch_dev(ref_arr) if ref_arr is not None else None
+    if dev is not None:
+        kwargs = {'device': dev}
+        if dtype is not None:
+            kwargs['dtype'] = dtype
+        return xp.asarray(data, **kwargs)
+    if dtype is not None:
+        return xp.asarray(data, dtype=dtype)
+    return xp.asarray(data)
+
+
+def xp_empty(shape, dtype, xp, ref_arr=None):
+    """Device-aware ``xp.empty``.  *ref_arr* provides the target device."""
+    dev = _torch_dev(ref_arr) if ref_arr is not None else None
+    if dev is not None:
+        return xp.empty(shape, dtype=dtype, device=dev)
+    return xp.empty(shape, dtype=dtype)
+
+
+def xp_arange(n, dtype=None, xp=None, ref_arr=None):
+    """Device-aware ``xp.arange``.  *ref_arr* provides the target device."""
+    dev = _torch_dev(ref_arr) if ref_arr is not None else None
+    if dev is not None:
+        kwargs = {'device': dev}
+        if dtype is not None:
+            kwargs['dtype'] = dtype
+        return xp.arange(n, **kwargs)
+    if dtype is not None:
+        return xp.arange(n, dtype=dtype)
+    return xp.arange(n)
+
+
+def xp_ones(shape, dtype, xp, ref_arr=None):
+    """Device-aware ``xp.ones``.  *ref_arr* provides the target device."""
+    dev = _torch_dev(ref_arr) if ref_arr is not None else None
+    if dev is not None:
+        return xp.ones(shape, dtype=dtype, device=dev)
+    return xp.ones(shape, dtype=dtype)
+
+
+def xp_maximum(arr, value, xp=None):
+    """Element-wise maximum that works for both numpy/cupy and torch.
+
+    Torch's ``maximum()`` requires both args to be tensors; numpy/cupy accept
+    scalars.  This helper wraps *value* as needed.
+    """
+    if _torch_dev(arr) is not None:
+        torch = _require_torch()
+        if not isinstance(value, torch.Tensor):
+            value = torch.tensor(value, dtype=arr.dtype, device=arr.device)
+        return torch.maximum(arr, value)
+    return xp.maximum(arr, value) if xp is not None else np.maximum(arr, value)
+
+
+def xp_copy(arr):
+    """Backend-safe copy (``.clone()`` for torch, ``.copy()`` otherwise)."""
+    if _torch_dev(arr) is not None:
+        return arr.clone()
+    return arr.copy()
+
+
+def xp_cholesky_solve(A, b, xp):
+    """Solve ``A @ x = b`` via Cholesky decomposition.
+
+    Works across numpy, cupy, and torch backends.  Handles the torch-specific
+    argument difference for ``solve_triangular`` (``upper=False`` vs ``lower=True``).
+    For cupy, uses general solve (no solve_triangular in cupy).
+    For numpy, uses scipy.linalg.solve_triangular.
+    """
+    if hasattr(A, 'get'):  # CuPy: no solve_triangular, use general solve directly
+        return xp.linalg.solve(A, b)
+    L = xp.linalg.cholesky(A)
+    if _torch_dev(L) is not None:
+        tmp = xp.linalg.solve_triangular(L, b, upper=False)
+        return xp.linalg.solve_triangular(L.T, tmp, upper=True)
+    # numpy: use scipy for solve_triangular
+    from scipy.linalg import solve_triangular
+    tmp = solve_triangular(L, b, lower=True)
+    return solve_triangular(L.T, tmp, lower=False)
+
+
+def torch_compile_supported():
+    """Check if torch.compile is safe to use (CUDA Capability >= 7.0)."""
+    try:
+        torch = _require_torch()
+        if torch.cuda.is_available():
+            cap = torch.cuda.get_device_capability()
+            return cap[0] >= 7
+    except ImportError:
+        return False  # torch not installed
+    except Exception:
+        pass
+    return False  # Can't verify — assume not supported

statgpu/core/__init__.py

@@ -0,0 +1,10 @@
+"""
+statgpu.core — Shared core utilities for statgpu.
+
+Provides common infrastructure used across all statgpu model modules:
+- ``core.formula``: R-style formula interface (``y ~ x1 + x2``).
+"""
+
+from . import formula
+
+__all__ = ["formula"]

statgpu/core/formula/__init__.py

@@ -0,0 +1,33 @@
+"""
+statgpu.core.formula – R-style formula interface for statgpu models.
+
+This module provides formula-based model fitting similar to statsmodels/patsy::
+
+    >>> import statgpu as sg
+    >>> model = sg.LinearRegression()
+    >>> model.fit(formula="y ~ x1 + x2 + C(cat)", data=df)
+    >>> model.summary()
+
+The formula syntax is parsed by `patsy` (optional dependency). Install with::
+
+    pip install statgpu[formula]
+
+Public API
+----------
+FormulaParser
+    Main class for parsing R-style formulas and building design matrices.
+parse_formula
+    Convenience function for one-shot formula evaluation.
+"""
+
+from ._parser import FormulaParser
+from ._design import parse_formula, parse_formula_safe
+from ._terms import make_surv_env, _surv
+
+__all__ = [
+    "FormulaParser",
+    "parse_formula",
+    "parse_formula_safe",
+    "make_surv_env",
+    "_surv",
+]

statgpu/core/formula/_design.py

@@ -0,0 +1,99 @@
+"""
+Design matrix building utilities.
+
+Provides convenience function for one-shot formula evaluation.
+"""
+
+from typing import Tuple, Optional, Any
+
+import numpy as np
+import pandas as pd
+
+from ._parser import FormulaParser
+
+
+def parse_formula(
+    formula: str,
+    data: pd.DataFrame,
+) -> Tuple[np.ndarray, np.ndarray, Any]:
+    """One-shot convenience function for formula parsing.
+
+    Parameters
+    ----------
+    formula : str
+        R-style formula string, e.g. ``"y ~ x1 + x2"``.
+    data : pd.DataFrame
+        DataFrame containing the referenced columns.
+
+    Returns
+    -------
+    y : ndarray
+        Response variable(s).
+    X : ndarray
+        Predictor design matrix.
+    design_info : patsy.DesignInfo
+        Metadata for the predictor design.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> df = pd.DataFrame({"y": [1, 2, 3], "x": [4, 5, 6]})
+    >>> y, X, info = parse_formula("y ~ x", df)
+    """
+    parser = FormulaParser(formula)
+    return parser.eval(data)
+
+
+def parse_formula_safe(
+    formula: Optional[str],
+    data: Optional[pd.DataFrame],
+    X: Optional[np.ndarray] = None,
+    y: Optional[np.ndarray] = None,
+) -> Tuple[np.ndarray, np.ndarray, Optional[Any]]:
+    """Safe formula parsing that falls back to raw arrays.
+
+    Used by model ``fit()`` methods to support both formula and array interfaces.
+
+    Parameters
+    ----------
+    formula : str or None
+        R-style formula string. If ``None``, ``X`` and ``y`` are used directly.
+    data : pd.DataFrame or None
+        DataFrame for formula parsing. Required when ``formula`` is given.
+    X : ndarray or None
+        Raw predictor matrix (used when ``formula`` is ``None``).
+    y : ndarray or None
+        Raw response vector (used when ``formula`` is ``None``).
+
+    Returns
+    -------
+    y : ndarray
+        Response variable(s).
+    X : ndarray
+        Predictor design matrix.
+    design_info : patsy.DesignInfo or None
+        Design metadata (``None`` when raw arrays are used).
+
+    Raises
+    ------
+    ValueError
+        If both formula and arrays are ``None``, or if formula is given without data.
+    """
+    if formula is not None:
+        if data is None:
+            raise ValueError(
+                "formula was provided but data (DataFrame) is None. "
+                "When using formula, pass data=your_dataframe."
+            )
+        return parse_formula(formula, data)
+
+    if X is None or y is None:
+        raise ValueError(
+            "Either formula+data or X+y must be provided. "
+            "Got formula=None and incomplete array input."
+        )
+
+    y = np.asarray(y)
+    if y.ndim == 2 and y.shape[1] == 1:
+        y = y.ravel()
+    return y, np.asarray(X), None