Refactor spline component

Author: jessegrabowski

pymc_experimental/model/modular/components.py

@@ -1,7 +1,8 @@
 from abc import ABC, abstractmethod
 
-import pandas as pd
+import numpy as np
 import pymc as pm
+import pytensor.tensor as pt
 
 from model.modular.utilities import (
     PRIOR_DEFAULT_KWARGS,
@@ -13,6 +14,7 @@
     select_data_columns,
 )
 from patsy import dmatrix
+from pytensor.graph import Apply, Op
 
 
 class GLMModel(ABC):
@@ -113,7 +115,7 @@ def __init__(
 
         self.prior = prior
         self.prior_params = prior_params if prior_params is not None else {}
-        self.pooling_columns = at_least_list(pooling_columns)
+        self.pooling_columns = pooling_columns
 
         name = name or f"Intercept(pooling_cols={pooling_columns})"
 
@@ -193,9 +195,9 @@ def __init__(
         prior_params:
             Additional keyword arguments to pass to the PyMC distribution specified by the prior argument.
         """
-        self.feature_columns = at_least_list(feature_columns)
+        self.feature_columns = feature_columns
         self.pooling = pooling
-        self.pooling_columns = at_least_list(pooling_columns)
+        self.pooling_columns = pooling_columns
 
         self.prior = prior
         self.prior_params = {} if prior_params is None else prior_params
@@ -210,7 +212,7 @@ def build(self, model=None):
         feature_dim = f"{self.name}_features"
 
         if feature_dim not in model.coords:
-            model.add_coord(feature_dim, self.feature_columns)
+            model.add_coord(feature_dim, at_least_list(self.feature_columns))
 
         with model:
             full_X = get_X_data(model)
@@ -237,17 +239,55 @@ def build(self, model=None):
             return regression_effect
 
 
-class Spline(Regression):
+class SplineTensor(Op):
+    def __init__(self, name, df=10, degree=3):
+        """
+        Thin wrapper around patsy dmatrix, allowing for the creation of spline basis functions given a symbolic input.
+
+        Parameters
+        ----------
+        name: str, optional
+            Name of the spline basis function.
+        df: int
+            Number of basis functions to generate
+        degree: int
+            Degree of the spline basis
+        """
+        self.name = name if name else ""
+        self.df = df
+        self.degree = degree
+
+    def make_node(self, x):
+        inputs = [pt.as_tensor(x)]
+        outputs = [pt.dmatrix(f"{self.name}_spline_basis")]
+
+        return Apply(self, inputs, outputs)
+
+    def perform(self, node: Apply, inputs: list[np.ndarray], outputs: list[list[None]]) -> None:
+        [x] = inputs
+
+        outputs[0][0] = np.asarray(
+            dmatrix(f"bs({self.name}, df={self.df}, degree={self.degree}) - 1", data={self.name: x})
+        )
+
+
+def pt_spline(x, name=None, df=10, degree=3) -> pt.Variable:
+    return SplineTensor(name=name, df=df, degree=degree)(x)
+
+
+class Spline(GLMModel):
     def __init__(
         self,
         name: str,
         *,
         feature_column: str | None = None,
         n_knots: int = 10,
-        prior: str = "Normal",
-        index_data: pd.Series | None = None,
+        spline_degree: int = 3,
         pooling: PoolingType = "complete",
-        **prior_params,
+        pooling_columns: ColumnType | None = None,
+        prior: str = "Normal",
+        prior_params: dict | None = None,
+        hierarchical_params: dict | None = None,
     ):
         """
         Class to represent a spline component in a GLM model.
@@ -263,25 +303,23 @@ def __init__(
         ----------
         name: str, optional
             Name of the intercept term. If None, a default name is generated based on the index_data.
-        n_knots: int, default 10
-            Number of knots to use in the spline basis.
         feature_column: str
             Column of the independent data to use in the spline.
-        index_data: Series or DataFrame, optional
-            Index data used to build hierarchical priors. If there are multiple columns, the columns are treated as
-            levels of a "telescoping" hierarchy, with the leftmost column representing the top level of the hierarchy,
-            and depth increasing to the right.
-
-            The index of the index_data must match the index of the observed data.
-        prior: str, optional
-            Name of the PyMC distribution to use for the intercept term. Default is "Normal".
+        n_knots: int, default 10
+            Number of knots to use in the spline basis.
+        spline_degree: int, default 3
+            Degree of the spline basis.
         pooling: str, one of ["none", "complete", "partial"], default "complete"
             Type of pooling to use for the intercept term. If "none", no pooling is applied, and each group in the
             index_data is treated as independent. If "complete", complete pooling is applied, and all data are treated
             as coming from the same group. If "partial", a hierarchical prior is constructed that shares information
             across groups in the index_data.
-        curve_type: str, one of ["log", "abc", "ns", "nss", "box-cox"]
-            Type of curve to build. For details, see the build_curve function.
+        pooling_columns: str or list of str, optional
+            Columns of the independent data to use as labels for pooling. These columns will be treated as categorical.
+            If None, no pooling is applied. If a list is provided, a "telescoping" hierarchy is constructed from left
+            to right, with the mean of each subsequent level centered on the mean of the previous level.
+        prior: str, optional
+            Name of the PyMC distribution to use for the intercept term. Default is "Normal".
         prior_params: dict, optional
             Additional keyword arguments to pass to the PyMC distribution specified by the prior argument.
         hierarchical_params: dict, optional
@@ -295,45 +333,49 @@ def __init__(
                 offset_dist: str, one of ["zerosum", "normal", "laplace"]
                     Name of the distribution to use for the offset distribution. Default is "zerosum"
         """
-        self.name = name if name else f"Spline({feature_column})"
         self.feature_column = feature_column
         self.n_knots = n_knots
+        self.spline_degree = spline_degree
+
         self.prior = prior
-        self.prior_params = prior_params
+        self.prior_params = {} if prior_params is None else prior_params
+        self.hierarchical_params = {} if hierarchical_params is None else hierarchical_params
 
+        self.pooling = pooling
+        self.pooling_columns = pooling_columns
+
+        name = name if name else f"Spline({feature_column}, df={n_knots}, degree={spline_degree})"
         super().__init__(name=name)
 
     def build(self, model: pm.Model | None = None):
         model = pm.modelcontext(model)
-        model.add_coord(f"{self.name}_spline", range(self.n_knots))
+        spline_dim = f"{self.name}_knots"
+        model.add_coord(spline_dim, range(self.n_knots))
 
         with model:
-            spline_data = {
-                self.feature_column: select_data_columns(
-                    get_X_data(model).get_value(), self.feature_column
-                )
-            }
-
-            X_spline = dmatrix(
-                f"bs({self.feature_column}, df={self.n_knots}, degree=3) - 1",
-                data=spline_data,
-                return_type="dataframe",
+            X_spline = pt_spline(
+                select_data_columns(self.feature_column, model),
+                name=self.feature_column,
+                df=self.n_knots,
+                degree=self.spline_degree,
             )
 
             if self.pooling == "complete":
-                beta = getattr(pm, self.prior)(
-                    f"{self.name}", **self.prior_params, dims=f"{self.feature_column}_spline"
-                )
+                prior_params = PRIOR_DEFAULT_KWARGS[self.prior].copy()
+                prior_params.update(self.prior_params)
+
+                beta = getattr(pm, self.prior)(f"{self.name}", **prior_params, dims=[spline_dim])
                 return X_spline @ beta
 
             elif self.pooling_columns is not None:
-                X = select_data_columns(self.pooling_columns, model)
                 beta = make_hierarchical_prior(
                     name=self.name,
-                    X=X,
+                    X=get_X_data(model),
+                    pooling=self.pooling,
+                    pooling_columns=self.pooling_columns,
                     model=model,
-                    dims=[f"{self.feature_column}_spline"],
-                    no_pooling=self.pooling == "none",
+                    dims=[spline_dim],
+                    **self.hierarchical_params,
                 )
 
             spline_effect = (X_spline * beta.T).sum(axis=-1)

pymc_experimental/model/modular/utilities.py

@@ -63,8 +63,7 @@ def select_data_columns(
     if cols is None:
         return
 
-    if isinstance(cols, str):
-        cols = [cols]
+    cols = at_least_list(cols)
 
     missing_cols = [col for col in cols if col not in model.coords["feature"]]
     if missing_cols:
@@ -407,6 +406,7 @@ def make_hierarchical_prior(
     **hierarchy_kwargs,
 ):
     model = pm.modelcontext(model)
+    pooling_columns = at_least_list(pooling_columns)
 
     if pooling == "none":
         return make_unpooled_hierarchy(

tests/model/modular/test_components.py

@@ -5,7 +5,7 @@
 
 from model.modular.utilities import at_least_list, encode_categoricals
 
-from pymc_experimental.model.modular.components import Intercept, PoolingType, Regression
+from pymc_experimental.model.modular.components import Intercept, PoolingType, Regression, Spline
 
 
 @pytest.fixture(scope="session")
@@ -91,3 +91,16 @@ def test_regression(pooling: PoolingType, prior, feature_columns, model):
         assert np.unique(beta_val).shape[0] == len(model.coords["city"]) * n_features
     else:
         assert np.unique(beta_val).shape[0] == n_features
+
+
+@pytest.mark.parametrize("pooling", ["partial", "none", "complete"], ids=str)
+@pytest.mark.parametrize("prior", ["Normal", "Laplace", "StudentT"], ids=str)
+def test_spline(pooling: PoolingType, prior, model):
+    spline = Spline(
+        name=None, feature_column="income", prior=prior, pooling=pooling, pooling_columns="city"
+    )
+
+    temp_model = model.copy()
+    xb = spline.build(temp_model)
+
+    assert "Spline(income, df=10, degree=3)_knots" in temp_model.coords.keys()