Merge pull request #396 from VincentRouvreau/tensorflow_wasserstein_test

Added tests for wasserstein distance with tensorflow

5 files changed, 63 insertions, 25 deletions

diff --git a/src/cmake/modules/GUDHI_third_party_libraries.cmake b/src/cmake/modules/GUDHI_third_party_libraries.cmake
index 1fbc4244..9dadac4f 100644
--- a/src/cmake/modules/GUDHI_third_party_libraries.cmake
+++ b/src/cmake/modules/GUDHI_third_party_libraries.cmake
@@ -155,6 +155,7 @@ if( PYTHONINTERP_FOUND )
   find_python_module("pykeops")
   find_python_module("eagerpy")
   find_python_module_no_version("hnswlib")
+  find_python_module("tensorflow")
 endif()
 
 if(NOT GUDHI_PYTHON_PATH)
diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt
index 4f26481e..c09996fe 100644
--- a/src/python/CMakeLists.txt
+++ b/src/python/CMakeLists.txt
@@ -103,6 +103,9 @@ if(PYTHONINTERP_FOUND)
     if(EAGERPY_FOUND)
       add_gudhi_debug_info("EagerPy version ${EAGERPY_VERSION}")
     endif()
+    if(TENSORFLOW_FOUND)
+      add_gudhi_debug_info("TensorFlow version ${TENSORFLOW_VERSION}")
+    endif()
 
     set(GUDHI_PYTHON_EXTRA_COMPILE_ARGS "${GUDHI_PYTHON_EXTRA_COMPILE_ARGS}'-DBOOST_RESULT_OF_USE_DECLTYPE', ")
     set(GUDHI_PYTHON_EXTRA_COMPILE_ARGS "${GUDHI_PYTHON_EXTRA_COMPILE_ARGS}'-DBOOST_ALL_NO_LIB', ")
@@ -496,11 +499,17 @@ if(PYTHONINTERP_FOUND)
 
     # Wasserstein
     if(OT_FOUND AND PYBIND11_FOUND)
-      if(TORCH_FOUND AND EAGERPY_FOUND)
+      # EagerPy dependency because of enable_autodiff=True
+      if(EAGERPY_FOUND)
         add_gudhi_py_test(test_wasserstein_distance)
       endif()
       add_gudhi_py_test(test_wasserstein_barycenter)
     endif()
+    if(OT_FOUND)
+      if(TORCH_FOUND AND TENSORFLOW_FOUND AND EAGERPY_FOUND)
+        add_gudhi_py_test(test_wasserstein_with_tensors)
+      endif()
+    endif()
 
     # Representations
     if(SKLEARN_FOUND AND MATPLOTLIB_FOUND)
diff --git a/src/python/doc/installation.rst b/src/python/doc/installation.rst
index 2f161d66..66efe45a 100644
--- a/src/python/doc/installation.rst
+++ b/src/python/doc/installation.rst
@@ -394,6 +394,11 @@ mathematics, science, and engineering.
 :class:`~gudhi.point_cloud.knn.KNearestNeighbors` can use the Python package
 `SciPy <http://scipy.org>`_ as a backend if explicitly requested.
 
+TensorFlow
+----------
+
+`TensorFlow <https://www.tensorflow.org>`_ is currently only used in some automatic differentiation tests.
+
 Bug reports and contributions
 *****************************
 
diff --git a/src/python/test/test_wasserstein_distance.py b/src/python/test/test_wasserstein_distance.py
index 90d26809..e3b521d6 100755
--- a/src/python/test/test_wasserstein_distance.py
+++ b/src/python/test/test_wasserstein_distance.py
@@ -97,27 +97,3 @@ def test_wasserstein_distance_pot():
 def test_wasserstein_distance_hera():
     _basic_wasserstein(hera_wrap(delta=1e-12), 1e-12, test_matching=False)
     _basic_wasserstein(hera_wrap(delta=.1), .1, test_matching=False)
-
-def test_wasserstein_distance_grad():
-    import torch
-
-    diag1 = torch.tensor([[2.7, 3.7], [9.6, 14.0], [34.2, 34.974]], requires_grad=True)
-    diag2 = torch.tensor([[2.8, 4.45], [9.5, 14.1]], requires_grad=True)
-    diag3 = torch.tensor([[2.8, 4.45], [9.5, 14.1]], requires_grad=True)
-    assert diag1.grad is None and diag2.grad is None and diag3.grad is None
-    dist12 = pot(diag1, diag2, internal_p=2, order=2, enable_autodiff=True)
-    dist30 = pot(diag3, torch.tensor([]), internal_p=2, order=2, enable_autodiff=True)
-    dist12.backward()
-    dist30.backward()
-    assert not torch.isnan(diag1.grad).any() and not torch.isnan(diag2.grad).any() and not torch.isnan(diag3.grad).any()
-    diag4 = torch.tensor([[0., 10.]], requires_grad=True)
-    diag5 = torch.tensor([[1., 11.], [3., 4.]], requires_grad=True)
-    dist45 = pot(diag4, diag5, internal_p=1, order=1, enable_autodiff=True)
-    assert dist45 == 3.
-    dist45.backward()
-    assert np.array_equal(diag4.grad, [[-1., -1.]])
-    assert np.array_equal(diag5.grad, [[1., 1.], [-1., 1.]])
-    diag6 = torch.tensor([[5., 10.]], requires_grad=True)
-    pot(diag6, diag6, internal_p=2, order=2, enable_autodiff=True).backward()
-    # https://github.com/jonasrauber/eagerpy/issues/6
-    # assert np.array_equal(diag6.grad, [[0., 0.]])
diff --git a/src/python/test/test_wasserstein_with_tensors.py b/src/python/test/test_wasserstein_with_tensors.py
new file mode 100755
index 00000000..e3f1411a
--- /dev/null
+++ b/src/python/test/test_wasserstein_with_tensors.py
@@ -0,0 +1,47 @@
+""" This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
+    See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
+    Author(s):       Mathieu Carriere
+
+    Copyright (C) 2020 Inria
+
+    Modification(s):
+      - YYYY/MM Author: Description of the modification
+"""
+
+from gudhi.wasserstein import wasserstein_distance as pot
+import numpy as np
+import torch
+import tensorflow as tf
+
+def test_wasserstein_distance_grad():
+    diag1 = torch.tensor([[2.7, 3.7], [9.6, 14.0], [34.2, 34.974]], requires_grad=True)
+    diag2 = torch.tensor([[2.8, 4.45], [9.5, 14.1]], requires_grad=True)
+    diag3 = torch.tensor([[2.8, 4.45], [9.5, 14.1]], requires_grad=True)
+    assert diag1.grad is None and diag2.grad is None and diag3.grad is None
+    dist12 = pot(diag1, diag2, internal_p=2, order=2, enable_autodiff=True)
+    dist30 = pot(diag3, torch.tensor([]), internal_p=2, order=2, enable_autodiff=True)
+    dist12.backward()
+    dist30.backward()
+    assert not torch.isnan(diag1.grad).any() and not torch.isnan(diag2.grad).any() and not torch.isnan(diag3.grad).any()
+    diag4 = torch.tensor([[0., 10.]], requires_grad=True)
+    diag5 = torch.tensor([[1., 11.], [3., 4.]], requires_grad=True)
+    dist45 = pot(diag4, diag5, internal_p=1, order=1, enable_autodiff=True)
+    assert dist45 == 3.
+    dist45.backward()
+    assert np.array_equal(diag4.grad, [[-1., -1.]])
+    assert np.array_equal(diag5.grad, [[1., 1.], [-1., 1.]])
+    diag6 = torch.tensor([[5., 10.]], requires_grad=True)
+    pot(diag6, diag6, internal_p=2, order=2, enable_autodiff=True).backward()
+    # https://github.com/jonasrauber/eagerpy/issues/6
+    # assert np.array_equal(diag6.grad, [[0., 0.]])
+
+def test_wasserstein_distance_grad_tensorflow():
+    with tf.GradientTape() as tape:
+        diag4 = tf.convert_to_tensor(tf.Variable(initial_value=np.array([[0., 10.]]),           trainable=True))
+        diag5 = tf.convert_to_tensor(tf.Variable(initial_value=np.array([[1., 11.], [3., 4.]]), trainable=True))
+        dist45 = pot(diag4, diag5, internal_p=1, order=1, enable_autodiff=True)
+        assert dist45 == 3.
+
+    grads = tape.gradient(dist45, [diag4, diag5])
+    assert np.array_equal(grads[0].values, [[-1., -1.]])
+    assert np.array_equal(grads[1].values, [[1., 1.], [-1., 1.]])
\ No newline at end of file