add tests with fixes

Jelmer Bot · Jelmer Bot · commit d1eff712eb83 · 2024-04-05T13:17:07.000+02:00
diff --git a/.gitignore b/.gitignore
@@ -65,3 +65,5 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+
+*.code-workspace
diff --git a/multi_mst/k_mst/k_mst.py b/multi_mst/k_mst/k_mst.py
@@ -1,6 +1,7 @@
 import numpy as np
 import warnings as warn
 from umap import UMAP
+from scipy.sparse import coo_matrix
 from sklearn.neighbors import KDTree
 from sklearn.utils import check_array
 from sklearn.base import BaseEstimator
@@ -29,7 +30,7 @@ def validate_parameters(data, num_neighbors, min_samples, epsilon):
     return data, epsilon
 
 
-def kMST(data, num_neighbors=5, min_samples=1, epsilon=None, umap_kwargs=None):
+def kMST(data, num_neighbors=3, min_samples=1, epsilon=None, umap_kwargs=None):
     """
     Computes a $k$-MST for the given data. Adapts the boruvka algorithm to look
     for $k$ candidate edges per point, of which the $k$ best per connected
@@ -43,7 +44,7 @@ def kMST(data, num_neighbors=5, min_samples=1, epsilon=None, umap_kwargs=None):
     data: array-like
         The data to construct a MST for.
     num_neighbors: int, optional
-        The number of edges to connect between each fragement. Default is 5.
+        The number of edges to connect between each fragement. Default is 3.
     min_samples: int, optional
         The number of neighbors for computing the mutual reachability distance.
         Value must be lower or equal to the number of neighbors. `epsilon`
@@ -98,7 +99,7 @@ class KMST(BaseEstimator):
     Parameters
     ----------
     num_neighbors: int, optional
-        The number of edges to connect between each fragement. Default is 5.
+        The number of edges to connect between each fragement. Default is 35.
     min_samples: int, optional
         The number of neighbors for computing the mutual reachability distance.
         Value must be lower or equal to the number of neighbors. `epsilon`
@@ -130,11 +131,12 @@ class KMST(BaseEstimator):
         missing values. Use the graph_ and embedding_ attributes instead!
     """
 
-    def __init__(self, *, num_neighbors=5, min_samples=1, epsilon=None, **umap_kwargs):
+    def __init__(self, *, num_neighbors=3, min_samples=1, epsilon=None, **umap_kwargs):
         self.num_neighbors = num_neighbors
         self.min_samples = min_samples
         self.epsilon = epsilon
         self.umap_kwargs = umap_kwargs
+        # TODO: Sklearn does not support **kwargs in __init__ for BaseEstimator...
 
     def fit(self, X, y=None, **fit_params):
         """
@@ -170,21 +172,29 @@ def fit(self, X, y=None, **fit_params):
             clean_data = X
 
         kwargs = self.get_params()
-        self.mst_indices_, self.mst_distances_, self._umap = kMST(clean_data, **kwargs)
+        print(kwargs, self.umap_kwargs)
+        self.mst_indices_, self.mst_distances_, self._umap = kMST(
+            clean_data, umap_kwargs=self.umap_kwargs, **kwargs
+        )
         self.graph_ = self._umap.graph_.copy()
         self.embedding_ = (
-            self._umap.embedding_.copy() if self._umap.embedding_ is not None else None
+            self._umap.embedding_.copy() if hasattr(self._umap, "embedding_") else None
         )
 
         if not self._all_finite:
-            self.graph_.shape = (X.shape[0], X.shape[0])
+            self.graph_ = self.graph_.tocoo()
             for i in range(len(self.graph_.data)):
                 self.graph_.row[i] = internal_to_raw[self.graph_.row[i]]
                 self.graph_.col[i] = internal_to_raw[self.graph_.col[i]]
+            self.graph_ = coo_matrix(
+                (self.graph_.data, (self.graph_.row, self.graph_.col)),
+                shape=(X.shape[0], X.shape[0]),
+            )
 
-            new_embedding = np.full((X.shape[0], self.num_components), np.nan)
-            new_embedding[finite_index] = self.embedding_
-            self.embedding_ = new_embedding
+            if self.embedding_ is not None:
+                new_embedding = np.full((X.shape[0], self.embedding_.shape[1]), np.nan)
+                new_embedding[finite_index] = self.embedding_
+                self.embedding_ = new_embedding
 
             new_indices = np.full(
                 (X.shape[0], self.mst_indices_.shape[1]), -1, dtype=np.int32
diff --git a/multi_mst/noisy_mst/noisy_mst.py b/multi_mst/noisy_mst/noisy_mst.py
@@ -1,6 +1,7 @@
 import numpy as np
 import warnings as warn
 from umap import UMAP
+from scipy.sparse import coo_matrix
 from sklearn.neighbors import KDTree
 from sklearn.utils import check_array
 from sklearn.base import BaseEstimator
@@ -29,7 +30,7 @@ def validate_parameters(data, num_trees, noise_fraction, min_samples):
     return data
 
 
-def noisyMST(data, num_trees=5, noise_fraction=0.1, min_samples=1, umap_kwargs=None):
+def noisyMST(data, num_trees=3, noise_fraction=0.1, min_samples=1, umap_kwargs=None):
     """
     Computes a union of $k$ noisy MSTs for the given data. Adapts the boruvka
     algorithm construct multiple noisy miminum spanning trees.
@@ -45,7 +46,7 @@ def noisyMST(data, num_trees=5, noise_fraction=0.1, min_samples=1, umap_kwargs=N
         Adds Gaussian noise with scale=noise_fraction * distance to every computed
         distance value.
     num_trees: int, optional
-        The number of noisy MSTS to create. Default is 5.
+        The number of noisy MSTS to create. Default is 3.
     min_samples: int, optional
         The number of neighbors for computing the mutual reachability distance.
         Value must be lower or equal to the number of neighbors. `epsilon`
@@ -96,7 +97,7 @@ class NoisyMST(BaseEstimator):
     Parameters
     ----------
     num_trees: int, optional
-        The number of minimum spanning trees created. Default is 5.
+        The number of minimum spanning trees created. Default is 3.
     noise_fraction:
         Adds Gaussian noise with scale=noise_fraction * distance to every computed
         distance value.
@@ -129,12 +130,13 @@ class NoisyMST(BaseEstimator):
     """
 
     def __init__(
-        self, *, num_trees=5, noise_fraction=0.1, min_samples=1, **umap_kwargs
+        self, *, num_trees=3, noise_fraction=0.1, min_samples=1, **umap_kwargs
     ):
         self.num_trees = num_trees
         self.noise_fraction = noise_fraction
         self.min_samples = min_samples
         self.umap_kwargs = umap_kwargs
+        # TODO: Sklearn does not support **kwargs in __init__ for BaseEstimator...
 
     def fit(self, X, y=None, **fit_params):
         """
@@ -171,22 +173,27 @@ def fit(self, X, y=None, **fit_params):
 
         kwargs = self.get_params()
         self.mst_indices_, self.mst_distances_, self._umap = noisyMST(
-            clean_data, **kwargs
+            clean_data, umap_kwargs=self.umap_kwargs, **kwargs
         )
         self.graph_ = self._umap.graph_.copy()
         self.embedding_ = (
-            self._umap.embedding_.copy() if self._umap.embedding_ is not None else None
+            self._umap.embedding_.copy() if hasattr(self._umap, "embedding_") else None
         )
 
         if not self._all_finite:
-            self.graph_.shape = (X.shape[0], X.shape[0])
+            self.graph_ = self.graph_.tocoo()
             for i in range(len(self.graph_.data)):
                 self.graph_.row[i] = internal_to_raw[self.graph_.row[i]]
                 self.graph_.col[i] = internal_to_raw[self.graph_.col[i]]
+            self.graph_ = coo_matrix(
+                (self.graph_.data, (self.graph_.row, self.graph_.col)),
+                shape=(X.shape[0], X.shape[0]),
+            )
 
-            new_embedding = np.full((X.shape[0], self.num_components), np.nan)
-            new_embedding[finite_index] = self.embedding_
-            self.embedding_ = new_embedding
+            if self.embedding_ is not None:
+                new_embedding = np.full((X.shape[0], self.embedding_.shape[1]), np.nan)
+                new_embedding[finite_index] = self.embedding_
+                self.embedding_ = new_embedding
 
             new_indices = np.full(
                 (X.shape[0], self.mst_indices_.shape[1]), -1, dtype=np.int32
diff --git a/tests/test_k_mst.py b/tests/test_k_mst.py
@@ -0,0 +1,161 @@
+"""Tests for MultiMST"""
+import pytest
+import numpy as np
+from sklearn.datasets import make_blobs
+from sklearn.preprocessing import StandardScaler
+from sklearn.utils._testing import assert_raises
+from scipy.sparse.csgraph import connected_components
+
+from umap import UMAP
+from multi_mst.k_mst import kMST, KMST
+
+
+def generate_noisy_data():
+    blobs, yBlobs = make_blobs(
+        n_samples=50,
+        centers=[(-0.75, 2.25), (2.0, -0.5)],
+        cluster_std=0.2,
+        random_state=3,
+    )
+    np.random.seed(5)
+    noise = np.random.uniform(-1.0, 3.0, (50, 2))
+    yNoise = np.full(50, -1)
+    return (
+        np.vstack((blobs, noise)),
+        np.concatenate((yBlobs, yNoise)),
+    )
+
+
+X, y = generate_noisy_data()
+X = StandardScaler().fit_transform(X)
+
+X_missing_data = X.copy()
+X_missing_data[0] = [np.nan, 1]
+X_missing_data[6] = [np.nan, np.nan]
+
+
+def test_badargs():
+    """Tests parameter validation."""
+    assert_raises(ValueError, kMST, X, num_neighbors=1.0)
+    assert_raises(ValueError, kMST, X, num_neighbors=-1)
+    assert_raises(ValueError, kMST, X, epsilon=1)
+    assert_raises(ValueError, kMST, X, epsilon=0.6)
+    assert_raises(ValueError, kMST, X, epsilon=-0.4)
+    assert_raises(ValueError, kMST, X, min_samples=1.0)
+    assert_raises(ValueError, kMST, X, num_neighbors=3, min_samples=4)
+    assert_raises(ValueError, kMST, X, min_samples=0)
+    assert_raises(ValueError, kMST, X, min_samples=-1)
+
+
+def test_defaults():
+    """Tests with default parameters."""
+    p = KMST()
+    
+    embedding = p.fit_transform(X)
+    assert embedding.shape[0] == X.shape[0]
+    assert embedding.shape[1] == 2 # Default num_components
+    assert np.issubdtype(embedding.dtype, np.floating)
+
+    assert p.mst_indices_.shape[0] == X.shape[0]
+    assert p.mst_indices_.shape[1] >= 3 # Default num_neighbors
+    assert np.issubdtype(p.mst_indices_.dtype, np.integer)
+    assert p.mst_distances_.shape[0] == X.shape[0]
+    assert p.mst_distances_.shape[1] >= 3 # Default num_neighbors
+    assert np.issubdtype(p.mst_distances_.dtype, np.floating)
+    assert p.graph_.shape[0] == X.shape[0]
+    assert p.graph_.shape[1] == X.shape[0]
+    assert p.embedding_.shape[0] == X.shape[0]
+    assert p.embedding_.shape[1] == 2 # Default num_components
+    assert np.issubdtype(p.embedding_.dtype, np.floating)
+    assert isinstance(p._umap, UMAP)
+    assert connected_components(p._umap.graph_, directed=False, return_labels=False) == 1
+
+
+def test_with_missing_data():
+    """Tests with nan data."""
+    clean_indices = list(range(1, 6)) + list(range(7, X.shape[0]))
+    model = KMST().fit(X_missing_data)
+    clean_model = KMST().fit(X_missing_data[clean_indices])
+    
+    assert np.all(model.mst_indices_[0, :] == -1)
+    assert np.isinf(model.mst_distances_[0, :]).all()
+    assert np.isnan(model.embedding_[0, :]).all()
+    assert np.all(model.graph_.row != 0) & np.all(model.graph_.col != 0)
+
+    assert np.all(model.mst_indices_[6, :] == -1)
+    assert np.isinf(model.mst_distances_[6, :]).all()
+    assert np.isnan(model.embedding_[6, :]).all()
+    assert np.all(model.graph_.row != 6) & np.all(model.graph_.col != 6)
+
+    assert np.allclose(clean_model.graph_.data, model.graph_.data)
+    assert np.allclose(clean_model.mst_indices_, model.mst_indices_[clean_indices])
+    assert np.allclose(clean_model.mst_distances_, model.mst_distances_[clean_indices])
+
+
+def test_with_missing_data_graph_mode():
+    """Tests with nan data."""
+    clean_indices = list(range(1, 6)) + list(range(7, X.shape[0]))
+    model = KMST(transform_mode='graph').fit(X_missing_data)
+    clean_model = KMST(transform_mode='graph').fit(X_missing_data[clean_indices])
+    
+    assert np.all(model.mst_indices_[0, :] == -1)
+    assert np.isinf(model.mst_distances_[0, :]).all()
+    assert np.all(model.graph_.row != 0) & np.all(model.graph_.col != 0)
+
+    assert np.all(model.mst_indices_[6, :] == -1)
+    assert np.isinf(model.mst_distances_[6, :]).all()
+    assert np.all(model.graph_.row != 6) & np.all(model.graph_.col != 6)
+
+    assert np.allclose(clean_model.graph_.data, model.graph_.data)
+    assert np.allclose(clean_model.mst_indices_, model.mst_indices_[clean_indices])
+    assert np.allclose(clean_model.mst_distances_, model.mst_distances_[clean_indices])
+
+
+def test_graph_mode():
+    """Tests with transform_mode='graph'."""
+    p = KMST(transform_mode='graph')
+    embedding = p.fit_transform(X)
+    
+    assert embedding is None
+    assert p.embedding_ is None
+    
+
+def test_min_samples():
+    """Tests with higher min_samples."""
+    p = KMST(min_samples=3, transform_mode='graph').fit(X)
+    
+    assert p.mst_indices_.shape[1] >= 3 # Max(min_samples, num_neighbors)
+    assert p.mst_distances_.shape[1] >= 3 # Max(min_samples, num_neighbors)
+    assert connected_components(p._umap.graph_, directed=False, return_labels=False) == 1
+
+
+def test_num_neighbors():
+    """Tests with lower num_neighbors."""
+    p = KMST(num_neighbors=1, transform_mode='graph').fit(X)
+    
+    assert p.mst_indices_.shape[1] >= 1 # Max(min_samples, num_neighbors)
+    assert p.mst_distances_.shape[1] >= 1 # Max(min_samples, num_neighbors)
+    assert connected_components(p._umap.graph_, directed=False, return_labels=False) == 1
+
+
+def test_epsilon():
+    """Tests with epsilon."""
+    base = KMST(transform_mode='graph').fit(X)
+    p = KMST(epsilon=1.1, transform_mode='graph').fit(X)
+
+    assert connected_components(p._umap.graph_, directed=False, return_labels=False) == 1
+    assert p.graph_.nnz < base.graph_.nnz
+
+
+def test_num_components():
+    """Tests with higher num_components."""
+    p = KMST(n_components=3)
+    
+    embedding = p.fit_transform(X)
+    assert embedding.shape[0] == X.shape[0]
+    assert embedding.shape[1] == 3
+    assert np.issubdtype(embedding.dtype, np.floating)
+
+    assert p.embedding_.shape[0] == X.shape[0]
+    assert p.embedding_.shape[1] == 3
+    assert np.issubdtype(p.embedding_.dtype, np.floating)
diff --git a/tests/test_noisy_mst.py b/tests/test_noisy_mst.py
