Merge pull request #253 from NNPDF/harmonic_sum_cache_2

Refactor harmonic cache

benchmarks/ekore/benchmark_ad.py

@@ -136,9 +136,9 @@ def check_gamma_1_pegasus(N, NF):
     P1NSP = CF * ((CF - CA / 2.0) * PNPA + CA * PNSB + TR * NF * PNSC)
     P1NSM = CF * ((CF - CA / 2.0) * PNMA + CA * PNSB + TR * NF * PNSC)
 
-    sx = h.sx(N, 2)
-    np.testing.assert_allclose(ad_as2.gamma_nsp(N, NF, sx), -P1NSP)
-    np.testing.assert_allclose(ad_as2.gamma_nsm(N, NF, sx), -P1NSM)
+    cache = h.cache.reset()
+    np.testing.assert_allclose(ad_as2.gamma_nsp(N, NF, cache), -P1NSP)
+    np.testing.assert_allclose(ad_as2.gamma_nsm(N, NF, cache), -P1NSM)
 
     NS = N * N
     NT = NS * N
@@ -255,7 +255,7 @@ def check_gamma_1_pegasus(N, NF):
     P1Sgq = (CF * CF * PGQA + CF * CA * PGQB + TR * NF * CF * PGQC) * 4.0
     P1Sgg = (CA * CA * PGGA + TR * NF * (CA * PGGB + CF * PGGC)) * 4.0
 
-    gS1 = ad_as2.gamma_singlet(N, NF, sx)
+    gS1 = ad_as2.gamma_singlet(N, NF, cache)
     np.testing.assert_allclose(gS1[0, 0], -P1Sqq)
     np.testing.assert_allclose(gS1[0, 1], -P1Sqg)
     np.testing.assert_allclose(gS1[1, 0], -P1Sgq)

extras/ad_n3lo/parametrised_ad/translate_singlet_replicas.py

@@ -44,17 +44,17 @@
         name = gamma_name[1:]
         if "PS" in name:
             name = "ps"
-        oo.write(f"def gamma_{name}_{nf}(n, sx, variation):\n")
+        oo.write(f"def gamma_{name}_{nf}(n, cache, variation):\n")
         if "S1" in nls[0]:
-            oo.write("    S1 = sx[0][0]\n")
+            oo.write("    S1 = c.get(c.S1, cache, n)\n")
         if "S2" in nls[0]:
-            oo.write("    S2 = sx[1][0]\n")
+            oo.write("    S2 = c.get(c.S2, cache, n)\n")
         if "S3" in nls[0]:
-            oo.write("    S3 = sx[2][0]\n")
+            oo.write("    S3 = c.get(c.S3, cache, n)\n")
         if "S4" in nls[0]:
-            oo.write("    S4 = sx[3][0]\n")
+            oo.write("    S4 = c.get(c.S4, cache, n)\n")
         if "S5" in nls[0]:
-            oo.write("    S5 = sx[4][0]\n")
+            oo.write("    S5 = c.get(c.S5, cache, n)\n")
         if "S3m2" in nls[1]:
             oo.write(
                 "    S3m2 = (-(((-1 + 2 * n) * (1 - n + n**2))/((-1 + n)**3 * n**3)) + S3)/n\n"

extras/ad_n3lo/singlet_nf3/translate_singlet_nf3.py

@@ -36,9 +36,10 @@
         oo.write('"""\n')
         oo.write("import numba as nb\n")
         oo.write("import numpy as np\n")
+        oo.write("from .....harmonics import cache as c\n")
         oo.write("\n\n")
         oo.write("@nb.njit(cache=True)\n")
-        oo.write(f"def gamma_{gamma_name[1:]}_nf3(n, sx):\n")
+        oo.write(f"def gamma_{gamma_name[1:]}_nf3(n, cache):\n")
         oo.write("\treturn ")
         for l in nls:
             oo.write(l)

src/eko/evolution_operator/operator_matrix_element.py

@@ -134,36 +134,27 @@ def quad_ker(
     integrand = ker_base.integrand(areas)
     if integrand == 0.0:
         return 0.0
-
-    max_weight_dict = {1: 2, 2: 3, 3: 5}
-    sx = harmonics.compute_cache(
-        ker_base.n, max_weight_dict[order[0]], ker_base.is_singlet
-    )
     # compute the ome
     if ker_base.is_singlet or ker_base.is_QEDsinglet:
         indices = {21: 0, 100: 1, 90: 2}
         if is_polarized:
             if is_time_like:
                 raise NotImplementedError("Polarized, time-like is not implemented")
-            else:
-                A = ome_ps.A_singlet(order, ker_base.n, sx, nf, L, is_msbar)
+            A = ome_ps.A_singlet(order, ker_base.n, nf, L, is_msbar)
         else:
             if is_time_like:
                 A = ome_ut.A_singlet(order, ker_base.n, L)
-            else:
-                A = ome_us.A_singlet(order, ker_base.n, sx, nf, L, is_msbar)
+            A = ome_us.A_singlet(order, ker_base.n, nf, L, is_msbar)
     else:
         indices = {200: 0, 91: 1}
         if is_polarized:
             if is_time_like:
                 raise NotImplementedError("Polarized, time-like is not implemented")
-            else:
-                A = ome_ps.A_non_singlet(order, ker_base.n, sx, nf, L)
+            A = ome_ps.A_non_singlet(order, ker_base.n, nf, L)
         else:
             if is_time_like:
                 A = ome_ut.A_non_singlet(order, ker_base.n, L)
-            else:
-                A = ome_us.A_non_singlet(order, ker_base.n, sx, nf, L)
+            A = ome_us.A_non_singlet(order, ker_base.n, nf, L)
 
     # correct for scale variations
     if sv_mode == sv.Modes.exponentiated:

src/ekore/anomalous_dimensions/polarized/space_like/__init__.py

@@ -11,32 +11,10 @@
 import numba as nb
 import numpy as np
 
-from .... import harmonics
+from ....harmonics import cache as c
 from . import as1, as2, as3
 
 
-@nb.njit(cache=True)
-def compute_cache(n, pto):
-    """Compute the harmonic cache for polarized anomalous dimension.
-
-    Parameters
-    ----------
-    n : complex
-        Mellin variable
-    pto : int
-        perturbative order
-
-    Returns
-    -------
-    list
-        harmonics cache
-
-    """
-    max_weight = pto if pto != 3 else 4
-    cache = harmonics.sx(n, max_weight)
-    return cache
-
-
 @nb.njit(cache=True)
 def gamma_ns(order, mode, n, nf):
     r"""Compute the tower of the non-singlet anomalous dimensions.
@@ -58,28 +36,26 @@ def gamma_ns(order, mode, n, nf):
         non-singlet anomalous dimensions
 
     """
-    # cache the s-es
-    sx = compute_cache(n, order[0] + 1)
-    # now combine
+    cache = c.reset()
     gamma_ns = np.zeros(order[0], np.complex_)
-    gamma_ns[0] = as1.gamma_ns(n, sx[0])
+    gamma_ns[0] = as1.gamma_ns(n, cache)
     # NLO and beyond
     if order[0] >= 2:
         if mode == 10101:
-            gamma_ns_1 = as2.gamma_nsp(n, nf, sx)
+            gamma_ns_1 = as2.gamma_nsp(n, nf, cache)
         # To fill the full valence vector in NNLO we need to add gamma_ns^1 explicitly here
         elif mode in [10201, 10200]:
-            gamma_ns_1 = as2.gamma_nsm(n, nf, sx)
+            gamma_ns_1 = as2.gamma_nsm(n, nf, cache)
         else:
             raise NotImplementedError("Non-singlet sector is not implemented")
         gamma_ns[1] = gamma_ns_1
     if order[0] >= 3:
         if mode == 10101:
-            gamma_ns_2 = as3.gamma_nsp(n, nf, sx)
+            gamma_ns_2 = as3.gamma_nsp(n, nf, cache)
         elif mode == 10201:
-            gamma_ns_2 = as3.gamma_nsm(n, nf, sx)
+            gamma_ns_2 = as3.gamma_nsm(n, nf, cache)
         elif mode == 10200:
-            gamma_ns_2 = as3.gamma_nsv(n, nf, sx)
+            gamma_ns_2 = as3.gamma_nsv(n, nf, cache)
         gamma_ns[2] = gamma_ns_2
     if order[0] >= 4:
         raise NotImplementedError("Polarized beyond NNLO is not yet implemented")
@@ -105,15 +81,13 @@ def gamma_singlet(order, n, nf):
         singlet anomalous dimensions matrices
 
     """
-    # cache the s-es
-    sx = compute_cache(n, order[0] + 1)
-
+    cache = c.reset()
     gamma_s = np.zeros((order[0], 2, 2), np.complex_)
-    gamma_s[0] = as1.gamma_singlet(n, sx[0], nf)
+    gamma_s[0] = as1.gamma_singlet(n, cache, nf)
     if order[0] >= 2:
-        gamma_s[1] = as2.gamma_singlet(n, nf, sx)
+        gamma_s[1] = as2.gamma_singlet(n, nf, cache)
     if order[0] >= 3:
-        gamma_s[2] = as3.gamma_singlet(n, nf, sx)
+        gamma_s[2] = as3.gamma_singlet(n, nf, cache)
     if order[0] >= 4:
         raise NotImplementedError("Polarized beyond NNLO is not yet implemented")
     return gamma_s

src/ekore/anomalous_dimensions/polarized/space_like/as1.py

@@ -5,6 +5,7 @@
 
 from eko import constants
 
+from ....harmonics import cache as c
 from ...unpolarized.space_like.as1 import gamma_ns
 
 
@@ -51,15 +52,15 @@ def gamma_gq(N):
 
 
 @nb.njit(cache=True)
-def gamma_gg(N, s1, nf):
+def gamma_gg(N, cache, nf):
     r"""Compute the |LO| polarized gluon-gluon anomalous dimension :cite:`Gluck:1995yr` (eq A.1).
 
     Parameters
     ----------
     N : complex
       Mellin moment
-    s1 : complex
-      harmonic sum :math:`S_{1}`
+    cache: numpy.ndarray
+        Harmonic sum cache
     nf : int
       Number of active flavors
 
@@ -69,13 +70,14 @@ def gamma_gg(N, s1, nf):
       |LO| gluon-gluon anomalous dimension :math:`\\gamma_{gg}^{(0)}(N)`
 
     """
-    gamma = -s1 + 2 / N / (N + 1)
+    S1 = c.get(c.S1, cache, N)
+    gamma = -S1 + 2 / N / (N + 1)
     result = constants.CA * (-4.0 * gamma - 11.0 / 3.0) + 4.0 / 3.0 * constants.TR * nf
     return result
 
 
 @nb.njit(cache=True)
-def gamma_singlet(N, s1, nf):
+def gamma_singlet(N, cache, nf):
     r"""Compute the |LO| polarized singlet anomalous dimension matrix.
 
       .. math::
@@ -88,8 +90,8 @@ def gamma_singlet(N, s1, nf):
     ----------
     N : complex
       Mellin moment
-    s1 : complex
-      harmonic sum :math:`S_{1}`
+    cache: numpy.ndarray
+        Harmonic sum cache
     nf : int
       Number of active flavors
 
@@ -99,9 +101,9 @@ def gamma_singlet(N, s1, nf):
       |LO| singlet anomalous dimension matrix :math:`\gamma_{S}^{(0)}(N)`
 
     """
-    gamma_qq = gamma_ns(N, s1)
+    gamma_qq = gamma_ns(N, cache)
     gamma_S_0 = np.array(
-        [[gamma_qq, gamma_qg(N, nf)], [gamma_gq(N), gamma_gg(N, s1, nf)]],
+        [[gamma_qq, gamma_qg(N, nf)], [gamma_gq(N), gamma_gg(N, cache, nf)]],
         np.complex_,
     )
     return gamma_S_0