Repository: NNPDF/eko
Issue category: Feature Request

Issue:
Allow for parallelization in `evolven3fit_new`
Currently `evolven3fit_new` evolves the PDFs sequentially both in Q2 and replicas. For many replicas this means that evolution can take up to several hours while even for the usual 100 replica fits it generally will take over 20 minutes for the evolution to complete. 

Whenever resources are available (which they are on any clusters) it would be nice to have the option to parallelize parts of the the evolution to speed up the computation.


Hints:
As per the discussion w @felixhekhorn, the necessary change here is for eko to allow for more than one replica at once. To first approximation it should be trivial (the EKO x PDF contraction should have one extra axis for the number of replicas).

Then of course the changes will need to be propagated to evolven3fit and there might also be collateral damage somewhere. 

Patch:
diff --git a/benchmarks/eko/benchmark_inverse_matching.py b/benchmarks/eko/benchmark_inverse_matching.py
index d63b776f0..1004a5481 100644
--- a/benchmarks/eko/benchmark_inverse_matching.py
+++ b/benchmarks/eko/benchmark_inverse_matching.py
@@ -90,14 +90,10 @@ def benchmark_inverse_matching():
     with pytest.raises(AssertionError):
         np.testing.assert_allclose(op1_nf3.operator, op2_nf3.operator)
 
-    pdf1 = apply.apply_pdf(eko_output1, toy.mkPDF("ToyLH", 0))
-    pdf2 = apply.apply_pdf(eko_output2, toy.mkPDF("ToyLH", 0))
+    pdf1, _ = apply.apply_pdf(eko_output1, toy.mkPDF("ToyLH", 0))
+    pdf2, _ = apply.apply_pdf(eko_output2, toy.mkPDF("ToyLH", 0))
 
     # test that different PTO matching is applied correctly
-    np.testing.assert_allclose(
-        pdf1[(MC**2, 4)]["pdfs"][C_PID], pdf2[(MC**2, 4)]["pdfs"][C_PID]
-    )
+    np.testing.assert_allclose(pdf1[(MC**2, 4)][C_PID], pdf2[(MC**2, 4)][C_PID])
     with pytest.raises(AssertionError):
-        np.testing.assert_allclose(
-            pdf1[(MC**2, 3)]["pdfs"][C_PID], pdf2[(MC**2, 3)]["pdfs"][C_PID]
-        )
+        np.testing.assert_allclose(pdf1[(MC**2, 3)][C_PID], pdf2[(MC**2, 3)][C_PID])
diff --git a/src/ekobox/apply.py b/src/ekobox/apply.py
index 8d92f9a7b..09dad2d0f 100644
--- a/src/ekobox/apply.py
+++ b/src/ekobox/apply.py
@@ -1,131 +1,181 @@
-"""Apply operator evolution to PDF set."""
+"""Apply evolution operator to a PDF."""
 
-from dataclasses import dataclass
-from typing import Dict, Optional, Union
+from collections.abc import Sequence
+from typing import Optional
 
 import numpy as np
+import numpy.typing as npt
 
 from eko import basis_rotation as br
 from eko import interpolation
 from eko.io import EKO
 from eko.io.types import EvolutionPoint
 
+RawPdfResult = dict[EvolutionPoint, npt.ArrayLike]
+"""Evolved PDFs as raw grids.
+
+The key is given by the associated evolution point. The values are
+tensors sorted by (replica, flavor, xgrid).
+"""
+
+RawErrorResult = dict[EvolutionPoint, Optional[npt.ArrayLike]]
+"""Integration errors for evolved PDFs as raw grids.
+
+The key is given by the associated evolution point. The values are
+tensors sorted by (replica, flavor, xgrid).
+"""
+
+
+LabeledPdfResult = dict[EvolutionPoint, dict[int, npt.ArrayLike]]
+"""Evolved PDFs labeled by their PDF identifier.
+
+The outer key is given by the associated evolution point. The inner key
+is the |PID|. The inner values are the values for along the xgrid.
+"""
+LabeledErrorResult = dict[EvolutionPoint, dict[int, Optional[npt.ArrayLike]]]
+"""Integration errors for evolved PDFs labeled by their PDF identifier.
+
+The outer key is given by the associated evolution point. The inner key
+is the |PID|. The inner values are the values for along the xgrid.
+"""
+
 
 def apply_pdf(
     eko: EKO,
     lhapdf_like,
-    targetgrid=None,
-    rotate_to_evolution_basis=False,
-):
-    """Apply all available operators to the input PDFs.
+    targetgrid: npt.ArrayLike = None,
+    rotate_to_evolution_basis: bool = False,
+) -> tuple[LabeledPdfResult, LabeledErrorResult]:
+    """Apply all available operators to the input PDF.
 
     Parameters
     ----------
-        output : eko.output.EKO
-            eko output object containing all informations
-        lhapdf_like : object
-            object that provides an xfxQ2 callable (as `lhapdf <https://lhapdf.hepforge.org/>`_
-            and :class:`ekomark.toyLH.toyPDF` do) (and thus is in flavor basis)
-        targetgrid : list
-            if given, interpolates to the pdfs given at targetgrid (instead of xgrid)
-        rotate_to_evolution_basis : bool
-            if True rotate to evoluton basis
+    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)
+    targetgrid :
+        if given, interpolates to the targetgrid (instead of xgrid)
+    rotate_to_evolution_basis :
+        if True rotate to evoluton basis
 
     Returns
     -------
-        out_grid : dict
-            output PDFs and their associated errors for the computed mu2grid
+    pdfs :
+        PDFs for the computed evolution points
+    errors :
+        Integration errors for PDFs for the computed evolution points
     """
+    # prepare post-process
     qed = eko.theory_card.order[1] > 0
+    flavor_rotation = None
+    labels = br.flavor_basis_pids
     if rotate_to_evolution_basis:
         if not qed:
-            rotate_flavor_to_evolution = br.rotate_flavor_to_evolution
-            labels = br.evol_basis
+            flavor_rotation = br.rotate_flavor_to_evolution
+            labels = br.evol_basis_pids
         else:
-            rotate_flavor_to_evolution = br.rotate_flavor_to_unified_evolution
-            labels = br.unified_evol_basis
-        return apply_pdf_flavor(
-            eko, lhapdf_like, targetgrid, rotate_flavor_to_evolution, labels=labels
-        )
-    return apply_pdf_flavor(eko, lhapdf_like, targetgrid)
-
-
-CONTRACTION = "ajbk,bk"
-
-_PdfLabel = Union[int, str]
-"""PDF identifier: either PID or label."""
-
-
-@dataclass
-class PdfResult:
-    """Helper class to collect PDF results."""
-
-    pdfs: Dict[_PdfLabel, float]
-    errors: Optional[Dict[_PdfLabel, float]] = None
+            flavor_rotation = br.rotate_flavor_to_unified_evolution
+            labels = br.unified_evol_basis_pids
+    return apply_pdf_flavor(eko, lhapdf_like, labels, targetgrid, flavor_rotation)
 
 
 def apply_pdf_flavor(
-    eko: EKO, lhapdf_like, targetgrid=None, flavor_rotation=None, labels=None
-):
-    """Apply all available operators to the input PDFs.
+    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
     ----------
-        output : eko.output.EKO
-            eko output object containing all informations
-        lhapdf_like : object
-            object that provides an xfxQ2 callable (as `lhapdf <https://lhapdf.hepforge.org/>`_
-            and :class:`ekomark.toyLH.toyPDF` do) (and thus is in flavor basis)
-        targetgrid : list
-            if given, interpolates to the pdfs given at targetgrid (instead of xgrid)
-        flavor_rotation : np.ndarray
-            Rotation matrix in flavor space
-        labels : list
-            list of labels
+    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
     -------
-        out_grid : dict
-            output PDFs and their associated errors for the computed mu2grid
+    pdfs :
+        PDFs for the computed evolution points
+    errors :
+        Integration errors for PDFs for the computed evolution points
     """
     # create pdfs
-    pdfs = np.zeros((len(br.flavor_basis_pids), len(eko.xgrid)))
+    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
-        pdfs[j] = np.array(
+        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 = {}
+    for ep, pdf in new_grids.items():
+        pdfs[ep] = {
+            lab: grid[0] if grid is not None else None for lab, grid in pdf.items()
+        }
+        errors[ep] = {
+            lab: (grid[0] if grid is not None else None)
+            for lab, grid in new_errors[ep].items()
+        }
+    return pdfs, errors
 
-    # build output
-    out_grid: Dict[EvolutionPoint, PdfResult] = {}
-    for ep, elem in eko.items():
-        pdf_final = np.einsum(CONTRACTION, elem.operator, pdfs, optimize="optimal")
-        if elem.error is not None:
-            error_final = np.einsum(CONTRACTION, elem.error, pdfs, optimize="optimal")
-        else:
-            error_final = None
-        out_grid[ep] = PdfResult(dict(zip(br.flavor_basis_pids, pdf_final)))
-        if error_final is not None:
-            out_grid[ep].errors = dict(zip(br.flavor_basis_pids, error_final))
 
-    qed = eko.theory_card.order[1] > 0
+def rotate_result(
+    eko: EKO,
+    grids: RawErrorResult,
+    flavor_labels: Sequence[int],
+    targetgrid: Optional[npt.ArrayLike] = None,
+    flavor_rotation: Optional[npt.ArrayLike] = None,
+) -> LabeledErrorResult:
+    """Rotate and relabel PDFs.
+
+    Parameters
+    ----------
+    eko :
+        eko output object containing all informations
+    grids :
+        Raw grids coming from evolution
+    flavor_labels :
+        flavors names
+    targetgrid :
+        if given, interpolates to the targetgrid (instead of xgrid)
+    flavor_rotation :
+        if given, rotates in flavor space
+
+    Returns
+    -------
+    pdfs :
+        relabeled and, if requested rotated, PDFs
+    """
     # rotate to evolution basis
     if flavor_rotation is not None:
-        for q2, op in out_grid.items():
-            pdf = flavor_rotation @ np.array(
-                [op.pdfs[pid] for pid in br.flavor_basis_pids]
+        new_grids = {}
+        for ep, pdf_grid in grids.items():
+            new_grids[ep] = (
+                np.einsum("ab,rbk->rak", flavor_rotation, pdf_grid, optimize="optimal")
+                if pdf_grid is not None
+                else None
             )
-            if not qed:
-                evol_basis = br.evol_basis
-            else:
-                evol_basis = br.unified_evol_basis
-            op.pdfs = dict(zip(evol_basis, pdf))
-            if op.errors is not None:
-                errors = flavor_rotation @ np.array(
-                    [op.errors[pid] for pid in br.flavor_basis_pids]
-                )
-                op.errors = dict(zip(evol_basis, errors))
+        grids = new_grids
 
     # rotate/interpolate to target grid
     if targetgrid is not None:
@@ -135,14 +185,74 @@ def apply_pdf_flavor(
             mode_N=False,
         )
 
-        rot = b.get_interpolation(targetgrid)
-        for evpdf in out_grid.values():
-            for pdf_label in evpdf.pdfs:
-                evpdf.pdfs[pdf_label] = np.matmul(rot, evpdf.pdfs[pdf_label])
-                if evpdf.errors is not None:
-                    evpdf.errors[pdf_label] = np.matmul(rot, evpdf.errors[pdf_label])
-    # cast back to be backward compatible
-    real_out_grid = {}
-    for ep, res in out_grid.items():
-        real_out_grid[ep] = {"pdfs": res.pdfs, "errors": res.errors}
-    return real_out_grid
+        x_rotation = b.get_interpolation(targetgrid)
+        new_grids = {}
+        for ep, pdf_grid in grids.items():
+            new_grids[ep] = (
+                np.einsum("jk,rbk->rbj", x_rotation, pdf_grid, optimize="optimal")
+                if pdf_grid is not None
+                else None
+            )
+        grids = new_grids
+
+    # relabel
+    new_grids = {}
+    for ep, pdf_grid in grids.items():
+        new_grids[ep] = dict(
+            zip(
+                flavor_labels,
+                np.swapaxes(pdf_grid, 0, 1)
+                if pdf_grid is not None
+                else [None] * len(flavor_labels),
+            )
+        )
+    grids = new_grids
+
+    return grids
+
+
+_EKO_CONTRACTION = "ajbk,rbk->raj"
+"""Contract eko for all replicas."""
+
+
+def apply_grids(
+    eko: EKO, input_grids: npt.ArrayLike
+) -> tuple[RawPdfResult, RawErrorResult]:
+    """Apply all available operators to the input grids.
+
+    Parameters
+    ----------
+    eko :
+        eko output object
+    input_grids :
+        3D PDF grids evaluated at the inital scale. The axis have to be (replica, flavor, xgrid)
+
+    Returns
+    -------
+    pdfs :
+        output PDFs for the computed evolution points
+    pdfs :
+        associated integration errors for the computed evolution points
+    """
+    # sanity check
+    if len(input_grids.shape) != 3 or input_grids.shape[1:] != (
+        len(br.flavor_basis_pids),
+        len(eko.xgrid),
+    ):
+        raise ValueError(
+            "input grids have to be sorted by replica, flavor, xgrid!"
+            f"The shape has to be (r,{len(br.flavor_basis_pids)},{len(eko.xgrid)})"
+        )
+    # iterate
+    pdfs: RawPdfResult = {}
+    errors: RawErrorResult = {}
+    for ep, elem in eko.items():
+        pdfs[ep] = np.einsum(
+            _EKO_CONTRACTION, elem.operator, input_grids, optimize="optimal"
+        )
+        errors[ep] = (
+            np.einsum(_EKO_CONTRACTION, elem.error, input_grids, optimize="optimal")
+            if elem.error is not None
+            else None
+        )
+    return pdfs, errors
diff --git a/src/ekobox/evol_pdf.py b/src/ekobox/evol_pdf.py
index 15e837736..b629dd249 100644
--- a/src/ekobox/evol_pdf.py
+++ b/src/ekobox/evol_pdf.py
@@ -78,11 +78,11 @@ def evolve_pdfs(
     with EKO.read(eko_path) as eko_output:
         for initial_PDF in initial_PDF_list:
             evolved_PDF_list.append(
-                apply.apply_pdf(eko_output, initial_PDF, targetgrid)
+                apply.apply_pdf(eko_output, initial_PDF, targetgrid)[0]
             )
 
         # separate by nf the evolgrid (and order per nf/q)
-        q2block_per_nf = regroup_evolgrid(eko_output.evolgrid)  # pylint: disable=E1101
+        q2block_per_nf = regroup_evolgrid(eko_output.evolgrid)
 
     # update info file
     if targetgrid is None:
@@ -128,7 +128,7 @@ def collect_blocks(evolved_PDF: dict, q2block_per_nf: dict, xgrid: list):
     # fake xfxQ2
     def pdf_xq2(pid, x, Q2):
         x_idx = xgrid.index(x)
-        return x * evolved_PDF[(Q2, nf)]["pdfs"][pid][x_idx]
+        return x * evolved_PDF[(Q2, nf)][pid][x_idx]
 
     # loop on nf patches
     for nf, q2grid in q2block_per_nf.items():
diff --git a/src/ekomark/benchmark/runner.py b/src/ekomark/benchmark/runner.py
index 4b8e94c4d..a1d79c67a 100644
--- a/src/ekomark/benchmark/runner.py
+++ b/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):


Test patch:
diff --git a/tests/ekobox/test_apply.py b/tests/ekobox/test_apply.py
index 70eb74adc..522c4dc6d 100644
--- a/tests/ekobox/test_apply.py
+++ b/tests/ekobox/test_apply.py
@@ -5,34 +5,31 @@
 from tests.conftest import EKOFactory
 
 
-class TestApply:
-    def test_apply(self, eko_factory: EKOFactory, fake_pdf):
-        eko = eko_factory.get()
-        ep_out = eko.evolgrid[0]
-        # fake pdfs
-        pdf_grid = apply.apply_pdf(eko, fake_pdf)
-        assert len(pdf_grid) == len(eko.evolgrid)
-        pdfs = pdf_grid[ep_out]["pdfs"]
-        assert list(pdfs.keys()) == list(br.flavor_basis_pids)
-        # rotate to target_grid
-        target_grid = [0.75]
-        pdf_grid = apply.apply_pdf(eko, fake_pdf, target_grid)
-        assert len(pdf_grid) == 1
-        pdfs = pdf_grid[ep_out]["pdfs"]
-        assert list(pdfs.keys()) == list(br.flavor_basis_pids)
+def test_apply(eko_factory: EKOFactory, fake_pdf):
+    eko = eko_factory.get()
+    ep_out = eko.evolgrid[0]
+    # base application
+    pdfs, errors = apply.apply_pdf(eko, fake_pdf)
+    assert len(pdfs) == len(eko.evolgrid)
+    assert len(pdfs) == len(errors)
+    ep_pdfs = pdfs[ep_out]
+    assert list(ep_pdfs.keys()) == list(br.flavor_basis_pids)
+    # rotate to target_grid
+    for target_grid in ([0.75], [0.5, 0.6, 0.7]):
+        pdfs, _errors = apply.apply_pdf(eko, fake_pdf, target_grid)
+        ep_pdfs = pdfs[ep_out]
+        assert list(ep_pdfs.keys()) == list(br.flavor_basis_pids)
+        assert len(ep_pdfs[21]) == len(target_grid)
+    # rotate flavor
+    pdfs, _errors = apply.apply_pdf(eko, fake_pdf, rotate_to_evolution_basis=True)
+    ep_pdfs = pdfs[ep_out]
+    assert list(ep_pdfs.keys()) == list(br.evol_basis_pids)
 
-    def test_apply_flavor(self, eko_factory: EKOFactory, fake_pdf, monkeypatch):
-        eko = eko_factory.get()
-        ep_out = eko.evolgrid[0]
-        # fake pdfs
-        monkeypatch.setattr(
-            "eko.basis_rotation.rotate_flavor_to_evolution", np.ones((14, 14))
-        )
-        fake_evol_basis = tuple(
-            ["a", "b"] + [str(x) for x in range(len(br.flavor_basis_pids) - 2)]
-        )
-        monkeypatch.setattr("eko.basis_rotation.evol_basis", fake_evol_basis)
-        pdf_grid = apply.apply_pdf(eko, fake_pdf, rotate_to_evolution_basis=True)
-        assert len(pdf_grid) == len(eko.evolgrid)
-        pdfs = pdf_grid[ep_out]["pdfs"]
-        assert list(pdfs.keys()) == list(fake_evol_basis)
+
+def test_apply_grids(eko_factory: EKOFactory):
+    eko = eko_factory.get()
+    # since everything is random, we can only test the tensor shapes here
+    input_grids = np.random.rand(3, len(br.flavor_basis_pids), len(eko.xgrid))
+    pdfs, errors = apply.apply_grids(eko, input_grids)
+    for ep, res in pdfs.items():
+        assert res.shape == input_grids.shape
diff --git a/tests/ekobox/test_evol_pdf.py b/tests/ekobox/test_evol_pdf.py
index 7d788277b..baf70726d 100644
--- a/tests/ekobox/test_evol_pdf.py
+++ b/tests/ekobox/test_evol_pdf.py
@@ -113,11 +113,7 @@ def test_collect_blocks():
     def mk(eps):
         f = {}
         for ep in eps:
-            f[ep] = {
-                "pdfs": {
-                    pid: np.random.rand(len(xgrid)) for pid in br.flavor_basis_pids
-                }
-            }
+            f[ep] = {pid: np.random.rand(len(xgrid)) for pid in br.flavor_basis_pids}
         return f
 
     # basic