Add optional install options with pip (#627)

* add extra optional dependency options

* check for dependencies, more informative errors: da, dr

* update RELEASES.md

* update README.md

* Update setup.py

Co-authored-by: Rémi Flamary <[email protected]>

* change filename to requirements_all in all files

* woops some of those were for docs

---------

Co-authored-by: Rémi Flamary <[email protected]>

.circleci/config.yml

@@ -46,7 +46,7 @@ jobs:
             command: |
               python -m pip install --user --upgrade --progress-bar off pip
               python -m pip install --user -e .
-              python -m pip install --user --upgrade --no-cache-dir --progress-bar off -r requirements.txt 
+              python -m pip install --user --upgrade --no-cache-dir --progress-bar off -r requirements_all.txt 
               python -m pip install --user --upgrade --progress-bar off -r docs/requirements.txt
               python -m pip install --user --upgrade --progress-bar off ipython sphinx-gallery memory_profiler              
           # python -m pip install --user --upgrade --progress-bar off ipython "https://api.github.com/repos/sphinx-gallery/sphinx-gallery/zipball/master" memory_profiler

.github/workflows/build_doc.yml

@@ -24,7 +24,7 @@ jobs:
     - name: Get Python running
       run: |
         python -m pip install --user --upgrade --progress-bar off pip
-        python -m pip install --user --upgrade --progress-bar off -r requirements.txt
+        python -m pip install --user --upgrade --progress-bar off -r requirements_all.txt
         python -m pip install --user --upgrade --progress-bar off -r docs/requirements.txt
         python -m pip install --user --upgrade --progress-bar off ipython "https://api.github.com/repos/sphinx-gallery/sphinx-gallery/zipball/master" memory_profiler
         python -m pip install --user -e .

.github/workflows/build_tests.yml

@@ -36,7 +36,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -r requirements.txt
+        pip install -r requirements_all.txt
         pip install pytest pytest-cov
     - name: Run tests
       run: |
@@ -108,7 +108,7 @@ jobs:
      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
-         pip install -r requirements.txt
+         pip install -r requirements_all.txt
          pip install pytest 
      - name: Run tests
        run: |

README.md

@@ -113,6 +113,12 @@ or get the very latest version by running:
 pip install -U https://github.com/PythonOT/POT/archive/master.zip # with --user for user install (no root)
 ```
 
+Optional dependencies may be installed with
+```console
+pip install POT[all]
+```
+Note that this installs `cvxopt`, which is licensed under GPL 3.0. Alternatively, if you cannot use GPL-licensed software, the specific optional dependencies may be installed individually, or per-submodule. The available optional installations are `backend-jax, backend-tf, backend-torch, cvxopt, dr, gnn, all`.
+
 #### Anaconda installation with conda-forge
 
 If you use the Anaconda python distribution, POT is available in [conda-forge](https://conda-forge.org). To install it and the required dependencies:

RELEASES.md

@@ -10,6 +10,7 @@
 + New general unbalanced solvers for `ot.solve` and BFGS solver and illustrative example (PR #620)
 + Add gradient computation with envelope theorem to sinkhorn solver of `ot.solve` with `grad='envelope'` (PR #605).
 + Added support for [Low rank Gromov-Wasserstein](https://proceedings.mlr.press/v162/scetbon22b/scetbon22b.pdf) with `ot.gromov.lowrank_gromov_wasserstein_samples` (PR #614)
++ Optional dependencies may now be installed with `pip install POT[all]` The specific backends or submodules' dependencies may also be installed individually. The pip options are: `backend-jax, backend-tf, backend-torch, cvxopt, dr, gnn, all`. The installation of the `cupy` backend should be done with conda.
 
 #### Closed issues
 - Fix gpu compatibility of sr(F)GW solvers when `G0 is not None`(PR #596)

ot/da.py

@@ -376,8 +376,10 @@ def emd_laplace(a, b, xs, xt, M, sim='knn', sim_param=None, reg='pos', eta=1, al
     elif sim == 'knn':
         if sim_param is None:
             sim_param = 3
-
-        from sklearn.neighbors import kneighbors_graph
+        try:
+            from sklearn.neighbors import kneighbors_graph
+        except ImportError:
+            raise ValueError('scikit-learn must be installed to use knn similarity. Install with `$pip install scikit-learn`.')
 
         sS = nx.from_numpy(kneighbors_graph(
             X=nx.to_numpy(xs), n_neighbors=int(sim_param)

ot/dr.py

@@ -17,12 +17,15 @@
 # License: MIT License
 
 from scipy import linalg
-import autograd.numpy as np
-from sklearn.decomposition import PCA
-
-import pymanopt
-import pymanopt.manifolds
-import pymanopt.optimizers
+try:
+    import autograd.numpy as np
+    from sklearn.decomposition import PCA
+
+    import pymanopt
+    import pymanopt.manifolds
+    import pymanopt.optimizers
+except ImportError:
+    raise ImportError("Missing dependency for ot.dr. Requires autograd, pymanopt, scikit-learn. You can install with install with 'pip install POT[dr]', or 'conda install autograd pymanopt scikit-learn'")
 
 from .bregman import sinkhorn as sinkhorn_bregman
 from .utils import dist as dist_utils, check_random_state

setup.py

@@ -46,6 +46,9 @@
     sdk_path = subprocess.check_output(['xcrun', '--show-sdk-path'])
     os.environ['CFLAGS'] = '-isysroot "{}"'.format(sdk_path.rstrip().decode("utf-8"))
 
+with open('requirements_all.txt') as f:
+    optional_requirements = f.read().splitlines()
+
 setup(
     name='POT',
     version=__version__,
@@ -70,6 +73,17 @@
     scripts=[],
     data_files=[],
     install_requires=["numpy>=1.16", "scipy>=1.6"],
+    extras_require={
+        'backend-numpy': [], # in requirements.
+        'backend-jax': ['jax<=0.4.24', 'jaxlib<=0.4.24'],
+        'backend-cupy': [], # should be installed with conda, not pip, or figure out what CUDA version above.
+        'backend-tf': ['tensorflow'],
+        'backend-torch': ['torch'],
+        'cvxopt': ['cvxopt'], # on it's own to prevent accidental GPL violations
+        'dr': ['scikit-learn', 'pymanopt', 'autograd'],
+        'gnn': ['torch', 'torch_geometric'],
+        'all': optional_requirements
+    },
     python_requires=">=3.7",
     classifiers=[
         'Development Status :: 5 - Production/Stable',