more metrics

alexgiving · May 25, 2024 · 2f57829 · 2f57829
1 parent 050719c
commit 2f57829
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/lkmeans/examples/main.py b/lkmeans/examples/main.py
@@ -5,7 +5,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
-from sklearn.metrics import accuracy_score, adjusted_mutual_info_score, adjusted_rand_score
+from sklearn.metrics import accuracy_score, adjusted_mutual_info_score, adjusted_rand_score, completeness_score, homogeneity_score, normalized_mutual_info_score, v_measure_score
 from tap import Tap
 
 from lkmeans.clustering import HardSSLKMeans, LKMeans, SoftSSLKMeans
@@ -46,6 +46,10 @@ def calculate_metrics(labels: NDArray, generated_labels: NDArray) -> Dict[str, f
     return {
         'ari': float(adjusted_rand_score(labels, generated_labels)),
         'ami': float(adjusted_mutual_info_score(labels, generated_labels)),
+        'completeness': float(completeness_score(labels, generated_labels)),
+        'homogeneity': float(homogeneity_score(labels, generated_labels)),
+        'nmi': float(normalized_mutual_info_score(labels, generated_labels)),
+        'v_measure': float(v_measure_score(labels, generated_labels)),
         'accuracy': float(accuracy_score(labels, generated_labels)),
     }