Create GUDHI_PYTHON_MODULES_EXTRA without auto-import

Put Wasserstein in it.

4 files changed, 25 insertions, 26 deletions

diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt
index 1b1684e1..dbef7183 100644
--- a/src/python/CMakeLists.txt
+++ b/src/python/CMakeLists.txt
@@ -49,7 +49,8 @@ if(PYTHONINTERP_FOUND)
     set(GUDHI_PYTHON_MODULES "${GUDHI_PYTHON_MODULES}'alpha_complex', ")
     set(GUDHI_PYTHON_MODULES "${GUDHI_PYTHON_MODULES}'euclidean_witness_complex', ")
     set(GUDHI_PYTHON_MODULES "${GUDHI_PYTHON_MODULES}'euclidean_strong_witness_complex', ")
-    set(GUDHI_PYTHON_MODULES "${GUDHI_PYTHON_MODULES}'wasserstein', ")
+    # Modules that should not be auto-imported in __init__.py
+    set(GUDHI_PYTHON_MODULES_EXTRA "${GUDHI_PYTHON_MODULES_EXTRA}'wasserstein', ")
 
     add_gudhi_debug_info("Python version ${PYTHON_VERSION_STRING}")
     add_gudhi_debug_info("Cython version ${CYTHON_VERSION}")
diff --git a/src/python/doc/wasserstein_distance_user.rst b/src/python/doc/wasserstein_distance_user.rst
index bcb7f19d..fc7bd82c 100644
--- a/src/python/doc/wasserstein_distance_user.rst
+++ b/src/python/doc/wasserstein_distance_user.rst
@@ -13,7 +13,7 @@ This implementation is based on ideas from "Large Scale Computation of Means and
 
 Function
 --------
-.. autofunction:: gudhi.wasserstein_distance
+.. autofunction:: gudhi.wasserstein.wasserstein_distance
 
 
 Basic example
@@ -24,13 +24,13 @@ Note that persistence diagrams must be submitted as (n x 2) numpy arrays and mus
 
 .. testcode::
 
-    import gudhi
+    import gudhi.wasserstein
     import numpy as np
 
     diag1 = np.array([[2.7, 3.7],[9.6, 14.],[34.2, 34.974]])
     diag2 = np.array([[2.8, 4.45],[9.5, 14.1]])
 
-    message = "Wasserstein distance value = " + '%.2f' % gudhi.wasserstein_distance(diag1, diag2, q=2., p=1.)
+    message = "Wasserstein distance value = " + '%.2f' % gudhi.wasserstein.wasserstein_distance(diag1, diag2, q=2., p=1.)
     print(message)
 
 The output is:
diff --git a/src/python/gudhi/__init__.py.in b/src/python/gudhi/__init__.py.in
index 28bab0e1..4039c08f 100644
--- a/src/python/gudhi/__init__.py.in
+++ b/src/python/gudhi/__init__.py.in
@@ -21,13 +21,13 @@ __debug_info__ = @GUDHI_PYTHON_DEBUG_INFO@
 from sys import exc_info
 from importlib import import_module
 
-__all__ = [@GUDHI_PYTHON_MODULES@]
+__all__ = [@GUDHI_PYTHON_MODULES@ @GUDHI_PYTHON_MODULES_EXTRA@]
 
 __available_modules = ''
 __missing_modules = ''
 
 # try to import * from gudhi.__module_name
-for __module_name in __all__:
+for __module_name in [@GUDHI_PYTHON_MODULES@]:
     try:
         __module = import_module('gudhi.' + __module_name)
         try:
diff --git a/src/python/test/test_wasserstein_distance.py b/src/python/test/test_wasserstein_distance.py
index c1b568e2..a6bf9901 100755
--- a/src/python/test/test_wasserstein_distance.py
+++ b/src/python/test/test_wasserstein_distance.py
@@ -1,4 +1,4 @@
-import gudhi
+from gudhi.wasserstein import wasserstein_distance
 import numpy as np
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
@@ -23,28 +23,26 @@ def test_basic_wasserstein():
     diag4 = np.array([[0, 3], [4, 8]])
     emptydiag = np.array([[]])
 
-    assert gudhi.wasserstein_distance(emptydiag, emptydiag, q=2., p=1.) == 0.
-    assert gudhi.wasserstein_distance(emptydiag, emptydiag, q=np.inf, p=1.) == 0.
-    assert gudhi.wasserstein_distance(emptydiag, emptydiag, q=np.inf, p=2.) == 0.
-    assert gudhi.wasserstein_distance(emptydiag, emptydiag, q=2., p=2.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, q=2., p=1.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, q=np.inf, p=1.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, q=np.inf, p=2.) == 0.
+    assert wasserstein_distance(emptydiag, emptydiag, q=2., p=2.) == 0.
 
-    assert gudhi.wasserstein_distance(diag3, emptydiag, q=np.inf, p=1.) == 2.
-    assert gudhi.wasserstein_distance(diag3, emptydiag, q=1., p=1.) == 4.
+    assert wasserstein_distance(diag3, emptydiag, q=np.inf, p=1.) == 2.
+    assert wasserstein_distance(diag3, emptydiag, q=1., p=1.) == 4.
 
-    assert gudhi.wasserstein_distance(diag4, emptydiag, q=1., p=2.) == 5.  # thank you Pythagorician triplets
-    assert gudhi.wasserstein_distance(diag4, emptydiag, q=np.inf, p=2.) == 2.5
-    assert gudhi.wasserstein_distance(diag4, emptydiag, q=2., p=2.) == 3.5355339059327378
+    assert wasserstein_distance(diag4, emptydiag, q=1., p=2.) == 5.  # thank you Pythagorician triplets
+    assert wasserstein_distance(diag4, emptydiag, q=np.inf, p=2.) == 2.5
+    assert wasserstein_distance(diag4, emptydiag, q=2., p=2.) == 3.5355339059327378
 
-    assert gudhi.wasserstein_distance(diag1, diag2, q=2., p=1.) == 1.4453593023967701
-    assert gudhi.wasserstein_distance(diag1, diag2, q=2.35, p=1.74) == 0.9772734057168739
+    assert wasserstein_distance(diag1, diag2, q=2., p=1.) == 1.4453593023967701
+    assert wasserstein_distance(diag1, diag2, q=2.35, p=1.74) == 0.9772734057168739
 
-    assert gudhi.wasserstein_distance(diag1, emptydiag, q=2.35, p=1.7863) == 3.141592214572228
+    assert wasserstein_distance(diag1, emptydiag, q=2.35, p=1.7863) == 3.141592214572228
 
-    assert gudhi.wasserstein_distance(diag3, diag4, q=1., p=1.) == 3.
-    assert gudhi.wasserstein_distance(diag3, diag4, q=np.inf, p=1.) == 3.  # no diag matching here
-    assert gudhi.wasserstein_distance(diag3, diag4, q=np.inf, p=2.) == np.sqrt(5)
-    assert gudhi.wasserstein_distance(diag3, diag4, q=1., p=2.) == np.sqrt(5)
-    assert gudhi.wasserstein_distance(diag3, diag4, q=4.5, p=2.) == np.sqrt(5)
-
-    
+    assert wasserstein_distance(diag3, diag4, q=1., p=1.) == 3.
+    assert wasserstein_distance(diag3, diag4, q=np.inf, p=1.) == 3.  # no diag matching here
+    assert wasserstein_distance(diag3, diag4, q=np.inf, p=2.) == np.sqrt(5)
+    assert wasserstein_distance(diag3, diag4, q=1., p=2.) == np.sqrt(5)
+    assert wasserstein_distance(diag3, diag4, q=4.5, p=2.) == np.sqrt(5)