from __future__ import annotations
import abc
from collections.abc import Callable, Iterable, Iterator
from typing import TYPE_CHECKING
import jax
import jax.numpy as jnp
from flax import nnx
from jaxtyping import Array, Bool
from evermore.binned.effect import H, OffsetAndScale
from evermore.parameters.parameter import V
from evermore.util import tree_stack
if TYPE_CHECKING:
from evermore.binned.effect import BaseEffect
__all__ = [
"BooleanMask",
"Compose",
"Modifier",
"ModifierBase",
"Transform",
"TransformOffset",
"TransformScale",
"Where",
]
def __dir__():
return __all__
[docs]
class ModifierBase(nnx.Module):
"""Base class for modules that modify histogram templates.
Subclasses implement :meth:`offset_and_scale` and automatically gain a
callable interface as well as support for composition via the matrix
multiplication operator.
Examples:
>>> import jax.numpy as jnp
>>> import evermore as evm
>>> modifier = evm.Parameter(1.0).scale()
>>> modifier(jnp.array([10.0, 20.0]))
Array([10., 20.], dtype=float32)
"""
@abc.abstractmethod
def offset_and_scale(self: ModifierBase, hist: H) -> OffsetAndScale: ...
def __call__(self: ModifierBase, hist: H) -> H:
os = self.offset_and_scale(hist=hist)
return os.scale * (hist + os.offset) # ty:ignore[invalid-return-type]
def __matmul__(self: ModifierBase, other: ModifierBase) -> Compose:
return Compose(self, other)
[docs]
class Modifier(ModifierBase):
"""Pairs a parameter with an effect to build a modifier.
Args:
parameter: Parameter instance that provides the nuisance strength.
effect: Effect describing how the histogram is altered.
Examples:
>>> import jax.numpy as jnp
>>> import evermore as evm
>>> modifier = evm.Modifier(
... value=1.1,
... effect=evm.effect.Linear(offset=0.0, slope=1.0),
... )
>>> modifier(jnp.array([10, 20, 30]))
Array([11., 22., 33.], dtype=float32)
"""
def __init__(self, value: V, effect: BaseEffect) -> None:
self.value = value
self.effect = effect
def offset_and_scale(self, hist: H) -> OffsetAndScale:
return self.effect(value=self.value, hist=hist)
[docs]
class Where(ModifierBase):
"""Chooses between two modifiers based on a boolean condition.
Args:
condition: Boolean array indicating where to apply ``modifier_true``.
modifier_true: Modifier evaluated where ``condition`` is ``True``.
modifier_false: Modifier evaluated where ``condition`` is ``False``.
Examples:
>>> import jax.numpy as jnp
>>> import evermore as evm
>>> hist = jnp.array([5, 20, 30])
>>> syst = evm.NormalParameter(value=0.1)
>>> norm = syst.scale_log_asymmetric(up=jnp.array([1.1]), down=jnp.array([0.9]))
>>> shape = syst.morphing(
... up_template=jnp.array([7, 22, 31]),
... down_template=jnp.array([4, 16, 27]),
... )
>>> modifier = evm.modifier.Where(hist < 10, norm, shape)
>>> modifier(hist)
Array([ 5.049494, 20.281374, 30.181376], dtype=float32)
"""
def __init__(
self,
condition: Bool[Array, ...],
modifier_true: ModifierBase,
modifier_false: ModifierBase,
):
self.condition = condition
self.modifier_true = modifier_true
self.modifier_false = modifier_false
def offset_and_scale(self, hist: H) -> OffsetAndScale:
true_os = self.modifier_true.offset_and_scale(hist)
false_os = self.modifier_false.offset_and_scale(hist)
def _where(
true: Bool[Array, ...],
false: Bool[Array, ...],
) -> Bool[Array, ...]:
return jnp.where(self.condition, true, false)
return jax.tree.map(_where, true_os, false_os)
[docs]
class BooleanMask(ModifierBase):
"""Applies a modifier only to bins selected by a boolean mask.
Args:
mask: Boolean array indicating which bins receive the modifier.
modifier: Modifier that provides the offsets and scales.
Examples:
>>> import jax.numpy as jnp
>>> import evermore as evm
>>> hist = jnp.array([5, 20, 30])
>>> syst = evm.NormalParameter(value=0.1)
>>> norm = syst.scale_log_asymmetric(up=1.1, down=0.9)
>>> mask = jnp.array([True, False, True])
>>> modifier = evm.modifier.BooleanMask(mask, norm)
>>> modifier(hist)
Array([ 5.049494, 20. , 30.296963], dtype=float32)
"""
def __init__(
self,
mask: Bool[Array, ...],
modifier: ModifierBase,
):
self.mask = mask
self.modifier = modifier
def offset_and_scale(self, hist: H) -> OffsetAndScale:
os = self.modifier.offset_and_scale(hist)
def _mask(
true: Bool[Array, ...],
false: Bool[Array, ...],
) -> Bool[Array, ...]:
return jnp.where(self.mask, true, false)
return OffsetAndScale(
offset=_mask(os.offset, jnp.zeros_like(os.offset)),
scale=_mask(os.scale, jnp.ones_like(os.offset)),
).broadcast()
[docs]
class Compose(ModifierBase):
"""Combines multiple modifiers and applies them in parallel.
Modifiers are grouped by their NNX graph structure and vectorized using XLA.
The combined effect multiplies all scale factors together and adds all offsets
to the input histogram. This corresponds to the standard HistFactory-style
combination of independent systematic effects.
Note: This is *parallel* composition, not sequential chaining. For pure
scale modifiers the two are equivalent, but when offsets are involved the
order-independent parallel combination is used.
Args:
*modifiers: Modifiers to compose. They are flattened if nested ``Compose`` instances are provided.
Examples:
>>> import jax.numpy as jnp
>>> import evermore as evm
>>> mu = evm.Parameter(value=1.1)
>>> syst = evm.NormalParameter(value=0.1)
>>> hist = jnp.array([10, 20, 30])
>>> composition = evm.modifier.Compose(
... mu.scale(offset=0, slope=1),
... syst.scale_log_asymmetric(up=1.1, down=0.9),
... )
>>> composition(hist)
Array([11.155, 22.237, 33.318], dtype=float32)
"""
def __init__(self, *modifiers: ModifierBase) -> None:
if not modifiers:
msg = "At least one modifier must be provided to Compose."
raise ValueError(msg)
# unpack the modifiers, if they are already a Compose instance
# this allows for nested compositions, e.g.:
# `Compose(Modifier1, Compose(Modifier2, Modifier3))`
# will flatten to `Compose(Modifier1, Modifier2, Modifier3)`
self.modifiers = nnx.List(unroll(modifiers))
def __len__(self) -> int:
return len(self.modifiers)
def offset_and_scale(self, hist: H) -> OffsetAndScale:
from collections import defaultdict
# initial scale and offset
scale = jnp.ones_like(hist)
offset = jnp.zeros_like(hist)
groups = defaultdict(list)
# first group modifiers into same NNX graph structures
for mod in self.modifiers:
graphdef, state = nnx.split(mod)
groups[graphdef].append(state)
# then do the `vmap` trick to calculate the scale factors in parallel per group.
for graphdef, states in groups.items():
# skip empty groups
if not states:
continue
# Essentially we are turning an array of modifiers (AOS) into a single modifier of stacked leaves (SOA).
# Then we can use XLA's vectorization/loop constructs (e.g.: `jax.vmap` or `jax.lax.scan`) to calculate
# the scale factors without having to compile the fully unrolled loop.
dynamic_stack = tree_stack(states, broadcast_leaves=True)
def calc_sf(graphdef, _hist, dynamic_stack):
stack = nnx.merge(graphdef, dynamic_stack)
os = stack.offset_and_scale(_hist)
# return plain arrays instead of an nnx.Pytree to avoid tracer leaks
# when the vmapped function returns a graph node type on some Python/JAX
# combinations (see evermore#XXX).
return os.offset, os.scale
# Vectorize over the first axis of the stack.
# Using `jax.vmap` is the most efficient way to do this,
# however it needs `hist` and `dynamic_stack` to fit into memory.
# If this is not the case, we should consider using `jax.lax.scan` instead.
# See: https://github.com/jax-ml/jax/discussions/19114#discussioncomment-7996283
vec_calc_sf = jax.vmap(
calc_sf,
in_axes=(
None,
None,
0,
), # vectorize over the batch axis of the dynamic_stack
out_axes=(0, 0), # return trees of offset/scale arrays
)
os_offsets, os_scales = vec_calc_sf(graphdef, hist, dynamic_stack)
scale *= jnp.prod(os_scales, axis=0)
offset += jnp.sum(os_offsets, axis=0)
return OffsetAndScale(offset=offset, scale=scale).broadcast()
def unroll(modifiers: Iterable[ModifierBase]) -> Iterator[ModifierBase]:
# Helper to recursively flatten nested Compose instances into a single list
for mod in modifiers:
if isinstance(mod, Compose):
# recursively yield from the modifiers of the Compose instance
yield from unroll(mod.modifiers)
elif isinstance(mod, ModifierBase):
# yield the modifier if it is a ModifierBase instance
yield mod
else:
# raise an error if the modifier is not a ModifierBase instance
msg = f"Modifier {mod} is not a ModifierBase instance."
raise TypeError(msg)
