Skip to main content

A Python package for BayesForge modeling and diagnostics.

Project description

BayesForge for Python

::: {align="center"} A unified probabilistic programming library, bringing JAX-powered Bayesian inference 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 | +==================================================================================================================================================================================================================================================================================================================+==============================================================================================================================================================================================================================================================================================================================+========================================================================================================================================================================================================================================================================================================================+ | ```python 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) ``` | ```julia @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 ``` | ```r 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 BayesForge import BayesForge
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 BayesForge import BayesForge

# 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

  • BFR - R implementation
  • BFJ - Julia implementation

::: {align="center"} BayesForge
Based on "The BayesForge library for Python, R, Julia" by Sosa, McElreath, & Ross (2026).

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.49.tar.gz (418.8 kB 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.49-py3-none-any.whl (525.9 kB view details)

Uploaded Python 3

File details

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

File metadata

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

File hashes

Hashes for bayesforge-0.0.49.tar.gz
Algorithm Hash digest
SHA256 ee7a6a361076ea647caf4bc0c272c84808503e1cb0961422903ce6b473d67bbf
MD5 ded7c1464db03305999376589fc72d69
BLAKE2b-256 175f491202c431f69ebee63d078ae82e39298ff0b6c5e6d3525e80041143bbc6

See more details on using hashes here.

File details

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

File metadata

  • Download URL: bayesforge-0.0.49-py3-none-any.whl
  • Upload date:
  • Size: 525.9 kB
  • 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.49-py3-none-any.whl
Algorithm Hash digest
SHA256 798db819ba8db162da823e2f0eb8af596e293638c36475cfb92547ae36d48729
MD5 632a96771816f133a26405e23fa1de1b
BLAKE2b-256 51ea96a4cf3dc3b148df6091e6be7b61950e1f8ed84e2a98bebb658e95330235

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