feat: add numba.jit support

ComPWA · Feb 18, 2021 · 6da0e0b · 6da0e0b
1 parent 81bf402
commit 6da0e0b
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 0 deletions.
diff --git a/cspell.json b/cspell.json
@@ -155,6 +155,7 @@
         "ncalls",
         "ndarray",
         "noqa",
+        "numba",
         "pandoc",
         "phasespace",
         "phsp",

diff --git a/reqs/3.6/requirements-dev.txt b/reqs/3.6/requirements-dev.txt
@@ -117,6 +117,7 @@ nbstripout==0.3.9
 nest-asyncio==1.5.1
 nodeenv==1.5.0
 notebook==6.2.0
+numba==0.52.0
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0

diff --git a/reqs/3.7/requirements-dev.txt b/reqs/3.7/requirements-dev.txt
@@ -113,6 +113,7 @@ nbstripout==0.3.9
 nest-asyncio==1.5.1
 nodeenv==1.5.0
 notebook==6.2.0
+numba==0.52.0
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0

diff --git a/reqs/3.8/requirements-dev.txt b/reqs/3.8/requirements-dev.txt
@@ -113,6 +113,7 @@ nbstripout==0.3.9
 nest-asyncio==1.5.1
 nodeenv==1.5.0
 notebook==6.2.0
+numba==0.52.0
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0

diff --git a/setup.cfg b/setup.cfg
@@ -83,6 +83,8 @@ ignore_errors = True
 ignore_missing_imports = True
 [mypy-matplotlib.*]
 ignore_missing_imports = True
+[mypy-numba.*]
+ignore_missing_imports = True
 [mypy-numpy.*]
 ignore_missing_imports = True
 [mypy-pandas.*]

diff --git a/src/tensorwaves/physics/amplitude.py b/src/tensorwaves/physics/amplitude.py
@@ -61,6 +61,17 @@ def jax_lambdify() -> Callable:
     if isinstance(backend, str):
         if backend == "jax":
             return jax_lambdify()
+        if backend == "numba":
+            from numba import jit
+
+            return jit(
+                sympy.lambdify(
+                    variables,
+                    expression,
+                    modules="numpy",
+                ),
+                parallel=True,
+            )
     if isinstance(backend, tuple):
         if any("jax" in x.__name__ for x in backend):
             return jax_lambdify()