Use EvoMethods in kernels

NNPDF · Jul 19, 2024 · 5636b4a · 5636b4a
1 parent 7a6c1cc
commit 5636b4a
Show file tree

Hide file tree

Showing 10 changed files with 87 additions and 82 deletions.
diff --git a/src/eko/kernels/non_singlet.py b/src/eko/kernels/non_singlet.py
@@ -4,6 +4,7 @@
 import numpy as np
 
 from .. import beta
+from . import EvoMethods
 from . import as4_evolution_integrals as as4_ei
 from . import evolution_integrals as ei
 
@@ -355,7 +356,7 @@ def dispatcher(
     ----------
     order : tuple(int,int)
         perturbation order
-    method : str
+    method : int
         method
     gamma_ns : numpy.ndarray
         non-singlet anomalous dimensions
@@ -378,38 +379,38 @@ def dispatcher(
     # use always exact in LO
     if order[0] == 1:
         return lo_exact(gamma_ns, a1, a0, betalist)
-    if method == "ordered-truncated":
+    if method is EvoMethods.ORDERED_TRUNCATED:
         return eko_ordered_truncated(
             gamma_ns, a1, a0, betalist, order, ev_op_iterations
         )
-    if method == "truncated":
+    if method is EvoMethods.TRUNCATED:
         return eko_truncated(gamma_ns, a1, a0, betalist, order, ev_op_iterations)
 
     # NLO
     if order[0] == 2:
         if method in [
-            "iterate-expanded",
-            "decompose-expanded",
-            "perturbative-expanded",
+            EvoMethods.ITERATE_EXPANDED,
+            EvoMethods.DECOMPOSE_EXPANDED,
+            EvoMethods.PERTURBATIVE_EXPANDED,
         ]:
             return nlo_expanded(gamma_ns, a1, a0, betalist)
-        # if method in ["iterate-exact", "decompose-exact", "perturbative-exact"]:
+        # exact is left
         return nlo_exact(gamma_ns, a1, a0, betalist)
     # NNLO
     if order[0] == 3:
         if method in [
-            "iterate-expanded",
-            "decompose-expanded",
-            "perturbative-expanded",
+            EvoMethods.ITERATE_EXPANDED,
+            EvoMethods.DECOMPOSE_EXPANDED,
+            EvoMethods.PERTURBATIVE_EXPANDED,
         ]:
             return nnlo_expanded(gamma_ns, a1, a0, betalist)
         return nnlo_exact(gamma_ns, a1, a0, betalist)
     # N3LO
     if order[0] == 4:
         if method in [
-            "iterate-expanded",
-            "decompose-expanded",
-            "perturbative-expanded",
+            EvoMethods.ITERATE_EXPANDED,
+            EvoMethods.DECOMPOSE_EXPANDED,
+            EvoMethods.PERTURBATIVE_EXPANDED,
         ]:
             return n3lo_expanded(gamma_ns, a1, a0, betalist)
         return n3lo_exact(gamma_ns, a1, a0, betalist)

diff --git a/src/eko/kernels/non_singlet_qed.py b/src/eko/kernels/non_singlet_qed.py
@@ -146,7 +146,7 @@ def as4_exact(gamma_ns, a1, a0, beta):
 @nb.njit(cache=True)
 def dispatcher(
     order,
-    method,
+    _method,
     gamma_ns,
     as_list,
     aem_half,
@@ -164,7 +164,7 @@ def dispatcher(
     ----------
     order : tuple(int,int)
         perturbation order
-    method : str
+    method : int
         method
     gamma_ns : numpy.ndarray
         non-singlet anomalous dimensions

diff --git a/src/eko/kernels/singlet.py b/src/eko/kernels/singlet.py
@@ -6,6 +6,7 @@
 from ekore import anomalous_dimensions as ad
 
 from .. import beta
+from . import EvoMethods
 from . import as4_evolution_integrals as as4_ei
 from . import evolution_integrals as ei
 
@@ -579,7 +580,7 @@ def dispatcher(  # pylint: disable=too-many-return-statements
     ----------
     order :  tuple(int,int)
         perturbative order
-    method : str
+    method : int
         method
     gamma_singlet : numpy.ndarray
         singlet anomalous dimensions matrices
@@ -609,9 +610,9 @@ def dispatcher(  # pylint: disable=too-many-return-statements
         return lo_exact(gamma_singlet, a1, a0, betalist)
 
     # Common method for NLO and NNLO
-    if method in ["iterate-exact", "iterate-expanded"]:
+    if method in [EvoMethods.ITERATE_EXACT, EvoMethods.ITERATE_EXPANDED]:
         return eko_iterate(gamma_singlet, a1, a0, betalist, order, ev_op_iterations)
-    if method == "perturbative-exact":
+    if method is EvoMethods.PERTURBATIVE_EXACT:
         return eko_perturbative(
             gamma_singlet,
             a1,
@@ -622,7 +623,7 @@ def dispatcher(  # pylint: disable=too-many-return-statements
             ev_op_max_order,
             True,
         )
-    if method == "perturbative-expanded":
+    if method is EvoMethods.PERTURBATIVE_EXPANDED:
         return eko_perturbative(
             gamma_singlet,
             a1,
@@ -633,19 +634,19 @@ def dispatcher(  # pylint: disable=too-many-return-statements
             ev_op_max_order,
             False,
         )
-    if method in ["truncated", "ordered-truncated"]:
+    if method in [EvoMethods.TRUNCATED, EvoMethods.ORDERED_TRUNCATED]:
         return eko_truncated(gamma_singlet, a1, a0, betalist, order, ev_op_iterations)
     # These methods are scattered for nlo and nnlo
-    if method == "decompose-exact":
+    if method is EvoMethods.DECOMPOSE_EXACT:
         if order[0] == 2:
             return nlo_decompose_exact(gamma_singlet, a1, a0, betalist)
-        elif order[0] == 3:
+        if order[0] == 3:
             return nnlo_decompose_exact(gamma_singlet, a1, a0, betalist)
         return n3lo_decompose_exact(gamma_singlet, a1, a0, nf)
-    if method == "decompose-expanded":
+    if method is EvoMethods.DECOMPOSE_EXPANDED:
         if order[0] == 2:
             return nlo_decompose_expanded(gamma_singlet, a1, a0, betalist)
-        elif order[0] == 3:
+        if order[0] == 3:
             return nnlo_decompose_expanded(gamma_singlet, a1, a0, betalist)
         return n3lo_decompose_expanded(gamma_singlet, a1, a0, nf)
     raise NotImplementedError("Selected method is not implemented")
diff --git a/src/eko/kernels/singlet_qed.py b/src/eko/kernels/singlet_qed.py
@@ -6,6 +6,7 @@
 from ekore import anomalous_dimensions as ad
 
 from .. import beta
+from . import EvoMethods
 
 
 @nb.njit(cache=True)
@@ -66,15 +67,15 @@ def dispatcher(
     a_half,
     nf,
     ev_op_iterations,
-    ev_op_max_order,
+    _ev_op_max_order,
 ):
     """Determine used kernel and call it.
 
     Parameters
     ----------
     order : tuple(int,int)
         perturbative order
-    method : str
+    method : int
         method
     gamma_singlet : numpy.ndarray
         singlet anomalous dimensions matrices
@@ -96,7 +97,7 @@ def dispatcher(
     e_s : numpy.ndarray
         singlet EKO
     """
-    if method in ["iterate-exact", "iterate-expanded"]:
+    if method in [EvoMethods.ITERATE_EXACT, EvoMethods.ITERATE_EXPANDED]:
         return eko_iterate(
             gamma_singlet, as_list, a_half, nf, order, ev_op_iterations, 4
         )

diff --git a/src/eko/kernels/valence_qed.py b/src/eko/kernels/valence_qed.py
@@ -2,6 +2,7 @@
 
 import numba as nb
 
+from . import EvoMethods
 from .singlet_qed import eko_iterate
 
 
@@ -14,7 +15,7 @@ def dispatcher(
     a_half,
     nf,
     ev_op_iterations,
-    ev_op_max_order,
+    _ev_op_max_order,
 ):
     """
     Determine used kernel and call it.
@@ -45,7 +46,7 @@ def dispatcher(
         e_v : numpy.ndarray
             singlet EKO
     """
-    if method in ["iterate-exact", "iterate-expanded"]:
+    if method in [EvoMethods.ITERATE_EXACT, EvoMethods.ITERATE_EXPANDED]:
         return eko_iterate(
             gamma_valence, as_list, a_half, nf, order, ev_op_iterations, 2
         )

diff --git a/tests/eko/kernels/test_kernels_QEDns.py b/tests/eko/kernels/test_kernels_QEDns.py
@@ -3,6 +3,7 @@
 
 from eko import basis_rotation as br
 from eko.couplings import Couplings
+from eko.kernels import EvoMethods
 from eko.kernels import non_singlet_qed as ns
 from eko.quantities.couplings import CouplingEvolutionMethod, CouplingsInfo
 from eko.quantities.heavy_quarks import QuarkMassScheme
@@ -14,7 +15,7 @@
     # "perturbative-expanded",
     # "truncated",
     # "ordered-truncated",
-    "iterate-exact",
+    EvoMethods.ITERATE_EXACT,
     # "decompose-exact",
     # "perturbative-exact",
 ]

diff --git a/tests/eko/kernels/test_kernels_QEDsinglet.py b/tests/eko/kernels/test_kernels_QEDsinglet.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from eko.kernels import EvoMethods
 from eko.kernels import singlet_qed as s
 from ekore.anomalous_dimensions.unpolarized import space_like as ad
 
@@ -10,7 +11,7 @@
     # "perturbative-expanded",
     # "truncated",
     # "ordered-truncated",
-    "iterate-exact",
+    EvoMethods.ITERATE_EXACT,
     # "decompose-exact",
     # "perturbative-exact",
 ]
@@ -125,5 +126,12 @@ def test_zero_true_gamma(monkeypatch):
 def test_error():
     with pytest.raises(NotImplementedError):
         s.dispatcher(
-            (3, 2), "AAA", np.random.rand(4, 3, 2, 2), [0.2, 0.1], [0.01], 3, 10, 10
+            (3, 2),
+            "iterate-exact",
+            np.random.rand(4, 3, 2, 2),
+            [0.2, 0.1],
+            [0.01],
+            3,
+            10,
+            10,
         )
diff --git a/tests/eko/kernels/test_kernels_QEDvalence.py b/tests/eko/kernels/test_kernels_QEDvalence.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from eko.kernels import EvoMethods
 from eko.kernels import valence_qed as val
 from ekore import anomalous_dimensions
 
@@ -10,13 +11,13 @@
     # "perturbative-expanded",
     # "truncated",
     # "ordered-truncated",
-    "iterate-exact",
+    EvoMethods.ITERATE_EXACT,
     # "decompose-exact",
     # "perturbative-exact",
 ]
 
 
-def test_zero(monkeypatch):
+def test_zero():
     """No evolution results in exp(0)"""
     nf = 3
     ev_op_iterations = 2
@@ -57,7 +58,7 @@ def test_zero(monkeypatch):
                 )
 
 
-def test_zero_true_gamma(monkeypatch):
+def test_zero_true_gamma():
     """No evolution results in exp(0)"""
     nf = 3
     ev_op_iterations = 2
@@ -101,5 +102,12 @@ def test_zero_true_gamma(monkeypatch):
 def test_error():
     with pytest.raises(NotImplementedError):
         val.dispatcher(
-            (3, 2), "AAA", np.random.rand(4, 3, 2, 2), [0.2, 0.1], [0.01], 3, 10, 10
+            (3, 2),
+            "iterate-exact",
+            np.random.rand(4, 3, 2, 2),
+            [0.2, 0.1],
+            [0.01],
+            3,
+            10,
+            10,
         )