A published Belle II likelihood: pyhf to FlatPPL to JAX

Import the full Belle II B+ to K+ nu nubar combined likelihood (ITA and HTA): five channels, thirty observed bins and 232 parameters. Check its expected counts and absolute log likelihood against pyhf, fit it with Optax, profile the signal strength, and sample a posterior with BlackJAX. pyhf, running its own fits on its own implementation, is the reference throughout.

Sources: the measurement (arXiv:2311.14647), the likelihood release (arXiv:2507.12393, HEPData 10.17182/hepdata.166082), and the author copy of the JSON at a fixed commit, checked by SHA-256.

The gallery runs a short sampling configuration; set FLATPPL_FULL=1 for four chains of 1,000 warmup and 2,000 retained draws.

Runtime budget: 240 s.

import hashlib
import json
import os
import urllib.request
from typing import Any

import blackjax
import jax
import jax.numpy as jnp
import numpy as np
import optax
import pyhf

from flatppl import from_pyhf

Read the published workspace

source = (
    "https://raw.githubusercontent.com/lorenzennio/knunu-reinterpretation-mini/"
    "710e304614ae47c044ec34f86c97b82b1f1ce956/scripts/data/combined_likelihood.json"
)
content = urllib.request.urlopen(source, timeout=60).read()
digest = "9ec0d563197912e03e31de202f04fea10babe39a10de99b9770d9292df285ab3"
assert hashlib.sha256(content).hexdigest() == digest
specification = json.loads(content)

pyhf.set_backend("numpy", precision="64b")
workspace = pyhf.Workspace(specification)
reference = workspace.model()
data = np.asarray(workspace.data(reference))
print(
    f"{len(reference.config.channels)} channels, {reference.config.nmaindata} bins, "
    f"{reference.config.npars} parameters, POI {reference.config.poi_name}"
)

5 channels, 30 bins, 232 parameters, POI mu

Import and compare in float64

model.pyhf lays the module out as pyhf does, so its compiled functions take pyhf’s flat parameter vector. A batch of points varies the signal strength, a histogram nuisance and a per-bin nuisance around the suggested initial values. Float64 is used because the comparisons and fits below need absolute log likelihoods to 1e-8; FlatPPL itself does not require it. jax.enable_x64 scopes it to each computation.

model = from_pyhf(json.dumps(specification))
view = model.pyhf
poi = view.poi_index
assert poi is not None
theta0 = np.asarray(view.suggested_init())
histogram = view.par_map["corr_KLeff_c1"].slice.start
points = np.repeat(theta0[None, :], 5, axis=0)
points[1, poi] -= 0.5
points[2, poi] += 0.5
points[3, histogram] -= 1.5
points[4, histogram] += 1.5
with jax.enable_x64(True):
    logpdf = view.logpdf(dtype="float64")
    expected = view.expected_data(dtype="float64")
    values = np.asarray(jax.jit(jax.vmap(logpdf))(points))
    counts = np.asarray(jax.jit(jax.vmap(expected))(points))
pyhf_values = np.array([reference.logpdf(point, data)[0] for point in points])
pyhf_counts = np.stack([reference.expected_actualdata(point) for point in points])
print("largest log likelihood difference:", np.max(np.abs(values - pyhf_values)))
np.testing.assert_allclose(values, pyhf_values, rtol=0, atol=1e-8)
np.testing.assert_allclose(counts, pyhf_counts, rtol=1e-11, atol=1e-10)

largest log likelihood difference: 3.922195901395753e-12

The cost of float32

The default precision is float32. Its absolute log likelihood differs from pyhf’s float64 by about 3e-5 here, from rounding in sums of large terms.

values32 = np.asarray(jax.vmap(view.logpdf())(points.astype(np.float32)))
print("float32 error:", np.max(np.abs(values32 - pyhf_values)))
assert np.max(np.abs(values32 - pyhf_values)) < 1e-4

float32 error: 3.235364943066088e-05

Maximum likelihood with Optax

logdensity("likelihood") takes the likelihood over its parameters as a record. The measurement’s fit bounds give each parameter a scaled logistic chart, and jacobian=False keeps the likelihood itself, so minimizing potential in free coordinates is the bounded maximum likelihood fit. L-BFGS starts at pyhf’s suggested initial values.

def lbfgs(potential: Any, start: Any, steps: int = 1500) -> Any:
    """Minimize ``potential`` from ``start`` with Optax L-BFGS."""
    solver = optax.lbfgs()
    value_and_grad = optax.value_and_grad_from_state(potential)

    def step(carry: Any, _: None) -> tuple[Any, Any]:
        params, state = carry
        value, grad = value_and_grad(params, state=state)
        updates, state = solver.update(
            grad, state, params, value=value, grad=grad, value_fn=potential
        )
        return (optax.apply_updates(params, updates), state), value

    (params, _), _ = jax.lax.scan(step, (start, solver.init(start)), length=steps)
    return params


with jax.enable_x64(True):
    fit = model.logdensity("likelihood", jacobian=False, dtype="float64")
    best = fit.constrain(jax.jit(lbfgs, static_argnums=0)(fit.potential, fit.init))
    flat_best = np.zeros(view.npars)
    for entry in view.par_map.values():
        flat_best[entry.slice] = np.atleast_1d(np.asarray(best[entry.name]))
    twice_nll = -2 * float(logpdf(flat_best))
    gradient = np.asarray(jax.grad(logpdf)(flat_best))

# pyhf's SciPy fit at a tight tolerance locates the minimum independently.
pyhf_best, pyhf_twice_nll = pyhf.infer.mle.fit(
    data, reference, return_fitted_val=True, tolerance=1e-10
)
print(
    f"mu: FlatPPL {flat_best[poi]:.4f}, pyhf {pyhf_best[poi]:.4f}; "
    f"-2 log L: FlatPPL {twice_nll:.6f}, pyhf {float(pyhf_twice_nll):.6f}"
)
assert abs(twice_nll - float(pyhf_twice_nll)) < 1e-5
# Away from the bounds, the gradient in model coordinates vanishes at the fit.
lower, upper = np.asarray(view.suggested_bounds()).T
interior = (flat_best > lower + 1e-6) & (flat_best < upper - 1e-6)
assert np.max(np.abs(gradient[interior])) < 1e-3

mu: FlatPPL 4.5991, pyhf 4.5995; -2 log L: FlatPPL -23.334280, pyhf -23.334278

Profile the signal strength

At each mu, refit every other parameter. optax.multi_transform holds mu at its free coordinate and fits the rest; unconstrain gives the free coordinate of the held value. pyhf’s fixed-POI fit is the reference.

assert view.poi is not None
mu_name = view.par_map[view.poi].name
labels = {name: ("held" if name == mu_name else "fit") for name in fit.names}
held = optax.multi_transform(
    {
        "fit": optax.adam(optax.cosine_decay_schedule(0.05, 3000, 1e-3)),
        "held": optax.set_to_zero(),
    },
    labels,
)


def profile(potential: Any, start: Any) -> jax.Array:
    """Minimize ``potential`` over every field but ``mu``."""

    def step(carry: Any, _: None) -> tuple[Any, None]:
        params, state = carry
        updates, state = held.update(jax.grad(potential)(params), state, params)
        return (optax.apply_updates(params, updates), state), None

    (params, _), _ = jax.lax.scan(step, (start, held.init(start)), length=3000)
    return potential(params)


grid = np.array([1.0, 3.0, 6.0])
with jax.enable_x64(True):
    starts = []
    for mu in grid:
        start = dict(best)
        start[mu_name] = jnp.float64(mu)
        starts.append(fit.unconstrain(start))
    stacked = jax.tree.map(lambda *x: jnp.stack(x), *starts)
    profiled = 2 * np.asarray(
        jax.jit(jax.vmap(lambda s: profile(fit.potential, s)))(stacked)
    )
pyhf_profiled = np.array(
    [
        float(
            pyhf.infer.mle.fixed_poi_fit(
                mu, data, reference, return_fitted_val=True, tolerance=1e-10
            )[1]
        )
        for mu in grid
    ]
)
print(
    "q(mu):",
    np.round(profiled - twice_nll, 4),
    "pyhf:",
    np.round(pyhf_profiled - float(pyhf_twice_nll), 4),
)
print("largest profile difference:", np.max(np.abs(profiled - pyhf_profiled)))
np.testing.assert_allclose(profiled, pyhf_profiled, atol=1e-5)

q(mu): [7.3771 1.4561 1.0773] pyhf: [7.3771 1.4561 1.0773]
largest profile difference: 1.303540386743407e-06

Sample a posterior with BlackJAX

A likelihood is not a posterior. Here every parameter gets a uniform prior on its published fit bounds: the default logdensity("likelihood") keeps the logistic chart and adds its log-Jacobian, which is exactly the density of the free coordinates under those uniform priors, up to a constant. The Gaussian auxiliary constraints are already part of the likelihood.

full = os.environ.get("FLATPPL_FULL") == "1"
chains, warmup, draws = (4, 1000, 2000) if full else (2, 300, 300)
density = model.logdensity("likelihood")


def run_chain(key: jax.Array, start: Any) -> Any:
    """Adapt one chain, then keep its draws and divergence flags."""
    warmup_key, sample_key = jax.random.split(key)
    adaptation = blackjax.window_adaptation(
        blackjax.nuts, density.fn, target_acceptance_rate=0.9
    )
    (state, parameters), _ = adaptation.run(warmup_key, start, num_steps=warmup)
    kernel = blackjax.nuts(density.fn, **parameters)

    def step(state: Any, key: jax.Array) -> tuple[Any, Any]:
        state, info = kernel.step(key, state)
        return state, (state.position, info.is_divergent)

    return jax.lax.scan(step, state, jax.random.split(sample_key, draws))[1]


starts = jax.tree.map(lambda x: jnp.stack([x] * chains), density.init)
free, divergent = jax.jit(jax.vmap(run_chain))(
    jax.random.split(jax.random.key(721), chains), starts
)
samples = jax.vmap(jax.vmap(density.constrain))(free)
mu_draws = np.asarray(samples[mu_name]).reshape(-1)
print(
    f"{chains} chains of {draws} draws, {int(divergent.sum())} divergences, "
    f"mu {mu_draws.mean():.3f} ± {mu_draws.std():.3f}, "
    f"68% interval {np.quantile(mu_draws, [0.16, 0.84]).round(3)}"
)

2 chains of 300 draws, 0 divergences, mu 4.618 ± 1.392, 68% interval [3.423 6.007]

The FlatPPL and pyhf likelihoods agree at the posterior draws too, in float64.

take = jax.tree.map(
    lambda x: np.asarray(x).reshape((-1, *np.shape(x)[2:]))[::50], samples
)
flat_draws = np.zeros((len(take[mu_name]), view.npars))
for entry in view.par_map.values():
    flat_draws[:, entry.slice] = np.asarray(take[entry.name]).reshape(
        len(flat_draws), -1
    )
with jax.enable_x64(True):
    draw_values = np.asarray(jax.vmap(logpdf)(flat_draws))
pyhf_draw_values = np.array([reference.logpdf(point, data)[0] for point in flat_draws])
np.testing.assert_allclose(draw_values, pyhf_draw_values, rtol=0, atol=1e-8)
print(f"{len(flat_draws)} draws scored; FlatPPL and pyhf agree to 1e-8")

12 draws scored; FlatPPL and pyhf agree to 1e-8