Add different metrics to compute KNN

pymde/preprocess/data_matrix.py

@@ -88,7 +88,7 @@ def distances(data, retain_fraction=1.0, verbose=False):
     return Graph.from_edges(edges, delta, n_items=n_items)
 
 
-def k_nearest_neighbors(data, k, max_distance=None, verbose=False):
+def k_nearest_neighbors(data, k, metric='euclidean', max_distance=None, verbose=False):
     """Compute k-nearest neighbors for each row in data matrix.
 
     Computes the k-nearest neighbor graph of data matrix, under
@@ -101,6 +101,8 @@ def k_nearest_neighbors(data, k, max_distance=None, verbose=False):
         The data matrix
     k: int
         The number of nearest neighbors per item
+    metric: str (optional)
+        Metric to be used to compute the KNN.
     max_distance: float (optional)
         If not None, neighborhoods are restricted to have a radius
         no greater than `max_distance`.
@@ -128,14 +130,17 @@ def k_nearest_neighbors(data, k, max_distance=None, verbose=False):
         if verbose:
             problem.LOGGER.info("Exact nearest neighbors by brute force ")
         nn = sklearn.neighbors.NearestNeighbors(
-            n_neighbors=k + 1, algorithm="brute"
+            n_neighbors=k + 1,
+            metric=metric,
+            algorithm="brute"
         )
         nn.fit(data)
         distances, neighbors = nn.kneighbors(data)
     else:
         # TODO default params (n_trees, max_candidates)
         index = pynndescent.NNDescent(
             data,
+            metric=metric,
             n_neighbors=k + 1,
             verbose=verbose,
             max_candidates=60,

pymde/preprocess/generic.py

@@ -57,7 +57,7 @@ def distances(data, retain_fraction=1.0, verbose=False):
         )
 
 
-def k_nearest_neighbors(data, k, max_distance=None, verbose=False):
+def k_nearest_neighbors(data, k, metric='euclidean', max_distance=None, verbose=False):
     """Compute k-nearest neighbors, given data matrix or graph.
 
     This function computes a k-nearest neighbor graph, given either
@@ -81,6 +81,9 @@ def k_nearest_neighbors(data, k, max_distance=None, verbose=False):
         A data matrix, shape ``(n_items, n_features)``, or a ``pymde.Graph``
     k: int
         The number of nearest neighbors per item
+    metric: str (optional)
+        Distance metric to be used to compute the KNN.
+        It is only used for data matrix data.
     max_distance: float (optional)
         If not ``None``, neighborhoods are restricted to have a radius
         no greater than `max_distance`.
@@ -96,7 +99,7 @@ def k_nearest_neighbors(data, k, max_distance=None, verbose=False):
 
     if _is_data_matrix(data):
         return data_matrix.k_nearest_neighbors(
-            data, k=k, max_distance=max_distance, verbose=verbose
+            data, k=k, metric=metric, max_distance=max_distance, verbose=verbose
         )
     else:
         return graph.k_nearest_neighbors(

pymde/recipes.py

@@ -226,6 +226,7 @@ def preserve_neighbors(
     constraint=None,
     n_neighbors=None,
     repulsive_fraction=None,
+    metric="euclidean",
     max_distance=None,
     init="quadratic",
     device="cpu",
@@ -272,6 +273,9 @@ def preserve_neighbors(
         Callable that constructs a distortion function, given negative
         weights. (If ``None``, only positive weights are used.) For example,
         ``pymde.penalties.Log`` or ``pymde.penalties.InversePower``.
+    metric: str
+        Distance metric to be used to compute the KNN.
+        It is only used when the data is a data matrix.
     constraint: pymde.constraints.Constraint (optional)
         Embedding constraint, like ``pymde.Standardized()`` or
         ``pymde.Anchored(anchors, values)`` (or ``None``). Defaults to no
@@ -340,13 +344,15 @@ def preserve_neighbors(
 
     if verbose:
         problem.LOGGER.info(
-            f"Computing {n_neighbors}-nearest neighbors, with "
+            f"Computing {n_neighbors}-nearest neighbors, "
+            f"using {metric} metric with "
             f"max_distance={max_distance}"
         )
 
     knn_graph = preprocess.generic.k_nearest_neighbors(
         data,
         k=n_neighbors,
+        metric=metric,
         max_distance=max_distance,
         verbose=verbose,
     )