Forward Map
Colibri provides a flexible platform for fitting parton distribution function (PDF) models to data involving at least one incoming proton. The data is modeled in the framework of collinear QCD factorization, where the scattering process is written as a convolution of the PDFs with perturbatively-computed hard-scattering cross sections.
In this context, inferring the PDFs from experimental measurements is an inverse problem: the unknowns are the PDFs, and the forward model consists of the hard-scattering cross section combined with PDF evolution kernels, commonly stored as FK tables. For more details on FK tables and evolution kernels, see [CHM22].
We distinguish two classes of forward maps based on whether the initial state involves one proton (Deep Inelastic Scattering, DIS) or two protons (hadron–hadron collisions).
DIS predictions
DIS data is the most abundant data type in the global PDF fits and is the most straightforward to model. For example, a measurement of the \(F_2\) structure function consisting of \(N_{\rm data}\) points, can be written as the contraction of two operators,
where the \(FK_{i,j,k}\) operator has shape \((N_{\rm data}, N_{\rm fl}, N_{\rm x})\) and \(f_{j,k}\) is the \((N_{\rm fl}, N_{\rm x})\) dimensional grid representing the PDF values at the input scale \(Q_0^2\).
In colibri, the forward modeling of DIS data is taken care of in the colibri.theory_predictions module
by the colibri.theory_predictions.make_dis_prediction() closure function which essentially takes in input an
FK-table object and returns a function that computes the DIS prediction for a given PDF and FK-array.
def make_dis_prediction(
fktable, FIT_XGRID, flavour_indices=None, fill_fk_xgrid_with_zeros=False
):
"""
Closure to compute the theory prediction for a DIS observable.
Parameters
----------
fktable : validphys.coredata.FKTableData
The fktable should be a validphys.coredata.FKTableData instance
and with cuts and masked flavours already applied.
FIT_XGRID: np.ndarray
xgrid of the theory, computed by a production rule by taking
the sorted union of the xgrids of the datasets entering the fit.
flavour_indices: list, default is None
fill_fk_xgrid_with_zeros: bool, default is False
If True, then the missing xgrid points in the FK table
will be filled with zeros. This is useful when the FK table
is needed as tensor of shape (Ndat, Nfl, Nfk_xgrid) with Nfk_xgrid and Nfl fixed
for all datasets.
Returns
-------
Callable
"""
lumi_indices = mask_luminosity_mapping(fktable, flavour_indices)
fk_xgrid_indices = fktable_xgrid_indices(
fktable, FIT_XGRID, fill_fk_xgrid_with_zeros=fill_fk_xgrid_with_zeros
)
def dis_prediction(pdf, fk_arr):
"""
Function to compute the theory prediction for a DIS observable.
Note that when running an ultranest fit this function gets compiled by jax.jit,
hence, ideally, this function should be pure.
However, luminosity indices and fk_xgrid_indices don't take much memory
and hence are left as global variables.
For more details on jax.jit issues with non pure functions see e.g.
https://github.com/google/jax/issues/5071
Parameters
----------
pdf: jnp.ndarray
pdf grid (shape is Nfl,Nx)
fk_arr: jnp.ndarray
fktable array
Returns
-------
jnp.ndarray
theory prediction for a hadronic observable (shape is Ndata, )
"""
return jnp.einsum(
"ijk, jk ->i", fk_arr, pdf[lumi_indices, :][:, fk_xgrid_indices]
)
return dis_prediction
Hadron-Hadron predictions
Hadron-hadron collisions are more complicated to model than DIS data, as they involve the convolution of two incoming partons, each with their own PDF. A \(N_{\rm data}\)-point measurement \(\sigma\) of a hadron-hadron cross section can be written as
where the \(FK_{i,j,k,l,m}\) operator has shape \((N_{\rm data}, N_{\rm fl}, N_{\rm fl}, N_{\rm x}, N_{\rm x})\).
In colibri, the forward modeling of hadron-hadron data is taken care of in the colibri.theory_predictions module
by the colibri.theory_predictions.make_had_prediction() closure function shown below.
def make_had_prediction(
fktable, FIT_XGRID, flavour_indices=None, fill_fk_xgrid_with_zeros=False
):
"""
Closure to compute the theory prediction for a Hadronic observable.
Parameters
----------
fktable : validphys.coredata.FKTableData
FIT_XGRID: np.ndarray
xgrid of the theory, computed by a production rule by taking
the sorted union of the xgrids of the datasets entering the fit.
flavour_indices: list, default is None
fill_fk_xgrid_with_zeros: bool, default is False
If True, then the missing xgrid points in the FK table
will be filled with zeros. This is useful when the FK table
is needed as tensor of shape (Ndat, Nfl, Nfk_xgrid) with Nfk_xgrid and Nfl fixed
for all datasets.
Returns
-------
Callable
"""
lumi_indices = mask_luminosity_mapping(fktable, flavour_indices)
first_lumi_indices = lumi_indices[0::2]
second_lumi_indices = lumi_indices[1::2]
fk_xgrid_indices = fktable_xgrid_indices(
fktable, FIT_XGRID, fill_fk_xgrid_with_zeros=fill_fk_xgrid_with_zeros
)
def had_prediction(pdf, fk_arr):
"""
Function to compute the theory prediction for a Hadronic observable.
Note that when running an ultranest fit this function gets compiled by jax.jit,
hence, ideally, this function should be pure.
However, luminosity indices and fk_xgrid_indices don't take much memory
and hence are left as global variables.
For more details on jax.jit issues with non pure functions see e.g.
https://github.com/google/jax/issues/5071
Parameters
----------
pdf: jnp.ndarray
pdf grid (shape is Nfl,Nx)
fk_arr: jnp.ndarray
fktable array
Returns
-------
jnp.ndarray
theory prediction for a hadronic observable (shape is Ndata, )
"""
return jnp.einsum(
"ijkl,jk,jl->i",
fk_arr,
pdf[first_lumi_indices, :][:, fk_xgrid_indices],
pdf[second_lumi_indices, :][:, fk_xgrid_indices],
)
return had_prediction