evermore 0.2.4

Differentiable (binned) likelihoods in JAX.

evermore

Differentiable (binned) likelihoods in JAX.

Installation

python -m pip install evermore

From source:

git clone https://github.com/pfackeldey/evermore
cd evermore
python -m pip install .

Example - Model and Loss Definition

See more in examples/

evermore in a nutshell:

import equinox as eqx
import jax
import jax.numpy as jnp
from jaxtyping import Array

import evermore as evm

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


# define a simple model with two processes and two parameters
class Model(eqx.Module):
    mu: evm.FreeFloating
    syst: evm.NormalConstrained

    def __call__(self, hists: dict[str, Array]) -> Array:
        mu_modifier = self.mu.unconstrained()
        syst_modifier = self.syst.log_normal(up=jnp.array([1.1]), down=jnp.array([0.9]))
        return mu_modifier(hists["signal"]) + syst_modifier(hists["bkg"])


nll = evm.loss.PoissonNLL()


def loss(model: Model, hists: dict[str, Array], observation: Array) -> Array:
    expectation = model(hists)
    # Poisson NLL of the expectation and observation
    log_likelihood = nll(expectation, observation)
    # Add parameter constraints from logpdfs
    constraints = evm.loss.get_log_probs(model)
    log_likelihood += evm.util.sum_leaves(constraints)
    return -jnp.sum(log_likelihood)


# setup model and data
hists = {"signal": jnp.array([3]), "bkg": jnp.array([10])}
observation = jnp.array([15])
model = Model(mu=evm.FreeFloating(), syst=evm.NormalConstrained())

# negative log-likelihood
loss_val = loss(model, hists, observation)
# gradients of negative log-likelihood w.r.t. model parameters
grads = eqx.filter_grad(loss)(model, hists, observation)
print(f"{grads.mu.value=}, {grads.syst.value=}")
# -> grads.mu.value=Array([-0.46153846]), grads.syst.value=Array([-0.15436207])

Contributing

See CONTRIBUTING.md for instructions on how to contribute.

License

Distributed under the terms of the BSD license.