bxai 0.1.1

Bayesian Feature Selection and Attribution Suite

bxai

Bayesian Feature Selection and Attribution Suite

A Python package implementing rigorous Bayesian methods for feature selection and model explainability. Where conjugate structure permits, bxai uses closed-form analytical updates (Beta-Binomial, Normal-Inverse-Gamma) to keep iteration fast. Where the model demands full posterior inference — Horseshoe GLMs, BART, and the optional MCMC path in BayLIME; it delegates to PyMC. The result is a unified toolkit that matches computational cost to statistical necessity.

Features

• Global Non-Linear Selection

  • BayesianBorutaSHAP: Tree-based feature selection using SHAP values, swapping frequentist p-values for Bayesian credible intervals, with dynamic pruning for fast performance. Supports discrete (Beta-Binomial) and continuous (Normal-Inverse-Gamma) modes.
  • BayesianPermutation: Model-agnostic importance tracking using paired validation loss drops updated via the Student-t continuous engine.

• Local Interpretability

  • BayLIME: Stable, prior-informed local explanations wrapping standard/custom perturbations in a Bayesian linear regression. The default backend uses a closed-form analytical posterior; an optional backend='mcmc' path (requires mcmc extra) adds heteroscedastic noise and Horseshoe priors for richer uncertainty quantification. Can be seeded with global SHAP weights from Phase 1.

• Parametric & Native Bayesian Importance

  • ShrinkagePIP: High-dimensional GLMs with Horseshoe and Lasso regularizing priors, tracking Posterior Inclusion Probabilities (PIP). Uses the posterior shrinkage factor κ_j = 1/(1+λ_j²τ²) for the Horseshoe prior (correct for continuous shrinkage distributions) and an auto-data-scaled |β_j| threshold for the Lasso (requires mcmc extra).
  • BARTImportance: Variable inclusion frequency (VIF) tracking from native Bayesian Additive Regression Trees (requires mcmc extra).

Installation

Install using uv:

uv add bxai

Or install with optional dependencies (e.g., SHAP/LightGBM for Boruta, PyMC for MCMC methods, catboost, xgboost):

uv add bxai --optional boruta --optional mcmc --optional xgboost --optional catboost

Quick Start

BayesianBorutaSHAP

import lightgbm as lgb
from sklearn.datasets import make_classification
from bxai.selection import BayesianBorutaSHAP

# Generate synthetic dataset
X, y = make_classification(n_samples=500, n_features=20, n_informative=5, random_state=42)

# Fit Bayesian BorutaSHAP
clf = lgb.LGBMClassifier(random_state=42)
selector = BayesianBorutaSHAP(model=clf, mode="discrete", max_iter=20, random_state=42)
selector.fit(X, y)

print("Confirmed Features:", selector.confirmed_)

BayesianPermutation

from bxai.selection import BayesianPermutation
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification

# Generate synthetic dataset
X_perm, y_perm = make_classification(n_samples=200, n_features=10, n_informative=3, random_state=42)

clf_perm = RandomForestClassifier(random_state=42).fit(X_perm, y_perm)

# Permutation feature selection with parallel jobs (n_jobs=2)
selector_perm = BayesianPermutation(
    model=clf_perm,
    scoring="accuracy",
    n_repeats=10,
    n_jobs=2,
    random_state=42
)
selector_perm.fit(X_perm, y_perm)

print("Confirmed Features:", selector_perm.confirmed_)
print(selector_perm.summary()[["feature", "mean", "hdi_lower", "hdi_upper", "selected"]])

BayLIME

from bxai.explanation import BayLIME
import numpy as np

# Instantiate BayLIME
explainer = BayLIME(
    training_data=X,
    feature_names=[f"feat_{i}" for i in range(20)]
)

# Explain a single instance
explanation = explainer.explain_instance(
    instance=X[0],
    predict_fn=clf.predict_proba
)

print(explanation.as_dataframe())

ShrinkagePIP

from bxai.parametric import ShrinkagePIP
from sklearn.datasets import make_regression

X, y = make_regression(n_samples=200, n_features=30, n_informative=5, random_state=42)

# Horseshoe prior — uses kappa-based PIP by default (pip_method='auto')
# PIP = P(κ_j < 0.5 | data), where κ_j = 1/(1 + λ_j² τ²)
selector = ShrinkagePIP(
    prior="horseshoe",
    kappa_threshold=0.5,   # κ < 0.5 → local scale dominates → signal
    pip_threshold=0.80,
    n_samples=500,
    random_state=42,
)
selector.fit(X, y)

print("Selected features:", selector.confirmed_)
print(selector.summary()[["feature", "pip", "kappa_mean", "selected"]])

# Lasso prior — uses auto-scaled |β| threshold (epsilon = std(y)/10)
selector_lasso = ShrinkagePIP(
    prior="lasso",
    pip_threshold=0.80,
    n_samples=500,
    random_state=42,
)
selector_lasso.fit(X, y)
print(f"Effective epsilon: {selector_lasso.epsilon_:.4f}")

BARTImportance

from bxai.parametric import BARTImportance
from sklearn.datasets import make_regression, make_classification

# Regression example
X_reg, y_reg = make_regression(n_samples=100, n_features=10, n_informative=3, random_state=42)
bart_reg = BARTImportance(model_type="regression", n_trees=20, n_samples=200, tune=100, chains=1, random_state=42)
bart_reg.fit(X_reg, y_reg)
print("Regression Selected features:", bart_reg.confirmed_)

# Classification example (Probit BART)
X_clf, y_clf = make_classification(n_samples=100, n_features=10, n_informative=3, random_state=42)
bart_clf = BARTImportance(model_type="classification", n_trees=20, n_samples=200, tune=100, chains=1, random_state=42)
bart_clf.fit(X_clf, y_clf)
print("Classification Selected features:", bart_clf.confirmed_)

License

MIT License