From ed7f1b1c91b3d595b7ab0984bab3b5bc90506ccc Mon Sep 17 00:00:00 2001 From: giacomomagni Date: Mon, 21 Oct 2024 10:22:19 +0200 Subject: [PATCH] Revert "Merge branch 'add_some_pol_tests' of https://github.com/N3PDF/eko into add_some_pol_tests" This reverts commit 74b092e8377c5da1dbff6e6ac35eb934b4e002c2, reversing changes made to 10986662fa4e1772f8882acd43d0e39dba7bbf6b. --- benchmarks/eko/benchmark_inverse_matching.py | 12 +- src/ekobox/apply.py | 298 ++++++------------- src/ekobox/evol_pdf.py | 6 +- src/ekomark/benchmark/runner.py | 9 +- tests/ekobox/test_apply.py | 57 ++-- tests/ekobox/test_evol_pdf.py | 6 +- 6 files changed, 145 insertions(+), 243 deletions(-) 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