Colibri PDF Models
In this section, we present colibri’s core functionality, which enables users to easily implement any PDF parametrisation. We then offer a comprehensive overview of the building blocks and general methodology that support Bayesian (and other) sampling of these parametrisations.
The Colibri PDF model class
In order to decouple the specifics of the implementation of parton distribution functions (PDFs) parametrisations from their numerical inference,
we introduce an abstract base class PDFModel in the colibri framework, see the code-block below for an overview of the class structure.
At its core, a PDFModel must provide:
A list of model parameters, representing the degrees of freedom used to describe the PDF.
A method that defines how these parameters are mapped to PDF values on a specified grid in momentum fraction
x, for each parton flavour.
This abstraction enables users to plug in a variety of model architectures ranging from simple parametric forms to more complex neural network-based approaches, while delegating performance-critical tasks such as convolution with pre‐tabulated perturbative kernels and Bayesian sampling of the parameters to optimized external engines.
In practice, each new PDF parametrisation is implemented as its own colibri application by subclassing PDFModel and completing the two required methods, as detailed in the Les Houches tutorial (TODO: reference).
class PDFModel(ABC):
"""An abstract class describing the key features of a PDF model."""
name = "Abstract PDFModel"
@property
@abstractmethod
def param_names(self) -> list:
"""This should return a list of names for the fitted parameters of the model.
The order of the names is important as it will be assumed to be the order of the parameters
fed to the model.
"""
pass
@abstractmethod
def grid_values_func(self, xgrid: ArrayLike) -> Callable[[jnp.array], jnp.ndarray]:
"""This function should produce a grid values function, which takes
in the model parameters, and produces the PDF values on the grid xgrid.
The grid values function should be a function of the parameters and return
an array of shape (N_fl, Nx). The first dimension is the number of flavours expected
by the FK tables belonging to the chosen theoryID.
The second dimension is the number of points in the xgrid, i.e. Nx = len(xgrid).
Example
-------
::
def grid_values_func(xgrid):
def func(params):
# Define expression for each flavour
fl_1 = params[0] + params[1] * xgrid
fl_2 = params[2] + params[3] * xgrid
# Combine the flavours into a single array
# This is just an example, the actual implementation will depend on the model
# and the number of flavours
return jnp.array([fl_1, fl_2])
return func
"""
pass
def pred_and_pdf_func(
self,
xgrid: ArrayLike,
forward_map: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray],
) -> Callable[[jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, jnp.ndarray]]:
"""Creates a function that returns a tuple of two arrays, given the model parameters and the fast kernel arrays as input.
The returned function produces:
- The first array: 1D vector of theory predictions for the data.
- The second array: PDF values evaluated on the x-grid, using `self.grid_values_func`, with shape (Nfl, Nx).
The `forward_map` is used to map the PDF values defined on the x-grid and the fast kernel arrays into the corresponding theory prediction vector.
"""
pdf_func = self.grid_values_func(xgrid)
def pred_and_pdf(params, fast_kernel_arrays):
"""
Parameters
----------
params: jnp.array
The model parameters.
fast_kernel_arrays: tuple
tuple of tuples of jnp.arrays
The FK tables to use.
Returns
-------
tuple
The predictions and the PDF values.
"""
pdf = pdf_func(params)
predictions = forward_map(pdf, fast_kernel_arrays)
return predictions, pdf
return pred_and_pdf
Parameter Specification
Each concrete PDF model must provide a parameter list via a param_names property. This list of strings defines the parameter names (e.g. normalizations, exponents, polynomial coefficients) in a fixed order.
Grid Evaluation Method
The core of the PDFModel class is the grid_values_func method, which returns a JAX-compatible function
mapping an \(N_{\rm p}\)-dimensional parameter vector \(\boldsymbol{\theta}\) into the PDF values (note: the function actually returns x*PDF values since this is the object convoluted with the FK-tables) for each parton flavour index evaluated on the user-provided \(x\) grid of length \(N_{\rm x}\). In practice, for a standard PDF fit, the user only needs to define the expression mapping the \(x\)-grid to the PDF values. The framework then automatically handles the construction of all the resources, such as theory predictions, needed for a PDF fit.
Prediction Construction
To compute physical observables (structure functions, cross sections, etc.), one must convolve the PDFs with perturbative coefficient functions.
In colibri this is handled via the pred_and_pdf_func method, which takes the x-grid and a forward map mapping the PDF to the physical observable,
and produces a function taking as input the PDF parameters and a tuple of fast-kernel arrays
that (i) evaluates the PDF on the grid via grid_values_func, and (ii) feeds the resulting \(N_{\rm fl} \times N_{\rm x}\)
array into the supplied forward_map to yield a 1D vector of theory predictions for all data points.
Note
The prediction function is already implemented, however the user is allowed to override it in its own PDF application if the specific model needs extra features.
Design Rationale and Benefits
Modularity: New PDF parametrisations can be added by defining only two methods, without touching the core fitting or convolution engines.
Performance: High-performance array computations and GPU compatibility thanks to the framework being written in JAX.
Universality: All PDF models share the same inference methods as well as data and theory predictions, allowing for more reliable comparison and studies of methodological differences.