"""
n3fit interface to validphys
Example
-------
>>> import numpy as np
>>> from n3fit.vpinterface import N3PDF
>>> from n3fit.model_gen import generate_pdf_model, ReplicaSettings
>>> from validphys.pdfgrids import xplotting_grid
>>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 's', 'sbar', 'g']]
>>> fake_x = np.linspace(1e-3,0.8,3)
>>> rps = [ReplicaSettings(nodes=[8], activations=["linear"], seed=4)]*4
>>> pdf_model = generate_pdf_model(rps, flav_info=fake_fl, fitbasis='FLAVOUR')
>>> n3pdf = N3PDF(pdf_model.split_replicas())
>>> res = xplotting_grid(n3pdf, 1.6, fake_x)
>>> res.grid_values.error_members().shape
(4, 8, 3)
# (nreplicas, flavours, x-grid)
"""
from collections.abc import Iterable
from dataclasses import dataclass
from functools import cached_property
import logging
import numpy as np
import pandas as pd
import scipy.linalg as la
from validphys.arclength import arc_lengths, integrability_number
from validphys.calcutils import calc_chi2, calc_phi
from validphys.convolution import central_predictions, predictions
from validphys.core import PDF, MCStats
from validphys.covmats import dataset_inputs_covmat_from_systematics
from validphys.lhapdfset import LHAPDFSet
from validphys.pdfbases import ALL_FLAVOURS, check_basis
log = logging.getLogger(__name__)
# Order of the evolution basis output from n3fit
EVOL_LIST = [
"photon",
"sigma",
"gluon",
"V",
"V3",
"V8",
"V15",
"V24",
"V35",
"T3",
"T8",
"T15",
"T24",
"T35",
]
[docs]@dataclass
class HyperoptMetrics:
chi2: float
phi2: float
logp: float
[docs]class N3Stats(MCStats):
"""The PDFs from n3fit are MC PDFs
however, since there is no grid, the CV has to be computed manually"""
[docs] def error_members(self):
return self.data
[docs] def central_value(self):
return np.mean(self.data, axis=0)
[docs]class N3LHAPDFSet(LHAPDFSet):
"""Extension of LHAPDFSet using n3fit models"""
def __init__(self, name, pdf_models, Q=1.65, is_t0=False):
log.debug("Creating LHAPDF-like n3fit PDF")
self._error_type = "replicas"
self._name = name
self._lhapdf_set = pdf_models
self._flavors = None
self._fitting_q = Q
self._is_t0 = is_t0
self.basis = check_basis("evolution", EVOL_LIST)["basis"]
[docs] def xfxQ(self, x, Q, n, fl):
"""Return the value of the PDF member for the given value in x"""
if Q != self._fitting_q:
log.warning(
"Querying N3LHAPDFSet at a value of Q=%f different from %f", Q, self._fitting_q
)
return self.grid_values([fl], [x]).squeeze()[n]
def _register_photon(self, xgrid):
"""If the PDF models contain photons, register the xgrid with them"""
for m in self._lhapdf_set:
pl = m.get_layer_re("add_photon")
# if pl is an empy list there's no photon
if not pl:
continue
pl[0].register_photon(xgrid)
# Recompile the model if necessary
if not pl[0].built:
m.compile()
def __call__(self, xarr, flavours=None, replica=None):
"""Uses the internal model to produce pdf values for the grid
The output is on the evolution basis.
Parameters
----------
xarr: numpy.ndarray
x-points with shape (xgrid_size,) (size-1 dimensions are removed)
flavours: list
list of flavours to output
replica: int
replica whose value must be returned (by default return all members)
replica 0 corresponds to the central value
Returns
-------
numpy.ndarray
(xgrid_size, flavours) pdf result
"""
if flavours is None:
flavours = EVOL_LIST
# Ensures that the input has the shape the model expect, no matter the input
# as the scaling is done by the model itself
mod_xgrid = xarr.reshape(1, -1, 1)
# Register the grid with the photon
self._register_photon(mod_xgrid)
if replica is None or replica == 0 or self._is_t0:
# We need generate output values for all replicas
result = np.concatenate(
[m.predict({"pdf_input": mod_xgrid}) for m in self._lhapdf_set], axis=0
)
if replica == 0 or self._is_t0:
# We want _only_ the central value
result = np.mean(result, axis=0, keepdims=True)
else:
result = self._lhapdf_set[replica - 1].predict({"pdf_input": mod_xgrid})
if flavours != "n3fit":
# Ensure that the result has its flavour in the basis-defined order
ii = self.basis._to_indexes(flavours)
result[:, :, ii] = result
return result
[docs] def grid_values(self, flavours, xarr, qmat=None):
"""
Parameters
----------
flavours: numpy.ndarray
flavours to compute
xarr: numpy.ndarray
x-points to compute, dim: (xgrid_size,)
qmat: numpy.ndarray
q-points to compute (not used by n3fit, used only for shaping purposes)
Returns
------
numpy.ndarray
array of shape (replicas, flavours, xgrid_size, qmat) with the values of
the ``pdf_model``(s) evaluated in ``xarr``
"""
n3fit_result = self(xarr.reshape(1, -1, 1))
# The results of n3fit are always in the 14-evolution basis used in fktables
# the calls to grid_values always assume the result will be LHAPDF flavours
# we need then to rotate them to the LHAPDF-flavour basis,
# we don't care that much for precision here
to_flav = la.inv(self.basis.from_flavour_mat)
to_flav[np.abs(to_flav) < 1e-12] = 0.0
flav_result = np.tensordot(n3fit_result, to_flav, axes=(-1, 1))
# Now drop the indices that are not requested
requested_flavours = [ALL_FLAVOURS.index(i) for i in flavours]
flav_result = flav_result[..., requested_flavours]
# Swap the flavour and xgrid axis for vp-compatibility
ret = flav_result.swapaxes(-2, -1)
# If given, insert as many copies of the grid as q values
ret = np.expand_dims(ret, -1)
if qmat is not None:
lq = len(qmat)
ret = ret.repeat(lq, -1)
return ret
[docs]class N3PDF(PDF):
"""
Creates a N3PDF object, extension of the validphys PDF object to perform calculation
with a n3fit generated model.
Parameters
----------
pdf_models: :py:class:`n3fit.backends.MetaModel` (or list thereof)
PDF trained with n3fit, x -> f(x)_{i} where i are the flavours in the evol basis
fit_basis: list(dict)
basis of the training, used for reporting
name: str
name of the N3PDF object
"""
def __init__(self, pdf_models, fit_basis=None, name="n3fit", Q=1.65):
self.fit_basis = fit_basis
if isinstance(pdf_models, Iterable):
self._models = pdf_models
else:
self._models = [pdf_models]
super().__init__(name)
self._stats_class = N3Stats
self._Q = Q
# Since there is no info file, create a fake `_info` dictionary
self._info = {"ErrorType": "replicas", "NumMembers": len(self._models)}
[docs] def select_models(self, indexes):
"""Given a list of indexes, return a ``N3PDF`` class with the same
settings as this one and the subset of models."""
# TODO: once the PDF class is moved to a dataclass as per a very old issue, this can simply use dataclass.replace
new_models = [self._models[i] for i in indexes]
return self.__class__(
new_models, fit_basis=self.fit_basis, name=f"{self.name}_reduced", Q=self._Q
)
@cached_property
def _lhapdf_set(self):
return N3LHAPDFSet(self.name, self._models, Q=self._Q)
[docs] def load(self):
"""If the function needs an LHAPDF object, return a N3LHAPDFSet"""
return self._lhapdf_set
[docs] def load_t0(self):
"""Load the central PDF object"""
return N3LHAPDFSet(self.name, self._models, Q=self._Q, is_t0=True)
[docs] def get_nn_weights(self):
"""Outputs all weights of the NN as numpy.ndarrays"""
return [model.get_weights() for model in self._models]
[docs] def get_preprocessing_factors(self, replica=None):
"""Loads the preprocessing alpha and beta arrays from the PDF trained model.
If a ``fit_basis`` given in the format of ``n3fit`` runcards is given it will be used
to generate a new dictionary with the names, the exponent and whether they are trainable
otherwise outputs a Nx2 array where [:,0] are alphas and [:,1] betas
"""
# If no replica is explicitly requested, get the preprocessing layer for the first model
# remember replicas start counting at one
if replica is None:
replica = 1
# Keep the import here to avoid loading the backend when it is not necessary
from n3fit.backends import PREPROCESSING_LAYER_ALL_REPLICAS
preprocessing_layer = self._models[replica - 1].get_layer(PREPROCESSING_LAYER_ALL_REPLICAS)
alphas_and_betas = None
if self.fit_basis is not None:
output_dictionaries = []
for d in self.fit_basis:
flavour = d["fl"]
alpha = preprocessing_layer.get_weight_by_name(f"alpha_{flavour}")
beta = preprocessing_layer.get_weight_by_name(f"beta_{flavour}")
if alpha is not None:
alpha = float(alpha.numpy())
if beta is not None:
beta = float(beta.numpy())
output_dictionaries.append(
{
"fl": flavour,
"smallx": alpha,
"largex": beta,
"trainable": d.get("trainable", True),
}
)
alphas_and_betas = output_dictionaries
return alphas_and_betas
def __call__(self, xarr, flavours=None, replica=None):
"""Uses the internal model to produce pdf values for the grid
The output is on the evolution basis.
Parameters
----------
xarr: numpy.ndarray
x-points with shape (xgrid_size,) (size-1 dimensions are removed)
flavours: list
list of flavours to output
replica: int
replica whose value must be returned (by default return all members)
replica 0 corresponds to the central value
Returns
-------
numpy.ndarray
(xgrid_size, flavours) pdf result
"""
return self._lhapdf_set(xarr, flavours=flavours, replica=replica)
# Utilities and wrapper to avoid having to pass around unnecessary information
[docs]def integrability_numbers(n3pdf, q0=1.65, flavours=None):
"""Compute the integrability numbers for the current PDF
using the corresponding validphys action
Parameters
----------
q0: float
energy at which the integrability is computed
flavours: list
flavours for which the integrability is computed
Returns
-------
np.array(float)
Value for the integrability for each of the flavours
Example
-------
>>> from n3fit.vpinterface import N3PDF, integrability_numbers
>>> from n3fit.model_gen import pdfNN_layer_generator
>>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'g', 's', 'sbar']]
>>> pdf_model = pdfNN_layer_generator(nodes=[8], activations=['linear'], seed=0, flav_info=fake_fl, fitbasis="FLAVOUR")
>>> n3pdf = N3PDF(pdf_model)
>>> res = integrability_numbers(n3pdf)
"""
if flavours is None:
flavours = ["V", "T3", "V3", "T8", "V8"]
return integrability_number(n3pdf, [q0], flavours=flavours)
[docs]def compute_arclength(self, q0=1.65, basis="evolution", flavours=None):
"""
Given the layer with the fit basis computes the arc length
using the corresponding validphys action
Parameters
----------
pdf_function: function
pdf function has received by the writer or ``pdf_model``
q0: float
energy at which the arc length is computed
basis: str
basis in which to compute the arc length
flavours: list
output flavours
Example
-------
>>> from n3fit.vpinterface import N3PDF, compute_arclength
>>> from n3fit.model_gen import pdfNN_layer_generator
>>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'g', 's', 'sbar']]
>>> pdf_model = pdfNN_layer_generator(nodes=[8], activations=['linear'], seed=0, flav_info=fake_fl, fitbasis="FLAVOUR")
>>> n3pdf = N3PDF(pdf_model)
>>> res = compute_arclength(n3pdf)
"""
if flavours is None:
flavours = ["sigma", "gluon", "V", "V3", "V8"]
ret = arc_lengths(self, [q0], basis, flavours)
return ret.stats.central_value()
[docs]def compute_hyperopt_metrics(n3pdf, experimental_data) -> HyperoptMetrics:
"""Compute the different hyperopt quantities from which one defines
the hyperopt metric.
Parameters
----------
n3pdfs: :class:`n3fit.vpinterface.N3PDF`
`N3PDF` instance defining the n3fitted multi-replica PDF
experimental_data: List[validphys.core.DataGroupSpec]
List of experiment group datasets as `DataGroupSpec` instances
Returns
-------
HyperoptMetrics: :class:`n3fit.vpinterface.HyperoptMetrics`
dataclass holding the values of chi2, phi2 and logp
Example
-------
>>> from n3fit.vpinterface import N3PDF, compute_hyperopt_metrics
>>> from n3fit.model_gen import generate_pdf_model, ReplicaSettings
>>> from validphys.loader import Loader
>>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'g', 's', 'sbar']]
>>> rps = [ReplicaSettings(nodes=[8], activations=["linear"], seed=i) for i in [0,1]]
>>> pdf_model = generate_pdf_model(rps, flav_info=fake_fl, fitbasis="FLAVOUR")
>>> n3pdf = N3PDF(pdf_model.split_replicas())
>>> ds = Loader().check_dataset("NMC_NC_NOTFIXED_P_EM-SIGMARED", theoryid=40_000_000, cuts="internal", variant="legacy")
>>> data_group_spec = Loader().check_experiment("My DataGroupSpec", [ds])
>>> hyperopt_losses = compute_hyperopt_metrics(n3pdf, [data_group_spec])
"""
exp_cv = []
th_cvs = []
th_rep = []
cds_list = []
# Loop over the list of `DataGroupSpec` objects
for datagroupspec in experimental_data:
# datagroupspec is an instance of `DataGroupSpec`
# Loop over `DataGroupSpec` datasets
for datasetspec in datagroupspec.datasets:
# datasetspec is an instance of `DataSetSpec`
# update list of CommonData and corresponding central values
cd = datasetspec.load_commondata()
cds_list.append(cd)
exp_cv.append(cd.central_values)
# update list of th pred, for the central value and for each replica
th_cvs.append(central_predictions(datasetspec, n3pdf))
th_rep.append(predictions(datasetspec, n3pdf))
pred_cvs = pd.concat(th_cvs, axis=0, ignore_index=True)
pred_rep = pd.concat(th_rep, axis=0, ignore_index=True)
expr_cvs = pd.concat(exp_cv, axis=0, ignore_index=True)
diffs = pred_cvs.values.flatten() - expr_cvs.values.flatten()
diffs_reps = pred_rep.values - expr_cvs.values[:, np.newaxis]
exp_cov = dataset_inputs_covmat_from_systematics(cds_list, use_weights_in_covmat=False)
exp_covmat_col = la.cholesky(exp_cov, lower=True)
# If there is only one replica, we don't account for PDF covmat
if pred_rep.shape[1] == 1:
total_covmat = exp_cov
else:
pdf_cov = np.cov(pred_rep.values)
assert exp_cov.shape == pdf_cov.shape
total_covmat = exp_cov + pdf_cov
# Compute the log_det
# Normalize the total covmat to central values of experimental data
norm_total_covmat = total_covmat / np.outer(
expr_cvs.values.flatten(), expr_cvs.values.flatten()
)
norm_total_covmat_chol = la.cholesky(norm_total_covmat, lower=True)
log_det_total_cov = 2 * np.sum(np.log(np.diag(norm_total_covmat_chol)))
# Compute the chi2
total_covmat_chol = la.cholesky(total_covmat, lower=True)
chi2 = calc_chi2(sqrtcov=total_covmat_chol, diffs=diffs)
# Compute phi2
phi2 = calc_phi(sqrtcov=exp_covmat_col, diffs=diffs_reps)
ndat = len(diffs)
logp = -0.5 * (len(diffs) * np.log(2 * np.pi) + log_det_total_cov + chi2)
return HyperoptMetrics(chi2=chi2 / ndat, phi2=phi2, logp=-logp / ndat)