diff --git a/src/smefit/chi2.py b/src/smefit/chi2.py index 24ebc859..117ecc6e 100644 --- a/src/smefit/chi2.py +++ b/src/smefit/chi2.py @@ -6,8 +6,7 @@ import numpy as np import scipy.optimize as opt from rich.progress import track -import jax.numpy as jnp - +from . import compute_theory_kappa as prk from . import compute_theory as pr from .coefficients import CoefficientManager from .loader import DataTuple, load_datasets @@ -44,9 +43,10 @@ def compute_chi2( """ # compute theory prediction for each point in the dataset - theory_predictions = pr.make_predictions( + + theory_predictions = prk.make_predictions( dataset, coefficients_values, use_quad, use_multiplicative_prescription - ) + ) #SIMONE: Here I chagned the imported library callinkg prk.make_predictions. I would add a label kappa_frm in the run_card that is going to be read here, is it is true make prediction will make them in the kappa-framework, as explained in the compute_theory_kappa.py file # compute experimental central values - theory if use_replica: @@ -56,7 +56,7 @@ def compute_chi2( invcovmat = dataset.InvCovMat # note @ is slower when running with mpiexec - return jnp.einsum("i,ij,j->", diff, invcovmat, diff) + return np.einsum("i,ij,j->", diff, invcovmat, diff) class Scanner: diff --git a/src/smefit/compute_theory_kappa.py b/src/smefit/compute_theory_kappa.py new file mode 100644 index 00000000..265a43af --- /dev/null +++ b/src/smefit/compute_theory_kappa.py @@ -0,0 +1,73 @@ +# -*- coding: utf-8 -*- + +""" +Module for the generation of theory predictions +""" +import numpy as np +import jax.numpy as jnp + + +def flatten(quad_mat, axis=0): + """ + Delete lower triangular part of a quadratic matrix + and flatten it into an array + + Parameters + ---------- + quad_mat: numpy.ndarray + tensor to flatten + axis: int + axis along which the triangular part is selected + """ + size = quad_mat.shape[axis] + return quad_mat[np.triu_indices(size)] + + +def make_predictions( + dataset, coefficients_values, use_quad, use_multiplicative_prescription +): + """ + Generate the corrected theory predictions for dataset + given a set of |SMEFT| coefficients. + + Parameters + ---------- + dataset : DataTuple + dataset tuple + coefficients_values : numpy.ndarray + |EFT| coefficients values + use_quad: bool + if True include also |HO| corrections + Returns + ------- + corrected_theory : numpy.ndarray + SM + EFT theory predictions + """ + + coefficients_values = jnp.array(coefficients_values) + + # Compute total linear correction + # note @ is slower when running with mpiexec + summed_corrections = jnp.einsum( + "ij,j->i", dataset.LinearCorrections, coefficients_values + ) + + # Compute total quadratic correction + coeff2=jnp.power(coefficients_values,2) + """ + if use_quad: INSTEAD OF USE_QUAD we should put a kappa_frm as label, that would also enter the run_card + """ + #each kappa is entering squared + coeff_outer_coeff = jnp.outer(coeff2, coeff2) + # note @ is slower when running with mpiexec + summed_quad_corrections = jnp.einsum( + "ij,j->i", dataset.QuadraticCorrections, flatten(coeff_outer_coeff) + ) + summed_corrections += summed_quad_corrections + + # Sum of SM theory + SMEFT corrections + # define the total k_H + k_weights=jnp.array([0.57,0.06,0.029,0.22,0.027,0.086,0.0023,0.0016,0.00022,0.]) #If this is going to become a module we should put this array somewhere else + k_H=jnp.sum(k_weights*coeff2) + return summed_corrections/k_H + diff --git a/src/smefit/writetheorycardkappa/kappa.py b/src/smefit/writetheorycardkappa/kappa.py new file mode 100644 index 00000000..21a412fe --- /dev/null +++ b/src/smefit/writetheorycardkappa/kappa.py @@ -0,0 +1,21 @@ +ATLAS_SSinc_RunII={"I":{"kb2":[0 for _ in range(16)], + "ktau2":[0 for _ in range(16)], + "kc2":[0 for _ in range(16)], + "kW2":[0,0,0,0,0.73,0.73,0.73,0.73,0.73,0.66,0.66,0.66,0,0,0,0], + "kZ2":[0,0,0,0,0.27,0.27,0.27,0.27,0.27,0.34,0.34,0.34,0,0,0,0], + "kg2":[1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0], + "kgamma2":[0 for _ in range(16)], + "kZgamma2":[0 for _ in range(16)], + "kmu2":[0 for _ in range(16)], + "kt2":[0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,1]}, + "F":{"kb2":[0,0,0,0,0,0,0,0,1,0,0,1,0,0,0,1], + "ktau2":[0,0,0,1,0,0,0,1,0,0,0,0,0,0,1,0], + "kc2":[0 for _ in range(16)], + "kW2":[0,0,1,0,0,0,1,0,0,0,0,0,0,0.88,0,0], + "kZ2":[0,1,0,0,0,1,0,0,0,0,1,0,0,0.12,0,0], + "kg2":[0 for _ in range(16)], + "kgamma2":[1,0,0,0,1,0,0,0,0,1,0,0,1,0,0,0], + "kZgamma2":[0 for _ in range(16)], + "kmu2":[0 for _ in range(16)] , + "kt2":[0 for _ in range(16)] + }} \ No newline at end of file diff --git a/src/smefit/writetheorycardkappa/write_kappa.py b/src/smefit/writetheorycardkappa/write_kappa.py new file mode 100644 index 00000000..840b72bc --- /dev/null +++ b/src/smefit/writetheorycardkappa/write_kappa.py @@ -0,0 +1,49 @@ +import json +import kappa +#function to combine the different kappas labels +"""def kappa_combiner(arr): + #arr is an array of string of kappas + mult_array=[] + len_=len(arr) + for i in range(len_): + for j in range(i,len_): + mult_array.append(kappas0[i]+"*"+kappas0[j]) + return mult_array""" +n=16 #number of point in the dataset +EMPTY=[0. for _ in range(n)] + +#create the different kappas combination +kappas0 = ["kb2","ktau2","kc2","kW2","kZ2","kg2","kgamma2","kZgamma2","kmu2","kt2"] +#kappas0_combination=kappa_combiner(kappas0) + +#We read the array from an external file, is a dictionary with I (initial) and F (final) states. +dict_dataset=kappa.ATLAS_SSinc_RunII #change this line for different database + + +kappacombination_dict={} +I=dict_dataset["I"] +F=dict_dataset["F"] + +for i_,istate in enumerate(I.keys()): + for fstate in list(F.keys())[i_:]: + if not(istate==fstate): + arr1=[i_val*f_val for i_val,f_val in zip(I[istate],F[fstate])] + arr2=[i_val*f_val for i_val,f_val in zip(I[fstate],F[istate])] #swith to consider them all if not the same label + kappacombination_dict[istate+"*"+fstate]=[a1+a2 for a1,a2 in zip(arr1,arr2)] + else: + kappacombination_dict[istate+"*"+fstate]=[i_val*f_val for i_val,f_val in zip(I[istate],F[fstate])] + + +#create the dictionary and print it in a file + +kappa_dict={} +LO_dict={} +kappa_dict["best_sm"]=EMPTY +kappa_dict["theory_cov"]=[EMPTY for _ in range(n)] +LO_dict["SM"]=EMPTY +for k in kappas0: + LO_dict[k]=EMPTY +LO_dict.update(kappacombination_dict) +kappa_dict["LO"]=LO_dict +with open("ATLAS_SSinc_RunII.json", "w") as outfile: + json.dump(kappa_dict, outfile,indent="") \ No newline at end of file