everwillow 0.1.2

Statistical inference steps for high-energy physics analyses with JAX.

everwillow is a statistical inference library for high-energy physics built on JAX pytrees and optimistix optimizers. It provides tools for fitting, profiling, and hypothesis testing with flexible parameter handling and parameter bounds via transformations. It works with any JAX-based statistical model.

Installation

pip install everwillow

or with uv:

uv add everwillow

From source:

git clone https://github.com/MoAly98/everwillow.git
cd everwillow
uv sync

Example

A Poisson counting experiment: define a model, fit it, compute CLs p-values, and find a 95% CL upper limit — both with asymptotic formulas and toys.

import jax
import jax.numpy as jnp

import everwillow as ew
import everwillow.statelib as sl
from everwillow.hypotest.calculators import AsymptoticCalculator, HypoTestCalculator
from everwillow.hypotest.distributions import (
    QTildeAsymptotic,
    SimpleEmpiricalDistribution,
)
from everwillow.hypotest.test_statistics import QTilde
from everwillow.hypotest.toys import ToyGenerator
from everwillow.hypotest.upper_limit import upper_limit, upper_limit_toys

jax.config.update("jax_enable_x64", True)

# --- Model ---

signal, background = 10.0, 5.0


def nll(params, observation):
    mu = params["mu"]
    expected = mu * signal + background
    return expected - observation["n"] * jnp.log(expected)


def predict(params_state):
    mu = params_state.to_pytree()["mu"]
    return {"n": mu * signal + background}


params = sl.State.from_pytree({"mu": 1.0})
observed = {"n": 12.0}

# --- Fit ---

result = ew.fit(nll_fn=nll, params=params, observation=observed)
print(result.params.to_pytree())
# {'mu': Array(0.7, dtype=float64)}

# --- Asymptotic hypothesis test ---

calc = AsymptoticCalculator(
    nll_fn=nll,
    params=params,
    observation=observed,
    poi_key="mu",
    predict_fn=predict,
    test_statistic=QTilde(),
    distribution=QTildeAsymptotic(),
)
result = calc.test(1.0)
print(f"CLs: {calc.cls(result):.4f}")
# CLs: 0.2140


def cls_objective(poi):
    return calc.cls(calc.test(poi))


limit = upper_limit(cls_objective, bounds=(0.0, 5.0), level=0.05)
print(f"95% CL upper limit (asymptotic): {float(limit):.4f}")
# 95% CL upper limit (asymptotic): 1.3673

# --- Toy-based upper limit ---

toy_gen = ToyGenerator(test_statistic=QTilde(), ntoys=5000)


def cls_from_toys(poi, key):
    """Generate toys and compute CLs at a given POI value."""
    toys = toy_gen.generate(
        nll,
        params,
        observed,
        "mu",
        poi_null=poi,
        poi_alt=0.0,
        key=key,
        predict_fn=predict,
    )
    dist = SimpleEmpiricalDistribution.from_toys(toys)
    toy_calc = HypoTestCalculator(
        nll_fn=nll,
        params=params,
        observation=observed,
        poi_key="mu",
        test_statistic=QTilde(),
        distribution=dist,
    )
    return toy_calc.cls(toy_calc.test(poi))


toy_limit = upper_limit_toys(
    cls_from_toys, bounds=(0.0, 5.0), key=jax.random.key(0), tol=0.01
)
print(f"95% CL upper limit (toys): {float(toy_limit):.4f}")

Documentation

• Quickstart
• API Reference

Contributing

We welcome contributions! See CONTRIBUTING.md for setup and workflow.

License

everwillow is distributed under the BSD-3-Clause License.