feat(spatial.models): design-matrix helper (formulaic-backed)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

src/statspai/spatial/models/_base.py

@@ -0,0 +1,31 @@
+"""Shared plumbing for spatial regression models."""
+from __future__ import annotations
+
+from typing import List, Tuple
+
+import numpy as np
+import pandas as pd
+
+try:
+    import formulaic
+except ImportError as e:                     # pragma: no cover
+    raise ImportError(
+        "formulaic is required for spatial regression formulas. "
+        "It is already a core StatsPAI dependency — "
+        "please reinstall with `pip install statspai`."
+    ) from e
+
+
+def build_design_matrix(
+    formula: str, data: pd.DataFrame
+) -> Tuple[np.ndarray, np.ndarray, List[str]]:
+    """Parse `y ~ x1 + x2` into (y, X, column_names).
+
+    Column names come from the formulaic model matrix header and may include
+    "Intercept" (when the intercept is present).
+    """
+    y_mat, X_mat = formulaic.Formula(formula).get_model_matrix(data)
+    y = np.asarray(y_mat).ravel().astype(float)
+    X = np.asarray(X_mat, dtype=float)
+    names = list(X_mat.columns)
+    return y, X, names

tests/spatial/test_models_base.py

@@ -0,0 +1,19 @@
+import numpy as np
+import pandas as pd
+from statspai.spatial.models._base import build_design_matrix
+
+
+def test_build_design_matrix_adds_intercept():
+    df = pd.DataFrame({"y": [1.0, 2.0, 3.0], "x1": [4.0, 5.0, 6.0]})
+    y, X, names = build_design_matrix("y ~ x1", df)
+    assert y.shape == (3,)
+    assert X.shape == (3, 2)
+    assert names[0] in {"Intercept", "(Intercept)", "1"}  # formulaic convention
+    assert "x1" in names
+
+
+def test_build_design_matrix_no_intercept():
+    df = pd.DataFrame({"y": [1.0, 2.0, 3.0], "x1": [4.0, 5.0, 6.0]})
+    y, X, names = build_design_matrix("y ~ x1 - 1", df)
+    assert X.shape == (3, 1)
+    assert names == ["x1"]