ttbayes

Tensor-train Bayesian inference — compile a black-box Bayesian posterior into a tensor-train (TT) representation, with gravitational-wave parameter estimation as the first application.

Status: name placeholder. 0.0.1 reserves the name on PyPI. None of the API below is implemented yet — do not install this expecting a working library. Development follows the checkpoint plan in plan.md.

What it will do

Given a black-box batched log-likelihood and a separable prior over continuous bounded parameters, approximate the square root of the unnormalized posterior with a tensor train:

$$\sqrt{\pi(\theta)} \approx \Psi_{\mathrm{TT}}(\theta), \qquad \widehat{\pi}{\mathrm{TT}}(\theta) = \Psi{\mathrm{TT}}(\theta)^2$$

Squaring keeps the approximate posterior nonnegative by construction. Building the TT with tempered TT-cross in physically informed (Fisher-whitened) coordinates is intended to reduce both the TT ranks and the number of expensive likelihood evaluations needed to reconstruct the posterior.

Once compiled, the TT posterior supports evaluation, log-evidence estimation, 1-D and 2-D marginals, moments, and approximate sampling — with an optional independence Metropolis–Hastings correction that restores the exact target.

Intended usage

import ttbayes as ttb

model = ttb.BlackBoxModel(
    parameter_names=("chirp_mass", "mass_ratio", "luminosity_distance", "cos_inclination"),
    log_likelihood=log_likelihood,
    log_prior=log_prior,
)

posterior = ttb.compile(
    model=model,
    domain={
        "chirp_mass": (28.0, 32.0, 64),
        "mass_ratio": (0.4, 1.0, 64),
        "luminosity_distance": (100.0, 1000.0, 64),
        "cos_inclination": (-1.0, 1.0, 64),
    },
    transform="fisher",
    positivity="square_root",
    tempering=[0.1, 0.25, 0.5, 0.75, 1.0],
    tolerance=1e-4,
    max_rank=32,
)

print(posterior.log_evidence())
samples = posterior.sample_corrected(model=model, n=10_000, method="independence_mh")

Installation

pip install ttbayes                # core: numpy, torch, tntorch
pip install "ttbayes[bilby]"       # + bilby, dynesty for the GW adapter
pip install -e ".[dev]"            # development

Requires Python 3.11+.

Scope

The first prototype is deliberately limited to continuous bounded parameters, tensor-product grids, black-box batched log-likelihoods, separable priors, and a tntorch TT-cross backend. Bilby is an optional adapter; Dynesty through Bilby is the reference sampler.

The public API is kept independent of tntorch — tntorch.Tensor is not exposed.

License

MIT