Cache harmonics sums for the as2 anomalous dimensions

benchmarks/eko/benchmark_ad.py

@@ -138,8 +138,9 @@ def check_gamma_1_pegasus(N, NF):
     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)
+    sx_new = h.compute_additional_sx_cache(N, False, False, sx)
+    np.testing.assert_allclose(ad_as2.gamma_nsp(N, NF, sx, sx_new), -P1NSP)
+    np.testing.assert_allclose(ad_as2.gamma_nsm(N, NF, sx, sx_new), -P1NSM)
 
     NS = N * N
     NT = NS * N
@@ -256,7 +257,8 @@ 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)
+    sx_new = h.compute_additional_sx_cache(N, True, False, sx)
+    gS1 = ad_as2.gamma_singlet(N, NF, sx, sx_new)
     np.testing.assert_allclose(gS1[0, 0], -P1Sqq)
     np.testing.assert_allclose(gS1[0, 1], -P1Sqg)
     np.testing.assert_allclose(gS1[1, 0], -P1Sgq)

src/eko/anomalous_dimensions/__init__.py

@@ -135,7 +135,9 @@ def gamma_ns(order, mode, n, nf):
     """
     # cache the s-es
     if order[0] >= 4:
-        full_sx_cache = harmonics.compute_cache(n, 5, is_singlet=False)
+        full_sx_cache = harmonics.compute_cache(
+            n, 5, is_singlet=False, additional=True, qed=False
+        )
         sx = np.array(
             [
                 full_sx_cache[0][0],
@@ -144,18 +146,20 @@ def gamma_ns(order, mode, n, nf):
                 full_sx_cache[3][0],
             ]
         )
+        sx_new = full_sx_cache[-1]
     else:
         sx = harmonics.sx(n, max_weight=order[0] + 1)
+        sx_new = harmonics.compute_additional_sx_cache(n, False, False, sx)
     # now combine
     gamma_ns = np.zeros(order[0], np.complex_)
     gamma_ns[0] = as1.gamma_ns(n, sx[0])
     # 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, sx, sx_new)
         # 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, sx, sx_new)
         else:
             raise NotImplementedError("Non-singlet sector is not implemented")
         gamma_ns[1] = gamma_ns_1
@@ -208,7 +212,9 @@ def gamma_singlet(order, n, nf):
     """
     # cache the s-es
     if order[0] >= 4:
-        full_sx_cache = harmonics.compute_cache(n, 5, is_singlet=False)
+        full_sx_cache = harmonics.compute_cache(
+            n, 5, is_singlet=False, additional=True, qed=False
+        )
         sx = np.array(
             [
                 full_sx_cache[0][0],
@@ -217,16 +223,18 @@ def gamma_singlet(order, n, nf):
                 full_sx_cache[3][0],
             ]
         )
+        sx_new = full_sx_cache[-1]
     elif order[0] >= 3:
         # here we need only S1,S2,S3,S4
         sx = harmonics.sx(n, max_weight=order[0] + 1)
+        sx_new = harmonics.compute_additional_sx_cache(n, True, False, sx)
     else:
         sx = harmonics.sx(n, max_weight=order[0])
-
+        sx_new = harmonics.compute_additional_sx_cache(n, True, False, sx)
     gamma_s = np.zeros((order[0], 2, 2), np.complex_)
     gamma_s[0] = as1.gamma_singlet(n, sx[0], nf)
     if order[0] >= 2:
-        gamma_s[1] = as2.gamma_singlet(n, nf, sx)
+        gamma_s[1] = as2.gamma_singlet(n, nf, sx, sx_new)
     if order[0] >= 3:
         gamma_s[2] = as3.gamma_singlet(n, nf, sx)
     if order[0] >= 4:
@@ -279,21 +287,21 @@ def gamma_ns_qed(order, mode, n, nf):
         sx = harmonics.sx(n, max_weight=max_weight + 1)
     else:
         sx = harmonics.sx(n, max_weight=3)
-    sx_ns_qed = harmonics.compute_qed_ns_cache(n, sx[0])
+    sx_new = harmonics.compute_additional_sx_cache(n, False, True, sx)
     # now combine
     gamma_ns = np.zeros((order[0] + 1, order[1] + 1), np.complex_)
     gamma_ns[1, 0] = as1.gamma_ns(n, sx[0])
     gamma_ns[0, 1] = choose_ns_ad_aem1(mode, n, sx)
-    gamma_ns[1, 1] = choose_ns_ad_as1aem1(mode, n, sx, sx_ns_qed)
+    gamma_ns[1, 1] = choose_ns_ad_as1aem1(mode, n, sx, sx_new)
     # NLO and beyond
     if order[0] >= 2:
         if mode in [10102, 10103]:
-            gamma_ns[2, 0] = as2.gamma_nsp(n, nf, sx)
+            gamma_ns[2, 0] = as2.gamma_nsp(n, nf, sx, sx_new)
         # To fill the full valence vector in NNLO we need to add gamma_ns^1 explicitly here
         else:
-            gamma_ns[2, 0] = as2.gamma_nsm(n, nf, sx)
+            gamma_ns[2, 0] = as2.gamma_nsm(n, nf, sx, sx_new)
     if order[1] >= 2:
-        gamma_ns[0, 2] = choose_ns_ad_aem2(mode, n, nf, sx, sx_ns_qed)
+        gamma_ns[0, 2] = choose_ns_ad_aem2(mode, n, nf, sx, sx_new)
     # NNLO and beyond
     if order[0] >= 3:
         if mode in [10102, 10103]:
@@ -331,7 +339,7 @@ def choose_ns_ad_aem1(mode, n, sx):
 
 
 @nb.njit(cache=True)
-def choose_ns_ad_as1aem1(mode, n, sx, sx_ns_qed):
+def choose_ns_ad_as1aem1(mode, n, sx, sx_new):
     r"""
     Select the non-singlet anomalous dimension at O(as1aem1) with the correct charge factor.
 
@@ -350,17 +358,17 @@ def choose_ns_ad_as1aem1(mode, n, sx, sx_ns_qed):
             non-singlet anomalous dimensions
     """
     if mode == 10102:
-        return constants.eu2 * as1aem1.gamma_nsp(n, sx, sx_ns_qed)
+        return constants.eu2 * as1aem1.gamma_nsp(n, sx, sx_new)
     elif mode == 10103:
-        return constants.ed2 * as1aem1.gamma_nsp(n, sx, sx_ns_qed)
+        return constants.ed2 * as1aem1.gamma_nsp(n, sx, sx_new)
     elif mode == 10202:
-        return constants.eu2 * as1aem1.gamma_nsm(n, sx, sx_ns_qed)
+        return constants.eu2 * as1aem1.gamma_nsm(n, sx, sx_new)
     elif mode == 10203:
-        return constants.ed2 * as1aem1.gamma_nsm(n, sx, sx_ns_qed)
+        return constants.ed2 * as1aem1.gamma_nsm(n, sx, sx_new)
 
 
 @nb.njit(cache=True)
-def choose_ns_ad_aem2(mode, n, nf, sx, sx_ns_qed):
+def choose_ns_ad_aem2(mode, n, nf, sx, sx_new):
     r"""
     Select the non-singlet anomalous dimension at O(aem2) with the correct charge factor.
 
@@ -379,13 +387,13 @@ def choose_ns_ad_aem2(mode, n, nf, sx, sx_ns_qed):
             non-singlet anomalous dimensions
     """
     if mode == 10102:
-        return constants.eu2 * aem2.gamma_nspu(n, nf, sx, sx_ns_qed)
+        return constants.eu2 * aem2.gamma_nspu(n, nf, sx, sx_new)
     elif mode == 10103:
-        return constants.ed2 * aem2.gamma_nspd(n, nf, sx, sx_ns_qed)
+        return constants.ed2 * aem2.gamma_nspd(n, nf, sx, sx_new)
     elif mode == 10202:
-        return constants.eu2 * aem2.gamma_nsmu(n, nf, sx, sx_ns_qed)
+        return constants.eu2 * aem2.gamma_nsmu(n, nf, sx, sx_new)
     elif mode == 10203:
-        return constants.ed2 * aem2.gamma_nsmd(n, nf, sx, sx_ns_qed)
+        return constants.ed2 * aem2.gamma_nsmd(n, nf, sx, sx_new)
 
 
 @nb.njit(cache=True)
@@ -423,15 +431,15 @@ def gamma_singlet_qed(order, n, nf):
         sx = harmonics.sx(n, max_weight=max_weight + 1)
     else:
         sx = harmonics.sx(n, max_weight=3)
-    sx_ns_qed = harmonics.compute_qed_ns_cache(n, sx[0])
+    sx_new = harmonics.compute_additional_sx_cache(n, False, True, sx)
     gamma_s = np.zeros((order[0] + 1, order[1] + 1, 4, 4), np.complex_)
     gamma_s[1, 0] = as1.gamma_singlet_qed(n, sx[0], nf)
     gamma_s[0, 1] = aem1.gamma_singlet(n, nf, sx)
-    gamma_s[1, 1] = as1aem1.gamma_singlet(n, nf, sx, sx_ns_qed)
+    gamma_s[1, 1] = as1aem1.gamma_singlet(n, nf, sx, sx_new)
     if order[0] >= 2:
-        gamma_s[2, 0] = as2.gamma_singlet_qed(n, nf, sx)
+        gamma_s[2, 0] = as2.gamma_singlet_qed(n, nf, sx, sx_new)
     if order[1] >= 2:
-        gamma_s[0, 2] = aem2.gamma_singlet(n, nf, sx, sx_ns_qed)
+        gamma_s[0, 2] = aem2.gamma_singlet(n, nf, sx, sx_new)
     if order[0] >= 3:
         gamma_s[3, 0] = as3.gamma_singlet_qed(n, nf, sx)
     return gamma_s
@@ -472,15 +480,15 @@ def gamma_valence_qed(order, n, nf):
         sx = harmonics.sx(n, max_weight=max_weight + 1)
     else:
         sx = harmonics.sx(n, max_weight=3)
-    sx_ns_qed = harmonics.compute_qed_ns_cache(n, sx[0])
+    sx_new = harmonics.compute_additional_sx_cache(n, False, True, sx)
     gamma_v = np.zeros((order[0] + 1, order[1] + 1, 2, 2), np.complex_)
     gamma_v[1, 0] = as1.gamma_valence_qed(n, sx[0])
     gamma_v[0, 1] = aem1.gamma_valence(n, nf, sx)
-    gamma_v[1, 1] = as1aem1.gamma_valence(n, nf, sx, sx_ns_qed)
+    gamma_v[1, 1] = as1aem1.gamma_valence(n, nf, sx, sx_new)
     if order[0] >= 2:
-        gamma_v[2, 0] = as2.gamma_valence_qed(n, nf, sx)
+        gamma_v[2, 0] = as2.gamma_valence_qed(n, nf, sx, sx_new)
     if order[1] >= 2:
-        gamma_v[0, 2] = aem2.gamma_valence(n, nf, sx, sx_ns_qed)
+        gamma_v[0, 2] = aem2.gamma_valence(n, nf, sx, sx_new)
     if order[0] >= 3:
         gamma_v[3, 0] = as3.gamma_valence_qed(n, nf, sx)
     return gamma_v

src/eko/anomalous_dimensions/aem2.py

@@ -151,7 +151,7 @@ def gamma_phd(N, nf, sx):
 
 
 @nb.njit(cache=True)
-def gamma_nspu(N, nf, sx, sx_ns_qed):
+def gamma_nspu(N, nf, sx, sx_new):
     r"""Compute the :math:`O(a_{em}^2)` singlet-like non-singlet anomalous dimension for up quarks.
 
     Implements sum of Eqs. (57-58) of :cite:`deFlorian:2016gvk` for q=u.
@@ -184,13 +184,11 @@ def gamma_nspu(N, nf, sx, sx_ns_qed):
         - 80.0 / 9.0 * S1
         + 16.0 / 3.0 * S2
     ) * eSigma2
-    return (
-        constants.eu2 * as1aem1.gamma_nsp(N, sx, sx_ns_qed) / constants.CF / 2.0 + tmp
-    )
+    return constants.eu2 * as1aem1.gamma_nsp(N, sx, sx_new) / constants.CF / 2.0 + tmp
 
 
 @nb.njit(cache=True)
-def gamma_nspd(N, nf, sx, sx_ns_qed):
+def gamma_nspd(N, nf, sx, sx_new):
     r"""Compute the :math:`O(a_{em}^2)` singlet-like non-singlet anomalous dimension for down quarks.
 
     Implements sum of Eqs. (57-58) of :cite:`deFlorian:2016gvk` for q=d.
@@ -223,13 +221,11 @@ def gamma_nspd(N, nf, sx, sx_ns_qed):
         - 80.0 / 9.0 * S1
         + 16.0 / 3.0 * S2
     ) * eSigma2
-    return (
-        constants.ed2 * as1aem1.gamma_nsp(N, sx, sx_ns_qed) / constants.CF / 2.0 + tmp
-    )
+    return constants.ed2 * as1aem1.gamma_nsp(N, sx, sx_new) / constants.CF / 2.0 + tmp
 
 
 @nb.njit(cache=True)
-def gamma_nsmu(N, nf, sx, sx_ns_qed):
+def gamma_nsmu(N, nf, sx, sx_new):
     r"""Compute the :math:`O(a_{em}^2)` valence-like non-singlet anomalous dimension for up quarks.
 
     Implements difference between Eqs. (57-58) of :cite:`deFlorian:2016gvk` for q=u.
@@ -262,13 +258,11 @@ def gamma_nsmu(N, nf, sx, sx_ns_qed):
         - 80.0 / 9.0 * S1
         + 16.0 / 3.0 * S2
     ) * eSigma2
-    return (
-        constants.eu2 * as1aem1.gamma_nsm(N, sx, sx_ns_qed) / constants.CF / 2.0 + tmp
-    )
+    return constants.eu2 * as1aem1.gamma_nsm(N, sx, sx_new) / constants.CF / 2.0 + tmp
 
 
 @nb.njit(cache=True)
-def gamma_nsmd(N, nf, sx, sx_ns_qed):
+def gamma_nsmd(N, nf, sx, sx_new):
     r"""Compute the :math:`O(a_{em}^2)` valence-like non-singlet anomalous dimension for down quarks.
 
     Implements difference between Eqs. (57-58) of :cite:`deFlorian:2016gvk` for q=d.
@@ -301,9 +295,7 @@ def gamma_nsmd(N, nf, sx, sx_ns_qed):
         - 80.0 / 9.0 * S1
         + 16.0 / 3.0 * S2
     ) * eSigma2
-    return (
-        constants.ed2 * as1aem1.gamma_nsm(N, sx, sx_ns_qed) / constants.CF / 2.0 + tmp
-    )
+    return constants.ed2 * as1aem1.gamma_nsm(N, sx, sx_new) / constants.CF / 2.0 + tmp
 
 
 @nb.njit(cache=True)
@@ -336,7 +328,7 @@ def gamma_ps(N, nf):
 
 
 @nb.njit(cache=True)
-def gamma_singlet(N, nf, sx, sx_ns_qed):
+def gamma_singlet(N, nf, sx, sx_new):
     r"""Compute the :math:`O(a_{em}^2)` singlet sector.
 
     Parameters
@@ -363,8 +355,8 @@ def gamma_singlet(N, nf, sx, sx_ns_qed):
     gamma_ph_d = gamma_phd(N, nf, sx)
     gamma_u_ph = gamma_uph(N, nf, sx)
     gamma_d_ph = gamma_dph(N, nf, sx)
-    gamma_ns_p_u = gamma_nspu(N, nf, sx, sx_ns_qed)
-    gamma_ns_p_d = gamma_nspd(N, nf, sx, sx_ns_qed)
+    gamma_ns_p_u = gamma_nspu(N, nf, sx, sx_new)
+    gamma_ns_p_d = gamma_nspd(N, nf, sx, sx_new)
     gamma_pure_singlet = gamma_ps(N, nf)
     gamma_S_02 = np.array(
         [
@@ -408,7 +400,7 @@ def gamma_singlet(N, nf, sx, sx_ns_qed):
 
 
 @nb.njit(cache=True)
-def gamma_valence(N, nf, sx, sx_ns_qed):
+def gamma_valence(N, nf, sx, sx_new):
     r"""Compute the :math:`O(a_{em}^2)` valence sector.
 
     Parameters
@@ -429,8 +421,8 @@ def gamma_valence(N, nf, sx, sx_ns_qed):
     nd = nf - nu
     vu = nu / nf
     vd = nd / nf
-    gamma_ns_m_u = gamma_nsmu(N, nf, sx, sx_ns_qed)
-    gamma_ns_m_d = gamma_nsmd(N, nf, sx, sx_ns_qed)
+    gamma_ns_m_u = gamma_nsmu(N, nf, sx, sx_new)
+    gamma_ns_m_d = gamma_nsmd(N, nf, sx, sx_new)
     gamma_V_02 = np.array(
         [
             [