subset-mixture-model 0.1.0__tar.gz

subset_mixture_model-0.1.0/PKG-INFO

@@ -0,0 +1,152 @@
+Metadata-Version: 2.4
+Name: subset-mixture-model
+Version: 0.1.0
+Summary: Interpretable empirical-Bayes aggregation of partition estimators for categorical regression
+Author-email: Aaron John Danielson <aaron.danielson@austin.utexas.edu>
+License: MIT
+Project-URL: Repository, https://github.com/aaronjdanielson/subset-mixture-model
+Keywords: interpretable machine learning,categorical features,empirical Bayes,uncertainty quantification,mixture model
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: numpy>=1.24
+Requires-Dist: pandas>=2.0
+Requires-Dist: scikit-learn>=1.3
+Requires-Dist: scipy>=1.11
+Provides-Extra: experiments
+Requires-Dist: matplotlib>=3.7; extra == "experiments"
+Requires-Dist: seaborn>=0.12; extra == "experiments"
+Requires-Dist: joblib>=1.3; extra == "experiments"
+Requires-Dist: lightgbm>=4.0; extra == "experiments"
+Requires-Dist: ngboost>=0.4; extra == "experiments"
+Requires-Dist: mapie>=0.6; extra == "experiments"
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+
+# Subset Mixture Model (SMM)
+
+**SMM** is an interpretable, empirical-Bayes method for regression on datasets with categorical features. It aggregates partition-based conditional-mean estimators over all non-empty feature subsets using learned simplex weights, adaptively balancing bias and variance across partition granularities.
+
+## Key idea
+
+Each feature subset $s$ induces a partition of the covariate space and a natural estimator of the conditional expectation — its empirical cell mean. SMM learns a convex combination of these estimators:
+
+$$\hat{f}(\mathbf{x}) = \sum_{s \in \mathcal{S}} \hat{\pi}_s \cdot \hat{\mu}_{m(s,\mathbf{x})}(s)$$
+
+The learned weights $\hat{\pi}_s$ are directly interpretable: they reveal which feature interactions drive predictions on average. Uncertainty is propagated from the MAP weight estimates via a Laplace approximation, yielding aleatoric/epistemic decompositions without post-hoc calibration.
+
+## Installation
+
+```bash
+pip install subset-mixture-model
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/aaronjdanielson/subset-mixture-model
+cd subset-mixture-model
+pip install -e .
+```
+
+## Quick start
+
+```python
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from smm import (
+    SubsetMaker, SubsetWeightsModel, SubsetDataset,
+    subset_mixture_neg_log_posterior, SubsetMixturePredictor,
+    compute_posterior_covariance, predict_with_uncertainty,
+)
+
+# --- your data (integer-coded categorical features) ---
+train_df = pd.read_csv("train.csv")
+val_df   = pd.read_csv("val.csv")
+test_df  = pd.read_csv("test.csv")
+
+cat_cols = ["feature_a", "feature_b", "feature_c"]
+target   = "y"
+
+# --- build lookup table ---
+subset_maker = SubsetMaker(train_df, cat_cols, [target])
+n_subsets = len(subset_maker.lookup)
+
+# --- train ---
+model     = SubsetWeightsModel(n_subsets)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
+loader    = DataLoader(SubsetDataset(train_df, cat_cols, [target]),
+                       batch_size=64, shuffle=True)
+
+for epoch in range(100):
+    for x, y in loader:
+        optimizer.zero_grad()
+        mus, variances, mask = subset_maker.batch_lookup(x)
+        loss = subset_mixture_neg_log_posterior(
+            model(), y, mus, variances, mask, alpha=1.1)
+        loss.backward()
+        optimizer.step()
+
+# --- predict with uncertainty ---
+pi_hat    = F.softmax(model.eta.detach(), dim=0)
+predictor = SubsetMixturePredictor(subset_maker, pi_hat)
+sigma_pi  = compute_posterior_covariance(
+    model, subset_maker, train_df, cat_cols, target, alpha=1.1)
+
+y_mean, y_std = predict_with_uncertainty(predictor, sigma_pi, test_df)
+# y_mean: point predictions
+# y_std:  total predictive standard deviation (aleatoric + epistemic)
+```
+
+## Interpretability
+
+```python
+import numpy as np
+
+subsets = list(subset_maker.lookup.keys())
+top_idx = np.argsort(pi_hat.numpy())[::-1][:10]
+
+for rank, i in enumerate(top_idx):
+    print(f"{rank+1:2d}. {subsets[i]}  π={pi_hat[i]:.4f}")
+```
+
+## Features
+
+- **Interpretable by construction**: learned weights reveal which feature interactions matter
+- **Principled uncertainty**: aleatoric/epistemic decomposition via Laplace approximation
+- **Efficient training**: only $2^D - 1$ logits optimized; lookup table precomputed once
+- **No post-hoc calibration**: well-calibrated predictive intervals out of the box
+- **Scalable to D ≤ 15** features; $k$-way truncation available for larger $D$
+
+## Datasets supported
+
+Any tabular dataset with integer-coded (or string, with encoding) categorical features and a continuous target.
+
+## Citation
+
+```bibtex
+@article{danielson2025smm,
+  title   = {Subset Mixture Model: Interpretable Empirical-Bayes Aggregation
+             of Partition Estimators for Categorical Regression},
+  author  = {Danielson, Aaron John},
+  journal = {Machine Learning},
+  year    = {2025},
+  note    = {Under review}
+}
+```
+
+## License
+
+MIT

subset_mixture_model-0.1.0/README.md

@@ -0,0 +1,118 @@
+# Subset Mixture Model (SMM)
+
+**SMM** is an interpretable, empirical-Bayes method for regression on datasets with categorical features. It aggregates partition-based conditional-mean estimators over all non-empty feature subsets using learned simplex weights, adaptively balancing bias and variance across partition granularities.
+
+## Key idea
+
+Each feature subset $s$ induces a partition of the covariate space and a natural estimator of the conditional expectation — its empirical cell mean. SMM learns a convex combination of these estimators:
+
+$$\hat{f}(\mathbf{x}) = \sum_{s \in \mathcal{S}} \hat{\pi}_s \cdot \hat{\mu}_{m(s,\mathbf{x})}(s)$$
+
+The learned weights $\hat{\pi}_s$ are directly interpretable: they reveal which feature interactions drive predictions on average. Uncertainty is propagated from the MAP weight estimates via a Laplace approximation, yielding aleatoric/epistemic decompositions without post-hoc calibration.
+
+## Installation
+
+```bash
+pip install subset-mixture-model
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/aaronjdanielson/subset-mixture-model
+cd subset-mixture-model
+pip install -e .
+```
+
+## Quick start
+
+```python
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from smm import (
+    SubsetMaker, SubsetWeightsModel, SubsetDataset,
+    subset_mixture_neg_log_posterior, SubsetMixturePredictor,
+    compute_posterior_covariance, predict_with_uncertainty,
+)
+
+# --- your data (integer-coded categorical features) ---
+train_df = pd.read_csv("train.csv")
+val_df   = pd.read_csv("val.csv")
+test_df  = pd.read_csv("test.csv")
+
+cat_cols = ["feature_a", "feature_b", "feature_c"]
+target   = "y"
+
+# --- build lookup table ---
+subset_maker = SubsetMaker(train_df, cat_cols, [target])
+n_subsets = len(subset_maker.lookup)
+
+# --- train ---
+model     = SubsetWeightsModel(n_subsets)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
+loader    = DataLoader(SubsetDataset(train_df, cat_cols, [target]),
+                       batch_size=64, shuffle=True)
+
+for epoch in range(100):
+    for x, y in loader:
+        optimizer.zero_grad()
+        mus, variances, mask = subset_maker.batch_lookup(x)
+        loss = subset_mixture_neg_log_posterior(
+            model(), y, mus, variances, mask, alpha=1.1)
+        loss.backward()
+        optimizer.step()
+
+# --- predict with uncertainty ---
+pi_hat    = F.softmax(model.eta.detach(), dim=0)
+predictor = SubsetMixturePredictor(subset_maker, pi_hat)
+sigma_pi  = compute_posterior_covariance(
+    model, subset_maker, train_df, cat_cols, target, alpha=1.1)
+
+y_mean, y_std = predict_with_uncertainty(predictor, sigma_pi, test_df)
+# y_mean: point predictions
+# y_std:  total predictive standard deviation (aleatoric + epistemic)
+```
+
+## Interpretability
+
+```python
+import numpy as np
+
+subsets = list(subset_maker.lookup.keys())
+top_idx = np.argsort(pi_hat.numpy())[::-1][:10]
+
+for rank, i in enumerate(top_idx):
+    print(f"{rank+1:2d}. {subsets[i]}  π={pi_hat[i]:.4f}")
+```
+
+## Features
+
+- **Interpretable by construction**: learned weights reveal which feature interactions matter
+- **Principled uncertainty**: aleatoric/epistemic decomposition via Laplace approximation
+- **Efficient training**: only $2^D - 1$ logits optimized; lookup table precomputed once
+- **No post-hoc calibration**: well-calibrated predictive intervals out of the box
+- **Scalable to D ≤ 15** features; $k$-way truncation available for larger $D$
+
+## Datasets supported
+
+Any tabular dataset with integer-coded (or string, with encoding) categorical features and a continuous target.
+
+## Citation
+
+```bibtex
+@article{danielson2025smm,
+  title   = {Subset Mixture Model: Interpretable Empirical-Bayes Aggregation
+             of Partition Estimators for Categorical Regression},
+  author  = {Danielson, Aaron John},
+  journal = {Machine Learning},
+  year    = {2025},
+  note    = {Under review}
+}
+```
+
+## License
+
+MIT

subset_mixture_model-0.1.0/pyproject.toml

@@ -0,0 +1,59 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "subset-mixture-model"
+version = "0.1.0"
+description = "Interpretable empirical-Bayes aggregation of partition estimators for categorical regression"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [
+  { name = "Aaron John Danielson", email = "aaron.danielson@austin.utexas.edu" },
+]
+keywords = [
+  "interpretable machine learning",
+  "categorical features",
+  "empirical Bayes",
+  "uncertainty quantification",
+  "mixture model",
+]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+requires-python = ">=3.9"
+dependencies = [
+  "torch>=2.0",
+  "numpy>=1.24",
+  "pandas>=2.0",
+  "scikit-learn>=1.3",
+  "scipy>=1.11",
+]
+
+[project.optional-dependencies]
+experiments = [
+  "matplotlib>=3.7",
+  "seaborn>=0.12",
+  "joblib>=1.3",
+  "lightgbm>=4.0",
+  "ngboost>=0.4",
+  "mapie>=0.6",
+]
+dev = [
+  "pytest>=7",
+  "pytest-cov",
+]
+
+[project.urls]
+Repository = "https://github.com/aaronjdanielson/subset-mixture-model"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["smm*"]

subset_mixture_model-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

subset_mixture_model-0.1.0/smm/__init__.py

@@ -0,0 +1,20 @@
+from .subset_maker import SubsetMaker
+from .model import SubsetWeightsModel, SubsetDataset, subset_mixture_neg_log_posterior, subset_mixture_mse
+from .predictor import SubsetMixturePredictor
+from .laplace import (
+    compute_posterior_covariance,
+    predict_with_uncertainty,
+    coverage,
+)
+
+__all__ = [
+    "SubsetMaker",
+    "SubsetWeightsModel",
+    "SubsetDataset",
+    "subset_mixture_neg_log_posterior",
+    "subset_mixture_mse",
+    "SubsetMixturePredictor",
+    "compute_posterior_covariance",
+    "predict_with_uncertainty",
+    "coverage",
+]

subset_mixture_model-0.1.0/smm/laplace.py

@@ -0,0 +1,266 @@
+"""
+Laplace approximation for posterior uncertainty in the Subset Mixture Model.
+
+After MAP estimation of eta (the logit weight vector), the Laplace approximation
+treats the negative log-posterior as a quadratic around the MAP:
+
+    q(eta) = N(eta | eta_hat, H^{-1})
+
+where H is the Hessian of the negative log-posterior at eta_hat.
+
+Transforming to the simplex via pi = softmax(eta) and applying the delta method:
+
+    Cov[pi | D] ≈ Sigma_pi = J H^{-1} J^T
+
+where J = d(pi)/d(eta) is the |S| x |S| softmax Jacobian.
+
+Predictive variance for a new point x_tilde then decomposes as:
+
+    Var[y | x_tilde] ≈  sum_s pi_s * sigma^2_s(x_tilde)        [aleatoric]
+                      +  mu_{x_tilde}^T Sigma_pi mu_{x_tilde}   [epistemic]
+
+References:
+    MacKay, D.J.C. (1992). A Practical Bayesian Framework for Backpropagation Networks.
+    Daxberger et al. (2021). Laplace Redux -- Effortless Bayesian Deep Learning.
+"""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from .model import SubsetDataset, subset_mixture_neg_log_posterior
+from .subset_maker import SubsetMaker
+from .predictor import SubsetMixturePredictor
+
+
+# ---------------------------------------------------------------------------
+# Step 1: Compute posterior covariance of pi
+# ---------------------------------------------------------------------------
+
+def softmax_jacobian(pi: torch.Tensor) -> torch.Tensor:
+    """
+    Jacobian of softmax(eta) w.r.t. eta, evaluated at pi = softmax(eta).
+
+    J_ij = pi_i * (delta_ij - pi_j)
+
+    Args:
+        pi: [S] mixture weights (must sum to 1).
+
+    Returns:
+        J: [S, S] Jacobian matrix.
+    """
+    return torch.diag(pi) - torch.outer(pi, pi)
+
+
+def compute_hessian(
+    model: torch.nn.Module,
+    subset_maker: SubsetMaker,
+    train_df,
+    cat_cols: list,
+    target: str,
+    alpha: float = 1.1,
+    batch_size: int = 512,
+) -> torch.Tensor:
+    """
+    Compute the exact Hessian of the negative log-posterior w.r.t. eta.
+
+    Because |S| is at most a few hundred, the Hessian is small and can be
+    computed exactly via torch.autograd.functional.hessian on the full
+    training set (loaded in batches if needed, accumulated via the empirical
+    Fisher for very large datasets).
+
+    For datasets where n is large we accumulate an approximation via the
+    diagonal outer-product (empirical Fisher). For the exact Hessian we need
+    the full batch, which is feasible when n < ~20K.
+
+    Args:
+        model:        Trained SubsetWeightsModel.
+        subset_maker: Fitted SubsetMaker.
+        train_df:     Training DataFrame (integer-encoded).
+        cat_cols:     Categorical feature column names.
+        target:       Target column name.
+        alpha:        Dirichlet concentration (must match training value).
+        batch_size:   Batch size for loading data (only affects memory).
+
+    Returns:
+        H: [S, S] Hessian tensor.
+    """
+    model.eval()
+    eta_hat = model.eta.detach().clone()
+
+    # Load all training data (full batch for exact Hessian)
+    ds = SubsetDataset(train_df, cat_cols, [target])
+    loader = DataLoader(ds, batch_size=len(ds), shuffle=False)
+    x_all, y_all = next(iter(loader))
+
+    mus, variances, mask = subset_maker.batch_lookup(x_all)
+    # Detach and keep on CPU
+    mus      = mus.detach()
+    variances = variances.detach()
+    mask     = mask.detach()
+
+    def loss_fn(eta):
+        return subset_mixture_neg_log_posterior(
+            eta, y_all, mus, variances, mask, alpha=alpha
+        )
+
+    eta_param = eta_hat.requires_grad_(True)
+    H = torch.autograd.functional.hessian(loss_fn, eta_param)  # [S, S]
+    return H.detach()
+
+
+def compute_posterior_covariance(
+    model: torch.nn.Module,
+    subset_maker: SubsetMaker,
+    train_df,
+    cat_cols: list,
+    target: str,
+    alpha: float = 1.1,
+    hessian_reg: float = 1e-4,
+) -> torch.Tensor:
+    """
+    Compute the Laplace approximation to the posterior covariance of pi.
+
+    Sigma_pi = J H^{-1} J^T
+
+    Args:
+        model:        Trained SubsetWeightsModel (provides eta_hat).
+        subset_maker: Fitted SubsetMaker.
+        train_df:     Training DataFrame (integer-encoded).
+        cat_cols:     Categorical feature column names.
+        target:       Target column name.
+        alpha:        Dirichlet concentration (must match training value).
+        hessian_reg:  Ridge added to H before inversion for numerical stability.
+
+    Returns:
+        sigma_pi: [S, S] posterior covariance of mixture weights.
+    """
+    eta_hat = model.eta.detach()
+    pi_hat  = F.softmax(eta_hat, dim=0)
+
+    S = len(eta_hat)
+    H = compute_hessian(model, subset_maker, train_df, cat_cols, target, alpha)
+
+    # Regularise to ensure positive definiteness
+    H_reg  = H + hessian_reg * torch.eye(S)
+    H_inv  = torch.linalg.inv(H_reg)
+
+    # Softmax Jacobian at MAP estimate
+    J = softmax_jacobian(pi_hat)                # [S, S]
+    sigma_pi = J @ H_inv @ J.T                  # [S, S]
+
+    return sigma_pi
+
+
+# ---------------------------------------------------------------------------
+# Step 2: Predictive distribution for new points
+# ---------------------------------------------------------------------------
+
+def predict_with_uncertainty(
+    predictor: SubsetMixturePredictor,
+    sigma_pi: torch.Tensor,
+    df,
+    return_components: bool = False,
+):
+    """
+    Compute predictive mean and uncertainty for a batch of test points.
+
+    Predictive mean:
+        y_hat(x) = pi_hat^T mu_x
+
+    Predictive variance:
+        Var[y|x] = sum_s pi_s * sigma_s^2(x)          [aleatoric]
+                 + mu_x^T Sigma_pi mu_x                 [epistemic]
+
+    For test points with no valid subset cell (all masked), the predictor
+    falls back to the global training mean with aleatoric variance equal to
+    the global training variance.
+
+    Args:
+        predictor:  SubsetMixturePredictor with learned weight vector.
+        sigma_pi:   [S, S] posterior covariance from compute_posterior_covariance.
+        df:         Test DataFrame (integer-encoded).
+        return_components: If True, also return aleatoric and epistemic stds.
+
+    Returns:
+        y_mean (np.ndarray):  [B] predicted means.
+        y_std  (np.ndarray):  [B] total predictive standard deviations.
+        (optional) aleatoric_std (np.ndarray): [B]
+        (optional) epistemic_std (np.ndarray): [B]
+    """
+    batch_tensor = torch.tensor(
+        df[predictor.subset_maker.subset_features].astype(float).values,
+        dtype=torch.float32,
+    )
+    mus, variances, mask = predictor.subset_maker.batch_lookup(batch_tensor)
+
+    B, S = mus.shape
+    pi    = predictor.weight_vector                         # [S]
+    weights = pi.unsqueeze(0).expand(B, S).clone()         # [B, S]
+
+    # Mask invalid cells
+    weights   = weights.masked_fill(~mask, 0.0)
+    mus_m     = mus.masked_fill(~mask, 0.0)
+    vars_m    = variances.masked_fill(~mask, 0.0)
+
+    weight_sums  = weights.sum(dim=1, keepdim=True)        # [B, 1]
+    norm_weights = weights / (weight_sums + 1e-9)          # [B, S]
+
+    # --- Predictive mean ---
+    fallback_mask = weight_sums.squeeze() < 1e-6           # [B]
+    y_mean_weighted = (mus_m * norm_weights).sum(dim=1)    # [B]
+    y_mean = torch.where(
+        fallback_mask,
+        torch.full_like(y_mean_weighted, predictor.subset_maker.fallback_mean),
+        y_mean_weighted,
+    )
+
+    # --- Aleatoric variance: E_pi[sigma^2(x)] = sum_s pi_s * sigma_s^2(x) ---
+    aleatoric_var = (norm_weights * vars_m).sum(dim=1)     # [B]
+    # Fallback: use global training variance
+    aleatoric_var = torch.where(
+        fallback_mask,
+        torch.full_like(aleatoric_var, predictor.subset_maker.fallback_var),
+        aleatoric_var,
+    )
+
+    # --- Epistemic variance: mu_x^T Sigma_pi mu_x ---
+    # mu_x is mus_m[i] — the vector of subset-conditional means for example i
+    # (zeroed for invalid cells, consistent with the masked prediction)
+    epistemic_var = torch.einsum("bi,ij,bj->b", mus_m, sigma_pi, mus_m)  # [B]
+    epistemic_var = epistemic_var.clamp(min=0.0)
+
+    total_var = aleatoric_var + epistemic_var
+    total_std = torch.sqrt(total_var.clamp(min=1e-12))
+
+    y_mean_np      = y_mean.detach().numpy()
+    total_std_np   = total_std.detach().numpy()
+
+    if return_components:
+        aleatoric_std_np = torch.sqrt(aleatoric_var.clamp(min=1e-12)).detach().numpy()
+        epistemic_std_np = torch.sqrt(epistemic_var.clamp(min=1e-12)).detach().numpy()
+        return y_mean_np, total_std_np, aleatoric_std_np, epistemic_std_np
+
+    return y_mean_np, total_std_np
+
+
+def coverage(y_true: np.ndarray, y_mean: np.ndarray, y_std: np.ndarray,
+             level: float = 0.95) -> float:
+    """
+    Empirical coverage at a given credible level.
+
+    Args:
+        y_true: [B] true targets.
+        y_mean: [B] predictive means.
+        y_std:  [B] predictive standard deviations.
+        level:  Nominal coverage (default 0.95).
+
+    Returns:
+        Scalar empirical coverage in [0, 1].
+    """
+    from scipy.stats import norm as scipy_norm
+    z = scipy_norm.ppf((1 + level) / 2)
+    lo = y_mean - z * y_std
+    hi = y_mean + z * y_std
+    return float(((y_true >= lo) & (y_true <= hi)).mean())

subset_mixture_model-0.1.0/smm/model.py

@@ -0,0 +1,125 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class SubsetWeightsModel(torch.nn.Module):
+    """
+    Learns a single global weight vector over all feature subsets.
+
+    A real-valued parameter vector eta of length |S| is initialized to zero and
+    passed through softmax to produce the mixture weights pi.
+
+    Args:
+        num_subsets (int): Number of subsets in the powerset (|S|).
+    """
+
+    def __init__(self, num_subsets: int):
+        super().__init__()
+        self.eta = torch.nn.Parameter(torch.zeros(num_subsets))
+
+    def forward(self, x=None):
+        # Return raw logits; softmax is applied inside the loss function.
+        return self.eta
+
+
+class SubsetDataset(torch.utils.data.Dataset):
+    """
+    PyTorch Dataset wrapping a DataFrame of integer-coded categorical features.
+
+    Args:
+        df (pd.DataFrame): Data split (train / val / test).
+        subset_features (list[str]): Categorical feature column names.
+        target (list[str]): Single-element list with the target column name.
+    """
+
+    def __init__(self, df, subset_features: list, target: list):
+        self.df = df
+        self.subset_features = subset_features
+        self.target = target
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        row = self.df.iloc[idx]
+        x = torch.tensor(row[self.subset_features].astype(np.float32).values)
+        y = torch.tensor(np.float32(row[self.target[0]]))
+        return x, y
+
+
+def subset_mixture_neg_log_posterior(
+    pi_logits: torch.Tensor,
+    y: torch.Tensor,
+    mus: torch.Tensor,
+    variances: torch.Tensor,
+    mask: torch.Tensor = None,
+    alpha: float = 1.1,
+) -> torch.Tensor:
+    """
+    Negative log-posterior of the subset mixture model.
+
+    Loss = -sum_i log( sum_s pi_s * N(y_i | mu_s(x_i), sigma^2_s(x_i)) )
+           - (alpha - 1) * sum_s log(pi_s)
+
+    Args:
+        pi_logits: [|S|] unnormalized logits (global parameter).
+        y:         [B] target values.
+        mus:       [B, |S|] empirical conditional means.
+        variances: [B, |S|] empirical conditional variances.
+        mask:      [B, |S|] bool tensor; True where the cell exists in training data.
+        alpha:     Dirichlet concentration parameter (>1 encourages non-degenerate weights).
+
+    Returns:
+        Scalar loss tensor.
+    """
+    pi = F.softmax(pi_logits, dim=0)          # [|S|]
+    log_pi = torch.log(pi + 1e-9)
+    log_pi_b = log_pi.unsqueeze(0)            # [1, |S|]
+
+    log_probs = (
+        -0.5 * torch.log(2 * torch.pi * variances)
+        - 0.5 * (y.unsqueeze(1) - mus) ** 2 / variances
+    )                                          # [B, |S|]
+
+    log_weighted = log_probs + log_pi_b        # [B, |S|]
+
+    if mask is not None:
+        log_weighted = log_weighted.masked_fill(~mask, float("-inf"))
+
+    log_likelihoods = torch.logsumexp(log_weighted, dim=1)  # [B]
+    nll = -log_likelihoods.sum()
+
+    batch_size = y.size(0)
+    log_prior = (alpha - 1.0) * log_pi.sum()
+
+    return nll - log_prior / batch_size
+
+
+def subset_mixture_mse(
+    pi_logits: torch.Tensor,
+    y: torch.Tensor,
+    mus: torch.Tensor,
+    mask: torch.Tensor = None,
+) -> torch.Tensor:
+    """
+    MSE loss for the subset mixture model (useful for warmup / debugging).
+
+    Args:
+        pi_logits: [|S|] unnormalized logits.
+        y:         [B] target values.
+        mus:       [B, |S|] empirical conditional means.
+        mask:      [B, |S|] bool tensor.
+
+    Returns:
+        Scalar MSE tensor.
+    """
+    pi = F.softmax(pi_logits, dim=0)           # [|S|]
+    weights = pi.unsqueeze(0).expand_as(mus)   # [B, |S|]
+
+    if mask is not None:
+        weights = weights.masked_fill(~mask, 0.0)
+        weights = weights / (weights.sum(dim=1, keepdim=True) + 1e-9)
+
+    preds = (mus * weights).sum(dim=1)
+    return F.mse_loss(preds, y, reduction="sum")

subset_mixture_model-0.1.0/smm/predictor.py

@@ -0,0 +1,60 @@
+import torch
+import pandas as pd
+import numpy as np
+from .subset_maker import SubsetMaker
+
+
+class SubsetMixturePredictor:
+    """
+    Inference wrapper for a trained Subset Mixture Model.
+
+    Args:
+        subset_maker (SubsetMaker): Trained SubsetMaker with lookup table.
+        weight_vector (torch.Tensor): Softmaxed mixture weights [|S|].
+    """
+
+    def __init__(self, subset_maker: SubsetMaker, weight_vector: torch.Tensor):
+        self.subset_maker = subset_maker
+        self.weight_vector = weight_vector
+
+    def predict(self, df: pd.DataFrame, return_debug: bool = False):
+        """
+        Predict target values for a new DataFrame.
+
+        For each example, weights are masked to valid subsets (those seen during
+        training) and re-normalized. Examples with no valid subsets fall back to
+        the global training mean.
+
+        Args:
+            df: Input DataFrame containing all subset_features columns.
+            return_debug: If True, also return per-example normalized weights
+                and a fallback mask.
+
+        Returns:
+            preds (np.ndarray): Predicted values [B].
+            (optional) norm_weights (np.ndarray): [B, |S|]
+            (optional) fallback_mask (np.ndarray): [B] bool
+        """
+        batch_tensor = torch.tensor(
+            df[self.subset_maker.subset_features].astype(np.float32).values
+        )
+        mus, _, mask = self.subset_maker.batch_lookup(batch_tensor)  # [B, S]
+
+        B, S = mus.shape
+        weights = self.weight_vector.unsqueeze(0).expand(B, S)       # [B, S]
+        weights = weights.masked_fill(~mask, 0.0)
+        mus = mus.masked_fill(~mask, 0.0)
+
+        weight_sums = weights.sum(dim=1, keepdim=True)               # [B, 1]
+        norm_weights = weights / (weight_sums + 1e-9)
+
+        fallback_mask = weight_sums.squeeze() < 1e-6                 # [B]
+        pred_weighted = (mus * norm_weights).sum(dim=1)
+        pred_fallback = torch.full_like(
+            pred_weighted, self.subset_maker.fallback_mean
+        )
+        preds = torch.where(fallback_mask, pred_fallback, pred_weighted)
+
+        if return_debug:
+            return preds.numpy(), norm_weights.numpy(), fallback_mask.numpy()
+        return preds.numpy()

subset_mixture_model-0.1.0/smm/subset_maker.py

@@ -0,0 +1,103 @@
+import torch
+import pandas as pd
+import numpy as np
+import itertools
+
+
+class SubsetMaker:
+    """
+    Builds a powerset lookup table from a training DataFrame.
+
+    For each non-empty subset of `subset_features`, groups training examples by
+    their unique value combination and stores the empirical mean and variance of
+    the target within each group.
+
+    Args:
+        df (pd.DataFrame): Training data. All `subset_features` columns must be
+            integer-coded categorical variables.
+        subset_features (list[str]): Names of the categorical feature columns.
+        target (list[str]): Single-element list with the target column name.
+    """
+
+    def __init__(self, df: pd.DataFrame, subset_features: list, target: list):
+        self.df = df
+        self.subset_features = subset_features
+        self.target = target
+        self.lookup = self._build_lookup()
+        valid_target = self.df[self.target[0]].dropna()
+        self.fallback_mean = float(valid_target.mean())
+        self.fallback_var = float(valid_target.var())
+
+    def _get_powerset(self) -> list:
+        powerset = []
+        for r in range(1, len(self.subset_features) + 1):
+            powerset += list(itertools.combinations(self.subset_features, r))
+        return [list(s) for s in powerset]
+
+    def _build_lookup(self, drop_missing_rows: bool = True) -> dict:
+        lookup = {}
+        for subset in self._get_powerset():
+            grouped = (
+                self.df[subset + self.target]
+                .groupby(subset)
+                .agg({self.target[0]: ["mean", "var"]})
+                .reset_index()
+            )
+            grouped.columns = [
+                "_".join(c).strip("_") if isinstance(c, tuple) else c
+                for c in grouped.columns
+            ]
+            if drop_missing_rows:
+                grouped = grouped.dropna()
+                grouped = grouped[grouped[f"{self.target[0]}_var"] > 1e-6]
+
+            mean_col = f"{self.target[0]}_mean"
+            var_col = f"{self.target[0]}_var"
+            group_dict = {
+                tuple(row[subset]): (row[mean_col], row[var_col])
+                for _, row in grouped.iterrows()
+            }
+            lookup[tuple(subset)] = (subset, grouped, group_dict)
+        return lookup
+
+    def batch_lookup(self, batch_tensor: torch.Tensor):
+        """
+        Look up empirical means/variances for a batch of integer-coded examples.
+
+        Args:
+            batch_tensor: [B, D] integer tensor matching the order of subset_features.
+
+        Returns:
+            means:     [B, |S|] float tensor
+            variances: [B, |S|] float tensor
+            mask:      [B, |S|] bool tensor (True where the cell exists in training data)
+        """
+        batch_df = pd.DataFrame(
+            batch_tensor.numpy(), columns=self.subset_features
+        ).astype(int)
+
+        all_means, all_vars, all_masks = [], [], []
+
+        for _, (subset_cols, _, group_dict) in self.lookup.items():
+            subset_vals = batch_df[subset_cols].apply(tuple, axis=1)
+            means, vars_, mask = [], [], []
+            for key in subset_vals:
+                if key in group_dict:
+                    m, v = group_dict[key]
+                    means.append(m)
+                    vars_.append(v)
+                    mask.append(True)
+                else:
+                    means.append(self.fallback_mean)
+                    vars_.append(self.fallback_var)
+                    mask.append(False)
+            all_means.append(means)
+            all_vars.append(vars_)
+            all_masks.append(mask)
+
+        # Transpose from [|S|, B] → [B, |S|]
+        return (
+            torch.tensor(all_means, dtype=torch.float32).T,
+            torch.tensor(all_vars, dtype=torch.float32).T,
+            torch.tensor(all_masks, dtype=torch.bool).T,
+        )

subset_mixture_model-0.1.0/subset_mixture_model.egg-info/PKG-INFO

@@ -0,0 +1,152 @@
+Metadata-Version: 2.4
+Name: subset-mixture-model
+Version: 0.1.0
+Summary: Interpretable empirical-Bayes aggregation of partition estimators for categorical regression
+Author-email: Aaron John Danielson <aaron.danielson@austin.utexas.edu>
+License: MIT
+Project-URL: Repository, https://github.com/aaronjdanielson/subset-mixture-model
+Keywords: interpretable machine learning,categorical features,empirical Bayes,uncertainty quantification,mixture model
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: numpy>=1.24
+Requires-Dist: pandas>=2.0
+Requires-Dist: scikit-learn>=1.3
+Requires-Dist: scipy>=1.11
+Provides-Extra: experiments
+Requires-Dist: matplotlib>=3.7; extra == "experiments"
+Requires-Dist: seaborn>=0.12; extra == "experiments"
+Requires-Dist: joblib>=1.3; extra == "experiments"
+Requires-Dist: lightgbm>=4.0; extra == "experiments"
+Requires-Dist: ngboost>=0.4; extra == "experiments"
+Requires-Dist: mapie>=0.6; extra == "experiments"
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+
+# Subset Mixture Model (SMM)
+
+**SMM** is an interpretable, empirical-Bayes method for regression on datasets with categorical features. It aggregates partition-based conditional-mean estimators over all non-empty feature subsets using learned simplex weights, adaptively balancing bias and variance across partition granularities.
+
+## Key idea
+
+Each feature subset $s$ induces a partition of the covariate space and a natural estimator of the conditional expectation — its empirical cell mean. SMM learns a convex combination of these estimators:
+
+$$\hat{f}(\mathbf{x}) = \sum_{s \in \mathcal{S}} \hat{\pi}_s \cdot \hat{\mu}_{m(s,\mathbf{x})}(s)$$
+
+The learned weights $\hat{\pi}_s$ are directly interpretable: they reveal which feature interactions drive predictions on average. Uncertainty is propagated from the MAP weight estimates via a Laplace approximation, yielding aleatoric/epistemic decompositions without post-hoc calibration.
+
+## Installation
+
+```bash
+pip install subset-mixture-model
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/aaronjdanielson/subset-mixture-model
+cd subset-mixture-model
+pip install -e .
+```
+
+## Quick start
+
+```python
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from smm import (
+    SubsetMaker, SubsetWeightsModel, SubsetDataset,
+    subset_mixture_neg_log_posterior, SubsetMixturePredictor,
+    compute_posterior_covariance, predict_with_uncertainty,
+)
+
+# --- your data (integer-coded categorical features) ---
+train_df = pd.read_csv("train.csv")
+val_df   = pd.read_csv("val.csv")
+test_df  = pd.read_csv("test.csv")
+
+cat_cols = ["feature_a", "feature_b", "feature_c"]
+target   = "y"
+
+# --- build lookup table ---
+subset_maker = SubsetMaker(train_df, cat_cols, [target])
+n_subsets = len(subset_maker.lookup)
+
+# --- train ---
+model     = SubsetWeightsModel(n_subsets)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
+loader    = DataLoader(SubsetDataset(train_df, cat_cols, [target]),
+                       batch_size=64, shuffle=True)
+
+for epoch in range(100):
+    for x, y in loader:
+        optimizer.zero_grad()
+        mus, variances, mask = subset_maker.batch_lookup(x)
+        loss = subset_mixture_neg_log_posterior(
+            model(), y, mus, variances, mask, alpha=1.1)
+        loss.backward()
+        optimizer.step()
+
+# --- predict with uncertainty ---
+pi_hat    = F.softmax(model.eta.detach(), dim=0)
+predictor = SubsetMixturePredictor(subset_maker, pi_hat)
+sigma_pi  = compute_posterior_covariance(
+    model, subset_maker, train_df, cat_cols, target, alpha=1.1)
+
+y_mean, y_std = predict_with_uncertainty(predictor, sigma_pi, test_df)
+# y_mean: point predictions
+# y_std:  total predictive standard deviation (aleatoric + epistemic)
+```
+
+## Interpretability
+
+```python
+import numpy as np
+
+subsets = list(subset_maker.lookup.keys())
+top_idx = np.argsort(pi_hat.numpy())[::-1][:10]
+
+for rank, i in enumerate(top_idx):
+    print(f"{rank+1:2d}. {subsets[i]}  π={pi_hat[i]:.4f}")
+```
+
+## Features
+
+- **Interpretable by construction**: learned weights reveal which feature interactions matter
+- **Principled uncertainty**: aleatoric/epistemic decomposition via Laplace approximation
+- **Efficient training**: only $2^D - 1$ logits optimized; lookup table precomputed once
+- **No post-hoc calibration**: well-calibrated predictive intervals out of the box
+- **Scalable to D ≤ 15** features; $k$-way truncation available for larger $D$
+
+## Datasets supported
+
+Any tabular dataset with integer-coded (or string, with encoding) categorical features and a continuous target.
+
+## Citation
+
+```bibtex
+@article{danielson2025smm,
+  title   = {Subset Mixture Model: Interpretable Empirical-Bayes Aggregation
+             of Partition Estimators for Categorical Regression},
+  author  = {Danielson, Aaron John},
+  journal = {Machine Learning},
+  year    = {2025},
+  note    = {Under review}
+}
+```
+
+## License
+
+MIT

subset_mixture_model-0.1.0/subset_mixture_model.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+README.md
+pyproject.toml
+smm/__init__.py
+smm/laplace.py
+smm/model.py
+smm/predictor.py
+smm/subset_maker.py
+subset_mixture_model.egg-info/PKG-INFO
+subset_mixture_model.egg-info/SOURCES.txt
+subset_mixture_model.egg-info/dependency_links.txt
+subset_mixture_model.egg-info/requires.txt
+subset_mixture_model.egg-info/top_level.txt
+tests/test_smm.py

subset_mixture_model-0.1.0/subset_mixture_model.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

subset_mixture_model-0.1.0/subset_mixture_model.egg-info/requires.txt

@@ -0,0 +1,17 @@
+torch>=2.0
+numpy>=1.24
+pandas>=2.0
+scikit-learn>=1.3
+scipy>=1.11
+
+[dev]
+pytest>=7
+pytest-cov
+
+[experiments]
+matplotlib>=3.7
+seaborn>=0.12
+joblib>=1.3
+lightgbm>=4.0
+ngboost>=0.4
+mapie>=0.6

subset_mixture_model-0.1.0/subset_mixture_model.egg-info/top_level.txt

@@ -0,0 +1 @@
+smm

subset_mixture_model-0.1.0/tests/test_smm.py

@@ -0,0 +1,259 @@
+"""
+Core tests for the Subset Mixture Model package.
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from smm import (
+    SubsetMaker,
+    SubsetWeightsModel,
+    SubsetDataset,
+    subset_mixture_neg_log_posterior,
+    SubsetMixturePredictor,
+    compute_posterior_covariance,
+    predict_with_uncertainty,
+    coverage,
+)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def synthetic_df():
+    """3-feature synthetic dataset with a known 3-way interaction."""
+    rng = np.random.default_rng(0)
+    n = 300
+    a = rng.integers(0, 4, n)
+    b = rng.integers(0, 3, n)
+    c = rng.integers(0, 2, n)
+    # target = interaction mean + noise
+    cell_means = {(i, j, k): rng.normal(0, 2)
+                  for i in range(4) for j in range(3) for k in range(2)}
+    y = np.array([cell_means[(a[i], b[i], c[i])] + rng.normal(0, 0.1)
+                  for i in range(n)])
+    return pd.DataFrame({"a": a, "b": b, "c": c, "y": y})
+
+
+@pytest.fixture
+def splits(synthetic_df):
+    from sklearn.model_selection import train_test_split
+    tr, te = train_test_split(synthetic_df, test_size=0.2, random_state=0)
+    tr, va = train_test_split(tr, test_size=0.15, random_state=0)
+    return tr.reset_index(drop=True), va.reset_index(drop=True), te.reset_index(drop=True)
+
+
+CAT_COLS = ["a", "b", "c"]
+TARGET   = "y"
+
+
+# ---------------------------------------------------------------------------
+# SubsetMaker
+# ---------------------------------------------------------------------------
+
+def test_subset_maker_powerset_size(splits):
+    tr, _, _ = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    # 3 features → 2^3 - 1 = 7 subsets
+    assert len(sm.lookup) == 7
+
+
+def test_subset_maker_batch_lookup_shapes(splits):
+    tr, _, _ = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    x = torch.tensor(tr[CAT_COLS].astype(np.float32).values[:16])
+    mus, variances, mask = sm.batch_lookup(x)
+    assert mus.shape == (16, 7)
+    assert variances.shape == (16, 7)
+    assert mask.shape == (16, 7)
+    assert mask.dtype == torch.bool
+
+
+def test_subset_maker_fallback(splits):
+    tr, _, _ = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    assert np.isfinite(sm.fallback_mean)
+    assert np.isfinite(sm.fallback_var)
+
+
+# ---------------------------------------------------------------------------
+# SubsetWeightsModel
+# ---------------------------------------------------------------------------
+
+def test_model_forward_returns_logits():
+    model = SubsetWeightsModel(7)
+    eta = model()
+    assert eta.shape == (7,)
+    assert not torch.isnan(eta).any()
+
+
+def test_softmax_sums_to_one():
+    model = SubsetWeightsModel(7)
+    pi = F.softmax(model(), dim=0)
+    assert torch.allclose(pi.sum(), torch.tensor(1.0), atol=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# Loss function
+# ---------------------------------------------------------------------------
+
+def test_loss_finite(splits):
+    tr, _, _ = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    model = SubsetWeightsModel(len(sm.lookup))
+    loader = DataLoader(SubsetDataset(tr, CAT_COLS, [TARGET]),
+                        batch_size=32, shuffle=False)
+    x, y = next(iter(loader))
+    mus, variances, mask = sm.batch_lookup(x)
+    loss = subset_mixture_neg_log_posterior(model(), y, mus, variances, mask)
+    assert torch.isfinite(loss)
+
+
+def test_loss_decreases_after_training(splits):
+    tr, va, _ = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    model = SubsetWeightsModel(len(sm.lookup))
+    opt = torch.optim.Adam(model.parameters(), lr=1e-2)
+    loader = DataLoader(SubsetDataset(tr, CAT_COLS, [TARGET]),
+                        batch_size=64, shuffle=True)
+
+    def epoch_loss():
+        total = 0.0
+        model.eval()
+        with torch.no_grad():
+            for x, y in loader:
+                mus, variances, mask = sm.batch_lookup(x)
+                total += subset_mixture_neg_log_posterior(
+                    model(), y, mus, variances, mask).item()
+        return total
+
+    loss_before = epoch_loss()
+    model.train()
+    for _ in range(5):
+        for x, y in loader:
+            opt.zero_grad()
+            mus, variances, mask = sm.batch_lookup(x)
+            subset_mixture_neg_log_posterior(
+                model(), y, mus, variances, mask).backward()
+            opt.step()
+    loss_after = epoch_loss()
+    assert loss_after < loss_before
+
+
+# ---------------------------------------------------------------------------
+# Predictor
+# ---------------------------------------------------------------------------
+
+def test_predictor_output_shape(splits):
+    tr, _, te = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    model = SubsetWeightsModel(len(sm.lookup))
+    pi = F.softmax(model.eta.detach(), dim=0)
+    predictor = SubsetMixturePredictor(sm, pi)
+    preds = predictor.predict(te)
+    assert preds.shape == (len(te),)
+    assert np.isfinite(preds).all()
+
+
+def test_predictor_fallback_for_unseen_cells():
+    """Test point with values not seen in training should fall back gracefully."""
+    rng = np.random.default_rng(1)
+    tr = pd.DataFrame({"a": [0, 1, 0, 1], "b": [0, 0, 1, 1], "y": [1.0, 2.0, 3.0, 4.0]})
+    sm = SubsetMaker(tr, ["a", "b"], ["y"])
+    model = SubsetWeightsModel(len(sm.lookup))
+    pi = F.softmax(model.eta.detach(), dim=0)
+    predictor = SubsetMixturePredictor(sm, pi)
+    # value 99 was never seen — should return fallback_mean
+    unseen = pd.DataFrame({"a": [99], "b": [99], "y": [0.0]})
+    preds = predictor.predict(unseen)
+    assert np.isfinite(preds).all()
+
+
+# ---------------------------------------------------------------------------
+# Laplace / uncertainty
+# ---------------------------------------------------------------------------
+
+def test_predict_with_uncertainty_shapes(splits):
+    tr, va, te = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    model = SubsetWeightsModel(len(sm.lookup))
+    opt = torch.optim.Adam(model.parameters(), lr=1e-2)
+    loader = DataLoader(SubsetDataset(tr, CAT_COLS, [TARGET]),
+                        batch_size=64, shuffle=True)
+    for _ in range(3):
+        for x, y in loader:
+            opt.zero_grad()
+            mus, variances, mask = sm.batch_lookup(x)
+            subset_mixture_neg_log_posterior(
+                model(), y, mus, variances, mask).backward()
+            opt.step()
+    pi = F.softmax(model.eta.detach(), dim=0)
+    predictor = SubsetMixturePredictor(sm, pi)
+    sigma_pi = compute_posterior_covariance(
+        model, sm, tr, CAT_COLS, TARGET, alpha=1.1)
+    y_mean, y_std = predict_with_uncertainty(predictor, sigma_pi, te)
+    assert y_mean.shape == (len(te),)
+    assert y_std.shape == (len(te),)
+    assert (y_std >= 0).all()
+
+
+def test_coverage_function():
+    y_true = np.array([0.0, 1.0, 2.0, 3.0])
+    y_mean = np.array([0.0, 1.0, 2.0, 3.0])
+    y_std  = np.array([1.0, 1.0, 1.0, 1.0])
+    cov = coverage(y_true, y_mean, y_std, level=0.95)
+    assert 0.0 <= cov <= 1.0
+
+
+# ---------------------------------------------------------------------------
+# Synthetic recovery (integration test)
+# ---------------------------------------------------------------------------
+
+def test_synthetic_weight_recovery(splits):
+    """
+    With the DGP driven purely by the 3-way interaction,
+    the full-powerset subset should receive the highest weight.
+    """
+    tr, va, te = splits
+    sm = SubsetMaker(tr, CAT_COLS, [TARGET])
+    model = SubsetWeightsModel(len(sm.lookup))
+    opt = torch.optim.Adam(model.parameters(), lr=5e-3)
+    loader = DataLoader(SubsetDataset(tr, CAT_COLS, [TARGET]),
+                        batch_size=64, shuffle=True)
+    best_state, best_val = None, float("inf")
+    val_loader = DataLoader(SubsetDataset(va, CAT_COLS, [TARGET]),
+                            batch_size=64, shuffle=False)
+    torch.manual_seed(0)
+    for epoch in range(60):
+        model.train()
+        for x, y in loader:
+            opt.zero_grad()
+            mus, variances, mask = sm.batch_lookup(x)
+            subset_mixture_neg_log_posterior(
+                model(), y, mus, variances, mask, alpha=1.1).backward()
+            opt.step()
+        model.eval()
+        vl = 0.0
+        with torch.no_grad():
+            for x, y in val_loader:
+                mus, variances, mask = sm.batch_lookup(x)
+                vl += subset_mixture_neg_log_posterior(
+                    model(), y, mus, variances, mask, alpha=1.1).item()
+        vl /= len(val_loader)
+        if vl < best_val:
+            best_val = vl
+            best_state = {k: v.clone() for k, v in model.state_dict().items()}
+    model.load_state_dict(best_state)
+    pi = F.softmax(model.eta.detach(), dim=0).numpy()
+    subsets = list(sm.lookup.keys())
+    # Full 3-way subset should be highest weight
+    full_subset = tuple(sorted(CAT_COLS))
+    idx = next(i for i, s in enumerate(subsets) if tuple(sorted(s)) == full_subset)
+    assert pi[idx] == pi.max(), (
+        f"3-way subset not top-weighted: pi={pi[idx]:.3f} vs max={pi.max():.3f}")