BASSA: Bayesian Analysis with Spike-and-Slab Arrangements

Overview

Most chemical datasets are small and high-dimensional, making deep learning impractical. Linear regression remains interpretable and effective, but feature selection is critical. Traditional methods pick a single “best” model, overlooking the fact that multiple plausible models may exist.

BASSA combines Bayesian spike-and-slab regression with a filtering method to efficiently discover and organize many valid regression models. This reveals diverse interpretations hidden in chemical data without overcommitting to a single solution.

---

Installation

pip install bassa-reg

Plots are generated through JavaScript/Vega renderers, so LaTeX is not required.

Example Use

import os

import numpy as np
import pandas as pd
from bassa_reg import Bassa
from bassa_reg.spike_and_slab.spike_and_slab import SpikeAndSlabConfigurations, SpikeAndSlabRegression
from bassa_reg.spike_and_slab.spike_and_slab_util_models import SpikeAndSlabPriors, TestSet

def generate_data(N, M, K, noise_level=0.1):
    X = pd.DataFrame(np.random.randn(N, M), columns=[fr's_{i}' for i in range(M)])
    coefficients = np.random.randn(K)
    Y = pd.Series(X.iloc[:, :K].dot(coefficients) + np.random.randn(N) * noise_level)
    return X, Y

x_train, y_train = generate_data(100, 20, 5)

priors = SpikeAndSlabPriors()
config = SpikeAndSlabConfigurations(sampler_iterations=5000)
abs_dir = os.path.dirname(os.path.abspath(__file__))

regression = SpikeAndSlabRegression(x=x_train,
                                    y=y_train,
                                    priors=priors,
                                    config=config,
                                    project_path=abs_dir,
                                    experiment_name="demo")

regression.run()
bassa = Bassa(model=regression)
bassa.run()

Results

After running both the spike-and-slab regression and BASSA, results are saved in the specified project directory.
The main output is the bassa_plot.png file, which represents the models chosen by BASSA. An interactive viewer is also saved at bassa_plot_viewer/index.html.

This chart visualizes the different models found by BASSA, with their feature combinations and performance metrics.
Key additional outputs include:

Markov Chain Visualization

The markov_chain_visualization.png file shows the exploration of the top 20 features over iterations. An HTML/SVG version is also saved as markov_chain_visualization.html.
It is sorted by feature inclusion frequency, highlighting the most commonly selected features.
Precise feature inclusion frequencies are also provided in a separate file named feature_stats.csv.

Survival Process Plot

The survival_plot.png file, accompanied by the survival_table.csv file, illustrates the survival process of models over iterations. An HTML/SVG version is also saved as survival_plot.html.

This is an auxiliary output that helps understand how models persist or change.

Additional Data

The meta_data.csv file contains information about the Spike-and-Slab regression run, including the number of iterations, and other configuration details. It also includes some metrics about the regression performance on the training data.

Convergence Diagnostics

Enable optional convergence tracking in the sampler configuration:

config = SpikeAndSlabConfigurations(
    sampler_iterations=5000,
    turn_on_convergence_tracking=True,
)

Each tracked run saves two convergence figures as standalone HTML (rendered as SVG) and PNG:

• sparsity_convergence.* shows the last 500 model-size samples and an autocorrelation curve for lags 0–20, calculated from the full sparsity chain.
• log_likelihood_convergence.* shows the last 500 log-density samples and an autocorrelation curve for lags 0–20, calculated from the full log-density chain.

The log-density is the log of the complete model density p(a) p(S|a) p(tau^-2) p(W|tau^-2) p(sigma^-2) p(Y|X,S,W,sigma^-2). Full sample values are saved in convergence_samples.csv, and full autocorrelation values are saved in convergence_autocorrelation.csv.

Prediction on New Data

In order to make predictions on new data, create a new TestSet object.

import os

import numpy as np
import pandas as pd
from bassa_reg.spike_and_slab.spike_and_slab import SpikeAndSlabConfigurations, SpikeAndSlabRegression
from bassa_reg.spike_and_slab.spike_and_slab_util_models import SpikeAndSlabPriors, TestSet

def generate_data(N, M, K, noise_level=0.1):
    X = pd.DataFrame(np.random.randn(N, M), columns=[fr's_{i}' for i in range(M)])
    coefficients = np.random.randn(K)
    Y = pd.Series(X.iloc[:, :K].dot(coefficients) + np.random.randn(N) * noise_level)
    x_test = pd.DataFrame(np.random.randn(int(N/2), M), columns=[fr's_{i}' for i in range(M)])
    return X, Y, x_test

x_train, y_train, x_test = generate_data(100, 20, 5)

priors = SpikeAndSlabPriors()
config = SpikeAndSlabConfigurations(sampler_iterations=5000)
abs_dir = os.path.dirname(os.path.abspath(__file__))

test_set = TestSet(x_test=x_test,
                   samples_per_y=100,
                   iterations=200)

regression = SpikeAndSlabRegression(x=x_train,
                                    y=y_train,
                                    priors=priors,
                                    config=config,
                                    test_set=test_set,
                                    project_path=abs_dir,
                                    experiment_name="prediction_demo")

regression.run()

The sampler will run an extra numbers of iterations set by the iterations parameter in the TestSet object.
In every iteration, the sampler will sample samples_per_y values of y for each sample in the test set.
The average of these samples will be the predicted value for each sample in the test set.

Continuing a Previous Run

In order to continue a previous run, you first need to set save_samples=True on the SpikeAndSlabConfigurations object.
Then, you can load the previous run using the SpikeAndSlabLoader object and pass it to the SpikeAndSlabRegression object.

import os

import numpy as np
import pandas as pd
from bassa_reg.spike_and_slab.spike_and_slab import SpikeAndSlabConfigurations, SpikeAndSlabRegression
from bassa_reg.spike_and_slab.spike_and_slab_util_models import SpikeAndSlabPriors, SpikeAndSlabLoader


def generate_data(N, M, K, noise_level=0.1):
    X = pd.DataFrame(np.random.randn(N, M), columns=[fr's_{i}' for i in range(M)])
    coefficients = np.random.randn(K)
    Y = pd.Series(X.iloc[:, :K].dot(coefficients) + np.random.randn(N) * noise_level)
    return X, Y

x_train, y_train = generate_data(100, 10, 6, noise_level=0.6)
priors = SpikeAndSlabPriors()
config = SpikeAndSlabConfigurations(sampler_iterations=5000,
                                    save_meta_data=True,
                                    save_samples=True)
abs_dir = os.path.dirname(os.path.abspath(__file__))
regression = SpikeAndSlabRegression(x=x_train,
                            y=y_train,
                            priors=priors,
                            config=config,
                            project_path=abs_dir,
                            experiment_name="example_run")
regression.run()
loader = SpikeAndSlabLoader(path = f"{abs_dir}/{regression.full_experiment_name}")
regression = SpikeAndSlabRegression(x=x_train,
                                    y=y_train,
                                    priors=priors,
                                    config=config,
                                    project_path=abs_dir,
                                    experiment_name="example_run",
                                    load_experiment=loader)
regression.run()

Choosing Priors For Spike-and-Slab

The spike-and-slab model has priors for three latent quantities: sparsity (a), signal magnitude (tau), and noise level (sigma). The recipe below is enough to choose the required SpikeAndSlabPriors values; see the paper for the motivation and derivations behind these choices.

1. Sparsity Prior

Choose the expected number of relevant features, k, out of the total number of features, d. Set:

a_alpha_prior = 1
a_beta_prior = (d - k) / k

It can be useful to set k a little higher than the desired final model size, roughly 1.2x-1.5x, because tau also adds regularization.

2. Noise Precision Prior

BASSA places a Gamma prior on the noise precision, sigma^-2, using:

sigma^-2 ~ Gamma(sigma_alpha_prior, sigma_beta_prior)

sigma_beta_prior is a Gamma rate parameter, not a scale parameter. Choose an expected R^2 and a confidence level from the table:

Expected R^2	Implied precision, lambda ~= 1 / sigma^2	Confidence level	sigma_alpha_prior	sigma_beta_prior
0.80	5.00	Low / broad	3	0.4
0.80	5.00	Medium	6	1.0
0.80	5.00	High / restrictive	11	2.0
0.70	3.33	Low / broad	3	0.6
0.70	3.33	Medium	6	1.5
0.70	3.33	High / restrictive	11	3.0
0.60	2.50	Low / broad	3	0.8
0.60	2.50	Medium	6	2.0
0.60	2.50	High / restrictive	11	4.0

When unsure, start with a higher expected R^2 and a low-confidence prior so the data can still drive the posterior.

3. Signal Magnitude Prior

Because BASSA standardizes both X and y, active coefficients are measured in standard-deviation units. A robust default is:

tau_alpha_prior = 1
tau_beta_prior = 1

Values of tau_beta_prior between 1 and 2 are often reasonable with tau_alpha_prior=1. If the results need tuning, adjust the sparsity and noise priors first, then revisit tau.

Example:

d = x_train.shape[1]
k = 8

priors = SpikeAndSlabPriors(
    a_alpha_prior=1,
    a_beta_prior=(d - k) / k,
    sigma_alpha_prior=3,
    sigma_beta_prior=0.4,
    tau_alpha_prior=1,
    tau_beta_prior=1,
)

BASSA Thresholds

BASSA thresholds mainly control how much is shown in the final plot and tables. They do not change the underlying spike-and-slab samples; they make the output smaller, more linear, and easier to interpret.

• Preferred metric (metric): choose the metric used to order and display models in the BASSA plot. Common choices are SurvivalMetric.R2 and SurvivalMetric.Q2.
• Metric threshold (metric_threshold): keep only models whose training R^2 is at least this value. This is a simple way to remove weak models before plotting.
• Minimum points (min_points): keep only models that survive across more than this number of restriction levels. Increasing this usually produces a smaller, cleaner result.
• Maximum restriction (max_restriction): the largest number of top features considered while building the survival curves. Lower values make the result more compact.
• Inclusion threshold (percentage_cutoff): keep only feature combinations that appear in at least this fraction of the sampled chain. A usual range is 0.005 to 0.02 (0.5%-2%).

In practice, start with the metric and a simple metric threshold, then tune min_points and percentage_cutoff until the output is neither too large nor too small. These settings mainly control what is included in the visualization, not the regression result itself.

Example:

from bassa_reg import BassaConfigurations
from bassa_reg.spike_and_slab.utils.bassa_enums import SurvivalMetric

bassa.run(BassaConfigurations(
    metric=SurvivalMetric.Q2,
    metric_threshold=0.70,
    min_points=2,
    max_restriction=20,
    percentage_cutoff=0.01,
))

Paper Reproducibility

The paper-specific experiment and figure scripts are kept in the repository, while their generated timestamped outputs are ignored by Git. Run each command with --help to inspect its configurable inputs and output directory.

python generate_paper_charts.py --help
python generate_synthetic_paper_panels.py
python expreiments/export_convergence_pdfs.py --help
python expreiments/synthetic_data/generate_clean_weight_plots.py
python expreiments/case_study_1/prediction_pipeline.py
python expreiments/case_study_3/prediction_pipeline.py --help
python expreiments/top_k_vs_bassa_comparison/generate_top_k_vs_bassa.py --help

The benchmark pipelines use the same AIC-based forward-selection and standardized LARS-Lasso helpers so their reported baselines are comparable.