Skip to main content

A Python package for BayesForge modeling and diagnostics.

Project description

BayesForge for Python

A unified probabilistic programming library, bringing JAX-powered BayesForge to the Python, R and Julia ecosystem.
Run bespoke models on CPU, GPU, or TPU with Julia's native syntax.

Website bioRxiv Python License: GPL (>= 3)


One Mental Model. Three Languages.

BayesForge (BF) provides a unified experience across Julia, Python, and R. Whether you work in R's formula syntax, Python's object-oriented approach, or Julia's mathematical elegance, the model logic remains consistent.

  • Zero Context Switching: Variable names, distribution signatures, and model logic remain consistent across all implementations.
  • NumPyro Power: All interfaces compile down to XLA via JAX for blazing fast inference.
  • Rich Diagnostics: Seamless integration with ArviZ for posterior analysis.

Compare the Syntax

Python Syntax Julia Syntax R Syntax
def model(height, weight):
    # Priors
    sigma = BF.dist.uniform(0, 50, name='sigma', shape=(1,))
    alpha = BF.dist.normal(178, 20, name='alpha', shape=(1,))
    beta  = BF.dist.normal(0, 1, name='beta', shape=(1,))

    # Likelihood
    mu = alpha + beta * weight
    BF.dist.normal(mu, sigma, obs=height)
@BF function model(weight, height)
    # Priors
    sigma = BF.dist.uniform(0, 50, name='sigma', shape=(1,))
    alpha = BF.dist.normal(178, 20, name='alpha', shape=(1,))
    beta  = BF.dist.normal(0, 1, name='beta', shape=(1,))

    # Likelihood
    mu = alpha + beta * weight
    BF.dist.normal(mu, sigma, obs=height)
end
model <- function(height, weight){
  # Priors
  sigma = BF.dist.uniform(0, 50, name='sigma', shape=c(1))
  alpha = BF.dist.normal(178, 20, name='alpha', shape=c(1))
  beta  = BF.dist.normal(0, 1, name='beta', shape=c(1))

  # Likelihood
  mu = alpha + beta * weight
  BF.dist.normal(mu, sigma, obs=height)
}

Built for Speed

Leveraging Just-In-Time (JIT) compilation via JAX, BF outperforms traditional engines on standard hardware and unlocks massive scalability on GPU clusters for large datasets.

Benchmark: Network Size 400 (Lower is Better)

Engine Execution Time Relative Performance
STAN (CPU) ████████████████████████████ Baseline
BF (CPU) ████████████ ~30x Faster
BF (GPU) ██ ~200x Faster

> Comparison of execution time for a Social Relations Model. Source: Sosa et al. (2026).


Installation & Setup

1. Install Python

Download and install Python 3.10 or later

2. Install Package

From pip

pip install BayesForge

Development Installation

pip install git+https://github.com/BGN-for-ASNA/BF.git

Or clone the repository and activate it locally:

git clone https://github.com/BGN-for-ASNA/BF.git
cd BF

Then in Python:

pip install -e .

3. Initialize Environment

from BF import BF
m = BF()

4. Select Backend

Choose "cpu", "gpu", or "tpu" when importing the library.

# Initialize on CPU (default)
m = BF(platform="cpu")

# Or on GPU (requires JAX GPU installation)
m = BF(platform="gpu")

Quick Start

from BF import BF

# Initialize BF
m = BF()

# Generate some data
x = m.dist.normal(0, 1, shape=(100,), sample=True)
y = m.dist.normal(0.2 + 0.6 * x, 1.2, sample=True)

# Define a Bayesian linear regression model
def linear_model(x, y):
    alpha = m.dist.normal(loc=0, scale=1, name="alpha")
    beta  = m.dist.normal(loc=0, scale=1, name="beta")
    sigma = m.dist.exponential(1, name="sigma")
    mu = alpha + beta * x
    m.dist.normal(mu, sigma, obs=y)


# Fit the model
m.fit(linear_model, num_warmup=1000, num_samples=1000, num_chains=1)

# Display results
m.summary()

# Plot results with @pyplot
m.plot_trace()

Features

Data Manipulation

  • One-hot encoding
  • Index variable conversion
  • Scaling and normalization

Modeling (via NumPyro)

  • Linear & Generalized Linear Models: Regression, Binomial, Poisson, Negative Binomial, etc.
  • Hierarchical/Multilevel Models: Varying intercepts and slopes.
  • Time Series & Processes: Gaussian Processes, Gaussian Random Walks, State Space Models.
  • Mixture Models: GMM, Dirichlet Process Mixtures.
  • Network Models: Network-based diffusion, Block models.
  • Bayesian Neural Networks (BNN).

Diagnostics (via ArviZ)

  • Posterior summary statistics and plots.
  • Trace plots, Density plots, Autocorrelation.
  • WAIC and LOO (ELPD) model comparison.
  • R-hat and Effective Sample Size (ESS).

Available Distributions

The package provides wrappers for a comprehensive set of distributions from NumPyro.

Continuous

  • m.dist.normal, m.dist.uniform, m.dist.student_t
  • m.dist.cauchy, m.dist.halfcauchy, m.dist.halfnormal
  • m.dist.gamma, m.dist.inverse_gamma, m.dist.exponential
  • m.dist.beta, m.dist.beta_proportion
  • m.dist.laplace, m.dist.asymmetric_laplace
  • m.dist.log_normal, m.dist.log_uniform
  • m.dist.pareto, m.dist.weibull, m.dist.gumbel
  • m.dist.chi2, m.dist.gompertz

Discrete

  • m.dist.bernoulli, m.dist.binomial
  • m.dist.poisson, m.dist.negative_binomial
  • m.dist.geometric, m.dist.discrete_uniform
  • m.dist.beta_binomial, m.dist.zero_inflated_poisson

Multivariate

  • m.dist.multivariate_normal, m.dist.multivariate_student_t
  • m.dist.dirichlet, m.dist.dirichlet_multinomial
  • m.dist.multinomial
  • m.dist.lkj, m.dist.lkj_cholesky
  • m.dist.wishart, m.dist.wishart_cholesky

Time Series & Stochastic Processes

  • m.dist.gaussian_random_walk
  • m.dist.gaussian_state_space
  • m.dist.euler_maruyama
  • m.dist.car (Conditional AutoRegressive)

Mixtures & Truncated

  • m.dist.mixture, m.dist.mixture_same_family
  • m.dist.truncated_normal, m.dist.truncated_cauchy
  • m.dist.lower_truncated_power_law

(See package documentation for the full list)


Documentation

For full documentation and examples:

# See the Quick Start guide
# QUICKSTART.md

# Explore example notebooks
# test/usage_example.ipynb

For help with specific functions in the underlying BF library, refer to the BayesForge documentation.


Platform Support

  • ✅ Linux
  • ✅ macOS
  • ✅ Windows

GPU support available on compatible systems with JAX GPU installation.


Related Packages

  • BIR - R implementation
  • BIJ - Julia implementation

**BayesForge**\ Based on "The BayesForge library for Python, R, Julia" by [Sosa, McElreath, & Ross (2026)](https://www.biorxiv.org/content/10.64898/2026.01.19.700318v1).

Official website | Issues | Quick Start

© 2026 BayesForge Team. Released under GPL-3.0.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

bayesforge-0.0.46.tar.gz (2.1 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

bayesforge-0.0.46-py3-none-any.whl (2.6 MB view details)

Uploaded Python 3

File details

Details for the file bayesforge-0.0.46.tar.gz.

File metadata

  • Download URL: bayesforge-0.0.46.tar.gz
  • Upload date:
  • Size: 2.1 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.10.12

File hashes

Hashes for bayesforge-0.0.46.tar.gz
Algorithm Hash digest
SHA256 413a4e49c707f14463962e6f6d619d1b2665a0b928ee6f2369e2eabc4359e87f
MD5 c8f6786c06a10374b56d45db028ec85f
BLAKE2b-256 34f15d45d30a62ec4993de49edbd5a5744ffaf4e8a7a78a8142ad49139c2cb27

See more details on using hashes here.

File details

Details for the file bayesforge-0.0.46-py3-none-any.whl.

File metadata

  • Download URL: bayesforge-0.0.46-py3-none-any.whl
  • Upload date:
  • Size: 2.6 MB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.10.12

File hashes

Hashes for bayesforge-0.0.46-py3-none-any.whl
Algorithm Hash digest
SHA256 61cc1f71c78091bf7d0ae7c53f6d0a116cc9dc5a0df6cffc84ce04ab88a398fd
MD5 3a407dc842ee285da7f5a53274d4b5b0
BLAKE2b-256 402689c04f68a8bd8d043dc6af06706da9275f531a9f0112bacf2da9dd57c523

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page