walnutpie

walnutpie is a Python package for Markov chain Monte Carlo (MCMC) sampling of differentiable target log densities.

• walnutpie can run models coded in Stan, PyMC, NumPyro, JAX, Numba, as well as models coded directly in Python.

• The underlying sampler is Walnuts, which adds dynamic step-size adaptation to the no-U-turn sampler (Nuts), which in turn is based on Hamiltonian Monte Carlo (HMC).

• The mass matrix and step-size adaptation scheme use an online variant of Nutpie

• For estimating the maximum step size during warmup, Walnuts uses Adam rather than dual averaging for stochastic gradient descent.

• Chain execution is multithreaded with optional convergence detection for warmup and sampling through lock-free buffers.

• Posterior analysis tools are included for the varying-length chains produced by asynchronous automatic stopping.

Python library

• Installation
  • Installing from source
• Python API
  • Sampling functions
  • Posterior Analysis Functions
• Using Walnuts from Python
  • Python
  • Numba

walnutpie accepts target log densities and gradients directly in Python, including those coded in NumPyro, PyMC, or JAX, even with foreign function calls. Stan models can be accessed directly at the C++ level through BridgeStan.

C++ interface documentation

walnutpie is implemented in multi-threaded C++20 with a stable client-facing API.

• C++ API
  • Sampling functionality
  • Configuration
  • Concepts

License

• walnutpie is distributed under the MIT License.

Bug reports and feature requests

Bug reports and feature requests are handled through GitHub.

• Walnutpie issue tracker

Developers and other contributors

We welcome new developers to the project and try to maintain a friendly and constructive environment. To get started, see the developers guide on GitHub:

• Contributing to walnutpie