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 `_
- 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 `_
+ 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 `_
- 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 `_
+ 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):
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