Adding _dtype to clustering algorithms

krstopro · Sep 16, 2024 · 68a4d49 · 68a4d49
1 parent 05d57aa
commit 68a4d49
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/unsupervised_bias_detection/clustering/_bahc.py b/unsupervised_bias_detection/clustering/_bahc.py
@@ -37,7 +37,7 @@ def fit(self, X, y):
             Fitted estimator.
         """
         X, y = self._validate_data(
-            X, y, reset=False, accept_large_sparse=False, dtype=[np.float32, np.float64], order="C"
+            X, y, reset=False, accept_large_sparse=False, dtype=self._dtype, order="C"
         )
         n_samples, _ = X.shape
         # We start with all samples in a single cluster

diff --git a/unsupervised_bias_detection/clustering/_kmeans.py b/unsupervised_bias_detection/clustering/_kmeans.py
@@ -1,3 +1,4 @@
+import numpy as np
 from ._bahc import BiasAwareHierarchicalClustering
 from sklearn.cluster import KMeans
 
@@ -41,6 +42,8 @@ class BiasAwareHierarchicalKMeans(BiasAwareHierarchicalClustering):
     array([ 10., -10.])
     """
 
+    _dtype = [np.float32, np.float64]
+
     def __init__(
         self,
         n_iter,

diff --git a/unsupervised_bias_detection/clustering/_kmodes.py b/unsupervised_bias_detection/clustering/_kmodes.py
@@ -41,6 +41,8 @@ class BiasAwareHierarchicalKModes(BiasAwareHierarchicalClustering):
     array([ 10., -10.])
     """
 
+    _dtype = None
+
     def __init__(self, n_iter, min_cluster_size, **kmodes_params):
         super().__init__(n_iter, min_cluster_size)