diff --git a/benchmarks/eko/benchmark_inverse_matching.py b/benchmarks/eko/benchmark_inverse_matching.py
index 1004a5481..d63b776f0 100644
--- a/benchmarks/eko/benchmark_inverse_matching.py
+++ b/benchmarks/eko/benchmark_inverse_matching.py
@@ -90,10 +90,14 @@ 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)][C_PID], pdf2[(MC**2, 4)][C_PID])
+ np.testing.assert_allclose(
+ pdf1[(MC**2, 4)]["pdfs"][C_PID], pdf2[(MC**2, 4)]["pdfs"][C_PID]
+ )
with pytest.raises(AssertionError):
- np.testing.assert_allclose(pdf1[(MC**2, 3)][C_PID], pdf2[(MC**2, 3)][C_PID])
+ np.testing.assert_allclose(
+ pdf1[(MC**2, 3)]["pdfs"][C_PID], pdf2[(MC**2, 3)]["pdfs"][C_PID]
+ )
diff --git a/src/ekobox/apply.py b/src/ekobox/apply.py
index 09dad2d0f..8d92f9a7b 100644
--- a/src/ekobox/apply.py
+++ b/src/ekobox/apply.py
@@ -1,181 +1,131 @@
-"""Apply evolution operator to a PDF."""
+"""Apply operator evolution to PDF set."""
-from collections.abc import Sequence
-from typing import Optional
+from dataclasses import dataclass
+from typing import Dict, Optional, Union
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: npt.ArrayLike = None,
- rotate_to_evolution_basis: bool = False,
-) -> tuple[LabeledPdfResult, LabeledErrorResult]:
- """Apply all available operators to the input PDF.
+ targetgrid=None,
+ rotate_to_evolution_basis=False,
+):
+ """Apply all available operators to the input PDFs.
Parameters
----------
- 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
+ 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
Returns
-------
- pdfs :
- PDFs for the computed evolution points
- errors :
- Integration errors for PDFs for the computed evolution points
+ out_grid : dict
+ output PDFs and their associated errors for the computed mu2grid
"""
- # 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:
- flavor_rotation = br.rotate_flavor_to_evolution
- labels = br.evol_basis_pids
+ rotate_flavor_to_evolution = br.rotate_flavor_to_evolution
+ labels = br.evol_basis
else:
- 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)
+ 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
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.
+ eko: EKO, lhapdf_like, targetgrid=None, flavor_rotation=None, labels=None
+):
+ """Apply all available operators to the input PDFs.
Parameters
----------
- 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
+ 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
Returns
-------
- pdfs :
- PDFs for the computed evolution points
- errors :
- Integration errors for PDFs for the computed evolution points
+ out_grid : dict
+ output PDFs and their associated errors for the computed mu2grid
"""
# create pdfs
- input_pdfs = np.zeros((len(br.flavor_basis_pids), len(eko.xgrid)))
+ 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(
+ 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
-
-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
+ # 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))
- Returns
- -------
- pdfs :
- relabeled and, if requested rotated, PDFs
- """
+ qed = eko.theory_card.order[1] > 0
# rotate to evolution basis
if flavor_rotation is not None:
- 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
+ for q2, op in out_grid.items():
+ pdf = flavor_rotation @ np.array(
+ [op.pdfs[pid] for pid in br.flavor_basis_pids]
)
- grids = new_grids
+ 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))
# rotate/interpolate to target grid
if targetgrid is not None:
@@ -185,74 +135,14 @@ def rotate_result(
mode_N=False,
)
- 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
+ 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
diff --git a/src/ekobox/evol_pdf.py b/src/ekobox/evol_pdf.py
index b629dd249..15e837736 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)[0]
+ apply.apply_pdf(eko_output, initial_PDF, targetgrid)
)
# separate by nf the evolgrid (and order per nf/q)
- q2block_per_nf = regroup_evolgrid(eko_output.evolgrid)
+ q2block_per_nf = regroup_evolgrid(eko_output.evolgrid) # pylint: disable=E1101
# 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)][pid][x_idx]
+ return x * evolved_PDF[(Q2, nf)]["pdfs"][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 a1d79c67a..4b8e94c4d 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, errors = apply.apply_pdf_flavor(
+ pdf_grid = 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,8 +243,9 @@ 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}")
- my_pdfs = pdf_grid[eps[0]]
- my_pdf_errs = errors[eps[0]]
+ res = pdf_grid[eps[0]]
+ my_pdfs = res["pdfs"]
+ my_pdf_errs = res["errors"]
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 522c4dc6d..70eb74adc 100644
--- a/tests/ekobox/test_apply.py
+++ b/tests/ekobox/test_apply.py
@@ -5,31 +5,34 @@
from tests.conftest import EKOFactory
-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)
+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_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
+ 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)
diff --git a/tests/ekobox/test_evol_pdf.py b/tests/ekobox/test_evol_pdf.py
index baf70726d..7d788277b 100644
--- a/tests/ekobox/test_evol_pdf.py
+++ b/tests/ekobox/test_evol_pdf.py
@@ -113,7 +113,11 @@ def test_collect_blocks():
def mk(eps):
f = {}
for ep in eps:
- f[ep] = {pid: np.random.rand(len(xgrid)) for pid in br.flavor_basis_pids}
+ f[ep] = {
+ "pdfs": {
+ pid: np.random.rand(len(xgrid)) for pid in br.flavor_basis_pids
+ }
+ }
return f
# basic