From 59cd2cdcbcb4cb5534f1ed93d06e7f9bcf6cb076 Mon Sep 17 00:00:00 2001 From: Felix Hekhorn Date: Tue, 22 Oct 2024 15:02:58 +0300 Subject: [PATCH] Change error propagation --- src/ekobox/apply.py | 74 ++++++++++++++------------------------ tests/ekobox/test_apply.py | 7 ++-- 2 files changed, 29 insertions(+), 52 deletions(-) diff --git a/src/ekobox/apply.py b/src/ekobox/apply.py index 09dad2d0f..510f3eb05 100644 --- a/src/ekobox/apply.py +++ b/src/ekobox/apply.py @@ -12,31 +12,20 @@ from eko.io.types import EvolutionPoint RawPdfResult = dict[EvolutionPoint, npt.ArrayLike] -"""Evolved PDFs as raw grids. +"""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). +tensors sorted by (replica, flavor, xgrid). It may be the PDF or the +associated integration error. """ LabeledPdfResult = dict[EvolutionPoint, dict[int, npt.ArrayLike]] -"""Evolved PDFs labeled by their PDF identifier. +"""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. +is the |PID|. The inner values are the values for along the xgrid. It +may be the PDF or the associated integration error. """ @@ -45,7 +34,7 @@ def apply_pdf( lhapdf_like, targetgrid: npt.ArrayLike = None, rotate_to_evolution_basis: bool = False, -) -> tuple[LabeledPdfResult, LabeledErrorResult]: +) -> tuple[LabeledPdfResult, LabeledPdfResult]: """Apply all available operators to the input PDF. Parameters @@ -87,7 +76,7 @@ def apply_pdf_flavor( flavor_labels: Sequence[int], targetgrid: npt.ArrayLike = None, flavor_rotation: npt.ArrayLike = None, -) -> tuple[LabeledPdfResult, LabeledErrorResult]: +) -> tuple[LabeledPdfResult, LabeledPdfResult]: """Apply all available operators to the input PDF. Parameters @@ -127,25 +116,21 @@ def apply_pdf_flavor( ) # unwrap the replica axis again pdfs: LabeledPdfResult = {} - errors: LabeledErrorResult = {} + errors: LabeledPdfResult = {} 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() - } + pdfs[ep] = {lab: grid[0] for lab, grid in pdf.items()} + if ep in new_errors: + errors[ep] = {lab: (grid[0]) for lab, grid in new_errors[ep].items()} return pdfs, errors def rotate_result( eko: EKO, - grids: RawErrorResult, + grids: RawPdfResult, flavor_labels: Sequence[int], targetgrid: Optional[npt.ArrayLike] = None, flavor_rotation: Optional[npt.ArrayLike] = None, -) -> LabeledErrorResult: +) -> LabeledPdfResult: """Rotate and relabel PDFs. Parameters @@ -170,10 +155,8 @@ def rotate_result( 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 + new_grids[ep] = np.einsum( + "ab,rbk->rak", flavor_rotation, pdf_grid, optimize="optimal" ) grids = new_grids @@ -188,10 +171,8 @@ def rotate_result( 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 + new_grids[ep] = np.einsum( + "jk,rbk->rbj", x_rotation, pdf_grid, optimize="optimal" ) grids = new_grids @@ -201,9 +182,7 @@ def rotate_result( new_grids[ep] = dict( zip( flavor_labels, - np.swapaxes(pdf_grid, 0, 1) - if pdf_grid is not None - else [None] * len(flavor_labels), + np.swapaxes(pdf_grid, 0, 1), ) ) grids = new_grids @@ -217,7 +196,7 @@ def rotate_result( def apply_grids( eko: EKO, input_grids: npt.ArrayLike -) -> tuple[RawPdfResult, RawErrorResult]: +) -> tuple[RawPdfResult, RawPdfResult]: """Apply all available operators to the input grids. Parameters @@ -231,7 +210,7 @@ def apply_grids( ------- pdfs : output PDFs for the computed evolution points - pdfs : + errors : associated integration errors for the computed evolution points """ # sanity check @@ -245,14 +224,13 @@ def apply_grids( ) # iterate pdfs: RawPdfResult = {} - errors: RawErrorResult = {} + errors: RawPdfResult = {} 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 - ) + if elem.error is not None: + errors[ep] = np.einsum( + _EKO_CONTRACTION, elem.error, input_grids, optimize="optimal" + ) return pdfs, errors diff --git a/tests/ekobox/test_apply.py b/tests/ekobox/test_apply.py index 522c4dc6d..5790a4822 100644 --- a/tests/ekobox/test_apply.py +++ b/tests/ekobox/test_apply.py @@ -9,9 +9,8 @@ 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) + 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 @@ -30,6 +29,6 @@ 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(): + pdfs, _errors = apply.apply_grids(eko, input_grids) + for _ep, res in pdfs.items(): assert res.shape == input_grids.shape