feat: support lambdify to tensorflow (#245)

ComPWA · Mar 24, 2021 · ea97355 · ea97355
1 parent 10ba92f
commit ea97355
Show file tree

Hide file tree

Showing 21 changed files with 51 additions and 30 deletions.
diff --git a/reqs/3.6/requirements-dev.txt b/reqs/3.6/requirements-dev.txt
@@ -84,7 +84,7 @@ jupyter-client==6.1.12
 jupyter-console==6.3.0
 jupyter-core==4.7.1
 jupyter-packaging==0.7.12
-jupyter-server==1.4.1
+jupyter-server==1.5.0
 jupyter-sphinx==0.3.1
 jupyter==1.0.0
 jupyterlab-code-formatter==1.4.5
@@ -94,7 +94,7 @@ jupyterlab==3.0.12
 keras-preprocessing==1.1.2
 kiwisolver==1.3.1
 labels==20.1.0
-lazy-object-proxy==1.5.2
+lazy-object-proxy==1.6.0
 livereload==2.6.3
 llvmlite==0.36.0
 markdown-it-py==0.6.2
@@ -126,7 +126,7 @@ packaging==20.9
 pandas==1.1.5
 pandocfilters==1.4.3
 parso==0.8.1
-particle==0.14.0
+particle==0.14.1
 pathspec==0.8.1
 pep517==0.10.0
 pep8-naming==0.11.1

diff --git a/reqs/3.6/requirements-doc.txt b/reqs/3.6/requirements-doc.txt
@@ -87,7 +87,7 @@ packaging==20.9
 pandas==1.1.5
 pandocfilters==1.4.3
 parso==0.8.1
-particle==0.14.0
+particle==0.14.1
 pexpect==4.8.0
 phasespace==1.2.0
 pickleshare==0.7.5

diff --git a/reqs/3.6/requirements-extras.txt b/reqs/3.6/requirements-extras.txt
@@ -37,7 +37,7 @@ numba==0.53.0
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 protobuf==3.15.6
 pyasn1-modules==0.2.8

diff --git a/reqs/3.6/requirements-sty.txt b/reqs/3.6/requirements-sty.txt
@@ -54,7 +54,7 @@ jaxlib==0.1.64
 jsonschema==3.2.0
 jupyter-core==4.7.1
 keras-preprocessing==1.1.2
-lazy-object-proxy==1.5.2
+lazy-object-proxy==1.6.0
 llvmlite==0.36.0
 markdown==3.3.4
 mccabe==0.6.1
@@ -69,7 +69,7 @@ numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
 packaging==20.9
-particle==0.14.0
+particle==0.14.1
 pathspec==0.8.1
 pep8-naming==0.11.1
 phasespace==1.2.0

diff --git a/reqs/3.6/requirements-test.txt b/reqs/3.6/requirements-test.txt
@@ -43,7 +43,7 @@ numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
 packaging==20.9
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 pluggy==0.13.1
 protobuf==3.15.6

diff --git a/reqs/3.6/requirements.txt b/reqs/3.6/requirements.txt
@@ -33,7 +33,7 @@ mpmath==1.2.1
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 protobuf==3.15.6
 pyasn1-modules==0.2.8

diff --git a/reqs/3.7/requirements-dev.txt b/reqs/3.7/requirements-dev.txt
@@ -80,7 +80,7 @@ jupyter-client==6.1.12
 jupyter-console==6.3.0
 jupyter-core==4.7.1
 jupyter-packaging==0.7.12
-jupyter-server==1.4.1
+jupyter-server==1.5.0
 jupyter-sphinx==0.3.1
 jupyter==1.0.0
 jupyterlab-code-formatter==1.4.5
@@ -90,7 +90,7 @@ jupyterlab==3.0.12
 keras-preprocessing==1.1.2
 kiwisolver==1.3.1
 labels==20.1.0
-lazy-object-proxy==1.5.2
+lazy-object-proxy==1.6.0
 livereload==2.6.3
 llvmlite==0.36.0
 markdown-it-py==0.6.2
@@ -122,7 +122,7 @@ packaging==20.9
 pandas==1.2.3
 pandocfilters==1.4.3
 parso==0.8.1
-particle==0.14.0
+particle==0.14.1
 pathspec==0.8.1
 pep517==0.10.0
 pep8-naming==0.11.1

diff --git a/reqs/3.7/requirements-doc.txt b/reqs/3.7/requirements-doc.txt
@@ -86,7 +86,7 @@ packaging==20.9
 pandas==1.2.3
 pandocfilters==1.4.3
 parso==0.8.1
-particle==0.14.0
+particle==0.14.1
 pexpect==4.8.0
 phasespace==1.2.0
 pickleshare==0.7.5

diff --git a/reqs/3.7/requirements-extras.txt b/reqs/3.7/requirements-extras.txt
@@ -36,7 +36,7 @@ numba==0.53.0
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 protobuf==3.15.6
 pyasn1-modules==0.2.8

diff --git a/reqs/3.7/requirements-sty.txt b/reqs/3.7/requirements-sty.txt
@@ -52,7 +52,7 @@ jaxlib==0.1.64
 jsonschema==3.2.0
 jupyter-core==4.7.1
 keras-preprocessing==1.1.2
-lazy-object-proxy==1.5.2
+lazy-object-proxy==1.6.0
 llvmlite==0.36.0
 markdown==3.3.4
 mccabe==0.6.1
@@ -67,7 +67,7 @@ numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
 packaging==20.9
-particle==0.14.0
+particle==0.14.1
 pathspec==0.8.1
 pep8-naming==0.11.1
 phasespace==1.2.0

diff --git a/reqs/3.7/requirements-test.txt b/reqs/3.7/requirements-test.txt
@@ -42,7 +42,7 @@ numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
 packaging==20.9
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 pluggy==0.13.1
 protobuf==3.15.6

diff --git a/reqs/3.7/requirements.txt b/reqs/3.7/requirements.txt
@@ -32,7 +32,7 @@ mpmath==1.2.1
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 protobuf==3.15.6
 pyasn1-modules==0.2.8

diff --git a/reqs/3.8/requirements-dev.txt b/reqs/3.8/requirements-dev.txt
@@ -80,7 +80,7 @@ jupyter-client==6.1.12
 jupyter-console==6.3.0
 jupyter-core==4.7.1
 jupyter-packaging==0.7.12
-jupyter-server==1.4.1
+jupyter-server==1.5.0
 jupyter-sphinx==0.3.1
 jupyter==1.0.0
 jupyterlab-code-formatter==1.4.5
@@ -90,7 +90,7 @@ jupyterlab==3.0.12
 keras-preprocessing==1.1.2
 kiwisolver==1.3.1
 labels==20.1.0
-lazy-object-proxy==1.5.2
+lazy-object-proxy==1.6.0
 livereload==2.6.3
 llvmlite==0.36.0
 markdown-it-py==0.6.2
@@ -122,7 +122,7 @@ packaging==20.9
 pandas==1.2.3
 pandocfilters==1.4.3
 parso==0.8.1
-particle==0.14.0
+particle==0.14.1
 pathspec==0.8.1
 pep517==0.10.0
 pep8-naming==0.11.1

diff --git a/reqs/3.8/requirements-doc.txt b/reqs/3.8/requirements-doc.txt
@@ -86,7 +86,7 @@ packaging==20.9
 pandas==1.2.3
 pandocfilters==1.4.3
 parso==0.8.1
-particle==0.14.0
+particle==0.14.1
 pexpect==4.8.0
 phasespace==1.2.0
 pickleshare==0.7.5

diff --git a/reqs/3.8/requirements-extras.txt b/reqs/3.8/requirements-extras.txt
@@ -35,7 +35,7 @@ numba==0.53.0
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 protobuf==3.15.6
 pyasn1-modules==0.2.8

diff --git a/reqs/3.8/requirements-sty.txt b/reqs/3.8/requirements-sty.txt
@@ -51,7 +51,7 @@ jaxlib==0.1.64
 jsonschema==3.2.0
 jupyter-core==4.7.1
 keras-preprocessing==1.1.2
-lazy-object-proxy==1.5.2
+lazy-object-proxy==1.6.0
 llvmlite==0.36.0
 markdown==3.3.4
 mccabe==0.6.1
@@ -66,7 +66,7 @@ numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
 packaging==20.9
-particle==0.14.0
+particle==0.14.1
 pathspec==0.8.1
 pep8-naming==0.11.1
 phasespace==1.2.0

diff --git a/reqs/3.8/requirements-test.txt b/reqs/3.8/requirements-test.txt
@@ -41,7 +41,7 @@ numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
 packaging==20.9
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 pluggy==0.13.1
 protobuf==3.15.6

diff --git a/reqs/3.8/requirements.txt b/reqs/3.8/requirements.txt
@@ -31,7 +31,7 @@ mpmath==1.2.1
 numpy==1.19.5
 oauthlib==3.1.0
 opt-einsum==3.3.0
-particle==0.14.0
+particle==0.14.1
 phasespace==1.2.0
 protobuf==3.15.6
 pyasn1-modules==0.2.8

diff --git a/setup.cfg b/setup.cfg
@@ -47,7 +47,7 @@ install_requires =
     phasespace >= 1.2.0
     PyYAML
     sympy
-    tensorflow >= 2.0
+    tensorflow >= 2.4  # tensorflow.experimental.numpy
     tqdm >= 4.24.0  # autonotebook
     typing-extensions; python_version < "3.8.0"
 packages = find:

diff --git a/src/tensorwaves/model.py b/src/tensorwaves/model.py
@@ -34,6 +34,11 @@ def get_backend_modules(
             return (jnp, jsp.special)
         if backend in {"numpy", "numba"}:
             return np.__dict__
+        if backend in {"tensorflow", "tf"}:
+            # pylint: disable=import-error
+            import tensorflow.experimental.numpy as tnp  # pyright: reportMissingImports=false
+
+            return tnp.__dict__
 
     return backend
 
@@ -119,7 +124,7 @@ def __init__(
 
     def lambdify(self, backend: Union[str, tuple, dict]) -> Callable:
         """Lambdify the model using `~sympy.utilities.lambdify.lambdify`."""
-        # pylint: disable=import-outside-toplevel
+        # pylint: disable=import-outside-toplevel,too-many-return-statements
         ordered_symbols = self.__argument_order
 
         def jax_lambdify() -> Callable:
@@ -147,17 +152,33 @@ def numba_lambdify() -> Callable:
                 parallel=True,
             )
 
+        def tensorflow_lambdify() -> Callable:
+            # pylint: disable=import-error
+            import tensorflow.experimental.numpy as tnp  # pyright: reportMissingImports=false
+
+            return sp.lambdify(
+                ordered_symbols,
+                self.__expression,
+                modules=tnp,
+            )
+
         backend_modules = get_backend_modules(backend)
         if isinstance(backend, str):
             if backend == "jax":
                 return jax_lambdify()
             if backend == "numba":
                 return numba_lambdify()
+            if backend in {"tensorflow", "tf"}:
+                return tensorflow_lambdify()
         if isinstance(backend, tuple):
             if any("jax" in x.__name__ for x in backend):
                 return jax_lambdify()
             if any("numba" in x.__name__ for x in backend):
                 return numba_lambdify()
+            if any("tensorflow" in x.__name__ for x in backend) or any(
+                "tf" in x.__name__ for x in backend
+            ):
+                return tensorflow_lambdify()
         return sp.lambdify(
             ordered_symbols,
             self.__expression,

diff --git a/src/tensorwaves/optimizer/minuit.py b/src/tensorwaves/optimizer/minuit.py
@@ -73,7 +73,7 @@ def wrapped_function(pars: list) -> float:
             update_parameters(pars)
             parameters = parameter_handler.unflatten(flattened_parameters)
             estimator_value = estimator(parameters)
-            progress_bar.set_postfix({"estimator": estimator_value})
+            progress_bar.set_postfix({"estimator": float(estimator_value)})
             progress_bar.update()
             logs = create_log(estimator_value, parameters)
             self.__callback.on_function_call_end(n_function_calls, logs)