Update build_tree function with SparseKmeans implementation

libmultilabel/linear/tree.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import scipy.sparse as sparse
-import sklearn.cluster
+from sparsekmeans import LloydKmeans, ElkanKmeans
 import sklearn.preprocessing
 from tqdm import tqdm
 import psutil
@@ -274,28 +274,29 @@ def _build_tree(label_representation: sparse.csr_matrix, label_map: np.ndarray,
     Returns:
         Node: Root of the (sub)tree built from label_representation.
     """
-    if d >= dmax or label_representation.shape[0] <= K:
-        return Node(label_map=label_map, children=[])
-
-    metalabels = (
-        sklearn.cluster.KMeans(
-            K,
-            random_state=np.random.randint(2**31 - 1),
-            n_init=1,
-            max_iter=300,
-            tol=0.0001,
-            algorithm="elkan",
+    children = []
+    if d < dmax and label_representation.shape[0] > K:
+        if label_representation.shape[0] > 10000:
+            kmeans_algo = ElkanKmeans
+        else:
+            kmeans_algo = LloydKmeans
+
+        kmeans = kmeans_algo(
+            n_clusters=K, max_iter=300, tol=0.0001, random_state=np.random.randint(2**31 - 1), verbose=True
         )
-        .fit(label_representation)
-        .labels_
-    )
+        metalabels = kmeans.fit(label_representation)
 
-    children = []
-    for i in range(K):
-        child_representation = label_representation[metalabels == i]
-        child_map = label_map[metalabels == i]
-        child = _build_tree(child_representation, child_map, d + 1, K, dmax)
-        children.append(child)
+        unique_labels = np.unique(metalabels)
+        if len(unique_labels) == K:
+            create_child_node = lambda i: _build_tree(
+                label_representation[metalabels == i], label_map[metalabels == i], d + 1, K, dmax
+            )
+        else:
+            create_child_node = lambda i: Node(label_map=label_map[metalabels == i], children=[])
+
+        for i in range(K):
+            child = create_child_node(i)
+            children.append(child)
 
     return Node(label_map=label_map, children=children)
 

requirements.txt

@@ -6,3 +6,4 @@ scikit-learn
 scipy
 tqdm
 psutil
+sparsekmeans

setup.cfg

@@ -33,6 +33,7 @@ install_requires =
     scipy
     tqdm
     psutil
+    sparsekmeans
 
 python_requires = >=3.10