Rework ekobox.apply

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

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

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():