Refactor harmonic cache

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

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

@@ -5,7 +5,7 @@
 
 from eko.constants import CA, CF, TR, zeta2, zeta3
 
-from .... import harmonics
+from ....harmonics import cache as c
 
 # Non Singlet sector is swapped
 from ...unpolarized.space_like.as2 import gamma_nsm as gamma_nsp
@@ -35,7 +35,7 @@ def gamma_ps(n, nf):
 
 
 @nb.njit(cache=True)
-def gamma_qg(n, nf, sx):
+def gamma_qg(n, nf, cache):
     r"""Compute the |NLO| polarized quark-gluon singlet anomalous dimension :cite:`Gluck:1995yr` (eq A.4).
 
     Parameters
@@ -44,18 +44,18 @@ def gamma_qg(n, nf, sx):
         Mellin moment
     nf : int
         number of active flavors
-    sx : numpy.ndarray
-        List of harmonic sums: :math:`S_{1},S_{2}`
+    cache: numpy.ndarray
+        Harmonic sum cache
 
     Returns
     -------
     complex
         |NLO| quark-gluon singlet anomalous dimension :math:`\\gamma_{qg}^{(1)}(n)`
 
     """
-    S1 = sx[0]
-    S2 = sx[1]
-    Sp2m = harmonics.S2((n - 1) / 2)
+    S1 = c.get(c.S1, cache, n)
+    S2 = c.get(c.S2, cache, n)
+    Sp2m = c.get(c.S2mh, cache, n)
     gqg1_nfca = (
         (S1**2 - S2 + Sp2m) * (n - 1) / (n * (n + 1))
         - 4 * S1 / (n * (1 + n) ** 2)
@@ -74,7 +74,7 @@ def gamma_qg(n, nf, sx):
 
 
 @nb.njit(cache=True)
-def gamma_gq(n, nf, sx):
+def gamma_gq(n, nf, cache):
     r"""Compute the |NLO| polarized gluon-quark singlet anomalous dimension :cite:`Gluck:1995yr` (eq A.5).
 
     Parameters
@@ -83,18 +83,18 @@ def gamma_gq(n, nf, sx):
         Mellin moment
     nf : int
         number of active flavors
-    sx : numpy.ndarray
-        List of harmonic sums: :math:`S_{1},S_{2}`
+    cache: numpy.ndarray
+        Harmonic sum cache
 
     Returns
     -------
     complex
         |NLO| gluon-quark singlet anomalous dimension :math:`\\gamma_{gq}^{(1)}(n)`
 
     """
-    S1 = sx[0]
-    S2 = sx[1]
-    Sp2m = harmonics.S2((n - 1) / 2)
+    S1 = c.get(c.S1, cache, n)
+    S2 = c.get(c.S2, cache, n)
+    Sp2m = c.get(c.S2mh, cache, n)
     ggq1_cfcf = (
         (2 * (S1**2 + S2) * (n + 2)) / (n * (n + 1))
         - (2 * S1 * (n + 2) * (1 + 3 * n)) / (n * (1 + n) ** 2)
@@ -116,7 +116,7 @@ def gamma_gq(n, nf, sx):
 
 
 @nb.njit(cache=True)
-def gamma_gg(n, nf, sx):
+def gamma_gg(n, nf, cache):
     r"""Compute the |NLO| polarized gluon-gluon singlet anomalous dimension :cite:`Gluck:1995yr` (eq A.6).
 
     Parameters
@@ -125,26 +125,22 @@ def gamma_gg(n, nf, sx):
         Mellin moment
     nf : int
         number of active flavors
-    sx : numpy.ndarray
-        List of harmonic sums: :math:`S_{1},S_{2}`
+    cache: numpy.ndarray
+        Harmonic sum cache
 
     Returns
     -------
     complex
         |NLO| gluon-quark singlet anomalous dimension :math:`\\gamma_{gq}^{(1)}(n)`
 
     """
-    S1 = sx[0]
-    Sp1m = harmonics.S1((n - 1) / 2)
-    Sp2m = harmonics.S2((n - 1) / 2)
-    Sp3m = harmonics.S3((n - 1) / 2)
-    S1h = harmonics.S1(n / 2)
-    SSCHLM = (
-        zeta2 / 2 * (+Sp1m - S1h + 2 / n)
-        - S1 / n**2
-        - harmonics.g_functions.mellin_g3(n, S1)
-        - 5 * zeta3 / 8
-    )
+    S1 = c.get(c.S1, cache, n)
+    Sp1m = c.get(c.S1mh, cache, n)
+    Sp2m = c.get(c.S2mh, cache, n)
+    Sp3m = c.get(c.S3mh, cache, n)
+    S1h = c.get(c.S1h, cache, n)
+    g3 = c.get(c.g3, cache, n)
+    SSCHLM = zeta2 / 2 * (+Sp1m - S1h + 2 / n) - S1 / n**2 - g3 - 5 * zeta3 / 8
     ggg1_caca = (
         -4 * S1 * Sp2m
         - Sp3m
@@ -178,7 +174,7 @@ def gamma_gg(n, nf, sx):
 
 
 @nb.njit(cache=True)
-def gamma_singlet(n, nf, sx):
+def gamma_singlet(n, nf, cache):
     r"""Compute the |NLO| polarized singlet anomalous dimension matrix.
 
         .. math::
@@ -193,18 +189,21 @@ def gamma_singlet(n, nf, sx):
         Mellin moment
     nf : int
         Number of active flavors
-    sx: list
-        harmonics cache
+    cache: numpy.ndarray
+        Harmonic sum cache
 
     Returns
     -------
     numpy.ndarray
         |NLO| singlet anomalous dimension matrix :math:`\gamma_{S}^{(1)}(N)`
 
     """
-    gamma_qq = gamma_nsp(n, nf, sx) + gamma_ps(n, nf)
+    gamma_qq = gamma_nsp(n, nf, cache) + gamma_ps(n, nf)
     gamma_S_0 = np.array(
-        [[gamma_qq, gamma_qg(n, nf, sx)], [gamma_gq(n, nf, sx), gamma_gg(n, nf, sx)]],
+        [
+            [gamma_qq, gamma_qg(n, nf, cache)],
+            [gamma_gq(n, nf, cache), gamma_gg(n, nf, cache)],
+        ],
         np.complex_,
     )
     return gamma_S_0