statgpu 0.1.0__py3-none-any.whl

statgpu/core/formula/_parser.py

@@ -0,0 +1,191 @@
+"""
+FormulaParser – R-style formula parser wrapping patsy.
+
+Provides the ``FormulaParser`` class that converts R-style formulas like
+``"y ~ x1 + x2 + C(sex)"`` into design matrices, using `patsy` internally.
+"""
+
+from typing import Optional, Tuple, List, Any
+
+import numpy as np
+import pandas as pd
+
+
+class FormulaParser:
+    """R-style formula parser that builds design matrices via patsy.
+
+    Parameters
+    ----------
+    formula : str
+        R-style formula string, e.g. ``"y ~ x1 + x2 + C(sex)"``.
+
+    Attributes
+    ----------
+    formula : str
+        The original formula string.
+    design_info : patsy.DesignInfo or None
+        Design matrix metadata (column names, term definitions).
+        Set after :meth:`eval` is called.
+    column_names : list[str] or None
+        Names of the predictor columns (excluding the response).
+        Set after :meth:`eval` is called.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> import numpy as np
+    >>> df = pd.DataFrame({
+    ...     "y": np.random.randn(100),
+    ...     "x1": np.random.randn(100),
+    ...     "x2": np.random.randn(100),
+    ... })
+    >>> parser = FormulaParser("y ~ x1 + x2")
+    >>> y, X, info = parser.eval(df)
+    >>> parser.column_names
+    ['x1', 'x2']
+    """
+
+    def __init__(self, formula: str):
+        self.formula = formula
+        self._design_info = None
+        self._y_names: Optional[List[str]] = None
+
+    @property
+    def design_info(self):
+        """Design matrix metadata, available after :meth:`eval`."""
+        return self._design_info
+
+    @property
+    def column_names(self) -> Optional[List[str]]:
+        """Predictor column names, available after :meth:`eval`."""
+        if self._design_info is None:
+            return None
+        return list(self._design_info.column_names)
+
+    def _require_patsy(self):
+        """Return patsy module or raise ImportError with guidance."""
+        try:
+            import patsy
+        except ImportError:
+            raise ImportError(
+                "The 'patsy' package is required for formula-based model fitting. "
+                "Install it with: pip install statgpu[formula] "
+                "or: pip install patsy"
+            )
+        return patsy
+
+    def eval(
+        self,
+        data: pd.DataFrame,
+        eval_env: int = 0,
+    ) -> Tuple[np.ndarray, np.ndarray, Any]:
+        """Parse formula and build design matrices from a DataFrame.
+
+        Parameters
+        ----------
+        data : pd.DataFrame
+            DataFrame containing the columns referenced in the formula.
+        eval_env : int, default=0
+            Evaluation frame depth for patsy name resolution.
+
+        Returns
+        -------
+        y : ndarray of shape (n_obs,) or (n_obs, n_responses)
+            Response variable(s).
+        X : ndarray of shape (n_obs, n_predictors)
+            Predictor design matrix.
+        design_info : patsy.DesignInfo
+            Metadata for the predictor design (column names, term info).
+        """
+        patsy = self._require_patsy()
+        data = data.copy()
+
+        y, X = patsy.dmatrices(
+            self.formula,
+            data,
+            eval_env=eval_env + 1,
+            return_type="matrix",
+        )
+
+        self._y_names = list(y.design_info.column_names)
+        self._design_info = X.design_info
+
+        y_arr = np.asarray(y)
+        if y_arr.ndim == 2 and y_arr.shape[1] == 1:
+            y_arr = y_arr.ravel()
+        X_arr = np.asarray(X)
+
+        return y_arr, X_arr, X.design_info
+
+    def transform(
+        self,
+        new_data: pd.DataFrame,
+        eval_env: int = 0,
+    ) -> np.ndarray:
+        """Build a design matrix for new data using the stored design_info.
+
+        Used during :meth:`predict` to ensure new data is encoded
+        with the same column structure (including categorical coding)
+        as the training data.
+
+        Parameters
+        ----------
+        new_data : pd.DataFrame
+            DataFrame with the same columns as the training data.
+        eval_env : int, default=0
+            Evaluation frame depth for patsy name resolution.
+
+        Returns
+        -------
+        X_new : ndarray of shape (n_new_obs, n_predictors)
+            Design matrix aligned with the training design.
+
+        Raises
+        ------
+        RuntimeError
+            If :meth:`eval` has not been called yet (no design_info available).
+        ValueError
+            If new_data has columns that don't match the training structure.
+        """
+        if self._design_info is None:
+            raise RuntimeError(
+                "Cannot transform: no design_info available. "
+                "Call eval() first on training data."
+            )
+
+        patsy = self._require_patsy()
+
+        X_new = patsy.build_design_matrices(
+            [self._design_info],
+            new_data,
+            return_type="matrix",
+        )[0]
+
+        return np.asarray(X_new)
+
+    def summary(self) -> str:
+        """Return a human-readable summary of the formula parsing.
+
+        Shows the formula string, response variables, predictor names,
+        and term definitions (useful for debugging categorical encoding).
+        """
+        lines = [f"Formula: {self.formula}"]
+
+        if self._design_info is None:
+            lines.append("(Not yet evaluated. Call eval() to parse.)")
+            return "\n".join(lines)
+
+        lines.append(f"Response: {self._y_names}")
+        lines.append(f"Predictors ({len(self.column_names)}):")
+        for name in self.column_names:
+            lines.append(f"  - {name}")
+
+        lines.append("\nTerms:")
+        for term in self._design_info.term_name_slices.keys():
+            lines.append(f"  {term}")
+
+        return "\n".join(lines)
+
+    def __repr__(self) -> str:
+        evaluated = "evaluated" if self._design_info is not None else "pending"
+        return f"FormulaParser({self.formula!r}, {evaluated})"

statgpu/core/formula/_terms.py

@@ -0,0 +1,70 @@
+"""
+Formula term helpers and custom evaluation environments.
+
+Patsy natively supports R-style formula terms:
+
+- ``C(var)`` — treat as categorical (one-hot encoding)
+- ``np.func(var)`` — apply numpy function (e.g. ``np.log(x)``)
+- ``x1:x2`` — interaction only
+- ``x1*x2`` — main effects + interaction
+- ``x1 + x2`` — additive
+- ``x1 + x2 - 1`` — additive without intercept
+- ``np.log(x)`` — transformations
+
+This module provides helper functions for constructing custom
+patsy evaluation environments, needed for model-specific syntax
+like ``Surv(time, event)`` in Cox PH models.
+"""
+
+from typing import Dict, Any, Optional
+
+import numpy as np
+
+
+def _surv(time, event):
+    """Survival function for patsy formula parsing.
+
+    Mimics R's survival::Surv() function for use in patsy formulas::
+
+        "Surv(time, event) ~ x1 + x2"
+
+    Parameters
+    ----------
+    time : array-like
+        Survival/follow-up times.
+    event : array-like
+        Event indicator (1 = event occurred, 0 = censored).
+
+    Returns
+    -------
+    result : ndarray of shape (n, 2)
+        Column 0: time, Column 1: event.
+    """
+    time = np.asarray(time, dtype=np.float64).ravel()
+    event = np.asarray(event, dtype=np.float64).ravel()
+
+    if len(time) != len(event):
+        raise ValueError(
+            f"time ({len(time)} elements) and event ({len(event)} elements) "
+            "must have the same length."
+        )
+
+    return np.column_stack([time, event])
+
+
+def make_surv_env() -> Dict[str, Any]:
+    """Create a patsy evaluation environment with ``Surv`` function.
+
+    Returns
+    -------
+    env : dict
+        Custom functions for patsy's ``EvalEnvironment``.
+
+    Examples
+    --------
+    >>> from statgpu.core.formula._terms import make_surv_env
+    >>> import patsy
+    >>> env = make_surv_env()
+    >>> # Then pass env to patsy.dmatrices or dmatrix
+    """
+    return {"Surv": _surv}

statgpu/core/formula/tests/test_parser.py

@@ -0,0 +1,194 @@
+"""
+Tests for statgpu.core.formula module.
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+
+@pytest.fixture
+def sample_df():
+    """Standard test DataFrame."""
+    np.random.seed(42)
+    n = 200
+    return pd.DataFrame({
+        "y": np.random.randn(n),
+        "x1": np.random.randn(n),
+        "x2": np.random.randn(n),
+        "cat": pd.Categorical(np.random.choice(["A", "B", "C"], n)),
+    })
+
+
+class TestFormulaParserBasic:
+    """Test basic formula parsing."""
+
+    def test_simple_formula(self, sample_df):
+        """Test simple y ~ x1 + x2 formula."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + x2")
+        y, X, info = parser.eval(sample_df)
+
+        assert y.shape == (200,)
+        assert X.shape == (200, 3)  # intercept + x1 + x2
+        assert parser.column_names == ["Intercept", "x1", "x2"]
+
+    def test_no_intercept(self, sample_df):
+        """Test y ~ x1 + x2 - 1 (no intercept)."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + x2 - 1")
+        y, X, info = parser.eval(sample_df)
+
+        assert X.shape == (200, 2)
+        assert parser.column_names == ["x1", "x2"]
+
+    def test_categorical_encoding(self, sample_df):
+        """Test C() for categorical variables."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + C(cat)")
+        y, X, info = parser.eval(sample_df)
+
+        # Intercept + x1 + cat[T.B] + cat[T.C] = 4 columns
+        assert X.shape[1] == 4
+        assert "x1" in parser.column_names
+        assert any("cat" in name for name in parser.column_names)
+
+    def test_interaction(self, sample_df):
+        """Test x1:x2 interaction."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + x2 + x1:x2")
+        y, X, info = parser.eval(sample_df)
+
+        assert X.shape[1] == 4  # intercept + x1 + x2 + x1:x2
+
+    def test_star_operator(self, sample_df):
+        """Test x1*x2 (main effects + interaction)."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 * x2")
+        y, X, info = parser.eval(sample_df)
+
+        assert X.shape[1] == 4  # intercept + x1 + x2 + x1:x2
+
+    def test_transform(self, sample_df):
+        """Test np() transformations."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ np.log(np.abs(x1)) + x2")
+        y, X, info = parser.eval(sample_df)
+
+        assert y.shape == (200,)
+        assert X.shape[1] == 3  # intercept + transformed_x1 + x2
+
+
+class TestFormulaParserTransform:
+    """Test transform (predict-time) functionality."""
+
+    def test_transform_new_data(self, sample_df):
+        """Test transform on new data with same structure."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + x2")
+        parser.eval(sample_df)
+
+        new_data = pd.DataFrame({
+            "x1": [0.5, -0.3],
+            "x2": [1.2, 0.8],
+        })
+        X_new = parser.transform(new_data)
+
+        assert X_new.shape == (2, 3)  # 2 rows, intercept + 2 cols
+
+    def test_transform_with_categorical(self, sample_df):
+        """Test transform handles categorical encoding from training."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + C(cat)")
+        parser.eval(sample_df)
+
+        new_data = pd.DataFrame({
+            "x1": [0.5],
+            "cat": pd.Categorical(["A"]),
+        })
+        X_new = parser.transform(new_data)
+
+        assert X_new.shape == (1, 4)  # intercept + x1 + cat[B] + cat[C]
+
+    def test_transform_no_design_info(self):
+        """Test transform raises when not yet evaluated."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1")
+        new_data = pd.DataFrame({"x1": [1.0]})
+
+        with pytest.raises(RuntimeError, match="no design_info available"):
+            parser.transform(new_data)
+
+
+class TestParseFormulaSafe:
+    """Test parse_formula_safe fallback logic."""
+
+    def test_formula_path(self, sample_df):
+        """Test formula path works."""
+        from statgpu.core.formula import parse_formula_safe
+
+        y, X, info = parse_formula_safe("y ~ x1", data=sample_df)
+        assert y.shape == (200,)
+        assert info is not None
+
+    def test_array_path(self, sample_df):
+        """Test array path when formula is None."""
+        from statgpu.core.formula import parse_formula_safe
+
+        X = sample_df[["x1", "x2"]].values
+        y = sample_df["y"].values
+        y_out, X_out, info = parse_formula_safe(None, None, X=X, y=y)
+
+        assert info is None
+        np.testing.assert_array_equal(y_out, y)
+        np.testing.assert_array_equal(X_out, X)
+
+    def test_formula_without_data_raises(self):
+        """Test that formula without data raises."""
+        from statgpu.core.formula import parse_formula_safe
+
+        with pytest.raises(ValueError, match="data"):
+            parse_formula_safe("y ~ x1", None)
+
+    def test_no_input_raises(self):
+        """Test that no input raises."""
+        from statgpu.core.formula import parse_formula_safe
+
+        with pytest.raises(ValueError, match="Either formula"):
+            parse_formula_safe(None, None)
+
+
+class TestFormulaParserSummary:
+    """Test FormulaParser.summary() output."""
+
+    def test_summary_before_eval(self, sample_df):
+        """Test summary shows pending state."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + x2")
+        s = parser.summary()
+
+        assert "y ~ x1 + x2" in s
+        assert "pending" in s.lower() or "Not yet evaluated" in s
+
+    def test_summary_after_eval(self, sample_df):
+        """Test summary shows parsed info."""
+        from statgpu.core.formula import FormulaParser
+
+        parser = FormulaParser("y ~ x1 + x2")
+        parser.eval(sample_df)
+        s = parser.summary()
+
+        assert "y ~ x1 + x2" in s
+        assert "x1" in s
+        assert "x2" in s
+        assert "Predictors (3)" in s

statgpu/covariance/__init__.py

@@ -0,0 +1,6 @@
+"""Covariance estimation with GPU acceleration."""
+
+from ._empirical import EmpiricalCovariance
+from ._shrinkage import LedoitWolf, OAS
+
+__all__ = ["EmpiricalCovariance", "LedoitWolf", "OAS"]