Adjust apply for ekomark

src/ekobox/apply.py

@@ -67,17 +67,7 @@ def apply_pdf(
     errors :
         Integration errors for PDFs for the computed evolution points
     """
-    # create pdfs
-    input_pdfs = np.zeros((len(br.flavor_basis_pids), len(eko.xgrid)))
-    for j, pid in enumerate(br.flavor_basis_pids):
-        if not lhapdf_like.hasFlavor(pid):
-            continue
-        input_pdfs[j] = np.array(
-            [lhapdf_like.xfxQ2(pid, x, eko.mu20) / x for x in eko.xgrid.raw]
-        )
-    # apply
-    grids, grid_errors = apply_grids(eko, input_pdfs[None, :])
-    # post-process
+    # prepare post-process
     qed = eko.theory_card.order[1] > 0
     flavor_rotation = None
     labels = br.flavor_basis_pids
@@ -88,8 +78,53 @@ def apply_pdf(
         else:
             flavor_rotation = br.rotate_flavor_to_unified_evolution
             labels = br.unified_evol_basis_pids
-    new_grids = rotate_result(eko, grids, labels, targetgrid, flavor_rotation)
-    new_errors = rotate_result(eko, grid_errors, labels, targetgrid, flavor_rotation)
+    return apply_pdf_flavor(eko, lhapdf_like, labels, targetgrid, flavor_rotation)
+
+
+def apply_pdf_flavor(
+    eko: EKO,
+    lhapdf_like,
+    flavor_labels: Sequence[int],
+    targetgrid: npt.ArrayLike = None,
+    flavor_rotation: npt.ArrayLike = None,
+) -> tuple[LabeledPdfResult, LabeledErrorResult]:
+    """Apply all available operators to the input PDF.
+
+    Parameters
+    ----------
+    eko :
+        eko output object containing all informations
+    lhapdf_like : Any
+        object that provides an `xfxQ2` callable (as `lhapdf <https://lhapdf.hepforge.org/>`_
+        and :class:`ekomark.toyLH.toyPDF` do) (and thus is in flavor basis)
+    flavor_labels :
+        flavor names
+    targetgrid :
+        if given, interpolates to the targetgrid (instead of xgrid)
+    flavor_rotation :
+        if give, rotate in flavor space
+
+    Returns
+    -------
+    pdfs :
+        PDFs for the computed evolution points
+    errors :
+        Integration errors for PDFs for the computed evolution points
+    """
+    # create pdfs
+    input_pdfs = np.zeros((len(br.flavor_basis_pids), len(eko.xgrid)))
+    for j, pid in enumerate(br.flavor_basis_pids):
+        if not lhapdf_like.hasFlavor(pid):
+            continue
+        input_pdfs[j] = np.array(
+            [lhapdf_like.xfxQ2(pid, x, eko.mu20) / x for x in eko.xgrid.raw]
+        )
+    # apply
+    grids, grid_errors = apply_grids(eko, input_pdfs[None, :])
+    new_grids = rotate_result(eko, grids, flavor_labels, targetgrid, flavor_rotation)
+    new_errors = rotate_result(
+        eko, grid_errors, flavor_labels, targetgrid, flavor_rotation
+    )
     # unwrap the replica axis again
     pdfs: LabeledPdfResult = {}
     errors: LabeledErrorResult = {}

src/ekomark/benchmark/runner.py

@@ -228,12 +228,12 @@ def log(self, theory, _, pdf, me, ext):
                 rotate_to_evolution[3, :] = [0, 0, 0, 0, 0, -1, -1, 0, 1, 1, 0, 0, 0, 0]
 
         with EKO.read(me) as eko:
-            pdf_grid = apply.apply_pdf_flavor(
+            pdf_grid, errors = apply.apply_pdf_flavor(
                 eko,
                 pdf,
+                labels,
                 xgrid,
                 flavor_rotation=rotate_to_evolution,
-                labels=labels,
             )
         for q2, ref_pdfs in ext["values"].items():
             log_tab = dfdict.DFdict()
@@ -243,9 +243,8 @@ def log(self, theory, _, pdf, me, ext):
                 if len(eps) == 0:
                     raise KeyError(f"PDF at Q2={q2} not found")
                 raise KeyError(f"More than one evolution point found for Q2={q2}")
-            res = pdf_grid[eps[0]]
-            my_pdfs = res["pdfs"]
-            my_pdf_errs = res["errors"]
+            my_pdfs = pdf_grid[eps[0]]
+            my_pdf_errs = errors[eps[0]]
 
             for key in my_pdfs:
                 if key in self.skip_pdfs(theory):