Merge pull request #108 from N3PDF/feature/nb-njit

Recover numba
NNPDF · Mar 29, 2022 · e6ee4bf · e6ee4bf
2 parents 606138f + d5f8bba
commit e6ee4bf
Show file tree

Hide file tree

Showing 54 changed files with 236 additions and 199 deletions.
diff --git a/benchmarks/eko/benchmark_evol_to_unity.py b/benchmarks/eko/benchmark_evol_to_unity.py
@@ -49,6 +49,7 @@ class BenchmarkBackwardForward:
         "ev_op_max_order": 1,
         "ev_op_iterations": 1,
         "backward_inversion": "exact",
+        "n_integration_cores": 1,
     }
     g = OperatorGrid.from_dict(
         theory_card,

diff --git a/benchmarks/lha_paper_bench.py b/benchmarks/lha_paper_bench.py
@@ -222,14 +222,14 @@ def benchmark_sv(self, pto):
 
     # Benchmark to LHA
     # obj = BenchmarkVFNS()
-    # obj = BenchmarkFFNS()
+    obj = BenchmarkFFNS()
 
-    # obj.benchmark_plain(2)
+    obj.benchmark_plain(0)
     # obj.benchmark_sv(2)
 
     # # VFNS benchmarks with LHA settings
-    programs = ["LHA", "pegasus", "apfel"]
-    for p in programs:
-        obj = BenchmarkRunner(p)
-        # obj.benchmark_plain(2)
-        obj.benchmark_sv(2)
+    # programs = ["LHA", "pegasus", "apfel"]
+    # for p in programs:
+    #     obj = BenchmarkRunner(p)
+    #     # obj.benchmark_plain(2)
+    #     obj.benchmark_sv(2)
diff --git a/pyproject.toml b/pyproject.toml
@@ -84,6 +84,7 @@ bench-run.env.NUMBA_DISABLE_JIT.default = "0"
 lint = "pylint src/**/*.py -E"
 lint-warnings = "pylint src/**/*.py --exit-zero"
 sandbox = "python benchmarks/sandbox.py"
+lha = "python benchmarks/lha_paper_bench.py"
 nav = "ekonav --config benchmarks/banana.yaml"
 navigator = "ekonav --config benchmarks/banana.yaml"
 docs = { "shell" = "cd doc; make html" }

diff --git a/src/eko/anomalous_dimensions/__init__.py b/src/eko/anomalous_dimensions/__init__.py
@@ -23,7 +23,7 @@
 from . import as1, as2, as3, harmonics
 
 
-@nb.njit("Tuple((c16[:,:],c16,c16,c16[:,:],c16[:,:]))(c16[:,:])", cache=True)
+@nb.njit(cache=True)
 def exp_singlet(gamma_S):
     r"""
     Computes the exponential and the eigensystem of the singlet anomalous dimension matrix
@@ -71,7 +71,7 @@ def exp_singlet(gamma_S):
     return exp, lambda_p, lambda_m, e_p, e_m
 
 
-@nb.njit("c16[:](u1,u2,c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_ns(order, mode, n, nf):
     r"""
     Computes the tower of the non-singlet anomalous dimensions
@@ -131,7 +131,7 @@ def gamma_ns(order, mode, n, nf):
     return gamma_ns
 
 
-@nb.njit("c16[:,:,:](u1,c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_singlet(order, n, nf):
     r"""
     Computes the tower of the singlet anomalous dimensions matrices

diff --git a/src/eko/anomalous_dimensions/aem1.py b/src/eko/anomalous_dimensions/aem1.py
@@ -5,7 +5,7 @@
 from . import as1
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def gamma_phq(N):
     """
     Computes the leading-order photon-quark anomalous dimension
@@ -26,7 +26,7 @@ def gamma_phq(N):
     return as1.gamma_gq(N) / constants.CF
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_qph(N, nf):
     """
     Computes the leading-order quark-photon anomalous dimension
@@ -50,7 +50,7 @@ def gamma_qph(N, nf):
     return as1.gamma_qg(N, nf) / constants.TR * constants.NC
 
 
-@nb.njit("c16(u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_phph(nf):
     """
     Computes the leading-order photon-photon anomalous dimension
@@ -71,7 +71,7 @@ def gamma_phph(nf):
     return 2 / 3 * constants.NC * 2 * nf
 
 
-@nb.njit("c16(c16,c16)", cache=True)
+@nb.njit(cache=True)
 def gamma_ns(N, s1):
     """
     Computes the leading-order non-singlet QED anomalous dimension.

diff --git a/src/eko/anomalous_dimensions/as1.py b/src/eko/anomalous_dimensions/as1.py
@@ -7,7 +7,7 @@
 from .. import constants
 
 
-@nb.njit("c16(c16,c16)", cache=True)
+@nb.njit(cache=True)
 def gamma_ns(N, s1):
     """
     Computes the leading-order non-singlet anomalous dimension.
@@ -31,7 +31,7 @@ def gamma_ns(N, s1):
     return result
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_qg(N, nf):
     """
     Computes the leading-order quark-gluon anomalous dimension
@@ -55,7 +55,7 @@ def gamma_qg(N, nf):
     return result
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def gamma_gq(N):
     """
     Computes the leading-order gluon-quark anomalous dimension
@@ -77,7 +77,7 @@ def gamma_gq(N):
     return result
 
 
-@nb.njit("c16(c16,c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_gg(N, s1, nf):
     """
     Computes the leading-order gluon-gluon anomalous dimension
@@ -103,7 +103,7 @@ def gamma_gg(N, s1, nf):
     return result
 
 
-@nb.njit("c16[:,:](c16,c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_singlet(N, s1, nf):
     r"""
       Computes the leading-order singlet anomalous dimension matrix

diff --git a/src/eko/anomalous_dimensions/as2.py b/src/eko/anomalous_dimensions/as2.py
@@ -14,7 +14,7 @@
 from . import harmonics
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_nsm(n, nf):
     """
     Computes the |NLO| valence-like non-singlet anomalous dimension.
@@ -57,7 +57,7 @@ def gamma_nsm(n, nf):
     return result
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_nsp(n, nf):
     """
     Computes the |NLO| singlet-like non-singlet anomalous dimension.
@@ -98,7 +98,7 @@ def gamma_nsp(n, nf):
     return result
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_ps(n, nf):
     """
     Computes the |NLO| pure-singlet quark-quark anomalous dimension.
@@ -125,7 +125,7 @@ def gamma_ps(n, nf):
     return result
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_qg(n, nf):
     """
     Computes the |NLO| quark-gluon singlet anomalous dimension.
@@ -158,7 +158,7 @@ def gamma_qg(n, nf):
     return result
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_gq(n, nf):
     """
     Computes the |NLO| gluon-quark singlet anomalous dimension.
@@ -194,7 +194,7 @@ def gamma_gq(n, nf):
     return result
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_gg(n, nf):
     """
     Computes the |NLO| gluon-gluon singlet anomalous dimension.
@@ -232,7 +232,7 @@ def gamma_gg(n, nf):
     return result
 
 
-@nb.njit("c16[:,:](c16,u1)", cache=True)
+@nb.njit(cache=True)
 def gamma_singlet(N, nf):
     r"""
       Computes the next-leading-order singlet anomalous dimension matrix

diff --git a/src/eko/anomalous_dimensions/as3.py b/src/eko/anomalous_dimensions/as3.py
@@ -16,7 +16,7 @@
 zeta3 = harmonics.zeta3
 
 
-@nb.njit("c16(c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_nsm(n, nf, sx):
     """
     Computes the |NNLO| valence-like non-singlet anomalous dimension.
@@ -93,7 +93,7 @@ def gamma_nsm(n, nf, sx):
     return result
 
 
-@nb.njit("c16(c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_nsp(n, nf, sx):
     """
     Computes the |NNLO| singlet-like non-singlet anomalous dimension.
@@ -170,7 +170,7 @@ def gamma_nsp(n, nf, sx):
     return result
 
 
-@nb.njit("c16(c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_nsv(n, nf, sx):
     """
     Computes the |NNLO| valence non-singlet anomalous dimension.
@@ -225,7 +225,7 @@ def gamma_nsv(n, nf, sx):
     return result
 
 
-@nb.njit("c16(c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_ps(n, nf, sx):
     """
     Computes the |NNLO| pure-singlet quark-quark anomalous dimension.
@@ -297,7 +297,7 @@ def gamma_ps(n, nf, sx):
     return result
 
 
-@nb.njit("c16(c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_qg(n, nf, sx):
     """
     Computes the |NNLO| quark-gluon singlet anomalous dimension.
@@ -371,7 +371,7 @@ def gamma_qg(n, nf, sx):
     return result
 
 
-@nb.njit("c16(c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_gq(n, nf, sx):
     """
     Computes the |NNLO| gluon-quark singlet anomalous dimension.
@@ -461,7 +461,7 @@ def gamma_gq(n, nf, sx):
     return result
 
 
-@nb.njit("c16(c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_gg(n, nf, sx):
     """
     Computes the |NNLO| gluon-gluon singlet anomalous dimension.
@@ -549,7 +549,7 @@ def gamma_gg(n, nf, sx):
     return result
 
 
-@nb.njit("c16[:,:](c16,u1,c16[:])", cache=True)
+@nb.njit(cache=True)
 def gamma_singlet(N, nf, sx):
     r"""
       Computes the |NNLO| singlet anomalous dimension matrix

diff --git a/src/eko/anomalous_dimensions/harmonics.py b/src/eko/anomalous_dimensions/harmonics.py
@@ -16,7 +16,7 @@
 zeta5 = scipy.special.zeta(5)
 
 
-@nb.njit("c16(c16,u1)", cache=True)
+@nb.njit(cache=True)
 def cern_polygamma(Z, K):  # pylint: disable=all
     """
     Computes the polygamma functions :math:`\\psi_k(z)`.
@@ -133,7 +133,7 @@ def cern_polygamma(Z, K):  # pylint: disable=all
     # fmt: on
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def harmonic_S1(N):
     r"""
     Computes the harmonic sum :math:`S_1(N)`.
@@ -161,7 +161,7 @@ def harmonic_S1(N):
     return cern_polygamma(N + 1.0, 0) + np.euler_gamma
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def harmonic_S2(N):
     r"""
     Computes the harmonic sum :math:`S_2(N)`.
@@ -189,7 +189,7 @@ def harmonic_S2(N):
     return -cern_polygamma(N + 1.0, 1) + zeta2
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def harmonic_S3(N):
     r"""
     Computes the harmonic sum :math:`S_3(N)`.
@@ -217,7 +217,7 @@ def harmonic_S3(N):
     return 0.5 * cern_polygamma(N + 1.0, 2) + zeta3
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def harmonic_S4(N):
     r"""
     Computes the harmonic sum :math:`S_4(N)`.
@@ -245,7 +245,7 @@ def harmonic_S4(N):
     return zeta4 - 1.0 / 6.0 * cern_polygamma(N + 1.0, 3)
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def harmonic_S5(N):
     r"""
     Computes the harmonic sum :math:`S_5(N)`.
@@ -273,7 +273,7 @@ def harmonic_S5(N):
     return zeta5 + 1.0 / 24.0 * cern_polygamma(N + 1.0, 4)
 
 
-@nb.njit("c16(c16)", cache=True)
+@nb.njit(cache=True)
 def mellin_g3(N):
     r"""
     Computes the Mellin transform of :math:`\text{Li}_2(x)/(1+x)`.

diff --git a/src/eko/beta.py b/src/eko/beta.py
@@ -11,7 +11,7 @@
 from .anomalous_dimensions.harmonics import zeta3
 
 
-@nb.njit("f8(u1)", cache=True)
+@nb.njit(cache=True)
 def beta_0(nf):
     """
     Computes the first coefficient of the QCD beta function.
@@ -32,7 +32,7 @@ def beta_0(nf):
     return beta_0
 
 
-@nb.njit("f8(u1)", cache=True)
+@nb.njit(cache=True)
 def beta_1(nf):
     """
     Computes the second coefficient of the QCD beta function.
@@ -57,7 +57,7 @@ def beta_1(nf):
     return beta_1
 
 
-@nb.njit("f8(u1)", cache=True)
+@nb.njit(cache=True)
 def beta_2(nf):
     """
     Computes the third coefficient of the QCD beta function
@@ -86,7 +86,7 @@ def beta_2(nf):
     return beta_2
 
 
-@nb.njit("f8(u1)", cache=True)
+@nb.njit(cache=True)
 def beta_3(nf):
     """
     Computes the fourth coefficient of the QCD beta function
@@ -113,7 +113,7 @@ def beta_3(nf):
     return beta_3
 
 
-@nb.njit("f8(u1,u1)", cache=True)
+@nb.njit(cache=True)
 def beta(k, nf):
     """
     Compute value of a beta coefficients
@@ -144,7 +144,7 @@ def beta(k, nf):
     return beta_
 
 
-@nb.njit("f8(u1,u1)", cache=True)
+@nb.njit(cache=True)
 def b(k, nf):
     """
     Compute b coefficient.