Merge pull request #138 from NNPDF/feature/nnlo-cc-hq

NNLO CC heavy quark

extras/heavy_cc_nnlo/.gitignore

@@ -0,0 +1,3 @@
+data/
+demoFelix/
+plots/

extras/heavy_cc_nnlo/reformat/__init__.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+import argparse
+import logging
+import pathlib
+
+from . import generate as gen
+from . import parse
+
+_logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("paths", nargs="*", type=pathlib.Path)
+    parser.add_argument("-a", "--actions", nargs="*", choices=["dump", "plot"])
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    for path in args.paths:
+        ch, blocks = parse.parse(path)
+        ds = gen.to_xarray(blocks)
+
+        if "dump" in args.actions:
+            dumpdest = pathlib.Path.cwd() / "data" / f"{path.stem}.nc"
+            dumpdest.parent.mkdir(exist_ok=True)
+            ds.to_netcdf(dumpdest)
+
+            _logger.info(
+                f"Saved xgrid array to {dumpdest.relative_to(pathlib.Path.cwd())}"
+            )
+
+        if "plot" in args.actions:
+            dest = pathlib.Path.cwd() / "plots" / f"{path.stem}.png"
+            dest.parent.mkdir(exist_ok=True)
+            gen.plot(ds, dest)

extras/heavy_cc_nnlo/reformat/__main__.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+import logging
+
+from rich.logging import RichHandler
+
+from . import main
+
+logging.basicConfig(
+    level=logging.INFO, format="%(message)s", datefmt="[%X]", handlers=[RichHandler()]
+)
+
+main()

extras/heavy_cc_nnlo/reformat/generate.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+import logging
+import os
+import pathlib
+
+import matplotlib.colors as clr
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+
+from . import parse
+
+_logger = logging.getLogger(__name__)
+
+
+def to_xarray(blocks: list[parse.Block]) -> xr.Dataset:
+    catdict = {
+        k: (["bl"], [getattr(dic, k) for dic in blocks]) for k in blocks[0].__dict__
+    }
+    catdict["grid"][0].extend(["x", "y"])
+    catdict["xgrid"][0].append("z")
+
+    return xr.Dataset(catdict)
+
+
+def plot(ds: xr.Dataset, path: os.PathLike):
+    path = pathlib.Path(path)
+
+    plt.figure()
+    ds["xgrid"].plot(norm=clr.LogNorm())
+    figpath = path.with_stem(f"{path.stem}-xgrid")
+    plt.savefig(figpath)
+
+    _logger.info(f"Saved xgrid plot to {figpath.relative_to(pathlib.Path.cwd())}")
+
+    plt.figure()
+    xbj = ds["xbj"]
+    nxbj = np.abs(np.log(xbj))
+    (nxbj / nxbj.max()).plot(label="xBj")
+    qc = ds["qcd_scale"]
+    (qc / qc.max()).plot(label="qcd")
+    en = ds["energy"]
+    nen = np.log(en)
+    (nen / nen.max()).plot(label="E")
+    plt.ylabel("logs and ratios")
+    plt.legend()
+    plt.title(
+        f"xBj = {xbj.min():0.3g} - {xbj.max():0.3g}\n"
+        f"qcd = {qc.min()} - {qc.max()}\n"
+        f"E = {en.min():0.3g} - {en.max():0.3g}"
+    )
+    plt.tight_layout()
+    figpath = path.with_stem(f"{path.stem}-xbj-scales")
+    plt.savefig(figpath)
+
+    _logger.info(f"Saved xgrid plot to {figpath.relative_to(pathlib.Path.cwd())}")

extras/heavy_cc_nnlo/reformat/parse.py

@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+import os
+import pathlib
+import re
+from dataclasses import dataclass
+
+import numpy as np
+
+DOUBLE = r"-?\d\.\d+D[-+]\d+"
+INT = r"-?\d+"
+
+BLOCK = r"^<Block([ \d]+)>(.*?)^<End([ \d]+)>"
+KINS = re.compile(
+    rf"^ Bjorken.*\n ({DOUBLE}) +({DOUBLE}) +({DOUBLE}) +({INT})", flags=re.MULTILINE
+)
+SCALES = re.compile(rf"^ Charm.*\n ({DOUBLE}) +({DOUBLE})", flags=re.MULTILINE)
+XGRID = re.compile(r"^ x-grid +(\d+)", flags=re.MULTILINE)
+INTERPOLATION = re.compile(
+    rf"^ Interpolation coefficients: +({INT}) +({INT}) +({INT}) +({INT})",
+    flags=re.MULTILINE,
+)
+
+
+@dataclass
+class Block:
+    xbj: float
+    ybj: float
+    energy: float
+    cpc: int
+    mc: float
+    qcd_scale: float
+    xgrid: np.ndarray
+    grid: np.ndarray
+
+
+def parse(path: os.PathLike) -> tuple[int, list[Block]]:
+    path = pathlib.Path(path)
+    content = path.read_text(encoding="utf-8")
+
+    charge = 1 if path.stem.startswith("nu") else -1
+
+    blocks = []
+    for bl in re.finditer(BLOCK, content, flags=re.MULTILINE | re.DOTALL):
+        blocks.append(block(bl))
+
+    return charge, blocks
+
+
+def dtof(double: str):
+    return float(double.replace("D", "e"))
+
+
+def block(matched: re.Match) -> Block:
+    if int(matched[1]) != int(matched[3]):
+        raise ValueError
+
+    content = matched[2]
+
+    parsed = {}
+
+    kins = re.search(KINS, content)
+    if kins is None:
+        raise ValueError()
+    parsed["xbj"] = dtof(kins[1])
+    parsed["ybj"] = dtof(kins[2])
+    parsed["energy"] = dtof(kins[3])
+    parsed["cpc"] = int(kins[4])
+
+    scales = re.search(SCALES, content)
+    if scales is None:
+        raise ValueError()
+    parsed["mc"] = dtof(scales[1])
+    parsed["qcd_scale"] = dtof(scales[2])
+
+    nx = int(re.search(XGRID, content)[1])
+    xgrid = [
+        dtof(x)
+        for x in content.split("x-grid")[1].split("Interpolation")[0].splitlines()[1:-1]
+    ]
+    parsed["xgrid"] = xgrid
+
+    if len(xgrid) != nx:
+        raise ValueError(
+            f"xgrid actual length, {len(xgrid)}, different from what declared, {nx}"
+        )
+
+    coeffs = re.search(INTERPOLATION, content)
+    if coeffs is None:
+        raise ValueError()
+    a, nx1, c, _ = int(coeffs[1]), int(coeffs[2]), int(coeffs[3]), int(coeffs[4])
+    grid = np.array(
+        [
+            [dtof(val) for val in line.strip().split()]
+            for line in content.split("Interpolation")[1].splitlines()[1:]
+        ]
+    )
+    parsed["grid"] = grid
+
+    if grid.shape != (a * nx1, c):
+        raise ValueError()
+
+    return Block(**parsed)