Bayesian Feature Selection

Scalable Bayesian feature selection using horseshoe priors with NumPyro/JAX. Select relevant features from high-dimensional data using a Bayesian GLM with the regularized horseshoe prior, which provides strong shrinkage for irrelevant features while preserving truly relevant signals.

• Free software: MIT license
• Documentation: https://macromagic.github.io/bayesian_feature_selection/

Features

• Horseshoe Prior GLM — Regularized horseshoe prior for automatic feature selection
• Multiple GLM Families — Gaussian (regression), Binomial (classification), and Poisson (count data)
• MCMC & SVI Inference — Full posterior via NUTS or fast approximation via Stochastic Variational Inference
• Flexible Feature Selection — Beta-based, lambda-based, or combined inclusion probabilities
• Data Loading & Preprocessing — Built-in CSV loading, train/test splitting, and standardization
• CLI Interface — Run experiments from the command line with YAML configuration files
• Visualization — Feature importance plots and MCMC diagnostic plots via ArviZ

Quick Start

import numpy as np
from bayesian_feature_selection import HorseshoeGLM, InferenceConfig

# Create synthetic data
rng = np.random.RandomState(42)
n, p = 100, 10
X = rng.randn(n, p)
true_beta = np.array([3.0, -2.0, 0, 0, 1.5, 0, 0, 0, 0, 0])
y = X @ true_beta + rng.randn(n) * 0.5

# Fit model
model = HorseshoeGLM(family="gaussian")
config = InferenceConfig(
    method="mcmc", num_warmup=500, num_samples=1000,
    num_chains=2, use_gpu=False, progress_bar=False
)
model.fit(X, y, config=config)

# Get selected features
importance = model.get_feature_importance(threshold=0.5)
selected = importance[importance["selected"]]
print(selected[["feature_idx", "beta_mean", "beta_inclusion_prob"]])

# Make predictions
predictions = model.predict(X)

CLI Usage

$ bayesian-fs -c configs/default.yaml

Installation

$ pip install bayesian_feature_selection

For development:

$ pip install -e ".[dev]"

Environment Setup (Python 3.12 + CUDA 12)

Python 3.12 (CPU only)

$ pip install bayesian_feature_selection

This installs JAX ≥ 0.7.0 and NumPyro ≥ 0.15.0 automatically.

Python 3.12 + CUDA 12 (GPU)

Install the package, then upgrade JAX with the CUDA 12 backend:

$ pip install bayesian_feature_selection
$ pip install "jax[cuda12]"

Verify the setup:

import jax
print(jax.devices())  # should show CudaDevice(id=0) when GPU is available

Important version notes

• NumPyro ≥ 0.15.0 requires JAX ≥ 0.7.0.
• JAX ≥ 0.10.0 removed internal symbols that NumPyro ≤ 0.20.1 depends on; use JAX 0.7.x – 0.9.x until NumPyro releases a compatible update.
• The gpu extra (pip install bayesian_feature_selection[gpu]) installs jax[cuda12] and is the recommended way to enable GPU support.

Credits

This package was created with Cookiecutter and the audreyr/cookiecutter-pypackage project template.