diff --git a/libmultilabel/linear/tree.py b/libmultilabel/linear/tree.py index fe6e94b4..ae4c89b1 100644 --- a/libmultilabel/linear/tree.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index fdff2e2e..814b4e18 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ scikit-learn scipy tqdm psutil +sparsekmeans \ No newline at end of file diff --git a/setup.cfg b/setup.cfg index a676ce91..0be32045 100644 --- a/setup.cfg +++ b/setup.cfg @@ -33,6 +33,7 @@ install_requires = scipy tqdm psutil + sparsekmeans python_requires = >=3.10