Merge branch 'num_fonll' of github.com:NNPDF/pineko into num_fonll

benchmarks/bench_fonll.py

@@ -23,4 +23,4 @@ def benchmark_produce_fonll_tcards(tmp_path, test_files, test_configs, theoryid)
         assert simFONLL_tcard["FNS"] == fns
         assert simFONLL_tcard["NfFF"] == nfff
         assert simFONLL_tcard["PTO"] == po
-        assert simFONLL_tcard["fonll-parts"] == part
+        assert simFONLL_tcard["FONLLParts"] == part

pyproject.toml

@@ -25,7 +25,7 @@ packages = [{ include = "pineko", from = "src" }]
 [tool.poetry.dependencies]
 python = ">=3.8,<3.12"
 eko = "^0.13.2"
-pineappl = { version = "^0.6.0a17", allow-prereleases = true }
+pineappl = "^0.6.0"
 PyYAML = "^6.0"
 numpy = "^1.21.0"
 pandas = "^1.4.1"

src/pineko/cli/fonll.py

@@ -1,7 +1,8 @@
 """CLI entry point to FONLL."""
 import click
+import rich
 
-from .. import configs, fonll, theory_card
+from .. import configs, fonll, parser, theory_card
 from ._base import command
 
 
@@ -12,32 +13,110 @@ class TheoryCardError(Exception):
 
 
 class InconsistentInputsError(Exception):
-    """Raised if the inputs are not consistent with FONLL accounting for only charm mass, or both charm and bottom masses."""
+    """Raised if the inputs are not consistent with FONLL."""
 
     pass
 
 
+def get_list_grids_name(yaml_file):
+    """Return the list of the grids in the yaml file."""
+    yaml_content = parser._load_yaml(yaml_file)
+    # Turn the operands and the members into paths (and check all of them exist)
+    ret = []
+    for operand in yaml_content["operands"]:
+        for member in operand:
+            ret.append(f"{member}.{parser.EXT}")
+    return ret
+
+
 @command.command("combine_fonll")
-@click.argument("FFNS3", type=click.Path(exists=True))
-@click.argument("FFN03", type=click.Path(exists=True))
-@click.argument("FFNS4til", type=click.Path(exists=True))
 @click.argument("theoryID", type=int)
-@click.option("--FFNS4bar", type=click.Path(exists=True))
-@click.option("--FFN04", type=click.Path(exists=True))
-@click.option("--FFNS5til", type=click.Path(exists=True))
-@click.option("--FFNS5bar", type=click.Path(exists=True))
-def subcommand(ffns3, ffn03, ffns4til, theoryid, ffns4bar, ffn04, ffns5til, ffns5bar):
+@click.argument("dataset", type=str)
+@click.option("--FFNS3", type=int, help="theoryID containing the ffns3 fktable")
+@click.option("--FFN03", type=int, help="theoryID containing the ffn03 fktable")
+@click.option("--FFNS4til", type=int, help="theoryID containing the ffns4til fktable")
+@click.option("--FFNS4bar", type=int, help="theoryID containing the ffns4bar fktable")
+@click.option("--FFN04", type=int, help="theoryID containing the ffn04 fktable")
+@click.option("--FFNS5til", type=int, help="theoryID containing the ffns5til fktable")
+@click.option("--FFNS5bar", type=int, help="theoryID containing the ffns5bar fktable")
+@click.option("--overwrite", is_flag=True, help="Allow files to be overwritten")
+def subcommand(
+    theoryid,
+    dataset,
+    ffns3,
+    ffn03,
+    ffns4til,
+    ffns4bar,
+    ffn04,
+    ffns5til,
+    ffns5bar,
+    overwrite,
+):
     """Combine the different FKs needed to produce the FONLL prescription."""
+    # Checks
     if (ffn04 is None and ffns5til is not None) or (
         ffn04 is not None and ffns5til is None
     ):
         print(
             "One between ffn04 and ffns5til has been provided without the other. Since they are both needed to construct FONLL, this does not make sense."
         )
         raise InconsistentInputsError
-    fonll.produce_combined_fk(
-        ffns3, ffn03, ffns4til, ffns4bar, ffn04, ffns5til, ffns5bar, theoryid
+    if (ffn03 is None and ffns4til is not None) or (
+        ffn03 is not None and ffns4til is None
+    ):
+        print(
+            "One between ffn03 and ffns4til has been provided without the other. Since they are both needed to construct FONLL, this does not make sense."
+        )
+        raise InconsistentInputsError
+    # Get theory info
+    configs.configs = configs.defaults(configs.load())
+    tcard = theory_card.load(theoryid)
+    if not "DAMPPOWER" in tcard:
+        if tcard["DAMP"] != 0:
+            raise InconsistentInputsError
+        tcard["DAMPPOWER"] = None
+    # Getting the paths to the grids
+    grids_name = get_list_grids_name(
+        configs.configs["paths"]["ymldb"] / f"{dataset}.yaml"
     )
+    for grid in grids_name:
+        # Checking if it already exists
+        new_fk_path = configs.configs["paths"]["fktables"] / str(theoryid) / grid
+        if new_fk_path.exists():
+            if not overwrite:
+                rich.print(
+                    f"[green]Success:[/] skipping existing FK Table {new_fk_path}"
+                )
+                return
+        fonll.produce_combined_fk(
+            configs.configs["paths"]["fktables"] / str(ffns3) / grid
+            if (configs.configs["paths"]["fktables"] / str(ffns3) / grid).exists()
+            else None,
+            configs.configs["paths"]["fktables"] / str(ffn03) / grid
+            if (configs.configs["paths"]["fktables"] / str(ffn03) / grid).exists()
+            else None,
+            configs.configs["paths"]["fktables"] / str(ffns4til) / grid
+            if (configs.configs["paths"]["fktables"] / str(ffns4til) / grid).exists()
+            else None,
+            configs.configs["paths"]["fktables"] / str(ffns4bar) / grid
+            if (configs.configs["paths"]["fktables"] / str(ffns4bar) / grid).exists()
+            else None,
+            configs.configs["paths"]["fktables"] / str(ffn04) / grid
+            if (configs.configs["paths"]["fktables"] / str(ffn04)).exists()
+            else None,
+            configs.configs["paths"]["fktables"] / str(ffns5til) / grid
+            if (configs.configs["paths"]["fktables"] / str(ffns5til) / grid).exists()
+            else None,
+            configs.configs["paths"]["fktables"] / str(ffns5bar) / grid
+            if (configs.configs["paths"]["fktables"] / str(ffns5bar) / grid).exists()
+            else None,
+            theoryid,
+            damp=(tcard["DAMP"], tcard["DAMPPOWER"]),
+        )
+        if new_fk_path.exists():
+            rich.print(f"[green]Success:[/] Wrote FK table to {new_fk_path}")
+        else:
+            rich.print(f"[red]Failure:[/]")
 
 
 @command.command("fonll_tcards")

src/pineko/fonll.py

@@ -29,48 +29,54 @@ def __init__(
     @property
     def fk_paths(self):
         """Returns the list of the fk paths needed to produce FONLL predictions."""
-        paths = [self.ffns3, self.ffn03, self.ffns4til]
-        if self.ffns4bar:
-            paths = paths + [self.ffns4bar]
-        if self.ffn04 and self.ffns5til:
-            paths = paths + [self.ffn04, self.ffns5til]
-            # It does not make sense to have ffns5bar without ffns5til
-            if self.ffns5bar:
-                paths = paths + [self.ffns5bar]
-        return [Path(p) for p in paths]
+        paths = {
+            "ffns3": self.ffns3,
+            "ffn03": self.ffn03,
+            "ffns4til": self.ffns4til,
+            "ffns4bar": self.ffns4bar,
+            "ffn04": self.ffn04,
+            "ffns5til": self.ffns5til,
+            "ffns5bar": self.ffns5bar,
+        }
+        return {p: Path(paths[p]) for p in paths if paths[p] is not None}
 
     @property
     def fks(self):
         """Returns the pineappl.Grid objects reading the paths provided by self.fk_paths."""
-        return [pineappl.grid.Grid.read(fk_path) for fk_path in self.fk_paths]
+        return {fk: pineappl.grid.Grid.read(self.fk_paths[fk]) for fk in self.fk_paths}
 
     @property
     def dataset_name(self):
         """Return the name of the dataset given by the fktables name, if all the fktables have the same name."""
-        if len({p.name for p in self.fk_paths}) == 1:
-            return self.fk_paths[0].name
+        if len({self.fk_paths[p].name for p in self.fk_paths}) == 1:
+            return self.fk_paths[list(self.fk_paths)[0]].name
         else:
             raise ValueError("not all fktables have the same name")
 
     @property
     def theorycard_no_fns_pto(self):
         """Return the common theory info between the different FONLL FKs."""
-        theorycards = [json.loads(fk.key_values()["theory"]) for fk in self.fks]
+        theorycards = [json.loads(self.fks[p].key_values()["theory"]) for p in self.fks]
         # Only these should differ
         for card in theorycards:
             del card["FNS"]
             del card["PTO"]
             del card["NfFF"]
             del card["ID"]
-            del card["fonll-parts"]
-        if not all([theorycards[0] == card in theorycards[1:]]):
+            del card["FONLLParts"]
+            del card["Comments"]
+        if len(theorycards) > 1 and not all(
+            [theorycards[0] == card in theorycards[1:]]
+        ):
             raise ValueError("theorycards not the same")
         return theorycards[0]
 
     @property
     def Q2grid(self):
         """Return the Q2grid of the fktables given by self.fks ."""
-        aa = json.loads(self.fks[0].raw.key_values()["runcard"])["observables"]
+        aa = json.loads(self.fks[list(self.fks)[0]].key_values()["runcard"])[
+            "observables"
+        ]
         bb = list(aa.values())[
             0
         ]  # there is only a single obseravble because it's a dis fktable
@@ -79,51 +85,61 @@ def Q2grid(self):
 
 
 # Notice we rely on the order defined by the FONLLInfo class
-FK_TO_DAMP = {"mc": [1, 2, 4, 5], "mb": [4, 5]}
-FK_WITH_MINUS = [1, 4]  # asy terms should be subtracted, therefore the sign
+FK_TO_DAMP = {
+    "mc": ["ffn03", "ffns4til", "ffn04", "ffns5til"],
+    "mb": ["ffn04", "ffns5til"],
+}
+FK_WITH_MINUS = ["ffn03", "ffn04"]  # asy terms should be subtracted, therefore the sign
 
 
 def produce_combined_fk(
-    ffns3, ffn03, ffns4til, ffns4bar, ffn04, ffns5til, ffns5bar, theoryid
+    ffns3,
+    ffn03,
+    ffns4til,
+    ffns4bar,
+    ffn04,
+    ffns5til,
+    ffns5bar,
+    theoryid,
+    damp=(0, None),
 ):
     """Combine the FONLL FKs to produce one single final FONLL FK."""
     fonll_info = FONLLInfo(ffns3, ffn03, ffns4til, ffns4bar, ffn04, ffns5til, ffns5bar)
-
     theorycard_constituent_fks = fonll_info.theorycard_no_fns_pto
-    if theorycard_constituent_fks["DAMP"] == 0:
+    fk_dict = fonll_info.fks
+    if damp[0] == 0:
         # then there is no damping, not even Heaviside only
-        combined_fk = fonll_info.fks[0]
-        for fk_path in fonll_info.fk_paths[1:]:
-            combined_fk.merge_from_file(fk_path)
+        combined_fk = fk_dict[list(fk_dict)[0]]
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            for fk in list(fk_dict)[1:]:
+                tmpfile_path = Path(tmpdirname) / f"{fk}.pineappl.lz4"
+                sign = -1 if fk in FK_WITH_MINUS else 1
+                fk_dict[fk].scale(sign)
+                fk_dict[fk].write_lz4(tmpfile_path)
+                combined_fk.merge_from_file(tmpfile_path)
     else:
         mc = theorycard_constituent_fks["mc"]
         mb = theorycard_constituent_fks["mb"]
         q2grid = fonll_info.Q2grid
         step_function_charm = mc**2 < q2grid
         step_function_bottom = mb**2 < q2grid
-        damping_factor_charm = (1 - q2grid / mc) ** theorycard_constituent_fks[
-            "DAMPPOWER"
-        ]
-        damping_factor_bottom = (1 - q2grid / mb) ** theorycard_constituent_fks[
-            "DAMPPOWER"
-        ]
+        damping_factor_charm = (1 - mc / q2grid) ** damp[1]
+        damping_factor_bottom = (1 - mb / q2grid) ** damp[1]
         damping_factor_charm *= step_function_charm
         damping_factor_bottom *= step_function_bottom
         dampings = {"mc": damping_factor_charm, "mb": damping_factor_bottom}
-        for i, fk in enumerate(fonll_info.fks):
-            sign = 1
-            if i in FK_WITH_MINUS:
-                sign = -1
-            for mass in FK_TO_DAMP:
-                if i in FK_TO_DAMP[mass]:
-                    fk.scale_by_bin(sign * dampings[mass])
         # pineappl does not support operating with two grids in memory:
         # https://github.com/NNPDF/pineappl/blob/8a672bef6d91b07a4edfdefbe4e30e4b1dd1f976/pineappl_py/src/grid.rs#L614-L617
         with tempfile.TemporaryDirectory() as tmpdirname:
-            combined_fk = fonll_info.fks[0]
-            for i, fk in enumerate(fonll_info.fks[1:]):
-                tmpfile_path = Path(tmpdirname) / f"{i}.pineappl.lz4"
-                fk.write_lz4(tmpfile_path)
+            combined_fk = fk_dict[list(fk_dict)[0]]
+            for fk in list(fk_dict)[1:]:
+                tmpfile_path = Path(tmpdirname) / f"{fk}.pineappl.lz4"
+                sign = -1 if fk in FK_WITH_MINUS else 1
+                fk_dict[fk].scale(sign)
+                for mass in FK_TO_DAMP:
+                    if fk in FK_TO_DAMP[mass]:
+                        fk_dict[fk].scale_by_bin(dampings[mass])
+                fk_dict[fk].write_lz4(tmpfile_path)
                 combined_fk.merge_from_file(tmpfile_path)
 
     # update theorycard entries for the combined fktable by reading the yamldb of the original theory
@@ -180,7 +196,7 @@ def produce_fonll_recipe(fns):
                 "FNS": fns,
                 "NfFF": nfff,
                 "PTO": po,
-                "fonll-parts": part,
+                "FONLLParts": part,
             }
         )
     return fonll_recipe

tests/test_fonll.py

@@ -22,7 +22,26 @@ def test_FONLLInfo():
     partialfonll_fake_info = pineko.fonll.FONLLInfo(
         full_list[0], full_list[1], full_list[2], None, None, None, None
     )
-    assert fullfonll_fake_info.fk_paths == [pathlib.Path(fk) for fk in full_list]
-    # In this case it ignores the ffns5bar fk
-    assert wrongfonll_fake_info.fk_paths == [pathlib.Path(fk) for fk in full_list[:4]]
-    assert partialfonll_fake_info.fk_paths == [pathlib.Path(fk) for fk in full_list[:3]]
+    name_list = [
+        "ffns3",
+        "ffn03",
+        "ffns4til",
+        "ffns4bar",
+        "ffn04",
+        "ffns5til",
+        "ffns5bar",
+    ]
+    assert fullfonll_fake_info.fk_paths == {
+        name: pathlib.Path(fk) for name, fk in zip(name_list, full_list)
+    }
+    assert wrongfonll_fake_info.fk_paths == {
+        name: pathlib.Path(fk)
+        for name, fk in zip(
+            name_list[:4] + [name_list[-1]], full_list[:4] + [full_list[-1]]
+        )
+    }
+    assert partialfonll_fake_info.fk_paths == {
+        name: pathlib.Path(fk)
+        for name, fk in zip(name_list[:3], full_list[:3])
+        if fk is not None
+    }