From f9a933862050ca95b3a96d7a8572d62f7f2205a9 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Sat, 11 Apr 2020 18:18:14 +0200
Subject: Use longer names

---
 src/python/gudhi/point_cloud/dtm.py | 10 ++++++----
 src/python/gudhi/point_cloud/knn.py |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/dtm.py b/src/python/gudhi/point_cloud/dtm.py
index 23c36b88..38368f29 100644
--- a/src/python/gudhi/point_cloud/dtm.py
+++ b/src/python/gudhi/point_cloud/dtm.py
@@ -7,10 +7,10 @@
 # Modification(s):
 #   - YYYY/MM Author: Description of the modification
 
-from .knn import KNN
+from .knn import KNearestNeighbors
 
 
-class DTM:
+class DistanceToMeasure:
     """
     Class to compute the distance to the empirical measure defined by a point set, as introduced in :cite:`dtm`.
     """
@@ -20,7 +20,7 @@ class DTM:
         Args:
             k (int): number of neighbors (possibly including the point itself).
             q (float): order used to compute the distance to measure. Defaults to 2.
-            kwargs: same parameters as :class:`~gudhi.point_cloud.knn.KNN`, except that metric="neighbors" means that :func:`transform` expects an array with the distances to the k nearest neighbors.
+            kwargs: same parameters as :class:`~gudhi.point_cloud.knn.KNearestNeighbors`, except that metric="neighbors" means that :func:`transform` expects an array with the distances to the k nearest neighbors.
         """
         self.k = k
         self.q = q
@@ -35,7 +35,9 @@ class DTM:
             X (numpy.array): coordinates for mass points.
         """
         if self.params.setdefault("metric", "euclidean") != "neighbors":
-            self.knn = KNN(self.k, return_index=False, return_distance=True, sort_results=False, **self.params)
+            self.knn = KNearestNeighbors(
+                self.k, return_index=False, return_distance=True, sort_results=False, **self.params
+            )
             self.knn.fit(X)
         return self
 
diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index 8369f1f8..6642a3c2 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -10,7 +10,7 @@
 import numpy
 
 
-class KNN:
+class KNearestNeighbors:
     """
     Class wrapping several implementations for computing the k nearest neighbors in a point set.
     """
-- 
cgit v1.2.3


From 83a1bc1fb6124a35d515f4836d2e830f3dbdf0e7 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Sun, 12 Apr 2020 21:57:51 +0200
Subject: Parallelize the "precomputed" case of knn

It is supposed to be possible to compile numpy with openmp, but it looks
like it isn't done in any of the usual packages.

It may be possible to refactor that code so there is less redundancy.
---
 src/python/gudhi/point_cloud/knn.py | 78 +++++++++++++++++++++++++++++--------
 src/python/test/test_dtm.py         |  3 ++
 src/python/test/test_knn.py         |  8 ++++
 3 files changed, 73 insertions(+), 16 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index 6642a3c2..f6870517 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -115,25 +115,71 @@ class KNearestNeighbors:
 
         if metric == "precomputed":
             # scikit-learn could handle that, but they insist on calling fit() with an unused square array, which is too unnatural.
-            X = numpy.array(X)
             if self.return_index:
-                neighbors = numpy.argpartition(X, k - 1)[:, 0:k]
-                if self.params.get("sort_results", True):
-                    X = numpy.take_along_axis(X, neighbors, axis=-1)
-                    ngb_order = numpy.argsort(X, axis=-1)
-                    neighbors = numpy.take_along_axis(neighbors, ngb_order, axis=-1)
+                n_jobs = self.params.get("n_jobs", 1)
+                # Supposedly numpy can be compiled with OpenMP and handle this, but nobody does that?!
+                if n_jobs == 1:
+                    neighbors = numpy.argpartition(X, k - 1)[:, 0:k]
+                    if self.params.get("sort_results", True):
+                        X = numpy.take_along_axis(X, neighbors, axis=-1)
+                        ngb_order = numpy.argsort(X, axis=-1)
+                        neighbors = numpy.take_along_axis(neighbors, ngb_order, axis=-1)
+                    else:
+                        ngb_order = neighbors
+                    if self.return_distance:
+                        distances = numpy.take_along_axis(X, ngb_order, axis=-1)
+                        return neighbors, distances
+                    else:
+                        return neighbors
                 else:
-                    ngb_order = neighbors
-                if self.return_distance:
-                    distances = numpy.take_along_axis(X, ngb_order, axis=-1)
-                    return neighbors, distances
-                else:
-                    return neighbors
+                    from joblib import Parallel, delayed, effective_n_jobs
+                    from sklearn.utils import gen_even_slices
+
+                    slices = gen_even_slices(len(X), effective_n_jobs(-1))
+                    parallel = Parallel(backend="threading", n_jobs=-1)
+                    if self.params.get("sort_results", True):
+
+                        def func(M):
+                            neighbors = numpy.argpartition(M, k - 1)[:, 0:k]
+                            Y = numpy.take_along_axis(M, neighbors, axis=-1)
+                            ngb_order = numpy.argsort(Y, axis=-1)
+                            return numpy.take_along_axis(neighbors, ngb_order, axis=-1)
+
+                    else:
+
+                        def func(M):
+                            return numpy.argpartition(M, k - 1)[:, 0:k]
+
+                    neighbors = numpy.concatenate(parallel(delayed(func)(X[s]) for s in slices))
+                    if self.return_distance:
+                        distances = numpy.take_along_axis(X, neighbors, axis=-1)
+                        return neighbors, distances
+                    else:
+                        return neighbors
             if self.return_distance:
-                distances = numpy.partition(X, k - 1)[:, 0:k]
-                if self.params.get("sort_results"):
-                    # partition is not guaranteed to sort the lower half, although it often does
-                    distances.sort(axis=-1)
+                n_jobs = self.params.get("n_jobs", 1)
+                if n_jobs == 1:
+                    distances = numpy.partition(X, k - 1)[:, 0:k]
+                    if self.params.get("sort_results"):
+                        # partition is not guaranteed to sort the lower half, although it often does
+                        distances.sort(axis=-1)
+                else:
+                    from joblib import Parallel, delayed, effective_n_jobs
+                    from sklearn.utils import gen_even_slices
+
+                    if self.params.get("sort_results"):
+
+                        def func(M):
+                            # Not partitioning in place, because we should not modify the user's array?
+                            r = numpy.partition(M, k - 1)[:, 0:k]
+                            r.sort(axis=-1)
+                            return r
+
+                    else:
+                        func = lambda M: numpy.partition(M, k - 1)[:, 0:k]
+                    slices = gen_even_slices(len(X), effective_n_jobs(-1))
+                    parallel = Parallel(backend="threading", n_jobs=-1)
+                    distances = numpy.concatenate(parallel(delayed(func)(X[s]) for s in slices))
                 return distances
             return None
 
diff --git a/src/python/test/test_dtm.py b/src/python/test/test_dtm.py
index 37934fdb..bc0d3698 100755
--- a/src/python/test/test_dtm.py
+++ b/src/python/test/test_dtm.py
@@ -33,6 +33,9 @@ def test_dtm_compare_euclidean():
     dtm = DistanceToMeasure(k, metric="precomputed")
     r4 = dtm.fit_transform(d)
     assert r4 == pytest.approx(r0)
+    dtm = DistanceToMeasure(k, metric="precomputed", n_jobs=2)
+    r4b = dtm.fit_transform(d)
+    assert r4b == pytest.approx(r0)
     dtm = DistanceToMeasure(k, implementation="keops")
     r5 = dtm.fit_transform(pts)
     assert r5 == pytest.approx(r0)
diff --git a/src/python/test/test_knn.py b/src/python/test/test_knn.py
index 6aac2006..6269df54 100755
--- a/src/python/test/test_knn.py
+++ b/src/python/test/test_knn.py
@@ -52,6 +52,14 @@ def test_knn_explicit():
     r = knn.fit_transform(dist)
     assert np.array_equal(r[0], [[0, 1], [1, 0], [2, 0]])
     assert np.array_equal(r[1], [[0, 3], [0, 1], [0, 1]])
+    # Second time in parallel
+    knn = KNearestNeighbors(2, metric="precomputed", return_index=True, return_distance=False, n_jobs=2)
+    r = knn.fit_transform(dist)
+    assert np.array_equal(r, [[0, 1], [1, 0], [2, 0]])
+    knn = KNearestNeighbors(2, metric="precomputed", return_index=True, return_distance=True, n_jobs=2)
+    r = knn.fit_transform(dist)
+    assert np.array_equal(r[0], [[0, 1], [1, 0], [2, 0]])
+    assert np.array_equal(r[1], [[0, 3], [0, 1], [0, 1]])
 
 
 def test_knn_compare():
-- 
cgit v1.2.3


From 280eb9d2323837619db1ae013b929adb9b45013b Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Mon, 13 Apr 2020 01:09:45 +0200
Subject: enable_autodiff with keops

This doesn't seem like the best way to handle it, we may want to handle
it like a wrapper that gets the indices from knn (whatever backend) and
then computes the distances.
---
 src/python/gudhi/point_cloud/knn.py | 33 +++++++++++++++++++++++++++++----
 src/python/test/test_dtm.py         |  8 ++++++++
 src/python/test/test_knn.py         |  6 ++++++
 3 files changed, 43 insertions(+), 4 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index f6870517..79362c09 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -36,6 +36,9 @@ class KNearestNeighbors:
             sort_results (bool): if True, then distances and indices of each point are
                 sorted on return, so that the first column contains the closest points.
                 Otherwise, neighbors are returned in an arbitrary order. Defaults to True.
+            enable_autodiff (bool): if the input is a torch.tensor, jax.numpy.array or similar, this instructs
+                the function to compute distances in a way that works with automatic differentiation.
+                This is experimental and not supported for all implementations.
             kwargs: additional parameters are forwarded to the backends.
         """
         self.k = k
@@ -202,13 +205,18 @@ class KNearestNeighbors:
         if self.params["implementation"] == "keops":
             import torch
             from pykeops.torch import LazyTensor
+            import eagerpy as ep
 
             # 'float64' is slow except on super expensive GPUs. Allow it with some param?
-            XX = torch.tensor(X, dtype=torch.float32)
-            if X is self.ref_points:
+            queries = X
+            X = ep.astensor(X)
+            XX = torch.as_tensor(X.numpy(), dtype=torch.float32)
+            if queries is self.ref_points:
+                Y = X
                 YY = XX
             else:
-                YY = torch.tensor(self.ref_points, dtype=torch.float32)
+                Y = ep.astensor(self.ref_points)
+                YY = torch.as_tensor(Y.numpy(), dtype=torch.float32)
 
             p = self.params["p"]
             if p == numpy.inf:
@@ -219,6 +227,24 @@ class KNearestNeighbors:
             else:
                 mat = ((LazyTensor(XX[:, None, :]) - LazyTensor(YY[None, :, :])).abs() ** p).sum(-1)
 
+            # pykeops does not support autodiff for kmin yet :-(
+            if self.params.get("enable_autodiff", False) and self.return_distance:
+                # Compute the indices of the neighbors, and recompute the relevant distances autodiff-friendly.
+                # Another strategy would be to compute the whole distance matrix with torch.cdist
+                # and use neighbors as indices into it.
+                neighbors = ep.astensor(mat.argKmin(k, dim=1)).numpy()
+                neighbor_pts = Y[neighbors]
+                diff = neighbor_pts - X[:, None, :]
+                if p == numpy.inf:
+                    distances = diff.abs().max(-1)
+                elif p == 2:
+                    distances = (diff ** 2).sum(-1) ** 0.5
+                else:
+                    distances = (diff.abs() ** p).sum(-1) ** (1.0 / p)
+                if self.return_index:
+                    return neighbors.raw, distances.raw
+                else:
+                    return distances.raw
             if self.return_index:
                 if self.return_distance:
                     distances, neighbors = mat.Kmin_argKmin(k, dim=1)
@@ -234,7 +260,6 @@ class KNearestNeighbors:
                     distances = distances ** (1.0 / p)
                 return distances
             return None
-        # FIXME: convert everything back to numpy arrays or not?
 
         if self.params["implementation"] == "ckdtree":
             qargs = {key: val for key, val in self.params.items() if key in {"p", "eps", "n_jobs"}}
diff --git a/src/python/test/test_dtm.py b/src/python/test/test_dtm.py
index bc0d3698..8709dd07 100755
--- a/src/python/test/test_dtm.py
+++ b/src/python/test/test_dtm.py
@@ -11,6 +11,7 @@
 from gudhi.point_cloud.dtm import DistanceToMeasure
 import numpy
 import pytest
+import torch
 
 
 def test_dtm_compare_euclidean():
@@ -39,6 +40,13 @@ def test_dtm_compare_euclidean():
     dtm = DistanceToMeasure(k, implementation="keops")
     r5 = dtm.fit_transform(pts)
     assert r5 == pytest.approx(r0)
+    pts2 = torch.tensor(pts, requires_grad=True)
+    assert pts2.grad is None
+    dtm = DistanceToMeasure(k, implementation="keops", enable_autodiff=True)
+    r6 = dtm.fit_transform(pts2)
+    assert r6.detach().numpy() == pytest.approx(r0)
+    r6.sum().backward()
+    assert pts2.grad is not None
 
 
 def test_dtm_precomputed():
diff --git a/src/python/test/test_knn.py b/src/python/test/test_knn.py
index 6269df54..415c9d48 100755
--- a/src/python/test/test_knn.py
+++ b/src/python/test/test_knn.py
@@ -32,6 +32,12 @@ def test_knn_explicit():
         .transform(query)
     )
     assert r == pytest.approx(np.array([[0.0, 1], [1, 1], [1, 2]]))
+    r = (
+        KNearestNeighbors(2, metric="chebyshev", return_distance=True, return_index=False, implementation="keops", enable_autodiff=True)
+        .fit(base)
+        .transform(query)
+    )
+    assert r == pytest.approx(np.array([[0.0, 1], [1, 1], [1, 2]]))
 
     knn = KNearestNeighbors(2, metric="minkowski", p=3, return_distance=False, return_index=True)
     knn.fit(base)
-- 
cgit v1.2.3


From 2f1576a23cf4ac055565875d384ca604c0ff6844 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Mon, 13 Apr 2020 15:01:51 +0200
Subject: Small autodiff tweaks

---
 src/python/gudhi/point_cloud/knn.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index 79362c09..ab3447d4 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -233,16 +233,17 @@ class KNearestNeighbors:
                 # Another strategy would be to compute the whole distance matrix with torch.cdist
                 # and use neighbors as indices into it.
                 neighbors = ep.astensor(mat.argKmin(k, dim=1)).numpy()
-                neighbor_pts = Y[neighbors]
+                # Work around https://github.com/pytorch/pytorch/issues/34452
+                neighbor_pts = Y[neighbors,]
                 diff = neighbor_pts - X[:, None, :]
                 if p == numpy.inf:
                     distances = diff.abs().max(-1)
                 elif p == 2:
-                    distances = (diff ** 2).sum(-1) ** 0.5
+                    distances = (diff ** 2).sum(-1).sqrt()
                 else:
                     distances = (diff.abs() ** p).sum(-1) ** (1.0 / p)
                 if self.return_index:
-                    return neighbors.raw, distances.raw
+                    return neighbors, distances.raw
                 else:
                     return distances.raw
             if self.return_index:
-- 
cgit v1.2.3


From 3a86402b733a48d9c25a4995325e72c7438c06c0 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Mon, 13 Apr 2020 15:21:06 +0200
Subject: Fix NaN gradient with pytorch

---
 src/python/gudhi/point_cloud/knn.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index ab3447d4..185a7764 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -236,12 +236,11 @@ class KNearestNeighbors:
                 # Work around https://github.com/pytorch/pytorch/issues/34452
                 neighbor_pts = Y[neighbors,]
                 diff = neighbor_pts - X[:, None, :]
-                if p == numpy.inf:
-                    distances = diff.abs().max(-1)
-                elif p == 2:
-                    distances = (diff ** 2).sum(-1).sqrt()
+                if isinstance(diff, ep.PyTorchTensor):
+                    # https://github.com/jonasrauber/eagerpy/issues/6
+                    distances = ep.astensor(diff.raw.norm(p, -1))
                 else:
-                    distances = (diff.abs() ** p).sum(-1) ** (1.0 / p)
+                    distances = diff.norms.lp(p, -1)
                 if self.return_index:
                     return neighbors, distances.raw
                 else:
-- 
cgit v1.2.3


From 3afce326428dddd638e22ab37ee4b2afe52eba75 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Mon, 13 Apr 2020 20:32:39 +0200
Subject: Generalize enable_autodiff to more implementations

Still limited to L^p
---
 src/python/gudhi/point_cloud/knn.py | 76 +++++++++++++++++++++++++++----------
 1 file changed, 55 insertions(+), 21 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index 185a7764..87b2798e 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -9,6 +9,7 @@
 
 import numpy
 
+# TODO: https://github.com/facebookresearch/faiss
 
 class KNearestNeighbors:
     """
@@ -67,6 +68,8 @@ class KNearestNeighbors:
                 self.params["implementation"] = "ckdtree"
             else:
                 self.params["implementation"] = "sklearn"
+        if not return_distance:
+            self.params["enable_autodiff"] = False
 
     def fit_transform(self, X, y=None):
         return self.fit(X).transform(X)
@@ -77,6 +80,10 @@ class KNearestNeighbors:
             X (numpy.array): coordinates for reference points.
         """
         self.ref_points = X
+        if self.params.get("enable_autodiff", False):
+            import eagerpy as ep
+            if self.params["implementation"] != "keops" or not isinstance(X, ep.PyTorchTensor):
+                X = ep.astensor(X).numpy()
         if self.params["implementation"] == "ckdtree":
             # sklearn could handle this, but it is much slower
             from scipy.spatial import cKDTree
@@ -113,6 +120,41 @@ class KNearestNeighbors:
         Args:
             X (numpy.array): coordinates for query points, or distance matrix if metric is "precomputed".
         """
+        if self.params.get("enable_autodiff", False):
+            # pykeops does not support autodiff for kmin yet, but when it does in the future,
+            # we may want a special path.
+            import eagerpy as ep
+            save_return_index = self.return_index
+            self.return_index = True
+            self.return_distance = False
+            self.params["enable_autodiff"] = False
+            try:
+                # FIXME: how do we test "X is ref_points" then?
+                newX = ep.astensor(X)
+                if self.params["implementation"] != "keops" or not isinstance(newX, ep.PyTorchTensor):
+                    newX = newX.numpy()
+                neighbors = self.transform(newX)
+            finally:
+                self.return_index = save_return_index
+                self.return_distance = True
+                self.params["enable_autodiff"] = True
+            # We can implement more later as needed
+            assert self.metric == "minkowski"
+            p = self.params["p"]
+            Y = ep.astensor(self.ref_points)
+            neighbor_pts = Y[neighbors,]
+            diff = neighbor_pts - X[:, None, :]
+            if isinstance(diff, ep.PyTorchTensor):
+                # https://github.com/jonasrauber/eagerpy/issues/6
+                distances = ep.astensor(diff.raw.norm(p, -1))
+            else:
+                distances = diff.norms.lp(p, -1)
+            if self.return_index:
+                return neighbors, distances.raw
+            else:
+                return distances.raw
+
+
         metric = self.metric
         k = self.k
 
@@ -207,16 +249,26 @@ class KNearestNeighbors:
             from pykeops.torch import LazyTensor
             import eagerpy as ep
 
-            # 'float64' is slow except on super expensive GPUs. Allow it with some param?
             queries = X
             X = ep.astensor(X)
-            XX = torch.as_tensor(X.numpy(), dtype=torch.float32)
+            if isinstance(X, ep.PyTorchTensor):
+                XX = X.raw
+            else:
+                # I don't know a clever way to reuse a GPU tensor from tensorflow in pytorch
+                # without copying to/from the CPU.
+                XX = X.numpy()
+            # 'float64' is slow except on super expensive GPUs. Allow it with some param?
+            XX = torch.as_tensor(XX, dtype=torch.float32)
             if queries is self.ref_points:
                 Y = X
                 YY = XX
             else:
                 Y = ep.astensor(self.ref_points)
-                YY = torch.as_tensor(Y.numpy(), dtype=torch.float32)
+                if isinstance(Y, ep.PyTorchTensor):
+                    YY = Y.raw
+                else:
+                    YY = Y.numpy()
+                YY = torch.as_tensor(YY, dtype=torch.float32)
 
             p = self.params["p"]
             if p == numpy.inf:
@@ -227,24 +279,6 @@ class KNearestNeighbors:
             else:
                 mat = ((LazyTensor(XX[:, None, :]) - LazyTensor(YY[None, :, :])).abs() ** p).sum(-1)
 
-            # pykeops does not support autodiff for kmin yet :-(
-            if self.params.get("enable_autodiff", False) and self.return_distance:
-                # Compute the indices of the neighbors, and recompute the relevant distances autodiff-friendly.
-                # Another strategy would be to compute the whole distance matrix with torch.cdist
-                # and use neighbors as indices into it.
-                neighbors = ep.astensor(mat.argKmin(k, dim=1)).numpy()
-                # Work around https://github.com/pytorch/pytorch/issues/34452
-                neighbor_pts = Y[neighbors,]
-                diff = neighbor_pts - X[:, None, :]
-                if isinstance(diff, ep.PyTorchTensor):
-                    # https://github.com/jonasrauber/eagerpy/issues/6
-                    distances = ep.astensor(diff.raw.norm(p, -1))
-                else:
-                    distances = diff.norms.lp(p, -1)
-                if self.return_index:
-                    return neighbors, distances.raw
-                else:
-                    return distances.raw
             if self.return_index:
                 if self.return_distance:
                     distances, neighbors = mat.Kmin_argKmin(k, dim=1)
-- 
cgit v1.2.3


From 521d8c17c2b7d71c46a51f0490ff2c13c809fc87 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Mon, 13 Apr 2020 21:13:19 +0200
Subject: Remove left-over code

eagerpy is only used with enable_autodiff
---
 src/python/gudhi/point_cloud/knn.py | 29 +++++++++--------------------
 1 file changed, 9 insertions(+), 20 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index 87b2798e..f2cddb38 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -82,8 +82,11 @@ class KNearestNeighbors:
         self.ref_points = X
         if self.params.get("enable_autodiff", False):
             import eagerpy as ep
+            X = ep.astensor(X)
             if self.params["implementation"] != "keops" or not isinstance(X, ep.PyTorchTensor):
-                X = ep.astensor(X).numpy()
+                # I don't know a clever way to reuse a GPU tensor from tensorflow in pytorch
+                # without copying to/from the CPU.
+                X = X.numpy()
         if self.params["implementation"] == "ckdtree":
             # sklearn could handle this, but it is much slower
             from scipy.spatial import cKDTree
@@ -133,6 +136,8 @@ class KNearestNeighbors:
                 newX = ep.astensor(X)
                 if self.params["implementation"] != "keops" or not isinstance(newX, ep.PyTorchTensor):
                     newX = newX.numpy()
+                else:
+                    newX = X
                 neighbors = self.transform(newX)
             finally:
                 self.return_index = save_return_index
@@ -247,29 +252,13 @@ class KNearestNeighbors:
         if self.params["implementation"] == "keops":
             import torch
             from pykeops.torch import LazyTensor
-            import eagerpy as ep
 
-            queries = X
-            X = ep.astensor(X)
-            if isinstance(X, ep.PyTorchTensor):
-                XX = X.raw
-            else:
-                # I don't know a clever way to reuse a GPU tensor from tensorflow in pytorch
-                # without copying to/from the CPU.
-                XX = X.numpy()
             # 'float64' is slow except on super expensive GPUs. Allow it with some param?
-            XX = torch.as_tensor(XX, dtype=torch.float32)
-            if queries is self.ref_points:
-                Y = X
+            XX = torch.as_tensor(X, dtype=torch.float32)
+            if X is self.ref_points:
                 YY = XX
             else:
-                Y = ep.astensor(self.ref_points)
-                if isinstance(Y, ep.PyTorchTensor):
-                    YY = Y.raw
-                else:
-                    YY = Y.numpy()
-                YY = torch.as_tensor(YY, dtype=torch.float32)
-
+                YY = torch.as_tensor(self.ref_points, dtype=torch.float32)
             p = self.params["p"]
             if p == numpy.inf:
                 # Requires pykeops 1.4 or later
-- 
cgit v1.2.3


From ce75f66da5a2d7ad2c479355112d48817c5ba68b Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Mon, 13 Apr 2020 21:38:24 +0200
Subject: Tweak to detect fit_transform

---
 src/python/gudhi/point_cloud/knn.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index f2cddb38..8b3cdb46 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -11,6 +11,7 @@ import numpy
 
 # TODO: https://github.com/facebookresearch/faiss
 
+
 class KNearestNeighbors:
     """
     Class wrapping several implementations for computing the k nearest neighbors in a point set.
@@ -82,6 +83,7 @@ class KNearestNeighbors:
         self.ref_points = X
         if self.params.get("enable_autodiff", False):
             import eagerpy as ep
+
             X = ep.astensor(X)
             if self.params["implementation"] != "keops" or not isinstance(X, ep.PyTorchTensor):
                 # I don't know a clever way to reuse a GPU tensor from tensorflow in pytorch
@@ -127,17 +129,19 @@ class KNearestNeighbors:
             # pykeops does not support autodiff for kmin yet, but when it does in the future,
             # we may want a special path.
             import eagerpy as ep
+
             save_return_index = self.return_index
             self.return_index = True
             self.return_distance = False
             self.params["enable_autodiff"] = False
             try:
-                # FIXME: how do we test "X is ref_points" then?
                 newX = ep.astensor(X)
-                if self.params["implementation"] != "keops" or not isinstance(newX, ep.PyTorchTensor):
+                if self.params["implementation"] != "keops" or (
+                    not isinstance(newX, ep.PyTorchTensor) and not isinstance(newX, ep.NumPyTensor)
+                ):
                     newX = newX.numpy()
                 else:
-                    newX = X
+                    newX = newX.raw
                 neighbors = self.transform(newX)
             finally:
                 self.return_index = save_return_index
@@ -159,7 +163,6 @@ class KNearestNeighbors:
             else:
                 return distances.raw
 
-
         metric = self.metric
         k = self.k
 
-- 
cgit v1.2.3


From 9518287cfa2a62948ede2e7d17d5c9f29092e0f4 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Tue, 14 Apr 2020 18:27:19 +0200
Subject: Doc improvements

---
 src/python/gudhi/point_cloud/dtm.py | 12 ++++++++++--
 src/python/gudhi/point_cloud/knn.py | 11 ++++++++---
 2 files changed, 18 insertions(+), 5 deletions(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/dtm.py b/src/python/gudhi/point_cloud/dtm.py
index 38368f29..58dec536 100644
--- a/src/python/gudhi/point_cloud/dtm.py
+++ b/src/python/gudhi/point_cloud/dtm.py
@@ -20,7 +20,9 @@ class DistanceToMeasure:
         Args:
             k (int): number of neighbors (possibly including the point itself).
             q (float): order used to compute the distance to measure. Defaults to 2.
-            kwargs: same parameters as :class:`~gudhi.point_cloud.knn.KNearestNeighbors`, except that metric="neighbors" means that :func:`transform` expects an array with the distances to the k nearest neighbors.
+            kwargs: same parameters as :class:`~gudhi.point_cloud.knn.KNearestNeighbors`, except that
+                metric="neighbors" means that :func:`transform` expects an array with the distances
+                to the k nearest neighbors.
         """
         self.k = k
         self.q = q
@@ -44,7 +46,13 @@ class DistanceToMeasure:
     def transform(self, X):
         """
         Args:
-            X (numpy.array): coordinates for query points, or distance matrix if metric is "precomputed", or distances to the k nearest neighbors if metric is "neighbors" (if the array has more than k columns, the remaining ones are ignored).
+            X (numpy.array): coordinates for query points, or distance matrix if metric is "precomputed",
+                or distances to the k nearest neighbors if metric is "neighbors" (if the array has more
+                than k columns, the remaining ones are ignored).
+
+        Returns:
+            numpy.array: a 1-d array with, for each point of X, its distance to the measure defined
+            by the argument of :func:`fit`.
         """
         if self.params["metric"] == "neighbors":
             distances = X[:, : self.k]
diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index 8b3cdb46..d7cf0b2a 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -38,9 +38,9 @@ class KNearestNeighbors:
             sort_results (bool): if True, then distances and indices of each point are
                 sorted on return, so that the first column contains the closest points.
                 Otherwise, neighbors are returned in an arbitrary order. Defaults to True.
-            enable_autodiff (bool): if the input is a torch.tensor, jax.numpy.array or similar, this instructs
-                the function to compute distances in a way that works with automatic differentiation.
-                This is experimental and not supported for all implementations.
+            enable_autodiff (bool): if the input is a torch.tensor, jax.numpy.ndarray or tensorflow.Tensor, this
+                instructs the function to compute distances in a way that works with automatic differentiation.
+                This is experimental and not supported for all metrics. Defaults to False.
             kwargs: additional parameters are forwarded to the backends.
         """
         self.k = k
@@ -124,6 +124,11 @@ class KNearestNeighbors:
         """
         Args:
             X (numpy.array): coordinates for query points, or distance matrix if metric is "precomputed".
+
+        Returns:
+            numpy.array: if return_index, an array of shape (len(X), k) with the indices (in the argument
+            of :func:`fit`) of the k nearest neighbors to the points of X. If return_distance, an array of the
+            same shape with the distances to those neighbors. If both, a tuple with the two arrays, in this order.
         """
         if self.params.get("enable_autodiff", False):
             # pykeops does not support autodiff for kmin yet, but when it does in the future,
-- 
cgit v1.2.3


From 1c1a99074049e4ff04fa28e7d6e1b6fc2067397a Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Mon, 20 Apr 2020 10:38:41 +0200
Subject: Add __license__

---
 src/python/gudhi/point_cloud/dtm.py | 4 ++++
 src/python/gudhi/point_cloud/knn.py | 8 +++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/gudhi/point_cloud/dtm.py b/src/python/gudhi/point_cloud/dtm.py
index 58dec536..13e16d24 100644
--- a/src/python/gudhi/point_cloud/dtm.py
+++ b/src/python/gudhi/point_cloud/dtm.py
@@ -9,6 +9,10 @@
 
 from .knn import KNearestNeighbors
 
+__author__ = "Marc Glisse"
+__copyright__ = "Copyright (C) 2020 Inria"
+__license__ = "MIT"
+
 
 class DistanceToMeasure:
     """
diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index d7cf0b2a..4017e498 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -11,6 +11,10 @@ import numpy
 
 # TODO: https://github.com/facebookresearch/faiss
 
+__author__ = "Marc Glisse"
+__copyright__ = "Copyright (C) 2020 Inria"
+__license__ = "MIT"
+
 
 class KNearestNeighbors:
     """
@@ -156,7 +160,9 @@ class KNearestNeighbors:
             assert self.metric == "minkowski"
             p = self.params["p"]
             Y = ep.astensor(self.ref_points)
-            neighbor_pts = Y[neighbors,]
+            neighbor_pts = Y[
+                neighbors,
+            ]
             diff = neighbor_pts - X[:, None, :]
             if isinstance(diff, ep.PyTorchTensor):
                 # https://github.com/jonasrauber/eagerpy/issues/6
-- 
cgit v1.2.3


From aa90b98bee73ab2aaf39ef91f39f5a750168e5d4 Mon Sep 17 00:00:00 2001
From: Marc Glisse <marc.glisse@inria.fr>
Date: Wed, 22 Apr 2020 13:04:15 +0200
Subject: Document several optional dependencies of knn

---
 src/python/doc/installation.rst     | 28 ++++++++++++++++++++++++++++
 src/python/gudhi/point_cloud/knn.py |  3 ++-
 2 files changed, 30 insertions(+), 1 deletion(-)

(limited to 'src/python/gudhi/point_cloud')

diff --git a/src/python/doc/installation.rst b/src/python/doc/installation.rst
index 48425d5e..09a843d5 100644
--- a/src/python/doc/installation.rst
+++ b/src/python/doc/installation.rst
@@ -211,6 +211,14 @@ The following examples requires CGAL version ≥ 4.11.0:
     * :download:`euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py <../example/euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py>`
     * :download:`euclidean_witness_complex_diagram_persistence_from_off_file_example.py <../example/euclidean_witness_complex_diagram_persistence_from_off_file_example.py>`
 
+EagerPy
+=======
+
+Some Python functions can handle automatic differentiation (possibly only when
+a flag `enable_autodiff=True` is used). In order to reduce code duplication, we
+use `EagerPy <https://eagerpy.jonasrauber.de/>`_ which wraps arrays from
+PyTorch, TensorFlow and JAX in a common interface.
+
 Eigen
 =====
 
@@ -229,6 +237,13 @@ The following examples require `Eigen <http://eigen.tuxfamily.org/>`_ version 
     * :download:`euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py <../example/euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py>`
     * :download:`euclidean_witness_complex_diagram_persistence_from_off_file_example.py <../example/euclidean_witness_complex_diagram_persistence_from_off_file_example.py>`
 
+Hnswlib
+=======
+
+:class:`~gudhi.point_cloud.knn.KNearestNeighbors` can use the Python package
+`Hnswlib <https://github.com/nmslib/hnswlib>`_ as a backend if explicitly
+requested, to speed-up queries.
+
 Matplotlib
 ==========
 
@@ -251,6 +266,13 @@ The following examples require the `Matplotlib <http://matplotlib.org>`_:
     * :download:`euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py <../example/euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py>`
     * :download:`euclidean_witness_complex_diagram_persistence_from_off_file_example.py <../example/euclidean_witness_complex_diagram_persistence_from_off_file_example.py>`
 
+PyKeOps
+=======
+
+:class:`~gudhi.point_cloud.knn.KNearestNeighbors` can use the Python package
+`PyKeOps <https://www.kernel-operations.io/keops/python/>`_ as a backend if
+explicitly requested, to speed-up queries using a GPU.
+
 Python Optimal Transport
 ========================
 
@@ -258,6 +280,12 @@ The :doc:`Wasserstein distance </wasserstein_distance_user>`
 module requires `POT <https://pot.readthedocs.io/>`_, a library that provides
 several solvers for optimization problems related to Optimal Transport.
 
+PyTorch
+=======
+
+`PyTorch <https://pytorch.org/>`_ is currently only used as a dependency of
+`PyKeOps`_, and in some tests.
+
 Scikit-learn
 ============
 
diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py
index 4017e498..07553d6d 100644
--- a/src/python/gudhi/point_cloud/knn.py
+++ b/src/python/gudhi/point_cloud/knn.py
@@ -44,7 +44,8 @@ class KNearestNeighbors:
                 Otherwise, neighbors are returned in an arbitrary order. Defaults to True.
             enable_autodiff (bool): if the input is a torch.tensor, jax.numpy.ndarray or tensorflow.Tensor, this
                 instructs the function to compute distances in a way that works with automatic differentiation.
-                This is experimental and not supported for all metrics. Defaults to False.
+                This is experimental, not supported for all metrics, and requires the package EagerPy.
+                Defaults to False.
             kwargs: additional parameters are forwarded to the backends.
         """
         self.k = k
-- 
cgit v1.2.3