Read histograms from ROOT files with uproot and hist

A HEP analysis usually starts from histograms in ROOT files. uproot reads them without ROOT, hist gives them axes, and their bin contents go into a FlatPPL model through flatppl.Module.set() as ordinary arrays.

The page writes a small ROOT file first, standing in for an analysis output: observed counts and an exposure template in ten bins. The model scales the template by a rate lam with a Gamma prior, a conjugate pair, so the posterior and its evidence are known exactly.

Runtime budget: 120 s.

import tempfile
from math import lgamma, log
from pathlib import Path

import hist
import jax
import jax.numpy as jnp
import numpy as np
import uproot

from flatppl import flatppl

Write a ROOT file

hist fills a histogram of simulated events, and uproot.recreate writes it with an exposure template. A real analysis reads its own file.

rng = np.random.default_rng(11)
axis = hist.axis.Regular(10, 0.0, 1.0, name="score")
observed = hist.Hist(axis)
observed.fill(score=rng.beta(2.0, 5.0, size=60))
exposure = hist.Hist(axis)
exposure[...] = np.linspace(2.0, 0.5, 10)

path = Path(tempfile.mkdtemp()) / "analysis.root"
with uproot.recreate(path) as file:
    file["observed"] = observed
    file["exposure"] = exposure

Read it back

to_hist() returns a hist.Hist with its axes; values() gives the bin contents as a NumPy array.

with uproot.open(path) as file:
    counts = file["observed"].to_hist().values().astype(np.int64)
    template = file["exposure"].to_hist().values()
print("counts:", counts.tolist())
print("exposure:", np.round(template, 3).tolist())

counts: [5, 14, 20, 9, 8, 3, 1, 0, 0, 0]
exposure: [2.0, 1.833, 1.667, 1.5, 1.333, 1.167, 1.0, 0.833, 0.667, 0.5]

Bind the bins to the model

Each bin’s count is Poisson with mean lam times its exposure. Module.set binds both arrays, so the compiled density holds them as constants.

model = flatppl(
    """
    counts = external(cartpow(nonnegintegers, 10))
    exposure = external(cartpow(posreals, 10))
    lam ~ Gamma(shape = 2.0, rate = 0.5)
    n ~ broadcast(Poisson, lam .* exposure)
    L = likelihoodof(kernelof(record(n = n), lam = lam), record(n = counts))
    posterior = bayesupdate(L, lawof(record(lam = lam)))
    """
).set(counts=counts, exposure=template)
density = model.logdensity("posterior")

Check against the closed form

With N = sum(counts) and T = sum(exposure) the posterior is Gamma(2 + N, 0.5 + T). Its mean is (2 + N) / (0.5 + T) and the evidence is prod(t^n / n!) 0.5^2 Gamma(2 + N) / (Gamma(2) (0.5 + T)^(2 + N)). A grid over the free coordinate u = log lam reproduces both.

a, b = 2.0, 0.5
total_counts, total_exposure = counts.sum(), template.sum()
log_evidence = (
    float((counts * np.log(template)).sum())
    - sum(lgamma(k + 1) for k in counts)
    + a * log(b) + lgamma(a + total_counts) - lgamma(a)
    - (a + total_counts) * log(b + total_exposure)
)  # fmt: skip
u = jnp.linspace(-4.0, 4.0, 4001)
values = np.asarray(jax.vmap(lambda v: density({"lam": v}))(u))
weights = np.exp(values)
evidence = np.trapezoid(weights, np.asarray(u))
mean = np.trapezoid(weights * np.exp(np.asarray(u)), np.asarray(u)) / evidence
exact_mean = (a + total_counts) / (b + total_exposure)
print(f"log evidence: grid {np.log(evidence):.6f}, closed form {log_evidence:.6f}")
print(f"posterior mean of lam: grid {mean:.6f}, closed form {exact_mean:.6f}")
assert abs(np.log(evidence) - log_evidence) < 1e-5
assert abs(mean - exact_mean) < 1e-5 * exact_mean

log evidence: grid -36.605385, closed form -36.605382
posterior mean of lam: grid 4.769231, closed form 4.769231