correctionlib-gradients 0.1.2

Automatic differentiation for high-energy physics correction factor calculations.

correctionlib-gradients

A JAX-friendly, auto-differentiable, Python-only implementation of correctionlib correction evaluations.

---

Table of Contents

• Installation
• Usage
• Supported types of corrections
• License

Installation

pip install correctionlib-gradients

Usage

1. construct a CorrectionWithGradient object from a correctionlib.schemav2.Correction
2. there is no point 2: you can use CorrectionWithGradient.evaluate as a normal JAX-friendly, auto-differentiable function

Example

import jax

from correctionlib import schemav2
from correctionlib_gradients import CorrectionWithGradient

# given a correctionlib schema:
formula_schema = schemav2.Correction(
    name="x squared",
    version=2,
    inputs=[schemav2.Variable(name="x", type="real")],
    output=schemav2.Variable(name="a scale", type="real"),
    data=schemav2.Formula(
        nodetype="formula",
        expression="x * x",
        parser="TFormula",
        variables=["x"],
    ),
)

# construct a CorrectionWithGradient
c = CorrectionWithGradient(formula_schema)

# use c.evaluate as a JAX-friendly, auto-differentiable function
value, grad = jax.value_and_grad(c.evaluate)(3.0)
assert jax.numpy.isclose(value, 9.0)
assert jax.numpy.isclose(grad, 6.0)

# jax.jit works too
value, grad = jax.jit(jax.value_and_grad(c.evaluate))(3.0)
assert jax.numpy.isclose(value, 9.0)
assert jax.numpy.isclose(grad, 6.0)

Supported types of corrections

Currently the following corrections from correctionlib.schemav2 are supported:

• Formula
• Binning with uniform and non-uniform bin edges, simple scalar bin values, and flow="clamp"
• scalar constants

License

correctionlib-gradients is distributed under the terms of the BSD 3-Clause license.