First commit in the branch

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:

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 
+

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)]                                               
+                         }}

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="")