From 36869374241231351da7399fb7bcc0d7abb7887b Mon Sep 17 00:00:00 2001 From: Chandan Singh Date: Sun, 3 Nov 2024 06:10:31 -0800 Subject: [PATCH] save all clustering models --- imodels/clustering/stableclustering.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/imodels/clustering/stableclustering.py b/imodels/clustering/stableclustering.py index aa1574d4..3fd78393 100644 --- a/imodels/clustering/stableclustering.py +++ b/imodels/clustering/stableclustering.py @@ -22,14 +22,14 @@ def __init__(self, k_values, n_repetitions=10, algorithm="k-means", metric="adju def fit(self, X): best_k = None best_score = -1 - best_model = None + self.models_ = [] for k in tqdm(self.k_values, desc="k"): clusters = [] for i_rep in tqdm(range(self.n_repetitions), desc='Repetitions', leave=False): if self.algorithm == "k-means": model = KMeans( - n_clusters=k, random_state=self.random_state + i_rep, init='random') + n_clusters=k, random_state=self.random_state + i_rep) # , init='random') labels = model.fit_predict(X) # elif self.algorithm == "nmf": # model = NMF(n_components=k, init='random', @@ -53,6 +53,8 @@ def fit(self, X): "Invalid metric: choose 'rand' or 'adjusted_rand'") scores.append(score) + self.models_.append(deepcopy(model)) + avg_score = np.mean(scores) # Store the average score for this k self.scores_[k] = float(avg_score) @@ -63,7 +65,7 @@ def fit(self, X): # Fit the final model on the whole data self.best_k_ = best_k - self.best_model_ = best_model.fit(X) + self.best_model_ = best_model return self def predict(self, X):