update documentation

7 files changed, 53 insertions, 17 deletions

diff --git a/README.md b/README.md
index 5f37ad6..e49c8da 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ This open source Python library provide several solvers for optimization problem
 It provides the following solvers:
 
 * OT Network Flow solver for the linear program/ Earth Movers Distance [1].
-* Entropic regularization OT solver with Sinkhorn Knopp Algorithm [2] and stabilized version [9][10] with optional GPU implementation (requires cudamat).
+* Entropic regularization OT solver with Sinkhorn Knopp Algorithm [2] and stabilized version [9][10] with optional GPU implementation (requires cupy).
 * Smooth optimal transport solvers (dual and semi-dual) for KL and squared L2 regularizations [17].
 * Non regularized Wasserstein barycenters [16] with LP solver (only small scale).
 * Bregman projections for Wasserstein barycenter [3] and unmixing [4].
@@ -83,12 +83,8 @@ Some sub-modules require additional dependences which are discussed below
 ```
 pip install pymanopt autograd
 ```
-* **ot.gpu** (GPU accelerated OT) depends on cudamat that have to be installed with:
-```
-git clone https://github.com/cudamat/cudamat.git
-cd cudamat
-python setup.py install --user # for user install (no root)
-```
+* **ot.gpu** (GPU accelerated OT) depends on cupy that have to be installed following instructions on [this page](https://docs-cupy.chainer.org/en/stable/install.html).
+
 
 obviously you need CUDA installed and a compatible GPU.
 
diff --git a/docs/source/all.rst b/docs/source/all.rst
index 9459023..1eaf3b1 100644
--- a/docs/source/all.rst
+++ b/docs/source/all.rst
@@ -48,6 +48,12 @@ ot.da
 
 .. automodule:: ot.da
   :members:
+  
+ot.gpu
+--------
+
+.. automodule:: ot.gpu
+  :members:
 
 ot.dr
 --------
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 114245d..433eca6 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -31,7 +31,7 @@ class Mock(MagicMock):
     @classmethod
     def __getattr__(cls, name):
         return MagicMock()
-MOCK_MODULES = ['ot.lp.emd_wrap','autograd','pymanopt','cudamat','autograd.numpy','pymanopt.manifolds','pymanopt.solvers']
+MOCK_MODULES = ['ot.lp.emd_wrap','autograd','pymanopt','cupy','autograd.numpy','pymanopt.manifolds','pymanopt.solvers']
 # 'autograd.numpy','pymanopt.manifolds','pymanopt.solvers',
 sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)
 # !!!!
diff --git a/ot/gpu/__init__.py b/ot/gpu/__init__.py
index de4825d..9de2c40 100644
--- a/ot/gpu/__init__.py
+++ b/ot/gpu/__init__.py
@@ -1,8 +1,28 @@
 # -*- coding: utf-8 -*-
+"""
+
+
+This module implement GPU ilmplementation for several OT solvers and utility 
+functions. The GPU backend in handled by `cupy 
+<https://cupy.chainer.org/>`_.
+
+By default, the functions in this module accept and return numpy arrays 
+in order to proide drop-in replacement for the other POT function but
+the transfer between CPU en GPU comes with a significant overhead.
+
+In order to get the best erformances, we recommend to given only cupy 
+arrays to the functions and desactivate the conversion to numpy of the 
+result of the function with parameter ``to_numpy=False``.
+
+
+
+
+"""
 
 from . import bregman
 from . import da
 from .bregman import sinkhorn
+from .da 
 
 from . import utils
 from .utils import dist, to_gpu, to_np
@@ -13,4 +33,4 @@ from .utils import dist, to_gpu, to_np
 #
 # License: MIT License
 
-__all__ = ["utils", "dist", "sinkhorn"]
+__all__ = ["utils", "dist", "sinkhorn", 'bregman', 'da', 'to_gpu', 'to_np']
diff --git a/ot/gpu/bregman.py b/ot/gpu/bregman.py
index 600ead4..978b307 100644
--- a/ot/gpu/bregman.py
+++ b/ot/gpu/bregman.py
@@ -16,7 +16,10 @@ from . import utils
 def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
                    verbose=False, log=False, to_numpy=True, **kwargs):
     """
-    Solve the entropic regularization optimal transport problem and return the OT matrix
+    Solve the entropic regularization optimal transport on GPU
+
+    If the input matrix are in numpy format, they will be uploaded to the
+    GPU first which can incur significant time overhead.
 
     The function solves the following optimization problem:
 
@@ -56,6 +59,8 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
         Print information along iterations
     log : bool, optional
         record log if True
+    to_numpy : boolean, optional (default True)
+        If true convert back the GPU array result to numpy format.
 
 
     Returns
diff --git a/ot/gpu/da.py b/ot/gpu/da.py
index 8c63870..6aba29c 100644
--- a/ot/gpu/da.py
+++ b/ot/gpu/da.py
@@ -22,7 +22,11 @@ def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
                      log=False, to_numpy=True):
     """
     Solve the entropic regularization optimal transport problem with nonconvex
-    group lasso regularization
+    group lasso regularization on GPU
+
+    If the input matrix are in numpy format, they will be uploaded to the
+    GPU first which can incur significant time overhead.
+
 
     The function solves the following optimization problem:
 
@@ -74,6 +78,8 @@ def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
         Print information along iterations
     log : bool, optional
         record log if True
+    to_numpy : boolean, optional (default True)
+        If true convert back the GPU array result to numpy format.
 
 
     Returns
diff --git a/ot/gpu/utils.py b/ot/gpu/utils.py
index d349a6d..41e168a 100644
--- a/ot/gpu/utils.py
+++ b/ot/gpu/utils.py
@@ -16,19 +16,22 @@ import cupy as cp  # cp used for cupy specific operations
 def euclidean_distances(a, b, squared=False, to_numpy=True):
     """
     Compute the pairwise euclidean distance between matrices a and b.
-     Parameters
+
+    If the input matrix are in numpy format, they will be uploaded to the
+    GPU first which can incur significant time overhead.
+
+    Parameters
     ----------
     a : np.ndarray (n, f)
         first matrix
     b : np.ndarray (m, f)
         second matrix
-    gpu : boolean, optional (default False)
-        if True and the module cupy is available, the computation is done
-        on the GPU and the type of the matrix returned is cupy.ndarray.
-        Otherwise, compute on the CPU and returns np.ndarray.
+    to_numpy : boolean, optional (default True)
+        If true convert back the GPU array result to numpy format.
     squared : boolean, optional (default False)
         if True, return squared euclidean distance matrix
-     Returns
+
+    Returns
     -------
     c : (n x m) np.ndarray or cupy.ndarray
         pairwise euclidean distance distance matrix