DEFER

DEnsity Function Estimation using Recursive partitioning

Approximate any black-box unnormalised density from point evaluations alone, and get samples, the normalising constant, and more. Runs on PyTorch or NumPy.

_{200k samples drawn from a DEFER approximation of a density concentrated on a 3-D spiral.}

Overview

DEFER approximates an arbitrary unnormalised (log-)density over a bounded box using only pointwise evaluations: no gradients, no closed form, no conjugacy. It recursively partitions the domain, spending evaluations where probability mass concentrates, and returns a single object that provides:

• fast, constant-time sampling (via the alias method),
• the normalising constant Z, i.e. the model evidence, which DEFER approximates,
• closed-form expectations under the approximation, including the mean, variance, and differential entropy.

It is robust and fast, and applies to problems of up to around 20 dimensions. The budget is a number of function evaluations: each evaluation adds one partition, so an approximation with N partitions costs N evaluations of your density and takes O(N log N) time and O(N) space. Every numerical step runs through a PyTorch backend (on a chosen device) or NumPy. The method is described in Black-box density function estimation using recursive partitioning (arXiv:2010.13632, ICML 2021).

Install

Managed with uv (Python 3.13):

uv sync

Quickstart

import numpy as np
from defer import defer

# Any unnormalised log-density you can evaluate pointwise (here a 2-D Gaussian).
def log_density(x):
    return -0.5 * (x**2).sum(-1)

approx = defer(
    log_density,
    lower=np.array([-5.0, -5.0]),
    upper=np.array([5.0, 5.0]),
    num_fn_calls=2000,       # density evaluations; one partition is added per evaluation
    use_numpy=True,          # or use_numpy=False, device="cpu" / "cuda"
)

approx.log_z                 # log normalising constant (the evidence)
approx.mean()                # posterior mean
samples = approx.sampler()(10_000)   # (10000, 2) samples, O(1) each

Set use_numpy=False, device="cuda" to run the same code on a GPU with PyTorch.

Examples

Runnable notebooks live in examples/ (uv sync installs the plotting extras):

Gaussian-process kernel selection Infer each kernel's hyperparameter posterior and use the evidence log Z to choose between a Matérn-3/2 and a sum of two Matérn-3/2. The data is generated from the sum kernel, and the evidence prefers it.

Black-box Gibbs density on a spiral A density concentrated on a 3-D spiral, defined only by distance to the curve. DEFER recovers a sampler, the normalising constant, and the shape itself.

Development

make lint        # ruff check + format
make typecheck   # pyright
make test        # pytest (numpy and torch backends)

License

MIT, see LICENSE.