Tests for evolven3fit_new and lhapdf .info files.

.git-blame-ignore-revs

@@ -12,3 +12,5 @@ b9084ce695c49b5645b82436b19ec6cc3bd30280
 a5d7385a79ef242e7eb1b1bf7b3746b6ebd6a0d0
 3e8cf70c2d3b36d19fcb764599f6dfa6eb924b77
 27f25d037854caa7e15496e4c6fa9d520f365ac9
+# Restyling of evolven3fit_new
+07d8d315aff44f98cb546e09088aa981e9778466

n3fit/src/evolven3fit_new/eko_utils.py

@@ -1,11 +1,12 @@
 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
 from eko.matchings import Atlas, nf_default
 from eko.quantities.heavy_quarks import MatchingScales
-from ekobox.cards import _operator as default_op_card
 from validphys.loader import Loader
 
 from . import utils
@@ -31,6 +32,7 @@
 NFREF_DEFAULT = 5
 NF0_DEFAULT = 4
 
+
 def construct_eko_cards(
     theoryID,
     q_fin,
@@ -51,6 +53,7 @@ def construct_eko_cards(
         theory_card_dict = {}
     if op_card_dict is None:
         op_card_dict = {}
+
     # theory_card construction
     theory = Loader().check_theoryID(theoryID).get_description()
     theory.pop("FNS")
@@ -59,22 +62,32 @@ def construct_eko_cards(
         theory["nfref"] = NFREF_DEFAULT
     if "nf0" not in theory:
         theory["nf0"] = NF0_DEFAULT
+
     # The Legacy function is able to construct a theory card for eko starting from an NNPDF theory
     legacy_class = runcards.Legacy(theory, {})
     theory_card = legacy_class.new_theory
+
+    # Prepare the thresholds according to MaxNfPdf
+    thresholds = {"c": theory["kcThr"], "b": theory["kbThr"], "t": theory["ktThr"]}
+    if theory["MaxNfPdf"] < 5:
+        thresholds["b"] = np.inf
+    if theory["MaxNfPdf"] < 6:
+        thresholds["t"] = np.inf
+
     # construct operator card
     q2_grid = utils.generate_q2grid(
         theory["Q0"],
         q_fin,
         q_points,
-        {theory["mb"]: theory["kbThr"], theory["mt"]: theory["ktThr"]},
+        {theory["mb"]: thresholds["b"], theory["mt"]: thresholds["t"]},
     )
     op_card = default_op_card
-    masses = np.array([theory["mc"],theory["mb"],theory["mt"]]) ** 2
-    thresholds_ratios=np.array([theory["kcThr"],theory["kbThr"],theory["ktThr"]]) ** 2
+    masses = np.array([theory["mc"], theory["mb"], theory["mt"]]) ** 2
+    thresholds_ratios = np.array([thresholds["c"], thresholds["b"], thresholds["t"]]) ** 2
+
     atlas = Atlas(
         matching_scales=MatchingScales(masses * thresholds_ratios),
-        origin=(theory["Q0"]**2, theory["nf0"])
+        origin=(theory["Q0"] ** 2, theory["nf0"]),
     )
     op_card.update(
         {
@@ -100,3 +113,22 @@ def construct_eko_cards(
 
     op_card = runcards.OperatorCard.from_dict(op_card)
     return theory_card, op_card
+
+
+def split_evolgrid(evolgrid):
+    """Split the evolgrid in blocks according to the number of flavors and repeating the last entry of one block in the first entry of the next."""
+    evolgrid_index_list = []
+    evolgrid.sort()
+    starting_nf = evolgrid[0][1]
+    for evo_point in evolgrid:
+        current_nf = evo_point[1]
+        if current_nf != starting_nf:
+            evolgrid_index_list.append(evolgrid.index(evo_point))
+            starting_nf = current_nf
+    start_index = 0
+    evolgrid_list = []
+    for index in evolgrid_index_list:
+        evolgrid_list.append(evolgrid[start_index : index + 1])
+        start_index = index
+    evolgrid_list.append(evolgrid[start_index:])
+    return evolgrid_list

n3fit/src/evolven3fit_new/evolve.py

@@ -2,11 +2,11 @@
 import pathlib
 import sys
 
+from ekobox import apply, genpdf, info_file
 import numpy as np
+
 import eko
-from eko import basis_rotation
-from eko import runner
-from ekobox import genpdf, info_file, apply
+from eko import basis_rotation, runner
 from reportengine.compat import yaml
 from validphys.loader import Loader
 
@@ -77,14 +77,9 @@ def evolve_fit(
 
     usr_path = pathlib.Path(fit_folder)
     initial_PDFs_dict = load_fit(usr_path)
-    x_grid = np.array(
-        initial_PDFs_dict[list(initial_PDFs_dict.keys())[0]]["xgrid"]
-    ).astype(float)
+    x_grid = np.array(initial_PDFs_dict[list(initial_PDFs_dict.keys())[0]]["xgrid"]).astype(float)
     theoryID = utils.get_theoryID_from_runcard(usr_path)
-    theory, op = eko_utils.construct_eko_cards(
-        theoryID, q_fin, q_points, x_grid, op_card_dict, theory_card_dict
-    )
-    qed = theory.order[1] > 0
+
     if eko_path is not None:
         eko_path = pathlib.Path(eko_path)
         _logger.info(f"Loading eko from : {eko_path}")
@@ -94,23 +89,38 @@ def evolve_fit(
             eko_path = (Loader().check_theoryID(theoryID).path) / "eko.tar"
         except FileNotFoundError:
             _logger.info(f"eko not found in theory {theoryID}, we will construct it")
+            theory, op = eko_utils.construct_eko_cards(
+                theoryID, q_fin, q_points, x_grid, op_card_dict, theory_card_dict
+            )
             runner.solve(theory, op, dump_eko)
             eko_path = dump_eko
+
     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)
-    info = info_file.build(theory, op, 1, info_update={})
-    info["NumMembers"] = "REPLACE_NREP"
-    info["ErrorType"] = "replicas"
-    info["XMin"] = float(x_grid[0])
-    info["XMax"] = float(x_grid[-1])
+
     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
+        op = eko_op.operator_card
+
+        qed = theory.order[1] > 0
+
+        # Modify the info file with the fit-specific info
+        info = info_file.build(theory, op, 1, info_update={})
+        info["NumMembers"] = "REPLACE_NREP"
+        info["ErrorType"] = "replicas"
+        info["XMin"] = float(x_grid[0])
+        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
         dump_info_file(usr_path, info)
-        for replica in initial_PDFs_dict.keys():
-            evolved_block = evolve_exportgrid(initial_PDFs_dict[replica], eko_op, x_grid, qed)
-            dump_evolved_replica(
-                evolved_block, usr_path, int(replica.removeprefix("replica_"))
-            )
+
+        for replica, pdf_data in initial_PDFs_dict.items():
+            evolved_blocks = evolve_exportgrid(pdf_data, eko_op, x_grid, qed)
+            dump_evolved_replica(evolved_blocks, usr_path, int(replica.removeprefix("replica_")))
+
     # remove folder:
     # The function dump_evolved_replica dumps the replica files in a temporary folder
     # We need then to remove it after fixing the position of those replica files
@@ -157,8 +167,8 @@ def evolve_exportgrid(exportgrid, eko, x_grid, qed):
             whether qed is activated or not
     Returns
     -------
-        : np.array
-        evolved block
+        : list(np.array)
+        list of evolved blocks
     """
     # construct LhapdfLike object
     pdf_grid = np.array(exportgrid["pdfgrid"]).transpose()
@@ -171,24 +181,28 @@ def evolve_exportgrid(exportgrid, eko, x_grid, qed):
     def ev_pdf(pid, x, Q2):
         return x * evolved_pdf[Q2]["pdfs"][pid][targetgrid.index(x)]
 
-    block = genpdf.generate_block(
-        ev_pdf,
-        xgrid=targetgrid,
-        evolgrid=eko.evolgrid,
-        pids=basis_rotation.flavor_basis_pids,
-    )
-    return block
+    evolgrid_list = eko_utils.split_evolgrid(eko.evolgrid)
+    blocks = []
+    for evgrid in evolgrid_list:
+        block = genpdf.generate_block(
+            ev_pdf,
+            xgrid=targetgrid,
+            evolgrid=evgrid,
+            pids=basis_rotation.flavor_basis_pids,
+        )
+        blocks.append(block)
+    return blocks
 
 
-def dump_evolved_replica(evolved_block, usr_path, replica_num):
+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
     the folder usr_path/nnfit/replica_<replica_num>/usr_path.stem.dat
 
     Parameters
     ----------
-        evolved_block: numpy.array
-            block of an evolved PDF
+        evolved_block: list(numpy.array)
+            list of blocks of an evolved PDF
         usr_path: pathlib.Path
             path of the fit folder
         replica_num: int
@@ -199,7 +213,7 @@ def dump_evolved_replica(evolved_block, usr_path, replica_num):
     path_where_dump.mkdir(exist_ok=True)
     to_write_in_head = f"PdfType: replica\nFromMCReplica: {replica_num}\n"
     genpdf.export.dump_blocks(
-        path_where_dump, replica_num, [evolved_block], pdf_type=to_write_in_head
+        path_where_dump, replica_num, evolved_blocks, pdf_type=to_write_in_head
     )
     # fixing_replica_path
     utils.fix_replica_path(usr_path, replica_num)