Forward Map
Colibri provides a flexible platform for fitting PDF models to data involving at least one incoming proton. The data is modeled in the framework of collinear QCD factorisation, 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 Ref. CHM22b.
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 takes as
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
Must match the value used to build the FK arrays passed to the
returned closure (see :func:`fast_kernel_arrays`). If True the FK
array is padded to the full fit x-grid and therefore has to be
masked back onto the FK x-grid before the convolution; if False the
FK array already lives on the FK x-grid and must not be masked.
Returns
-------
Callable
"""
lumi_indices = mask_luminosity_mapping(fktable, flavour_indices)
fk_xgrid_indices = fktable_xgrid_indices(fktable, FIT_XGRID)
# NOTE: the FK array is indexed on the fit x-grid only when it has been
# zero-padded onto it; otherwise its x-grid axis is already the FK x-grid.
if fill_fk_xgrid_with_zeros:
def mask_fk_arr(fk_arr):
return fk_arr[:, :, fk_xgrid_indices]
else:
def mask_fk_arr(fk_arr):
return fk_arr
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, )
"""
# NOTE: for computational efficiency, in the convolution, we only sum over xgrid points that are non-zero in the FK table.
return jnp.einsum(
"ijk, jk ->i",
mask_fk_arr(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
Must match the value used to build the FK arrays passed to the
returned closure (see :func:`fast_kernel_arrays`). If True the FK
array is padded to the full fit x-grid and therefore has to be
masked back onto the FK x-grid before the convolution; if False the
FK array already lives on the FK x-grid and must not be masked.
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)
# NOTE: the FK array is indexed on the fit x-grid only when it has been
# zero-padded onto it; otherwise its x-grid axes are already the FK x-grid.
if fill_fk_xgrid_with_zeros:
def mask_fk_arr(fk_arr):
return fk_arr[:, :, fk_xgrid_indices[:, None], fk_xgrid_indices[None, :]]
else:
def mask_fk_arr(fk_arr):
return fk_arr
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, )
"""
# NOTE: for computational efficiency, in the convolution, we only sum over xgrid points that are non-zero in the FK table.
return jnp.einsum(
"ijkl,jk,jl->i",
mask_fk_arr(fk_arr),
pdf[first_lumi_indices, :][:, fk_xgrid_indices],
pdf[second_lumi_indices, :][:, fk_xgrid_indices],
)
return had_prediction