PDMPFlux

Documentation	Workflows	Code Coverage	Quality Assurance

Overview

PDMPFlux.jl provides a fast and efficient implementation of Piecewise Deterministic Markov Process (PDMP) samplers using a grid-based Poisson thinning approach.

In this version v0.2.0, only Zig-Zag samplers are implemented. We will extend the functionality to include other PDMP samplers in the future.

Key Features

• To sample from a distribution $p(x)$, the only required inputs are its dimension $d$ and the negative log density $U(x)=-\log p(x)$ (up to constant).

Motivation

Markov Chain Monte Carlo (MCMC) methods are standard in sampling from distributions with unknown normalizing constants.

However, PDMPs offer a promising alternative due to their continuous and non-reversible dynamics, particularly in high-dimensional and big data contexts, as discussed in Bouchard-Côté et. al. (2018) and Bierkens et. al. (2019).

Despite their potential, practical applications of PDMPs remain limited by a lack of efficient and flexible implementations.

PDMPFlux.jl is my attempt to fill this gap, with the aid of the existing automatic differentiation engines.

Installation

Currently, julia >= 1.11 is required for compatibility.

To install the package, use Julia's package manager:

(@v1.11) pkg> add PDMPFlux

Usage

Basic

The following example demonstrates how to sample from a standard Gaussian distribution using a Zig-Zag sampler.

using PDMPFlux

function U_Gauss(x::Vector)
    return sum(x.^2) / 2
end

dim = 10
sampler = ZigZagAD(dim, U_Gauss)

N_sk, N, xinit, vinit = 1_000_000, 1_000_000, zeros(dim), ones(dim)
samples = sample(sampler, N_sk, N, xinit, vinit, seed=2024)

jointplot(samples)

Advanced

For more control, you can manually provide the gradient.

Also, by breaking down the sample() function into two steps: sample_skeleton() and sample_from_skeleton(), you can use plot_traj() and diagnostic() functions to diagnose the sampler:

using PDMPFlux
using Zygote

N_sk = 1_000_000 # number of skeleton points
N = 1_000_000 # number of samples

function ∇U_banana(x::Vector)
    mean_x2 = (x[1]^2 - 1)
    return - (- x[1] + -(x[2] - mean_x2) - sum(x[3:end]))  # don't forget the minus sign!
end

dim = 50
xinit = ones(dim)
vinit = ones(dim)
grid_size = 0  # use constant bounds

sampler = ZigZag(dim, ∇U_banana, grid_size=grid_size)  # manually providing the gradient
output = sample_skeleton(sampler, N_sk, xinit, vinit, verbose = true)  # simulate skeleton points
samples = sample_from_skeleton(sampler, N, output)  # get samples from the skeleton points

plot_traj(output, 10000)
diagnostic(output)

jointplot(samples)

Gallery

2D Funnel Distribution (Ground Truth)	2D Zig-Zag Trajectory (Tmax=10000)	2D Zig-Zag on Funnel	3D Zig-Zag on Funnel
2D Banana Density Contour (Ground Truth)	2D Zig-Zag Sample Jointplot	2D Zig-Zag on Banana	3D Zig-Zag on Banana

1D Zig-Zag on Cauchy	1D Zig-Zag on Gaussian	Cauchy vs. Gaussian Density Plot

Remarks

• The automatic Poisson thinning implementation in PDMPFlux.jl is based on the paper Andral and Kamatani (2024) Automated Techniques for Efficient Sampling of Piecewise-Deterministic Markov Processes and its accompanying Python package pdmp_jax.
• pdmp_jax has a jax based implementation, and typically about four times faster than current PDMPFlux.jl.
• Both ForwardDiff.jl and Zygote.jl are used for automatic differentiation, each with their own trade-offs.

References

• pdmp_jax by Charly Andral, on which this repository is strongly based on.
  • Andral and Kamatani (2024) Automated Techniques for Efficient Sampling of Piecewise-Deterministic Markov Processes
• ForwardDiff.jl and Zygote.jl are used for automatic differentiation.
  • Revels, Lubin, and Papamarkou (2016) Forward-Mode Automatic Differentiation rin Julia
  • Innes et. al. (2018) Fashionable Modelling with Flux
• Other PDMP packages:
  • Julia package ZigZagBoomerang.jl by Marcel Schauer
  • R package rjpdmp by Matthew Sutton