Merge pull request #2181 from NNPDF/use_eko_v15

Use eko v0.15

Author: scarlehoff

conda-recipe/meta.yaml

@@ -39,7 +39,7 @@ requirements:
         - prompt_toolkit
         - validobj
         - pineappl >=0.8.2
-        - eko >=0.14.2,<0.15
+        - eko >=0.15.1,<0.16
         - fiatlux
         - sphinx >=5.0.2
         - joblib

doc/sphinx/.gitignore

@@ -1,3 +1,4 @@
 # Ignore auto generated module references
 source/modules
 source/theories.csv
+source/theories_central.csv

doc/sphinx/Makefile

@@ -19,6 +19,7 @@ help:
 	@if test $@ != "clean"; then \
 		sphinx-apidoc -f -o ./$(SOURCEDIR)/modules/validphys ../../validphys2/src/validphys/ ; \
 		sphinx-apidoc -f -o ./$(SOURCEDIR)/modules/n3fit-code ../../n3fit/src/n3fit/ ; \
+		sphinx-apidoc -f -o ./$(SOURCEDIR)/modules/n3fit-code ../../n3fit/src/evolven3fit/ ; \
 	fi
 	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 

n3fit/src/evolven3fit/cli.py

@@ -2,7 +2,7 @@
 
 
 def cli_evolven3fit(
-    configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump, ncores
+    configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump
 ):
     """Evolves the fitted PDFs.
 
@@ -23,13 +23,5 @@ def cli_evolven3fit(
     """
     utils.check_is_a_fit(configuration_folder)
     return evolve.evolve_fit(
-        configuration_folder,
-        q_fin,
-        q_points,
-        op_card_info,
-        theory_card_info,
-        force,
-        load,
-        dump,
-        ncores,
+        configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump
     )

n3fit/src/evolven3fit/eko_utils.py

@@ -1,7 +1,6 @@
 import logging
 from typing import Any, Dict, Optional
 
-from ekobox.cards import _operator as default_op_card
 import numpy as np
 
 from eko.io import runcards
@@ -30,6 +29,11 @@
 NFREF_DEFAULT = 5
 
 
+def _eko_theory_from_nnpdf_theory(nnpdf_theory):
+    """Takes an NNPDF theory and returns an eko theory"""
+    return runcards.Legacy(nnpdf_theory, {}).new_theory
+
+
 def construct_eko_cards(
     nnpdf_theory,
     q_fin,
@@ -54,10 +58,7 @@ def construct_eko_cards(
     # eko needs a value for Qedref and for max nf alphas
     theory["Qedref"] = theory["Qref"]
     theory["MaxNfAs"] = theory["MaxNfPdf"]
-
-    # The Legacy function is able to construct a theory card for eko starting from a NNPDF theory
-    legacy_class = runcards.Legacy(theory, {})
-    theory_card = legacy_class.new_theory
+    theory_card = _eko_theory_from_nnpdf_theory(theory)
 
     # construct mugrid
 
@@ -127,10 +128,7 @@ def construct_eko_photon_cards(
     # Now make sure the Legacy class still gets a Qedref, which is equal to Qref
     theory["Qedref"] = theory["Qref"]
     theory["MaxNfAs"] = theory["MaxNfPdf"]
-
-    # The Legacy function is able to construct a theory card for eko starting from a NNPDF theory
-    legacy_class = runcards.Legacy(theory, {})
-    theory_card = legacy_class.new_theory
+    theory_card = _eko_theory_from_nnpdf_theory(theory)
 
     # The photon needs to be evolved down to Q0
     q_fin = theory["Q0"]
@@ -177,10 +175,31 @@ def build_opcard(op_card_dict, theory, x_grid, mu0, mugrid):
     if op_card_dict is None:
         op_card_dict = {}
 
-    op_card = default_op_card
-
-    op_card.update({"mu0": mu0, "mugrid": mugrid})
+    # Taken from cards.py https://github.com/NNPDF/eko/blob/master/src/ekobox/cards.py
+    # 7735fdb
+    op_card = dict(
+        init=(1.65, 4),
+        mugrid=[(100.0, 5)],
+        xgrid=np.geomspace(1e-7, 1.0, 50).tolist(),
+        configs=dict(
+            # These three values might be set by op_card_dict
+            ev_op_iterations=10,
+            n_integration_cores=1,
+            polarized=False,
+            # ... but these we fix here
+            ev_op_max_order=[10, 0],
+            interpolation_polynomial_degree=4,
+            interpolation_is_log=True,
+            scvar_method=None,
+            inversion_method=None,
+            evolution_method="iterate-exact",
+            time_like=False,
+        ),
+        debug=dict(skip_singlet=False, skip_non_singlet=False),
+    )
 
+    op_card["init"] = (mu0, theory["nf0"])
+    op_card["mugrid"] = mugrid
     op_card["xgrid"] = x_grid
 
     # Specify the evolution options and defaults differently from TRN / EXA

n3fit/src/evolven3fit/evolve.py

@@ -1,16 +1,17 @@
 from collections import defaultdict
+import dataclasses
 import json
 import logging
 import pathlib
 import sys
 
 from ekobox import apply, genpdf, info_file
-from joblib import Parallel, delayed
 import numpy as np
-import psutil
 
 import eko
 from eko import basis_rotation, runner
+from eko.interpolation import XGrid
+from eko.io import manipulate
 from validphys.utils import yaml_safe
 
 from . import eko_utils, utils
@@ -24,19 +25,9 @@
     "level": logging.DEBUG,
 }
 
-NUM_CORES = psutil.cpu_count(logical=False)
-
 
 def evolve_fit(
-    fit_folder,
-    q_fin,
-    q_points,
-    op_card_dict,
-    theory_card_dict,
-    force,
-    eko_path,
-    dump_eko=None,
-    ncores=1,
+    fit_folder, q_fin, q_points, op_card_dict, theory_card_dict, force, eko_path, dump_eko=None
 ):
     """
     Evolves all the fitted replica in fit_folder/nnfit
@@ -106,10 +97,7 @@ def evolve_fit(
         else:
             raise ValueError(f"dump_eko not provided and {eko_path=} not found")
 
-    with eko.EKO.edit(eko_path) as eko_op:
-        x_grid_obj = eko.interpolation.XGrid(x_grid)
-        eko.io.manipulate.xgrid_reshape(eko_op, targetgrid=x_grid_obj, inputgrid=x_grid_obj)
-
+    # Open the EKO in read-only mode, if it needs to be manipulated keep it in memory
     with eko.EKO.read(eko_path) as eko_op:
         # Read the cards directly from the eko to make sure they are consistent
         theory = eko_op.theory_card
@@ -118,6 +106,42 @@ def evolve_fit(
         _logger.debug(f"Theory card: {json.dumps(theory.raw)}")
         _logger.debug(f"Operator card: {json.dumps(op.raw)}")
 
+        eko_original_xgrid = eko_op.xgrid
+        if XGrid(x_grid) != eko_original_xgrid:
+            # If the xgrid of the eko is not directly usable, construct a copy in memory
+            # by replacing the internal inventory of operators in a readonly copy
+            new_xgrid = XGrid(x_grid)
+            new_metadata = dataclasses.replace(eko_op.metadata, xgrid=new_xgrid)
+
+            new_operators = {}
+            for target_key in eko_op.operators:
+                elem = eko_op[target_key.ep]
+
+                if eko_op.metadata.version == "0.13.4":
+                    # For eko 0.13.4 xgrid is the internal interpolation so we need to check
+                    # whether the rotation is truly needed
+                    # <in practice> this means checking whether the operator shape matches the grid
+                    oplen = elem.operator.shape[-1]
+                    if oplen != len(eko_original_xgrid):
+                        # The operator and its xgrid have different shape
+                        # either prepare an identity, or this EKO is not supported
+                        if oplen != len(x_grid):
+                            raise ValueError(
+                                f"The operator at {eko_path} is not usable, version not supported"
+                            )
+                        eko_original_xgrid = XGrid(x_grid)
+
+                new_operators[target_key] = manipulate.xgrid_reshape(
+                    elem,
+                    eko_original_xgrid,
+                    op.configs.interpolation_polynomial_degree,
+                    targetgrid=XGrid(x_grid),
+                    inputgrid=XGrid(x_grid),
+                )
+
+            new_inventory = dataclasses.replace(eko_op.operators, cache=new_operators)
+            eko_op = dataclasses.replace(eko_op, metadata=new_metadata, operators=new_inventory)
+
         # Modify the info file with the fit-specific info
         info = info_file.build(theory, op, 1, info_update={})
         info["NumMembers"] = "REPLACE_NREP"
@@ -126,21 +150,49 @@ def evolve_fit(
         info["XMax"] = float(x_grid[-1])
         # Save the PIDs in the info file in the same order as in the evolution
         info["Flavors"] = basis_rotation.flavor_basis_pids
-        info["NumFlavors"] = theory.heavy.num_flavs_max_pdf
+        info.setdefault("NumFlavors", 5)
         dump_info_file(usr_path, info)
 
-        def _wrap_evolve(pdf, replica):
-            evolved_blocks = evolve_exportgrid(pdf, eko_op, x_grid)
-            dump_evolved_replica(evolved_blocks, usr_path, int(replica.removeprefix("replica_")))
-
-        # Choose the number of cores to be the Minimal value
-        nb_cores = min(NUM_CORES, abs(ncores))
-        Parallel(n_jobs=nb_cores)(
-            delayed(_wrap_evolve)(pdf, r) for r, pdf in initial_PDFs_dict.items()
-        )
+        # Read the information from all the sorted replicas into what eko wants
+        n_replicas = len(initial_PDFs_dict)
+        all_replicas = []
+        for rep_idx in range(1, n_replicas + 1):
+            # swap photon position to match eko.basis_rotation.flavor_basis_pids
+            pdfgrid = np.array(initial_PDFs_dict[f"replica_{rep_idx}"]["pdfgrid"])
+            pdfgrid = np.append(pdfgrid[:, -1].reshape(x_grid.size, 1), pdfgrid[:, :-1], axis=1)
+            # and divide by x
+            all_replicas.append(pdfgrid.T / x_grid)
+
+        # output is {(Q2, nf): (replica, flavour, x)}
+        all_evolved, _ = apply.apply_grids(eko_op, np.array(all_replicas))
+
+        # Now, replica by replica, break into nf blocks
+        targetgrid = eko_op.xgrid.tolist()
+        by_nf = defaultdict(list)
+        for q2, nf in sorted(eko_op.evolgrid, key=lambda ep: ep[1]):
+            by_nf[nf].append(q2)
+        q2block_per_nf = {nf: sorted(q2s) for nf, q2s in by_nf.items()}
+
+        for replica in range(n_replicas):
+            blocks = []
+            for nf, q2grid in q2block_per_nf.items():
+
+                def pdf_xq2(pid, x, Q2):
+                    x_idx = targetgrid.index(x)
+                    pid_idx = info["Flavors"].index(pid)
+                    return x * all_evolved[(Q2, nf)][replica][pid_idx][x_idx]
+
+                block = genpdf.generate_block(
+                    pdf_xq2,
+                    xgrid=targetgrid,
+                    sorted_q2grid=q2grid,
+                    pids=basis_rotation.flavor_basis_pids,
+                )
+                blocks.append(block)
+            dump_evolved_replica(blocks, usr_path, replica + 1)
 
     # remove folder:
-    # The function dump_evolved_replica dumps the replica files in a temporary folder
+    # The function dump_evolved_replica uses a temporary folder
     # We need then to remove it after fixing the position of those replica files
     (usr_path / "nnfit" / usr_path.stem).rmdir()
 
@@ -169,52 +221,6 @@ def load_fit(usr_path):
     return pdf_dict
 
 
-def evolve_exportgrid(exportgrid, eko, x_grid):
-    """
-    Evolves the provided exportgrid for the desired replica with the eko and returns the evolved block
-
-    Parameters
-    ----------
-        exportgrid: dict
-            exportgrid of pdf at fitting scale
-        eko: eko object
-            eko operator for evolution
-        xgrid: list
-            xgrid to be used as the targetgrid
-    Returns
-    -------
-        : list(np.array)
-        list of evolved blocks
-    """
-    # construct LhapdfLike object
-    pdf_grid = np.array(exportgrid["pdfgrid"]).transpose()
-    pdf_to_evolve = utils.LhapdfLike(pdf_grid, exportgrid["q20"], x_grid)
-    # evolve pdf
-    evolved_pdf = apply.apply_pdf(eko, pdf_to_evolve)
-    # generate block to dump
-    targetgrid = eko.bases.targetgrid.tolist()
-
-    # Finally separate by nf block (and order per nf/q)
-    by_nf = defaultdict(list)
-    for q, nf in sorted(eko.evolgrid, key=lambda ep: ep[1]):
-        by_nf[nf].append(q)
-    q2block_per_nf = {nf: sorted(qs) for nf, qs in by_nf.items()}
-
-    blocks = []
-    for nf, q2grid in q2block_per_nf.items():
-
-        def pdf_xq2(pid, x, Q2):
-            x_idx = targetgrid.index(x)
-            return x * evolved_pdf[(Q2, nf)]["pdfs"][pid][x_idx]
-
-        block = genpdf.generate_block(
-            pdf_xq2, xgrid=targetgrid, sorted_q2grid=q2grid, pids=basis_rotation.flavor_basis_pids
-        )
-        blocks.append(block)
-
-    return blocks
-
-
 def dump_evolved_replica(evolved_blocks, usr_path, replica_num):
     """
     Dump the evolved replica given by evolved_block as the replica num "replica_num" in

n3fit/src/evolven3fit/q2grids.py

@@ -1,12 +1,13 @@
 """
-    Definition of default Q2 grids
+Definition of default Q2 grids
 
-    This file includes:
+This file includes:
 
-        ``Q2GRID_DEFAULT``: default NNPDF Q2 grid for evolution (55 points, starts at Q=1GeV)
-        ``Q2GRID_NNPDF40``: q2 grid used in the fits for the NNPDF4.0 release (49 points, starts at Q=1.65 GeV)
-        ``Q2GRID_Nf03``: q2 grid used in the perturvative charm fits for the NNPDF4.0 release (48 points, starts at Q=1GeV)
+    - ``Q2GRID_DEFAULT``: default NNPDF Q2 grid for evolution (55 points, starts at Q=1GeV)
+    - ``Q2GRID_NNPDF40``: q2 grid used in the fits for the NNPDF4.0 release (49 points, starts at Q=1.65 GeV)
+    - ``Q2GRID_Nf03``: q2 grid used in the perturvative charm fits for the NNPDF4.0 release (48 points, starts at Q=1GeV)
 """
+
 import numpy as np
 
 Q2GRID_NNPDF40 = (