Metadata-Version: 2.4
Name: everwillow
Version: 0.0.1
Summary: Statistical inference steps for high-energy physics analyses with JAX.
Project-URL: Homepage, https://github.com/MoAly98/everwillow
Project-URL: Bug Tracker, https://github.com/MoAly98/everwillow/issues
Project-URL: Discussions, https://github.com/MoAly98/everwillow/discussions
Project-URL: Changelog, https://github.com/MoAly98/everwillow/releases
Author-email: MoAly98 <ma2520@princeton.edu>, pfackeldey <peter.fackeldey@princeton.edu>, maxgalli <mg2708@princeton.edu>
License-Expression: BSD-3-Clause
License-File: LICENSE
Keywords: fitting,hep,hypothesis-testing,inference,jax,physics,statistics
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: equinox>=0.10.11
Requires-Dist: jax>=0.7.2
Requires-Dist: jaxtyping>=0.2.34
Requires-Dist: optimistix>=0.0.10
Requires-Dist: orbax-checkpoint>=0.6.0
Requires-Dist: tqdm>=4.60
Requires-Dist: treescope>=0.1.0
Description-Content-Type: text/markdown

<div align="center">
<tr>
<td width="200px">
<img src="https://raw.githubusercontent.com/MoAly98/everwillow/main/images/logo.svg" alt="everwillow logo" width="180">
</td>

</tr>
</div>

##
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

```bash
pip install everwillow
```

or with uv:

```bash
uv add everwillow
```

From source:

```bash
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.

```python
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](https://everwillow.readthedocs.io/en/latest/quickstart.html)
- [API Reference](https://everwillow.readthedocs.io/en/latest/api/)

## Contributing

We welcome contributions! See [CONTRIBUTING.md](CONTRIBUTING.md) for setup and workflow.

## License

everwillow is distributed under the [BSD-3-Clause License](LICENSE).
