statgpu 0.1.0__py3-none-any.whl

statgpu/_base.py

@@ -0,0 +1,544 @@
+"""
+Base classes for statgpu estimators.
+"""
+
+from __future__ import annotations
+
+__all__ = ["BaseEstimator"]
+
+from abc import ABC, abstractmethod
+from typing import Optional, Union, Any
+import numpy as np
+
+from statgpu._config import Device, get_device
+from statgpu.backends import (
+    get_backend,
+    BackendBase,
+    _get_torch_device_str,
+    _cupy_to_torch_dlpack,
+    _torch_to_cupy_dlpack,
+    _numpy_to_torch_tensor,
+    _move_torch_tensor,
+)
+
+
+class BaseEstimator(ABC):
+    """
+    Base class for all statgpu estimators.
+    
+    Provides common functionality for device management and input validation.
+    """
+    
+    def __init__(
+        self,
+        device: Union[str, Device] = Device.AUTO,
+        n_jobs: Optional[int] = None
+    ):
+        """
+        Initialize base estimator.
+        
+        Parameters
+        ----------
+        device : str or Device, default='auto'
+            Computation device: 'cpu', 'cuda', or 'auto'.
+        n_jobs : int, optional
+            Number of parallel jobs for CPU computation.
+            -1 means using all processors.
+        """
+        self.device = device if isinstance(device, Device) else Device(device)
+        self.n_jobs = n_jobs
+        self._fitted = False
+    
+    def _get_compute_device(self) -> Device:
+        """Resolve device for actual computation."""
+        if self.device == Device.AUTO:
+            return get_device()
+        return self.device
+
+    def _get_backend(self, backend: str = "auto") -> BackendBase:
+        """
+        Return the compute backend appropriate for this estimator's device.
+
+        Parameters
+        ----------
+        backend : {'auto', 'numpy', 'cupy', 'torch'}, default='auto'
+            Override which array library to use.  When ``'auto'``, the backend
+            is chosen based on :attr:`device` and GPU availability.
+
+        Returns
+        -------
+        BackendBase
+            A backend instance whose :attr:`~BackendBase.xp` attribute is the
+            underlying array module (NumPy, CuPy, or PyTorch).
+        """
+        compute_device = self._get_compute_device()
+        device_str = compute_device.value  # 'cpu', 'cuda', or 'torch'
+
+        if (
+            self.device != Device.AUTO
+            and compute_device == Device.CUDA
+            and backend == "auto"
+        ):
+            cupy_backend = get_backend(backend="cupy", device="cuda")
+            if not cupy_backend.is_available():
+                raise RuntimeError(
+                    "device='cuda' requires a working CuPy CUDA backend. "
+                    "Use device='auto' to allow automatic backend selection."
+                )
+            return cupy_backend
+
+        # Handle Device.TORCH explicitly - use torch backend
+        if compute_device == Device.TORCH:
+            try:
+                import torch
+            except Exception as exc:
+                raise RuntimeError(
+                    "device='torch' requires PyTorch with CUDA support."
+                ) from exc
+            if not torch.cuda.is_available():
+                raise RuntimeError(
+                    "device='torch' requires torch.cuda.is_available() to be True. "
+                    "Use device='auto' or device='cpu' if CUDA is unavailable."
+                )
+            return get_backend(backend="torch", device="cuda")
+
+        return get_backend(backend=backend, device=device_str)
+    
+    def _to_array(self, X, device: Optional[Device] = None, backend: Optional[str] = None) -> Any:
+        """
+        Convert input to appropriate array type for device.
+
+        Parameters
+        ----------
+        X : array-like
+            Input data.
+        device : Device, optional
+            Target device. If None, uses self._get_compute_device().
+        backend : {'numpy', 'cupy', 'torch'}, optional
+            Explicit backend selection. If None, uses device-based default.
+
+        Returns
+        -------
+        array
+            NumPy array (CPU), CuPy array (GPU), or Torch tensor.
+        """
+        target_device = device or self._get_compute_device()
+
+        # If backend is explicitly specified, use it
+        if backend == "torch":
+            torch_target = "cpu" if target_device == Device.CPU else "cuda"
+            return self._to_torch(X, device=torch_target)
+        elif backend == "cupy":
+            return self._to_cupy(X)
+        elif backend == "numpy":
+            if hasattr(X, "get"):
+                return X.get()
+            # Handle torch tensors that may be on CUDA — must move to CPU first
+            if hasattr(X, 'cpu') and hasattr(X, 'numpy'):
+                return X.detach().cpu().numpy()
+            return np.asarray(X)
+
+        # Otherwise, use device-based default
+        if target_device == Device.TORCH:
+            return self._to_torch(X, device="cuda")
+
+        if target_device == Device.CUDA:
+            # Strict CUDA means CuPy-backed arrays. Torch has its own explicit
+            # device mode and is not used as an implicit CUDA fallback.
+            try:
+                import cupy as cp
+                if isinstance(X, cp.ndarray):
+                    return X
+            except Exception:
+                pass
+
+        # Convert to numpy first for CPU paths or for non-CuPy inputs.
+        if hasattr(X, "get"):  # CuPy-like array
+            X_np = X.get()
+        elif hasattr(X, "cpu"):  # PyTorch tensor
+            X_cpu = X.detach().cpu() if hasattr(X, "detach") else X.cpu()
+            X_np = X_cpu.numpy() if hasattr(X_cpu, "numpy") else np.asarray(X_cpu)
+        else:
+            X_np = np.asarray(X)
+
+        if target_device == Device.CUDA:
+            try:
+                import cupy as cp
+                return cp.asarray(X_np)
+            except Exception as exc:
+                raise RuntimeError(
+                    "device='cuda' requires a working CuPy CUDA backend; "
+                    "no CPU/Torch fallback is performed for explicit CUDA."
+                ) from exc
+
+        return X_np
+
+    def _to_torch(self, X, device=None):
+        """Convert input to Torch tensor."""
+        import torch
+
+        if device == "cuda" and not torch.cuda.is_available():
+            raise RuntimeError(
+                "device='torch' requires torch.cuda.is_available() to be True; "
+                "no Torch CPU fallback is performed for explicit Torch GPU."
+            )
+
+        if isinstance(X, torch.Tensor):
+            return _move_torch_tensor(X, device=device) if device else X
+
+        if hasattr(X, "get"):  # CuPy
+            tensor = _cupy_to_torch_dlpack(X, device=device)
+            if tensor is not None:
+                return tensor
+            X_np = X.get()
+        elif hasattr(X, "cpu"):  # Tensor-like
+            X_cpu = X.detach().cpu() if hasattr(X, "detach") else X.cpu()
+            X_np = X_cpu.numpy() if hasattr(X_cpu, "numpy") else np.asarray(X_cpu)
+        else:
+            target_device = device or _get_torch_device_str()
+            return _numpy_to_torch_tensor(
+                X,
+                device=target_device,
+                pin_memory=str(target_device).startswith("cuda"),
+            )
+
+        target_device = device or _get_torch_device_str()
+        return _numpy_to_torch_tensor(
+            X_np,
+            device=target_device,
+            pin_memory=str(target_device).startswith("cuda"),
+        )
+
+    def _to_cupy(self, X):
+        """Convert input to CuPy array."""
+        import cupy as cp
+
+        if isinstance(X, cp.ndarray):
+            return X
+
+        if hasattr(X, "cpu"):  # PyTorch
+            arr = _torch_to_cupy_dlpack(X)
+            if arr is not None:
+                return arr
+            X_np = X.detach().cpu().numpy()
+        elif hasattr(X, "get"):  # CuPy (shouldn't happen, but handle it)
+            X_np = X.get()
+        else:
+            X_np = np.asarray(X)
+
+        return cp.asarray(X_np)
+    
+    def _to_numpy(self, X) -> np.ndarray:
+        """Convert array back to numpy."""
+        if hasattr(X, 'get'):  # CuPy
+            return X.get()
+        elif hasattr(X, 'cpu'):  # PyTorch
+            return X.detach().cpu().numpy()
+        return np.asarray(X)
+
+    def adjust_pvalues(
+        self,
+        pvalues=None,
+        method: str = "bh",
+        alpha: float = 0.05,
+        axis: Optional[int] = 0,
+        backend: str = "auto",
+    ):
+        """
+        Adjust p-values for multiple testing (FDR/FWER controls).
+
+        Parameters
+        ----------
+        pvalues : array-like, optional
+            Raw p-values. If omitted, uses this estimator's ``_pvalues``.
+        method : str, default='bh'
+            Adjustment method: ``bh``, ``by``, ``holm``, ``bonferroni``
+            (aliases accepted).
+        alpha : float, default=0.05
+            Rejection threshold in (0, 1).
+        axis : int or None, default=0
+            Axis along which to adjust. ``None`` flattens all entries.
+        backend : {'auto', 'numpy', 'cupy'}, default='auto'
+            Compute backend. ``'auto'`` uses CuPy when estimator device is CUDA.
+
+        Returns
+        -------
+        dict
+            Contains ``pvalues``, ``pvalues_adjusted``, ``reject``,
+            ``method``, ``alpha``, and ``axis``.
+        """
+        from statgpu.inference import adjust_pvalues as _adjust_pvalues
+
+        source = pvalues
+        if source is None:
+            source = getattr(self, "_pvalues", None)
+        if source is None:
+            raise RuntimeError(
+                "No p-values available. Fit with inference enabled or pass pvalues explicitly."
+            )
+
+        backend_name = str(backend).strip().lower()
+        if backend_name == "auto" and self._get_compute_device() == Device.CUDA:
+            backend_name = "cupy"
+
+        if backend_name == "cupy":
+            pvals = self._to_array(source, Device.CUDA)
+        else:
+            pvals = self._to_numpy(source)
+
+        reject, pvals_adj = _adjust_pvalues(
+            pvals,
+            method=method,
+            alpha=alpha,
+            axis=axis,
+            backend=backend_name,
+        )
+
+        return {
+            "method": method,
+            "alpha": float(alpha),
+            "axis": axis,
+            "backend": backend_name,
+            "pvalues": pvals,
+            "pvalues_adjusted": pvals_adj,
+            "reject": reject,
+        }
+
+    def combine_pvalues(
+        self,
+        pvalues=None,
+        method: str = "fisher",
+        weights=None,
+        axis: Optional[int] = None,
+        backend: str = "auto",
+    ):
+        """
+        Combine p-values into a global p-value.
+
+        Parameters
+        ----------
+        pvalues : array-like, optional
+            Raw p-values. If omitted, uses this estimator's ``_pvalues``.
+        method : str, default='fisher'
+            Combination method: ``fisher`` or ``cauchy`` (aliases accepted).
+        weights : array-like, optional
+            Optional non-negative weights for cauchy combination.
+        axis : int or None, default=None
+            Axis along which to combine p-values. ``None`` flattens input.
+        backend : {'auto', 'numpy', 'cupy'}, default='auto'
+            Compute backend. ``'auto'`` uses CuPy when estimator device is CUDA.
+
+        Returns
+        -------
+        dict
+            Contains ``pvalues``, ``statistic``, ``pvalue``,
+            ``method``, ``axis``, and ``backend``.
+        """
+        from statgpu.inference import combine_pvalues as _combine_pvalues
+
+        source = pvalues
+        if source is None:
+            source = getattr(self, "_pvalues", None)
+        if source is None:
+            raise RuntimeError(
+                "No p-values available. Fit with inference enabled or pass pvalues explicitly."
+            )
+
+        backend_name = str(backend).strip().lower()
+        if backend_name == "auto" and self._get_compute_device() == Device.CUDA:
+            backend_name = "cupy"
+
+        if backend_name == "cupy":
+            pvals = self._to_array(source, Device.CUDA)
+            w_cast = None if weights is None else self._to_array(weights, Device.CUDA)
+        elif backend_name == "numpy":
+            pvals = self._to_numpy(source)
+            w_cast = None if weights is None else self._to_numpy(weights)
+        else:
+            pvals = source
+            w_cast = weights
+
+        statistic, pvalue = _combine_pvalues(
+            pvals,
+            method=method,
+            weights=w_cast,
+            axis=axis,
+            backend=backend_name,
+        )
+
+        return {
+            "method": method,
+            "axis": axis,
+            "backend": backend_name,
+            "pvalues": pvals,
+            "weights": w_cast,
+            "statistic": statistic,
+            "pvalue": pvalue,
+        }
+
+    def bootstrap_statistic(
+        self,
+        statistic,
+        *arrays,
+        n_resamples: int = 200,
+        strategy: str = "iid",
+        strata=None,
+        clusters=None,
+        block_size: Optional[int] = None,
+        confidence_level: float = 0.95,
+        random_state: Optional[int] = None,
+        statistic_name: str = "statistic",
+        backend: str = "auto",
+    ):
+        """
+        Run unified bootstrap engine from model context.
+
+        This is a thin wrapper over ``statgpu.inference.bootstrap_statistic``.
+        """
+        from statgpu.inference import bootstrap_statistic as _bootstrap_statistic
+
+        arrays_use = arrays
+        if len(arrays_use) == 0:
+            X_cache = getattr(self, "_X_design", None)
+            y_cache = getattr(self, "_y", None)
+            if X_cache is None or y_cache is None:
+                raise RuntimeError(
+                    "No cached training arrays available. Pass arrays explicitly or fit first."
+                )
+            arrays_use = (X_cache, y_cache)
+
+        backend_name = str(backend).strip().lower()
+        if backend_name == "auto" and self._get_compute_device() == Device.CUDA:
+            backend_name = "cupy"
+
+        if backend_name == "cupy":
+            arrays_cast = tuple(self._to_array(a, Device.CUDA) for a in arrays_use)
+            strata_cast = None if strata is None else self._to_array(strata, Device.CUDA)
+            clusters_cast = None if clusters is None else self._to_array(clusters, Device.CUDA)
+        elif backend_name == "numpy":
+            arrays_cast = tuple(self._to_numpy(a) for a in arrays_use)
+            strata_cast = None if strata is None else self._to_numpy(strata)
+            clusters_cast = None if clusters is None else self._to_numpy(clusters)
+        else:
+            arrays_cast = arrays_use
+            strata_cast = strata
+            clusters_cast = clusters
+
+        return _bootstrap_statistic(
+            statistic,
+            *arrays_cast,
+            n_resamples=n_resamples,
+            strategy=strategy,
+            strata=strata_cast,
+            clusters=clusters_cast,
+            block_size=block_size,
+            confidence_level=confidence_level,
+            random_state=random_state,
+            statistic_name=statistic_name,
+            backend=backend_name,
+        )
+
+    def permutation_test(
+        self,
+        statistic,
+        X,
+        y,
+        n_resamples: int = 1000,
+        strategy: str = "iid",
+        strata=None,
+        groups=None,
+        alternative: str = "two-sided",
+        random_state: Optional[int] = None,
+        statistic_name: str = "statistic",
+        backend: str = "auto",
+    ):
+        """
+        Run unified permutation test engine from model context.
+
+        This is a thin wrapper over ``statgpu.inference.permutation_test``.
+        """
+        from statgpu.inference import permutation_test as _permutation_test
+
+        backend_name = str(backend).strip().lower()
+        if backend_name == "auto" and self._get_compute_device() == Device.CUDA:
+            backend_name = "cupy"
+
+        if backend_name == "cupy":
+            X_cast = self._to_array(X, Device.CUDA)
+            y_cast = self._to_array(y, Device.CUDA)
+            strata_cast = None if strata is None else self._to_array(strata, Device.CUDA)
+            groups_cast = None if groups is None else self._to_array(groups, Device.CUDA)
+        elif backend_name == "numpy":
+            X_cast = self._to_numpy(X)
+            y_cast = self._to_numpy(y)
+            strata_cast = None if strata is None else self._to_numpy(strata)
+            groups_cast = None if groups is None else self._to_numpy(groups)
+        else:
+            X_cast = X
+            y_cast = y
+            strata_cast = strata
+            groups_cast = groups
+
+        return _permutation_test(
+            statistic,
+            X_cast,
+            y_cast,
+            n_resamples=n_resamples,
+            strategy=strategy,
+            strata=strata_cast,
+            groups=groups_cast,
+            alternative=alternative,
+            random_state=random_state,
+            statistic_name=statistic_name,
+            backend=backend_name,
+        )
+    
+    @abstractmethod
+    def fit(self, X, y=None, **fit_params):
+        """Fit the estimator."""
+        pass
+    
+    @abstractmethod
+    def predict(self, X):
+        """Make predictions."""
+        pass
+    
+    def _check_is_fitted(self):
+        """Check if estimator has been fitted."""
+        if not self._fitted:
+            raise RuntimeError(
+                f"This {self.__class__.__name__} instance is not fitted yet. "
+                "Call 'fit' before using this method."
+            )
+    
+    def get_params(self, deep=True):
+        """Get parameters for this estimator.
+
+        Only returns parameters accepted by this class's own ``__init__``,
+        not parent class parameters. This matches sklearn's contract where
+        ``clone(est).__init__(**est.get_params())`` must work.
+        """
+        import inspect
+        params = {}
+        # Only look at the most specific __init__ (this class, not parents)
+        try:
+            sig = inspect.signature(type(self).__init__)
+        except (ValueError, TypeError):
+            return params
+        for name in sig.parameters:
+            if name == "self":
+                continue
+            if hasattr(self, name):
+                params[name] = getattr(self, name)
+            elif hasattr(self, f'_{name}'):
+                params[name] = getattr(self, f'_{name}')
+        return params
+    
+    def set_params(self, **params):
+        """Set parameters for this estimator."""
+        for key, value in params.items():
+            if key == 'device':
+                self.device = Device(value) if isinstance(value, str) else value
+            else:
+                setattr(self, key, value)
+        return self

statgpu/_config.py

@@ -0,0 +1,127 @@
+"""
+Device configuration and GPU detection.
+"""
+
+import os
+import warnings
+from enum import Enum
+from typing import Optional, Union
+
+
+class Device(Enum):
+    """Device types for computation."""
+    CPU = "cpu"
+    CUDA = "cuda"
+    TORCH = "torch"
+    AUTO = "auto"
+
+
+class _DeviceManager:
+    """Internal device state manager."""
+
+    def __init__(self):
+        self._current_device = Device.AUTO
+        self._cupy_available = None
+        self._torch_available = None
+        self._cuda_available = None
+
+    def _check_cupy(self) -> bool:
+        """Check if CuPy is available and working."""
+        if self._cupy_available is None:
+            try:
+                import cupy as cp
+                # Test actual CUDA functionality
+                cp.cuda.Device(0).use()
+                self._cupy_available = True
+            except Exception:
+                self._cupy_available = False
+        return self._cupy_available
+
+    def _check_torch(self) -> bool:
+        """Check if PyTorch CUDA is available."""
+        if self._torch_available is None:
+            try:
+                import torch
+                self._torch_available = torch.cuda.is_available()
+            except Exception:
+                self._torch_available = False
+        return self._torch_available
+
+    def cuda_available(self) -> bool:
+        """Check if any CUDA backend is available (CuPy or Torch)."""
+        if self._cuda_available is None:
+            self._cuda_available = self._check_cupy() or self._check_torch()
+        return self._cuda_available
+
+    def get_device(self) -> Device:
+        """Get current device setting."""
+        if self._current_device == Device.AUTO:
+            if self.cuda_available():
+                # Prefer CuPy if both are available for backward compatibility
+                return Device.CUDA if self._check_cupy() else Device.TORCH
+            return Device.CPU
+        return self._current_device
+
+    def set_device(self, device: Union[str, Device]) -> None:
+        """Set device for computation."""
+        if isinstance(device, str):
+            device = Device(device.lower())
+
+        if device in (Device.CUDA, Device.TORCH) and not self.cuda_available():
+            warnings.warn(
+                "CUDA requested but not available. statgpu keeps the explicit "
+                "device setting and model execution will raise unless a matching "
+                "GPU backend is installed; use device='auto' for automatic CPU selection.",
+                RuntimeWarning
+            )
+
+        self._current_device = device
+
+
+# Global device manager instance
+_device_manager = _DeviceManager()
+
+
+def get_device() -> Device:
+    """
+    Get the current computation device.
+    
+    Returns
+    -------
+    Device
+        Current resolved device. If the configured device is ``'auto'``, this
+        resolves to CuPy CUDA when available, then Torch CUDA, then CPU.
+    """
+    return _device_manager.get_device()
+
+
+def set_device(device: Union[str, Device]) -> None:
+    """
+    Set the computation device.
+    
+    Parameters
+    ----------
+    device : str or Device
+        Device to use: ``'cpu'``, ``'cuda'``, ``'torch'``, or ``'auto'``.
+        ``'cuda'`` and ``'torch'`` are explicit GPU requests and are kept as
+        configured; model execution raises if the matching backend is not
+        available. ``'auto'`` chooses CuPy CUDA when available, then Torch CUDA,
+        then CPU.
+    
+    Examples
+    --------
+    >>> import statgpu as sg
+    >>> sg.set_device('cuda')  # Force CuPy CUDA
+    >>> sg.set_device('torch') # Force Torch CUDA
+    >>> sg.set_device('auto')  # Auto-detect
+    """
+    _device_manager.set_device(device)
+
+
+def cuda_available() -> bool:
+    """Check if CUDA is available."""
+    return _device_manager.cuda_available()
+
+
+# Convenience imports
+__all__ = ['Device', 'get_device', 'set_device', 'cuda_available']

statgpu/anova/__init__.py

@@ -0,0 +1,5 @@
+"""One-way ANOVA with GPU acceleration."""
+
+from ._oneway import f_oneway, AnovaResult
+
+__all__ = ['f_oneway', 'AnovaResult']