Allow evolution and convolution with different ekos for a single grid

src/pineko/evolve.py

@@ -46,6 +46,25 @@ def sv_scheme(tcard):
     return modsv
 
 
+def get_ekos_convolution_type(kv):
+    """Temporary finction to retrive the ekos convolution type.
+
+    Parameters
+    ----------
+    kv: dict
+        pineappl grid metadata
+    """
+    if "convolution_type_1" in kv:
+        eko1 = kv["convolution_type_1"]
+    # TODO: this case is now deprecated and should be remved from yadism and pinefarm
+    elif "polarized" in kv:
+        eko1 = "polPDF"
+    else:
+        eko1 = "PDF"
+    eko2 = kv.get("convolution_type_2", "PDF")
+    return eko1, eko2
+
+
 def write_operator_card_from_file(
     pineappl_path: os.PathLike,
     default_card_path: os.PathLike,
@@ -155,8 +174,7 @@ def write_operator_card(pineappl_grid, default_card, card_path, tcard):
 
     # switch on polarization ?
     kv = pineappl_grid.key_values()
-    if "polarized" in kv:
-        operators_card["configs"]["polarized"] = kv["polarized"] == "True"
+    conv_type_a, conv_type_b = get_ekos_convolution_type(kv)
 
     # fragmentation function grid?
     if "timelike" in kv:
@@ -172,21 +190,37 @@ def write_operator_card(pineappl_grid, default_card, card_path, tcard):
             "are you sure that's what you want?"
         )
 
-    with open(card_path, "w", encoding="UTF-8") as f:
-        yaml.safe_dump(operators_card, f)
-        f.write(f"# {pineko_version=}")
+    def dump_card(card_path, operators_card, conv_type):
+        operators_card["configs"]["polarized"] = conv_type == "polPDF"
+        with open(card_path, "w", encoding="UTF-8") as f:
+            yaml.safe_dump(operators_card, f)
+            f.write(f"# {pineko_version=}")
+
+    # For hardonic obs we might need to dump 2 eko cards
+
+    if conv_type_a == conv_type_b:
+        dump_card(card_path, operators_card, conv_type_a)
+    else:
+        # dump eko1
+        card_name = f"{card_path.stem}_{conv_type_a}.yaml"
+        dump_card(card_path.parent / card_name, operators_card)
+
+        # dump eko2
+        card_name = f"{card_path.stem}_{conv_type_b}.yaml"
+        dump_card(card_path.parent / card_name, operators_card)
 
     return operators_card["xgrid"], q2_grid
 
 
 def evolve_grid(
     grid,
-    operators,
+    operators_a,
     fktable_path,
     max_as,
     max_al,
     xir,
     xif,
+    operators_b=None,
     assumptions="Nf6Ind",
     comparison_pdf=None,
     meta_data=None,
@@ -198,7 +232,7 @@ def evolve_grid(
     ----------
     grid : pineappl.grid.Grid
         unconvoluted grid
-    operators : eko.EKO
+    operators_a : eko.EKO
         evolution operator
     fktable_path : str
         target path for convoluted grid
@@ -210,6 +244,8 @@ def evolve_grid(
         renormalization scale variation
     xif : float
         factorization scale variation
+    operators_b: eko.EKO
+        additonal evolution operator if different from operators_a
     assumptions : str
         assumptions on the flavor dimension
     comparison_pdf : None or str
@@ -231,22 +267,29 @@ def evolve_grid(
     evol_info = grid.evolve_info(order_mask)
     x_grid = evol_info.x1
     mur2_grid = evol_info.ren1
-    xif = 1.0 if operators.operator_card.configs.scvar_method is not None else xif
-    tcard = operators.theory_card
-    opcard = operators.operator_card
+    xif = 1.0 if operators_a.operator_card.configs.scvar_method is not None else xif
+    tcard = operators_a.theory_card
+    opcard = operators_a.operator_card
     # rotate the targetgrid
     if "integrability_version" in grid.key_values():
         x_grid = np.append(x_grid, 1.0)
-    eko.io.manipulate.xgrid_reshape(
-        operators, targetgrid=eko.interpolation.XGrid(x_grid)
-    )
-    check.check_grid_and_eko_compatible(grid, operators, xif, max_as, max_al)
-    # rotate to evolution (if doable and necessary)
-    if np.allclose(operators.bases.inputpids, br.flavor_basis_pids):
-        eko.io.manipulate.to_evol(operators)
-    # Here we are checking if the EKO contains the rotation matrix (flavor to evol)
-    elif not np.allclose(operators.bases.inputpids, br.rotate_flavor_to_evolution):
-        raise ValueError("The EKO is neither in flavor nor in evolution basis.")
+
+    def xgrid_reshape(operators):
+        eko.io.manipulate.xgrid_reshape(
+            operators, targetgrid=eko.interpolation.XGrid(x_grid)
+        )
+        check.check_grid_and_eko_compatible(grid, operators, xif, max_as, max_al)
+        # rotate to evolution (if doable and necessary)
+        if np.allclose(operators.bases.inputpids, br.flavor_basis_pids):
+            eko.io.manipulate.to_evol(operators)
+        # Here we are checking if the EKO contains the rotation matrix (flavor to evol)
+        elif not np.allclose(operators.bases.inputpids, br.rotate_flavor_to_evolution):
+            raise ValueError("The EKO is neither in flavor nor in evolution basis.")
+
+    xgrid_reshape(operators_a)
+    if operators_b is not None:
+        xgrid_reshape(operators_b)
+
     # PineAPPL wants alpha_s = 4*pi*a_s
     # remember that we already accounted for xif in the opcard generation
     evmod = eko.couplings.couplings_mod_ev(opcard.configs.evolution_method)
@@ -274,18 +317,31 @@ def evolve_grid(
     ]
     # We need to use ekompatibility in order to pass a dictionary to pineappl
     fktable = grid.evolve(
-        ekompatibility.pineappl_layout(operators),
+        ekompatibility.pineappl_layout(operators_a),
         xir * xir * mur2_grid,
         alphas_values,
         "evol",
+        operators_b=(
+            ekompatibility.pineappl_layout(operators_b)
+            if operators_b is not None
+            else None
+        ),
         order_mask=order_mask,
         xi=(xir, xif),
     )
     rich.print(f"Optimizing for {assumptions}")
     fktable.optimize(assumptions)
-    fktable.set_key_value("eko_version", operators.metadata.version)
-    fktable.set_key_value("eko_theory_card", json.dumps(operators.theory_card.raw))
-    fktable.set_key_value("eko_operator_card", json.dumps(operators.operator_card.raw))
+    fktable.set_key_value("eko_version", operators_a.metadata.version)
+    fktable.set_key_value("eko_theory_card", json.dumps(operators_a.theory_card.raw))
+
+    fktable.set_key_value(
+        "eko_operator_card", json.dumps(operators_a.operator_card.raw)
+    )
+    if operators_b is not None:
+        fktable.set_key_value(
+            "eko_operator_card_b", json.dumps(operators_b.operator_card.raw)
+        )
+
     fktable.set_key_value("pineko_version", version.__version__)
     if meta_data is not None:
         for k, v in meta_data.items():

src/pineko/theory.py

@@ -398,8 +398,20 @@ def fk(self, name, grid_path, tcard, pdf):
         grid = pineappl.grid.Grid.read(grid_path)
         # remove zero subgrid
         grid.optimize()
+
+        # Do you need one or multiple ekos?
+        kv = grid.key_values()
+        conv_type_a, conv_type_b = evolve.get_ekos_convolution_type(kv)
+
         # setup data
-        eko_filename = self.ekos_path() / f"{name}.tar"
+        if conv_type_a == conv_type_b:
+            eko_filename = [self.ekos_path() / f"{name}.tar"]
+        else:
+            eko_filename = [
+                self.ekos_path() / f"{name}_{conv_type_a}.tar",
+                self.ekos_path() / f"{name}_{conv_type_b}.tar",
+            ]
+
         fk_filename = self.fks_path / f"{name}.{parser.EXT}"
         if fk_filename.exists():
             if not self.overwrite:
@@ -434,16 +446,28 @@ def fk(self, name, grid_path, tcard, pdf):
             if not np.isclose(xif, 1.0):
                 check_scvar_evolve(grid, max_as, max_al, check.Scale.FACT)
         # loading ekos to produce a tmp copy
-        with eko.EKO.read(eko_filename) as operators:
+        n_ekos = len(eko_filename)
+        with eko.EKO.read(eko_filename[0]) as operators_a:
+
             # Skip the computation of the fktable if the eko is empty
-            if len(operators.mu2grid) == 0 and check.is_num_fonll(tcard["FNS"]):
+            if len(operators_a.mu2grid) == 0 and check.is_num_fonll(tcard["FNS"]):
                 rich.print("[green] Skipping empty eko for nFONLL.")
                 return
-            eko_tmp_path = (
-                operators.paths.root.parent / f"eko-tmp-{name}-{np.random.rand()}.tar"
+
+            eko_tmp_path_a = (
+                operators_a.paths.root.parent / f"eko-tmp-{name}-{np.random.rand()}.tar"
             )
-            operators.deepcopy(eko_tmp_path)
-        with eko.EKO.edit(eko_tmp_path) as operators:
+            operators_a.deepcopy(eko_tmp_path_a)
+
+        if n_ekos > 1:
+            with eko.EKO.read(eko_filename[1]) as operators_b:
+                eko_tmp_path_b = (
+                    operators_a.paths.root.parent
+                    / f"eko-tmp-{name}-{np.random.rand()}.tar"
+                )
+            operators_b.deepcopy(eko_tmp_path_b)
+
+        with eko.EKO.edit(eko_tmp_path_a) as operators_a:
             # Obtain the assumptions hash
             assumptions = theory_card.construct_assumptions(tcard)
             # do it!
@@ -456,7 +480,6 @@ def fk(self, name, grid_path, tcard, pdf):
                 xif,
             )
             start_time = time.perf_counter()
-
             rich.print(
                 rich.panel.Panel.fit(
                     "Computing ...", style="magenta", box=rich.box.SQUARE
@@ -466,20 +489,40 @@ def fk(self, name, grid_path, tcard, pdf):
                 f"= {fk_filename}\n",
                 f"with max_as={max_as}, max_al={max_al}, xir={xir}, xif={xif}",
             )
-            _grid, _fk, comparison = evolve.evolve_grid(
-                grid,
-                operators,
-                fk_filename,
-                max_as,
-                max_al,
-                xir=xir,
-                xif=xif,
-                assumptions=assumptions,
-                comparison_pdf=pdf,
-                meta_data={"theory_card": json.dumps(tcard)},
-            )
+
+            if n_ekos == 1:
+                _grid, _fk, comparison = evolve.evolve_grid(
+                    grid,
+                    operators_a,
+                    fk_filename,
+                    max_as,
+                    max_al,
+                    xir=xir,
+                    xif=xif,
+                    assumptions=assumptions,
+                    comparison_pdf=pdf,
+                    meta_data={"theory_card": json.dumps(tcard)},
+                )
+            else:
+                with eko.EKO.edit(eko_tmp_path_b) as operators_b:
+                    _grid, _fk, comparison = evolve.evolve_grid(
+                        grid,
+                        operators_a,
+                        fk_filename,
+                        max_as,
+                        max_al,
+                        xir=xir,
+                        xif=xif,
+                        assumptions=assumptions,
+                        operators_b=operators_b,
+                        comparison_pdf=pdf,
+                        meta_data={"theory_card": json.dumps(tcard)},
+                    )
+                # Remove tmp ekos
+                eko_tmp_path_b.unlink()
+
         # Remove tmp ekos
-        eko_tmp_path.unlink()
+        eko_tmp_path_a.unlink()
 
         logger.info(
             "Finished computation of %s - took %f s",