"""
Module for the generation of theory predictions
"""
import numpy as np
[docs]def make_predictions(config, dataset, coeffs, labels):
""" Generate the corrected theory predictions for dataset `set`
given a set of SMEFT coefficients `coeffs`. Optionally a specific
operator may be selected with `iop` """
# Compute total linear correction
idx = np.where(dataset.CorrectionsDICT == labels)[0]
summed_corrections = np.einsum("ij,j->i", dataset.CorrectionsVAL, coeffs[idx])
# Compute total quadratic correction
if config["HOlambda"] == "HO":
idx1 = config["HOindex1"]
idx2 = config["HOindex2"]
coeffs_quad = coeffs[idx1] * coeffs[idx2]
summed_quad_corrections = np.einsum(
"ij,j->i", dataset.HOcorrectionsVAL, coeffs_quad
)
# Sum of SM theory + SMEFT corrections
if config["HOlambda"] == "HO":
corrected_theory = (
dataset.SMTheory + summed_corrections + summed_quad_corrections
)
else:
corrected_theory = dataset.SMTheory + summed_corrections
return corrected_theory