Allowing OOs to use precomputed RGE matrix with fixed scale.

external_chi2/optimal_observables/interface_oos_cepc.py

@@ -2,7 +2,7 @@
 import pathlib
 import jax.numpy as jnp
 from smefit import log
-from smefit.rge.rge import RGE
+from smefit.rge.rge import RGE, load_precomputed_rge_matrix, load_rge_mats_from_scales
 from smefit.coefficients import CoefficientManager
 
 current_file_path = pathlib.Path(__file__).resolve().parent
@@ -13,15 +13,24 @@
 _logger = log.logging.getLogger(__name__)
 
 
+# function to compute the RGE matrix at scale
 def compute_rge_matrix(coefficients, rge_dict, scale):
+    """
+    Compute the RGE solution at the scale `scale` and return it as an array.
+    It also returns an updated list of coefficients that includes the new operators generated by the RGE.
+    """
+    if type(rge_dict["obs_scale"]) == float:
+        low_scale = rge_dict["obs_scale"]
+    else:
+        low_scale = scale
     rge_runner = RGE(
         wc_names=coefficients.name,
         init_scale=rge_dict["init_scale"],
         accuracy=rge_dict.get("smeft_accuracy", "integrate"),
         adm_QCD=rge_dict.get("adm_QCD", False),
         yukawa=rge_dict.get("yukawa", "top"),
     )
-    rge_df = rge_runner.RGEmatrix(scale)
+    rge_df = rge_runner.RGEmatrix(low_scale)
     gen_operators = list(rge_df.index)
     _logger.info("The operators generated by the RGE are: ")
     _logger.info(gen_operators)
@@ -31,14 +40,80 @@ def compute_rge_matrix(coefficients, rge_dict, scale):
     return rge_df.values, new_coeffs
 
 
-class OptimalWWCEPC161:
-    def __init__(self, coefficients, rge_dict=None):
-        if rge_dict is not None:
-            rgemat, coefficients = compute_rge_matrix(
-                coefficients, rge_dict, scale=161.0
+# function that manages the inclusion og the RGE matrix in the optimal observables classes. It returns the RGE matrix and the updated coefficients.
+# It decides if the RGE matrix should be computed or loaded from a precomputed file based on the rge_dict input.
+def rge_matrix_manager(coefficients, rge_dict, scale):
+    """
+    Manage the inclusion of the RGE matrix in the optimal observables classes.
+    It returns the RGE matrix and the updated coefficients.
+    It decides if the RGE matrix should be computed or loaded from a precomputed file based on the rge_dict input.
+    If the RGE dictionary includes a float as the "obs_scale", it ignores the scale input and uses the "obs_scale"
+    as the scale at which to compute the RGE matrix.
+    If "obs_scale" is a float, it also checks if there is a precomputed RGE matrix for that scale and loads it if it exists, otherwise it computes the RGE matrix and returns it.
+    If "obs_scale" is not a float, it computes the RGE matrix at the scale input and returns it.
+    Future improvements should include the possibility to load precomputed RGE matrices for the case of dynamical scales.
+    """
+    if rge_dict is not None:
+        # need for the RGE. Initializes the runner.
+        rge_runner = RGE(
+            wc_names=coefficients.name,
+            init_scale=rge_dict["init_scale"],
+            accuracy=rge_dict["smeft_accuracy"],
+            adm_QCD=rge_dict["adm_QCD"],
+            yukawa=rge_dict["yukawa"],
+        )
+        # checks if the obs_scale is a float. If it is, it uses that as the scale at which to compute the RGE matrix
+        if type(rge_dict["obs_scale"]) == float:
+            # check if there is a precomputed RGE matrix address.
+            path_to_rge_mat = rge_dict.get("rg_matrix", False)
+            low_scale = rge_dict["obs_scale"]
+            rge_cache = {}
+            if path_to_rge_mat:
+                # if there is a precomputed RGE matrix, it loads it with the internal smefit function.
+                rge_cache = load_precomputed_rge_matrix(
+                    path_to_rge_mat,
+                    {
+                        "init_scale": rge_dict["init_scale"],
+                        "smeft_accuracy": rge_dict["smeft_accuracy"],
+                        "adm_QCD": rge_dict["adm_QCD"],
+                        "yukawa": rge_dict["yukawa"],
+                    },
+                )
+
+            # uses the internal smefit function to select from the (possibly) precomputed RGE matrix or compute it  if needed.
+            pre_rgemat = load_rge_mats_from_scales(
+                [low_scale], coefficients.name, rge_runner, rge_cache
             )
+
+            # array of values to return
+            rgemat = (pre_rgemat[0]).values
+
+            # updates the list of coefficients with the ones generated by the RGE.
+            # piece of code shared with the compute_rge_matrix function, should be refactored in a future update to avoid code duplication.
+            gen_operators = list((pre_rgemat[0]).index)
+            _logger.info("The operators generated by the RGE are: ")
+            _logger.info(gen_operators)
+            operators_dict = {k: {"max": 0.0, "min": 0.0} for k in gen_operators}
+            new_coeffs = CoefficientManager.from_dict(operators_dict)
+
         else:
-            rgemat = None
+            # assumes that the scale input is dynamica and hence the RGE should be computed at the right scale for the OOs, which is an input.
+            low_scale = scale
+            rgemat, new_coeffs = compute_rge_matrix(
+                coefficients, rge_dict, scale=low_scale
+            )
+    else:
+        rgemat = None
+        new_coeffs = coefficients
+
+    return rgemat, new_coeffs
+
+
+class OptimalWWCEPC161:
+    def __init__(self, coefficients, rge_dict=None):
+        self.rgemat, coefficients = rge_matrix_manager(
+            coefficients, rge_dict, scale=161.0
+        )
 
         self.project = np.zeros((len(oo_ww_wc_basis), coefficients.size))
         for i, op in enumerate(oo_ww_wc_basis):
@@ -57,8 +132,6 @@ def __init__(self, coefficients, rge_dict=None):
             incovs_reordered.append(temp)
         self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
 
-        self.rgemat = rgemat
-
         if self.rgemat is not None:
             # multiply the RGE matrix as well
             self.incov_tot = jnp.einsum(
@@ -77,12 +150,9 @@ def compute_chi2(self, coefficient_values):
 
 class OptimalWWCEPC240:
     def __init__(self, coefficients, rge_dict=None):
-        if rge_dict is not None:
-            rgemat, coefficients = compute_rge_matrix(
-                coefficients, rge_dict, scale=240.0
-            )
-        else:
-            rgemat = None
+        self.rgemat, coefficients = rge_matrix_manager(
+            coefficients, rge_dict, scale=240.0
+        )
 
         self.project = np.zeros((len(oo_ww_wc_basis), coefficients.size))
         for i, op in enumerate(oo_ww_wc_basis):
@@ -101,8 +171,6 @@ def __init__(self, coefficients, rge_dict=None):
             incovs_reordered.append(temp)
         self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
 
-        self.rgemat = rgemat
-
         if self.rgemat is not None:
             # multiply the RGE matrix as well
             self.incov_tot = jnp.einsum(
@@ -121,12 +189,9 @@ def compute_chi2(self, coefficient_values):
 
 class OptimalWWCEPC365:
     def __init__(self, coefficients, rge_dict=None):
-        if rge_dict is not None:
-            rgemat, coefficients = compute_rge_matrix(
-                coefficients, rge_dict, scale=365.0
-            )
-        else:
-            rgemat = None
+        self.rgemat, coefficients = rge_matrix_manager(
+            coefficients, rge_dict, scale=365.0
+        )
 
         self.project = np.zeros((len(oo_ww_wc_basis), coefficients.size))
         for i, op in enumerate(oo_ww_wc_basis):
@@ -145,8 +210,6 @@ def __init__(self, coefficients, rge_dict=None):
             incovs_reordered.append(temp)
         self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
 
-        self.rgemat = rgemat
-
         if self.rgemat is not None:
             # multiply the RGE matrix as well
             self.incov_tot = jnp.einsum(
@@ -165,12 +228,9 @@ def compute_chi2(self, coefficient_values):
 
 class OptimalttCEPC365:
     def __init__(self, coefficients, rge_dict=None):
-        if rge_dict is not None:
-            rgemat, coefficients = compute_rge_matrix(
-                coefficients, rge_dict, scale=365.0
-            )
-        else:
-            rgemat = None
+        self.rgemat, coefficients = rge_matrix_manager(
+            coefficients, rge_dict, scale=365.0
+        )
 
         self.project = np.zeros((len(oo_tt_wc_basis), coefficients.size))
         for i, op in enumerate(oo_tt_wc_basis):
@@ -186,8 +246,6 @@ def __init__(self, coefficients, rge_dict=None):
             incovs_reordered.append(temp)
         self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
 
-        self.rgemat = rgemat
-
         if self.rgemat is not None:
             # multiply the RGE matrix as well
             self.incov_tot = jnp.einsum(