scikit-network 0.28.3__cp39-cp39-macosx_12_0_arm64.whl

sknetwork/clustering/louvain_core.pyx

@@ -0,0 +1,134 @@
+# distutils: language = c++
+# cython: language_level=3
+# cython: linetrace=True
+# distutils: define_macros=CYTHON_TRACE_NOGIL=1
+from libcpp.set cimport set
+from libcpp.vector cimport vector
+cimport cython
+
+ctypedef fused int_or_long:
+    int
+    long
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def fit_core(float resolution, float tol, float[:] ou_node_probs, float[:] in_node_probs, float[:] self_loops,
+             float[:] data, int_or_long[:] indices, int_or_long[:] indptr):  # pragma: no cover
+    """Fit the clusters to the objective function.
+
+    Parameters
+    ----------
+    resolution :
+        Resolution parameter (positive).
+    tol :
+        Minimum increase in modularity to enter a new optimization pass.
+    ou_node_probs :
+        Distribution of node weights based on their out-edges (sums to 1).
+    in_node_probs :
+        Distribution of node weights based on their in-edges (sums to 1).
+    self_loops :
+        Weights of self loops.
+    data :
+        CSR format data array of the normalized adjacency matrix.
+    indices :
+        CSR format index array of the normalized adjacency matrix.
+    indptr :
+        CSR format index pointer array of the normalized adjacency matrix.
+
+    Returns
+    -------
+    labels :
+        Cluster index of each node.
+    total_increase :
+        Score of the clustering (total increase in modularity).
+    """
+    cdef int_or_long n = indptr.shape[0] - 1
+    cdef int_or_long increase = 1
+    cdef int_or_long cluster
+    cdef int_or_long cluster_best
+    cdef int_or_long cluster_node
+    cdef int_or_long i
+    cdef int_or_long j
+    cdef int_or_long j1
+    cdef int_or_long j2
+    cdef int_or_long label
+
+    cdef float increase_total = 0
+    cdef float increase_pass
+    cdef float delta
+    cdef float delta_best
+    cdef float delta_exit
+    cdef float delta_local
+    cdef float node_prob_in
+    cdef float node_prob_ou
+    cdef float ratio_in
+    cdef float ratio_ou
+
+    cdef vector[int_or_long] labels
+    cdef vector[float] neighbor_clusters_weights
+    cdef vector[float] ou_clusters_weights
+    cdef vector[float] in_clusters_weights
+    cdef set[int_or_long] unique_clusters = ()
+
+    for i in range(n):
+        labels.push_back(i)
+        neighbor_clusters_weights.push_back(0.)
+        ou_clusters_weights.push_back(ou_node_probs[i])
+        in_clusters_weights.push_back(in_node_probs[i])
+
+    while increase == 1:
+        increase = 0
+        increase_pass = 0
+
+        for i in range(n):
+            unique_clusters.clear()
+            cluster_node = labels[i]
+            j1 = indptr[i]
+            j2 = indptr[i + 1]
+
+            for j in range(j1, j2):
+                label = labels[indices[j]]
+                neighbor_clusters_weights[label] += data[j]
+                unique_clusters.insert(label)
+
+            unique_clusters.erase(cluster_node)
+
+            if not unique_clusters.empty():
+                node_prob_ou = ou_node_probs[i]
+                node_prob_in = in_node_probs[i]
+                ratio_ou = resolution * node_prob_ou
+                ratio_in = resolution * node_prob_in
+
+                delta_exit = 2 * (neighbor_clusters_weights[cluster_node] - self_loops[i])
+                delta_exit -= ratio_ou * (in_clusters_weights[cluster_node] - node_prob_in)
+                delta_exit -= ratio_in * (ou_clusters_weights[cluster_node] - node_prob_ou)
+
+                delta_best = 0
+                cluster_best = cluster_node
+
+                for cluster in unique_clusters:
+                    delta = 2 * neighbor_clusters_weights[cluster]
+                    delta -= ratio_ou * in_clusters_weights[cluster]
+                    delta -= ratio_in * ou_clusters_weights[cluster]
+
+                    delta_local = delta - delta_exit
+                    if delta_local > delta_best:
+                        delta_best = delta_local
+                        cluster_best = cluster
+
+                    neighbor_clusters_weights[cluster] = 0
+
+                if delta_best > 0:
+                    increase_pass += delta_best
+                    ou_clusters_weights[cluster_node] -= node_prob_ou
+                    in_clusters_weights[cluster_node] -= node_prob_in
+                    ou_clusters_weights[cluster_best] += node_prob_ou
+                    in_clusters_weights[cluster_best] += node_prob_in
+                    labels[i] = cluster_best
+
+            neighbor_clusters_weights[cluster_node] = 0
+
+        increase_total += increase_pass
+        if increase_pass > tol:
+            increase = 1
+    return labels, increase_total

sknetwork/clustering/metrics.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created in July 2018
+@author: Nathan de Lara <nathan.delara@polytechnique.org>
+@author: Thomas Bonald <bonald@enst.fr>
+"""
+from typing import Optional, Union, Tuple
+
+import numpy as np
+from scipy import sparse
+
+from sknetwork.utils.check import get_probs
+from sknetwork.utils.format import get_adjacency
+from sknetwork.utils.membership import get_membership
+
+
+def get_modularity(input_matrix: Union[sparse.csr_matrix, np.ndarray], labels: np.ndarray,
+                   labels_col: Optional[np.ndarray] = None, weights: str = 'degree',
+                   resolution: float = 1, return_all: bool = False) -> Union[float, Tuple[float, float, float]]:
+    """Modularity of a clustering.
+
+    The modularity of a clustering is
+
+    :math:`Q = \\dfrac{1}{w} \\sum_{i,j}\\left(A_{ij} - \\gamma \\dfrac{w_iw_j}{w}\\right)\\delta_{c_i,c_j}`
+    for graphs,
+
+    :math:`Q = \\dfrac{1}{w} \\sum_{i,j}\\left(A_{ij} - \\gamma \\dfrac{d^+_id^-_j}{w}\\right)\\delta_{c_i,c_j}`
+    for directed graphs,
+
+    where
+
+    * :math:`c_i` is the cluster of node :math:`i`,\n
+    * :math:`w_i` is the weight of node :math:`i`,\n
+    * :math:`w^+_i, w^-_i` are the out-weight, in-weight of node :math:`i` (for directed graphs),\n
+    * :math:`w = 1^TA1` is the total weight,\n
+    * :math:`\\delta` is the Kronecker symbol,\n
+    * :math:`\\gamma \\ge 0` is the resolution parameter.
+
+    Parameters
+    ----------
+    input_matrix :
+        Adjacency matrix or biadjacency matrix of the graph.
+    labels :
+        Labels of nodes.
+    labels_col :
+        Labels of column nodes (for bipartite graphs).
+    weights :
+        Weighting of nodes (``'degree'`` (default) or ``'uniform'``).
+    resolution:
+        Resolution parameter (default = 1).
+    return_all:
+        If ``True``, return modularity, fit, diversity.
+
+    Returns
+    -------
+    modularity : float
+    fit: float, optional
+    diversity: float, optional
+
+    Example
+    -------
+    >>> from sknetwork.clustering import get_modularity
+    >>> from sknetwork.data import house
+    >>> adjacency = house()
+    >>> labels = np.array([0, 0, 1, 1, 0])
+    >>> np.round(get_modularity(adjacency, labels), 2)
+    0.11
+    """
+    adjacency, bipartite = get_adjacency(input_matrix.astype(float))
+
+    if bipartite:
+        if labels_col is None:
+            raise ValueError('For bipartite graphs, you must specify the labels of both rows and columns.')
+        else:
+            labels = np.hstack((labels, labels_col))
+
+    if len(labels) != adjacency.shape[0]:
+        raise ValueError('Dimension mismatch between labels and input matrix.')
+
+    probs_row = get_probs(weights, adjacency)
+    probs_col = get_probs(weights, adjacency.T)
+    membership = get_membership(labels).astype(float)
+
+    fit = membership.T.dot(adjacency.dot(membership)).diagonal().sum() / adjacency.data.sum()
+    div = membership.T.dot(probs_col).dot(membership.T.dot(probs_row))
+    mod = fit - resolution * div
+    if return_all:
+        return mod, fit, div
+    else:
+        return mod

sknetwork/clustering/postprocess.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on July 10, 2019
+@author: Nathan de Lara <nathan.delara@polytechnique.org>
+@author: Thomas Bonald <bonald@enst.fr>
+"""
+from typing import Optional
+
+import numpy as np
+from scipy import sparse
+
+from sknetwork.utils.membership import get_membership
+
+
+def reindex_labels(labels: np.ndarray) -> np.ndarray:
+    """Reindex clusters in decreasing order of size.
+
+    Parameters
+    ----------
+    labels :
+        Label of each node.
+    Returns
+    -------
+    new_labels : np.ndarray
+        New label of each node.
+
+    Example
+    -------
+    >>> from sknetwork.clustering import reindex_labels
+    >>> labels = np.array([0, 1, 1])
+    >>> reindex_labels(labels)
+    array([1, 0, 0])
+    """
+    _, index, counts = np.unique(labels, return_inverse=True, return_counts=True)
+    _, new_index = np.unique(np.argsort(-counts), return_index=True)
+    return new_index[index]
+
+
+def aggregate_graph(input_matrix: sparse.csr_matrix, labels: Optional[np.ndarray] = None,
+                    labels_row: Optional[np.ndarray] = None, labels_col: Optional[np.ndarray] = None) \
+        -> sparse.csr_matrix:
+    """Aggregate graph per label. All nodes with the same label become a single node.
+    Negative labels are ignored (corresponding nodes are not discarded).
+
+    Parameters
+    ----------
+    input_matrix: sparse matrix
+        Adjacency or biadjacency matrix of the graph.
+    labels: np.ndarray
+        Labels of nodes.
+    labels_row: np.ndarray
+        Labels of rows (for bipartite graphs). Alias for labels.
+    labels_col: np.ndarray
+        Labels of columns (for bipartite graphs).
+    """
+    if labels_row is not None:
+        membership_row = get_membership(labels_row)
+    else:
+        membership_row = get_membership(labels)
+    if labels_col is not None:
+        membership_col = get_membership(labels_col)
+    else:
+        membership_col = membership_row
+    aggregate_matrix = membership_row.T.dot(input_matrix).dot(membership_col)
+    return aggregate_matrix

sknetwork/clustering/propagation_clustering.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+# coding: utf-8
+"""
+Created on May, 2020
+@author: Thomas Bonald <tbonald@enst.fr>
+"""
+from typing import Union
+
+import numpy as np
+from scipy import sparse
+
+from sknetwork.classification.propagation import Propagation
+from sknetwork.clustering.base import BaseClustering
+from sknetwork.utils.format import check_format, get_adjacency
+
+
+class PropagationClustering(BaseClustering, Propagation):
+    """Clustering by label propagation.
+
+    Parameters
+    ----------
+    n_iter : int
+        Maximum number of iterations (-1 for infinity).
+    node_order : str
+        * `'random'`: node labels are updated in random order.
+        * `'increasing'`: node labels are updated by increasing order of weight.
+        * `'decreasing'`: node labels are updated by decreasing order of weight.
+        * Otherwise, node labels are updated by index order.
+    weighted : bool
+        If ``True``, the vote of each neighbor is proportional to the edge weight.
+        Otherwise, all votes have weight 1.
+    sort_clusters :
+        If ``True``, sort labels in decreasing order of cluster size.
+    return_membership :
+        If ``True``, return the membership matrix of nodes to each cluster (soft clustering).
+    return_aggregate :
+        If ``True``, return the aggregate adjacency matrix or biadjacency matrix between clusters.
+
+    Attributes
+    ----------
+    labels_ : np.ndarray
+        Labels of the nodes.
+    labels_row_ : np.ndarray
+        Labels of the rows (for bipartite graphs).
+    labels_col_ : np.ndarray
+        Labels of the columns (for bipartite graphs).
+    membership_ : sparse.csr_matrix
+        Membership matrix of the nodes, shape (n_nodes, n_clusters).
+    membership_row_ : sparse.csr_matrix
+        Membership matrix of the rows (for bipartite graphs).
+    membership_col_ : sparse.csr_matrix
+        Membership matrix of the columns (for bipartite graphs).
+    aggregate_ : sparse.csr_matrix
+        Aggregate adjacency matrix or biadjacency matrix between clusters.
+
+    Example
+    -------
+    >>> from sknetwork.clustering import PropagationClustering
+    >>> from sknetwork.data import karate_club
+    >>> propagation = PropagationClustering()
+    >>> graph = karate_club(metadata=True)
+    >>> adjacency = graph.adjacency
+    >>> labels = propagation.fit_transform(adjacency)
+    >>> len(set(labels))
+    2
+
+    References
+    ----------
+    Raghavan, U. N., Albert, R., & Kumara, S. (2007).
+    `Near linear time algorithm to detect community structures in large-scale networks.
+    <https://arxiv.org/pdf/0709.2938.pdf>`_
+    Physical review E, 76(3), 036106.
+    """
+    def __init__(self, n_iter: int = 5, node_order: str = 'decreasing', weighted: bool = True,
+                 sort_clusters: bool = True, return_membership: bool = True, return_aggregate: bool = True):
+        Propagation.__init__(self, n_iter, node_order, weighted)
+        BaseClustering.__init__(self, sort_clusters, return_membership, return_aggregate)
+        self.bipartite = None
+
+    def fit(self, input_matrix: Union[sparse.csr_matrix, np.ndarray]) -> 'PropagationClustering':
+        """Clustering by label propagation.
+
+        Parameters
+        ----------
+        input_matrix :
+            Adjacency matrix or biadjacency matrix of the graph.
+
+        Returns
+        -------
+        self: :class:`PropagationClustering`
+        """
+        self._init_vars()
+
+        # input
+        input_matrix = check_format(input_matrix)
+        adjacency, bipartite = get_adjacency(input_matrix)
+
+        # propagation
+        Propagation.fit(self, adjacency)
+
+        # output
+        _, self.labels_ = np.unique(self.labels_, return_inverse=True)
+        if bipartite:
+            self._split_vars(input_matrix.shape)
+            self.bipartite = True
+        self._secondary_outputs(input_matrix)
+
+        return self

sknetwork/clustering/tests/__init__.py

@@ -0,0 +1 @@
+"""tests for clustering"""

sknetwork/clustering/tests/test_API.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Tests for clustering API"""
+import unittest
+
+from sknetwork.clustering import *
+from sknetwork.data import house
+from sknetwork.data.test_graphs import *
+from sknetwork.embedding.svd import GSVD
+
+
+class TestClusteringAPI(unittest.TestCase):
+
+    def test_regular(self):
+        for algo in [Louvain(return_aggregate=True), KMeans(embedding_method=GSVD(3), return_aggregate=True),
+                     PropagationClustering(return_aggregate=True)]:
+            for adjacency in [test_graph(), test_digraph(), test_graph_disconnect()]:
+                n = adjacency.shape[0]
+                labels = algo.fit_transform(adjacency)
+                n_labels = len(set(labels))
+                self.assertEqual(labels.shape, (n,))
+                self.assertEqual(algo.aggregate_.shape, (n_labels, n_labels))
+                adjacency_bool = adjacency.astype(bool)
+                labels = algo.fit_transform(adjacency_bool)
+                n_labels = len(set(labels))
+                self.assertEqual(labels.shape, (n,))
+                self.assertEqual(algo.aggregate_.shape, (n_labels, n_labels))
+
+    def test_bipartite(self):
+        biadjacency = test_bigraph()
+        n_row, n_col = biadjacency.shape
+        for algo in [Louvain(return_aggregate=True),
+                     KMeans(embedding_method=GSVD(3), co_cluster=True, return_aggregate=True),
+                     PropagationClustering(return_aggregate=True)]:
+            algo.fit_transform(biadjacency)
+            self.assertEqual(algo.labels_row_.shape, (n_row,))
+            self.assertEqual(algo.labels_col_.shape, (n_col,))

sknetwork/clustering/tests/test_kmeans.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created in October 2019
+@author: Nathan de Lara <nathan.delara@polytechnique.org>
+"""
+
+import unittest
+
+from sknetwork.clustering import KMeans
+from sknetwork.data.test_graphs import *
+from sknetwork.embedding import GSVD, Spectral
+
+
+class TestKMeans(unittest.TestCase):
+
+    def test_undirected(self):
+        n_clusters = 3
+        algo = KMeans(n_clusters, GSVD(2))
+        algo_options = KMeans(n_clusters, Spectral(3), co_cluster=True, sort_clusters=False)
+        for adjacency in [test_graph(), test_graph_disconnect(), test_digraph()]:
+            n = adjacency.shape[0]
+            labels = algo.fit_transform(adjacency)
+            self.assertEqual(len(set(labels)), n_clusters)
+            self.assertEqual(algo.membership_.shape, (n, n_clusters))
+            self.assertEqual(algo.aggregate_.shape, (n_clusters, n_clusters))
+            labels = algo_options.fit_transform(adjacency)
+            self.assertEqual(len(set(labels)), n_clusters)
+
+    def test_bipartite(self):
+        algo = KMeans(3, GSVD(2))
+        algo_options = KMeans(4, Spectral(3), co_cluster=True, sort_clusters=False)
+        for biadjacency in [test_bigraph(), test_bigraph_disconnect()]:
+            n_row, n_col = biadjacency.shape
+            algo.fit(biadjacency)
+            self.assertEqual(len(algo.labels_), n_row)
+            self.assertEqual(algo.membership_.shape, (n_row, 3))
+            self.assertEqual(algo.membership_row_.shape, (n_row, 3))
+            self.assertEqual(algo.membership_col_.shape, (n_col, 3))
+            self.assertEqual(algo.aggregate_.shape, (3, 3))
+            algo_options.fit(biadjacency)
+            labels = np.hstack((algo_options.labels_row_, algo_options.labels_col_))
+            self.assertEqual(len(set(labels)), 4)
+            self.assertEqual(algo_options.membership_.shape, (n_row, 4))
+            self.assertEqual(algo_options.membership_row_.shape, (n_row, 4))
+            self.assertEqual(algo_options.membership_col_.shape, (n_col, 4))
+            self.assertEqual(algo_options.aggregate_.shape, (4, 4))

sknetwork/clustering/tests/test_louvain.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Tests for Louvain"""
+import unittest
+
+from sknetwork.clustering import Louvain
+from sknetwork.data import karate_club, star_wars
+from sknetwork.data.test_graphs import *
+from sknetwork.utils import bipartite2undirected
+
+
+class TestLouvainClustering(unittest.TestCase):
+
+    def test_disconnected(self):
+        adjacency = test_graph_disconnect()
+        n = adjacency.shape[0]
+        labels = Louvain().fit_transform(adjacency)
+        self.assertEqual(len(labels), n)
+
+    def test_modularity(self):
+        adjacency = karate_club()
+        louvain_d = Louvain(modularity='dugue')
+        louvain_n = Louvain(modularity='newman')
+        labels_d = louvain_d.fit_transform(adjacency)
+        labels_n = louvain_n.fit_transform(adjacency)
+        self.assertTrue((labels_d == labels_n).all())
+
+        louvain_p = Louvain(modularity='potts')
+        louvain_p.fit_transform(adjacency)
+
+    def test_bilouvain(self):
+        biadjacency = star_wars()
+        adjacency = bipartite2undirected(biadjacency)
+        louvain = Louvain(modularity='newman')
+        labels1 = louvain.fit_transform(adjacency)
+        louvain.fit(biadjacency)
+        labels2 = np.concatenate((louvain.labels_row_, louvain.labels_col_))
+        self.assertTrue((labels1 == labels2).all())
+
+    def test_options(self):
+        adjacency = karate_club()
+
+        # resolution
+        louvain = Louvain(resolution=2)
+        labels = louvain.fit_transform(adjacency)
+        self.assertEqual(len(set(labels)), 7)
+
+        # tolerance
+        louvain = Louvain(resolution=2, tol_aggregation=0.1)
+        labels = louvain.fit_transform(adjacency)
+        self.assertEqual(len(set(labels)), 12)
+
+        # shuffling
+        louvain = Louvain(resolution=2, shuffle_nodes=True, random_state=42)
+        labels = louvain.fit_transform(adjacency)
+        self.assertEqual(len(set(labels)), 7)
+
+        # aggregate graph
+        louvain = Louvain(return_aggregate=True)
+        labels = louvain.fit_transform(adjacency)
+        n_labels = len(set(labels))
+        self.assertEqual(louvain.aggregate_.shape, (n_labels, n_labels))
+
+        # aggregate graph
+        Louvain(n_aggregations=1, sort_clusters=False).fit(adjacency)
+
+    def test_options_with_64_bit(self):
+        adjacency = karate_club()
+        # force 64-bit index
+        adjacency.indices = adjacency.indices.astype(np.int64)
+        adjacency.indptr = adjacency.indptr.astype(np.int64)
+
+        # resolution
+        louvain = Louvain(resolution=2)
+        labels = louvain.fit_transform(adjacency)
+        self.assertEqual(len(set(labels)), 7)
+
+        # tolerance
+        louvain = Louvain(resolution=2, tol_aggregation=0.1)
+        labels = louvain.fit_transform(adjacency)
+        self.assertEqual(len(set(labels)), 12)
+
+        # shuffling
+        louvain = Louvain(resolution=2, shuffle_nodes=True, random_state=42)
+        labels = louvain.fit_transform(adjacency)
+        self.assertEqual(len(set(labels)), 7)
+
+        # aggregate graph
+        louvain = Louvain(return_aggregate=True)
+        labels = louvain.fit_transform(adjacency)
+        n_labels = len(set(labels))
+        self.assertEqual(louvain.aggregate_.shape, (n_labels, n_labels))
+
+        # aggregate graph
+        Louvain(n_aggregations=1, sort_clusters=False).fit(adjacency)
+
+        # check if labels are 64-bit
+        self.assertEqual(labels.dtype, np.int64)
+
+    def test_invalid(self):
+        adjacency = karate_club()
+        louvain = Louvain(modularity='toto')
+        with self.assertRaises(ValueError):
+            louvain.fit(adjacency)

sknetwork/clustering/tests/test_metrics.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+# tests for metrics.py
+""""tests for clustering metrics"""
+import unittest
+
+import numpy as np
+
+from sknetwork.clustering import get_modularity, Louvain
+from sknetwork.data import star_wars, karate_club
+from sknetwork.data.test_graphs import test_graph
+
+
+class TestClusteringMetrics(unittest.TestCase):
+
+    def setUp(self):
+        """Basic graph for tests"""
+        self.adjacency = test_graph()
+        n = self.adjacency.shape[0]
+        labels = np.zeros(n)
+        labels[0] = 1
+        self.labels = labels.astype(int)
+        self.unique_cluster = np.zeros(n, dtype=int)
+
+    def test_api(self):
+        for metric in [get_modularity]:
+            _, fit, div = metric(self.adjacency, self.labels, return_all=True)
+            mod = metric(self.adjacency, self.labels, return_all=False)
+            self.assertAlmostEqual(fit - div, mod)
+            self.assertAlmostEqual(metric(self.adjacency, self.unique_cluster), 0.)
+
+            with self.assertRaises(ValueError):
+                metric(self.adjacency, self.labels[:3])
+
+    def test_modularity(self):
+        adjacency = karate_club()
+        labels = Louvain().fit_transform(adjacency)
+        self.assertAlmostEqual(get_modularity(adjacency, labels), 0.42, 2)
+
+    def test_bimodularity(self):
+        biadjacency = star_wars()
+        labels_row = np.array([0, 0, 1, 1])
+        labels_col = np.array([0, 1, 0])
+        self.assertAlmostEqual(get_modularity(biadjacency, labels_row, labels_col), 0.12, 2)
+
+        with self.assertRaises(ValueError):
+            get_modularity(biadjacency, labels_row)
+        with self.assertRaises(ValueError):
+            get_modularity(biadjacency, labels_row[:2], labels_col)
+        with self.assertRaises(ValueError):
+            get_modularity(biadjacency, labels_row, labels_col[:2])

sknetwork/clustering/tests/test_post_processing.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Tests for clustering post-processing"""
+
+import unittest
+
+import numpy as np
+
+from sknetwork.clustering import reindex_labels
+
+
+class TestClusteringPostProcessing(unittest.TestCase):
+
+    def test_reindex_clusters(self):
+        truth = np.array([1, 1, 2, 0, 0, 0])
+
+        labels = np.array([0, 0, 1, 2, 2, 2])
+        output = reindex_labels(labels)
+        self.assertTrue(np.array_equal(truth, output))
+
+        labels = np.array([0, 0, 5, 2, 2, 2])
+        output = reindex_labels(labels, consecutive=False)
+        self.assertTrue(np.array_equal(truth, output))