graphical-sampling 0.1.0__py3-none-any.whl

graphical_sampling/__init__.py

@@ -0,0 +1,14 @@
+from importlib import metadata
+
+from .design import Design
+from . import criteria
+from . import search
+from . import sampling
+from . import clustering
+from . import random
+from . import measure
+
+
+__version__ = metadata.version("geometric_sampling")
+
+__all__ = ["Design", "criteria", "search", "sampling", "clustering", "random", "measure"]

graphical_sampling/clustering/__init__.py

@@ -0,0 +1,13 @@
+from .soft_balanced_kmeans import SoftBalancedKMeans
+from .aggregate import AggregateBalancedKMeans
+from .one_boundary import OneBoundaryBalancedKMeans
+from .dubly_balanced_clustering import DublyBalancedKMeans
+# from .agg import Agg
+# from .agg_one import AggOne
+# from .aggregate import FinalAgg
+# from .final2 import FinalAgg2
+# from .swap import SwapAgg
+
+
+__all__ = ["SoftBalancedKMeans", "AggregateBalancedKMeans", "OneBoundaryBalancedKMeans", "DublyBalancedKMeans"]
+# __all__ = ["SoftBalancedKMeans", "AggregateBalancingKmeans", "Agg", "AggOne", "FinalAgg", "FinalAgg2", "SwapAgg"]

graphical_sampling/clustering/aggregate.py

@@ -0,0 +1,213 @@
+import numpy as np
+from numpy.typing import NDArray
+from sklearn.cluster import KMeans
+
+
+class AggregateBalancedKMeans:
+    def __init__(
+        self, k: int, *, initial_centroids: NDArray = None, tolerance: int = 5
+    ) -> None:
+        self.k = k
+        self.tolerance = tolerance
+        self.Y_features = None
+        self.X_features = None
+        self.weights = None
+        self.m = None
+        self.N = None
+        self.centroids = initial_centroids
+        self.labels: NDArray = None
+        self.membership: NDArray = None
+        self.Ti: NDArray = None
+        self.Tij: NDArray = None
+        self.goal: float = None
+        self.rng = np.random.default_rng()
+
+    def _generate_goal_j(self) -> float:
+        return self.X_features.sum(axis=0)/self.k
+
+    def _generate_goal(self) -> float:
+        return np.sum(self._generate_goal_j())
+
+    def _generate_membership(self):
+        membership = np.zeros((self.N, self.k))
+        for j in range(self.N):
+            membership[j, self.labels[j]] = 1
+        return membership
+
+    def _generate_Tij(self) -> NDArray:
+        return np.sum(
+            self.X_features[:, :, np.newaxis] * self.membership[:, np.newaxis, :],
+            axis=0,
+        ).T
+
+    def _generate_Ti(self) -> NDArray:
+        return np.sum(self._generate_Tij(), axis=1)
+
+    def _generate_Tij_cost(self):
+        return np.sum((self.weights*(self.Tij-self.goal_j))**2)
+
+    def _transfer_score(
+        self,
+        data_index: int,
+        old_cluster: int,
+        new_cluster: int,
+    ) -> float:
+        if (
+            self.Ti[old_cluster] - self.Ti[new_cluster]
+            > 10**-self.tolerance
+        ):
+            score = (
+                np.linalg.norm(
+                    self.Y_features[data_index] - self.centroids[new_cluster]
+                )
+                ** 2
+                - np.linalg.norm(
+                    self.Y_features[data_index] - self.centroids[old_cluster]
+                )
+                ** 2
+            ) / (self.Ti[old_cluster] - self.Ti[new_cluster] + 1e-9)
+            return score #if score > 0 else np.inf
+        else:
+            return np.inf
+
+    def _get_transfer_records(self, top_m: int):
+        transfer_records = []
+
+        for i in range(self.N):
+            for j_old in np.nonzero(self.membership[i])[0]:
+                j_new_min = np.argmin(
+                    [self._transfer_score(i, j_old, j_new) for j_new in range(self.k)]
+                )
+                cost = self._transfer_score(i, j_old, j_new_min)
+                transfer_records.append((cost, i, j_old, j_new_min))
+
+        transfer_records = np.array(transfer_records)
+        sorted_transfer_records = transfer_records[np.argsort(transfer_records[:, 0])]
+        best_cost = sorted_transfer_records[0, 0]
+        top_m_transfer_records = sorted_transfer_records[:top_m, 1:].astype(int)
+
+        # print("cost")
+        # print(sorted_transfer_records[:5])
+        # print()
+
+        return best_cost, top_m_transfer_records
+
+    def _transfer_percent(self, data_index: int, old_cluster: int, new_cluster: int):
+        if (
+            self.Ti[old_cluster] >= self.goal
+            and self.Ti[new_cluster] >= self.goal
+        ) or (
+            self.Ti[old_cluster] <= self.goal
+            and self.Ti[new_cluster] <= self.goal
+        ):
+            # print('case ++ or --')
+            return min(
+                self.membership[data_index, old_cluster],
+                ((self.Ti[old_cluster] - self.Ti[new_cluster]) / (2*np.sum(self.X_features[data_index])))
+            )
+        else:
+            # print('case +-')
+            return min(
+                self.membership[data_index, old_cluster],
+                ((self.Ti[old_cluster] - self.goal) / np.sum(self.X_features[data_index])),
+                ((self.goal - self.Ti[new_cluster]) / np.sum(self.X_features[data_index])),
+            )
+
+    def _transfer(self, data_index: int, old_cluster: int, new_cluster: int) -> None:
+        transfer_percent = self._transfer_percent(data_index, old_cluster, new_cluster)
+
+        # print(f'transfer {data_index} from {old_cluster} to {new_cluster}')
+        # print(f'BEFROE: T_old_cluster={round(self.Ti[old_cluster], 5)} and T_new_cluster={round(self.Ti[new_cluster], 5)}')
+        # print(f'transfer percent: {transfer_percent}')
+
+        self.membership[data_index, old_cluster] -= transfer_percent
+        self.membership[data_index, new_cluster] += transfer_percent
+
+    def _no_transfer_possible(self, best_cost: float) -> bool:
+        return best_cost == np.inf
+
+    def _is_transfer_possible(self, old_cluster: int, new_cluster: int) -> bool:
+        return (
+            self.Ti[old_cluster] - self.Ti[new_cluster]
+            > 10**-self.tolerance
+        )
+
+    def _stop_codition(self, tol) -> bool:
+        return np.all(np.abs(self.Ti - self.goal) < 10**-tol)
+
+    def _expected_num_transfers(self) -> float:
+        max_diff_sum = np.max(self.Ti - self.Ti[:, None])
+        possible_transfers = self.X_features.sum(axis=1)[:, np.newaxis]*self.membership
+        mean_nonzero_transfers = np.mean(possible_transfers[np.nonzero(possible_transfers)])
+        return max(int(np.floor(max_diff_sum / (2 * mean_nonzero_transfers))), 1)
+
+    def _update_centroids(self) -> None:
+        for i in range(self.k):
+            self.centroids[i] = np.mean(
+                self.Y_features[self.membership[:, i] > 0], axis=0
+            )
+
+    def fit(self, Y_features: NDArray, X_features: NDArray, weights: NDArray) -> None:
+        self.Y_features = Y_features
+        self.X_features = X_features
+        self.weights = weights
+        self.m = X_features.shape[1]
+        self.N = X_features.shape[0]
+        self.goal_j = self._generate_goal_j()
+        self.goal = self._generate_goal()
+
+        # print('Goal_j:', self.goal_j)
+        # print('Goal:', self.goal)
+        # print()
+
+        kmeans = KMeans(
+            n_clusters=self.k,
+            init=self.centroids if self.centroids is not None else "k-means++",
+            n_init=10,
+            tol=10**-self.tolerance,
+        )
+        kmeans.fit(self.Y_features)
+
+        self.centroids = kmeans.cluster_centers_
+        self.labels = kmeans.labels_
+        self.membership = self._generate_membership()
+        self.Tij = self._generate_Tij()
+        self.Ti = self._generate_Ti()
+        self.Tij_cost = self._generate_Tij_cost()
+        iter_ = 0
+
+        while not self._stop_codition(self.tolerance) and iter_ < 1000:
+            # print("================================================")
+            # print("iter:", iter_)
+            # print("\nTij", self.Tij)
+            # print("\nTij - goal_j", self.Tij - self.goal_j)
+            # print("\nTi", self.Ti)
+            # print("\nTij_cost", round(self.Tij_cost, 5))
+            # print()
+            best_cost, transfer_records = self._get_transfer_records(top_m=self._expected_num_transfers())
+            if self._no_transfer_possible(best_cost):
+                break
+            for data_index, old_cluster, new_cluster in transfer_records:
+                if self._is_transfer_possible(old_cluster, new_cluster):
+                    self._transfer(data_index, old_cluster, new_cluster)
+                    self.Tij = self._generate_Tij()
+                    self.Ti = self._generate_Ti()
+                    self.Tij_cost = self._generate_Tij_cost()
+                    # print(f'AFTER: T_old_cluster={round(self.Ti[old_cluster], 5)} and T_new_cluster={round(self.Ti[new_cluster], 5)}')
+                    # print()
+            self._update_centroids()
+            iter_ += 1
+
+    def get_clusters(self) -> NDArray:
+        clusters = []
+
+        for i in range(self.k):
+            probs = self.membership[:, i] * self.X_features.reshape(-1)
+            ids = np.nonzero(probs)[0]
+            units = np.concatenate(
+                [ids.reshape(-1, 1), self.Y_features[ids], probs[ids].reshape(-1, 1)],
+                axis=1,
+            )
+            clusters.append(units)
+
+        return clusters

graphical_sampling/clustering/dubly_balanced_clustering.py

@@ -0,0 +1,209 @@
+import numpy as np
+from numpy.typing import NDArray
+from k_means_constrained import KMeansConstrained
+from scipy.stats import mode
+
+
+class DublyBalancedKMeans:
+    def __init__(self, k, split_size=0.01):
+        self.k = k
+        self.split_size = split_size
+
+    def _generate_expanded_coords(self, coords, probs):
+        counts = (probs / self.split_size).round().astype(int)
+        expanded_coords = np.repeat(coords, counts, axis=0)
+        expanded_idx = np.repeat(np.arange(self.N), counts)
+        return expanded_coords, expanded_idx
+
+    def _generate_labels(self, extended_labels, expanded_idx, coords):
+        labels = np.zeros(self.N, dtype=int)
+        for i in range(self.N):
+            assigned_labels = extended_labels[expanded_idx == i]
+            if len(assigned_labels) == 0:
+                labels[i] = np.argmin(np.linalg.norm(self.centroids - coords[i], axis=1))
+            else:
+                labels[i] = mode(assigned_labels, keepdims=True)[0][0]
+        return labels
+
+    def fit(self, coords, probs):
+        self.N = coords.shape[0]
+        expanded_coords, expanded_idx = self._generate_expanded_coords(coords, probs)
+        cluster_size = len(expanded_idx) // self.k
+        kmeans = KMeansConstrained(
+            n_clusters=self.k,
+            size_min=cluster_size,
+            size_max=cluster_size+1
+        )
+        labels = kmeans.fit_predict(expanded_coords)
+        self.centroids = kmeans.cluster_centers_
+        self.labels = self._generate_labels(labels, expanded_idx, coords)
+
+        cb = ContinuesBalancing(self.k)
+        cb.fit(coords, probs, self.centroids, self.labels)
+
+        self.centroids = cb.centroids
+        self.labels = cb.labels
+        self.membership = cb.membership
+        self.Ti = cb.Ti
+        self.goal = cb.goal
+        self.clusters = cb.get_clusters()
+
+
+
+class ContinuesBalancing:
+    def __init__(
+        self, k: int, *, tolerance: int = 3
+    ) -> None:
+        self.k = k
+        self.tolerance = tolerance
+        self.Y_features = None
+        self.X_feature = None
+        self.labels: NDArray = None
+        self.membership: NDArray = None
+        self.Ti: NDArray = None
+        self.goal: float = None
+
+    def _generate_goal(self) -> float:
+        return self.X_feature.sum()/self.k
+
+    def _generate_membership(self):
+        membership = np.zeros((self.N, self.k))
+        for j in range(self.N):
+            membership[j, self.labels[j]] = 1
+        return membership
+
+    def _generate_Ti(self) -> NDArray:
+        return np.sum(
+            self.X_feature[:, np.newaxis] * self.membership,
+            axis=0,
+        )
+
+    def _transfer_score(
+        self,
+        data_index: int,
+        old_cluster: int,
+        new_cluster: int,
+    ) -> float:
+        if (
+            self.Ti[old_cluster] - self.Ti[new_cluster]
+            > 10**-self.tolerance
+        ):
+            score = (
+                np.linalg.norm(
+                    self.Y_features[data_index] - self.centroids[new_cluster]
+                )
+                ** 2
+                - np.linalg.norm(
+                    self.Y_features[data_index] - self.centroids[old_cluster]
+                )
+                ** 2
+            ) / (self.Ti[old_cluster] - self.Ti[new_cluster] + 1e-9)
+            return score
+        else:
+            return np.inf
+
+    def _get_transfer_records(self, top_m: int):
+        transfer_records = []
+
+        for i in range(self.N):
+            for j_old in np.nonzero(self.membership[i])[0]:
+                j_new_min = np.argmin(
+                    [self._transfer_score(i, j_old, j_new) for j_new in range(self.k)]
+                )
+                cost = self._transfer_score(i, j_old, j_new_min)
+                transfer_records.append((cost, i, j_old, j_new_min))
+
+        transfer_records = np.array(transfer_records)
+        sorted_transfer_records = transfer_records[np.argsort(transfer_records[:, 0])]
+        best_cost = sorted_transfer_records[0, 0]
+        top_m_transfer_records = sorted_transfer_records[:top_m, 1:].astype(int)
+
+        return best_cost, top_m_transfer_records
+
+    def _transfer_percent(self, data_index: int, old_cluster: int, new_cluster: int):
+        if (
+            self.Ti[old_cluster] >= self.goal
+            and self.Ti[new_cluster] >= self.goal
+        ) or (
+            self.Ti[old_cluster] <= self.goal
+            and self.Ti[new_cluster] <= self.goal
+        ):
+            return min(
+                self.membership[data_index, old_cluster],
+                ((self.Ti[old_cluster] - self.Ti[new_cluster]) / (2*np.sum(self.X_feature[data_index])))
+            )
+        else:
+            return min(
+                self.membership[data_index, old_cluster],
+                ((self.Ti[old_cluster] - self.goal) / np.sum(self.X_feature[data_index])),
+                ((self.goal - self.Ti[new_cluster]) / np.sum(self.X_feature[data_index])),
+            )
+
+    def _transfer(self, data_index: int, old_cluster: int, new_cluster: int) -> None:
+        transfer_percent = self._transfer_percent(data_index, old_cluster, new_cluster)
+        self.membership[data_index, old_cluster] -= transfer_percent
+        self.membership[data_index, new_cluster] += transfer_percent
+
+    def _no_transfer_possible(self, best_cost: float) -> bool:
+        return best_cost == np.inf
+
+    def _is_transfer_possible(self, old_cluster: int, new_cluster: int) -> bool:
+        return (
+            self.Ti[old_cluster] - self.Ti[new_cluster]
+            > 10**-self.tolerance
+        )
+
+    def _stop_codition(self, tol) -> bool:
+        return np.all(np.abs(self.Ti - self.goal) < 10**-tol)
+
+    def _expected_num_transfers(self) -> float:
+        max_diff_sum = np.max(self.Ti - self.Ti[:, None])
+        possible_transfers = self.X_feature[:, np.newaxis] * self.membership
+        mean_nonzero_transfers = np.mean(possible_transfers[np.nonzero(possible_transfers)])
+        return max(int(np.floor(max_diff_sum / (2 * mean_nonzero_transfers))), 1)
+
+    def _update_centroids(self) -> None:
+        for i in range(self.k):
+            self.centroids[i] = np.mean(
+                self.Y_features[self.membership[:, i] > 0], axis=0
+            )
+
+    def fit(self, Y_features: NDArray, X_feature: NDArray, centroids: NDArray, labels: NDArray, max_iteration: int = 1000) -> None:
+        self.N = len(X_feature)
+        self.Y_features = Y_features
+        self.X_feature = X_feature
+        self.centroids = centroids
+        self.labels = labels
+        self.goal = self._generate_goal()
+        self.membership = self._generate_membership()
+        self.Ti = self._generate_Ti()
+        self._update_centroids()
+        iter_ = 0
+
+        while not self._stop_codition(self.tolerance) and iter_ < max_iteration:
+            print(f"\nIteration {iter_}")
+            print(f"Ti: {self.Ti}")
+            print(f"Sum Ti: {np.sum(self.Ti)}")
+            best_cost, transfer_records = self._get_transfer_records(top_m=self._expected_num_transfers())
+            if self._no_transfer_possible(best_cost):
+                break
+            for data_index, old_cluster, new_cluster in transfer_records:
+                if self._is_transfer_possible(old_cluster, new_cluster):
+                    self._transfer(data_index, old_cluster, new_cluster)
+                    self.Ti = self._generate_Ti()
+            self._update_centroids()
+            iter_ += 1
+
+    def get_clusters(self) -> NDArray:
+        clusters = []
+
+        for i in range(self.k):
+            x = self.membership[:, i] * self.X_feature
+            ids = np.nonzero(x)[0]
+            units = np.concatenate(
+                [ids.reshape(-1, 1), self.Y_features[ids], x[ids].reshape(-1, 1)],
+                axis=1,
+            )
+            clusters.append(units)
+
+        return clusters

graphical_sampling/clustering/one_boundary.py

@@ -0,0 +1,233 @@
+import numpy as np
+from numpy.typing import NDArray
+from sklearn.cluster import KMeans
+import matplotlib.pyplot as plt
+
+
+class OneBoundaryBalancedKMeans:
+    def __init__(
+        self, k: int, *, initial_centroids: NDArray = None, tolerance: int = 5
+    ) -> None:
+        self.k = k
+        self.tolerance = tolerance
+        self.Y_features = None
+        self.x_feature = None
+        self.N = None
+        self.centroids = initial_centroids
+        self.labels: NDArray = None
+        self.membership: NDArray = None
+        self.Ti: NDArray = None
+        self.goal: float = None
+        self.rng = np.random.default_rng()
+
+    def _generate_goal(self) -> float:
+        return self.x_feature.sum()/self.k
+
+    def _generate_membership(self):
+        membership = np.zeros((self.N, self.k))
+        for j in range(self.N):
+            membership[j, self.labels[j]] = 1
+        return membership
+
+    def _generate_Ti(self) -> NDArray:
+        return np.sum(
+            self.x_feature[:, np.newaxis] * self.membership,
+            axis=0,
+        ).T
+
+    def _determine_points_to_gain(self) -> dict[int, list[tuple[int, float]]]:
+        """
+        For each cluster below goal, compute the points (and fractions) needed to fill its deficit,
+        then score each cluster by (mean distance of those points to its centroid) / deficit,
+        and finally return the gains for the single cluster with the lowest score.
+        Restriction: once any fraction of a point is earmarked for a cluster, that fraction is
+        'locked' and cannot be reused for another cluster in this round.  If 100% is locked,
+        the point is out entirely; if 50% is locked, only the remaining 50% can be used elsewhere.
+        """
+        deficits = self.goal - self.Ti
+        cluster_gains: dict[int, list[tuple[int, float]]] = {}
+        cluster_scores: dict[int, float] = {}
+        # tracks fraction of each point already locked in this round
+
+        for cluster_idx, deficit in enumerate(deficits):
+            if deficit <= 0:
+                continue
+
+            # build list of candidates not yet fully in this cluster and not fully locked
+            candidates = []
+            for j in range(self.N):
+                in_cluster = self.membership[j, cluster_idx]
+                locked = self.locked_frac.get(j, 0.0)
+                # available fraction of the point for this cluster:
+                avail_frac = (1 - in_cluster) - locked
+                if avail_frac > 0:
+                    candidates.append(j)
+            if not candidates:
+                continue
+
+            # distances to centroid
+            dists = np.linalg.norm(self.Y_features[candidates] - self.centroids[cluster_idx], axis=1)
+            sorted_idx = np.argsort(dists)
+            sorted_candidates = [candidates[i] for i in sorted_idx]
+            # print(f"centroid: {self.centroids[cluster_idx]}")
+            # print(dists[sorted_idx][:5], sorted_candidates[:5])
+
+            gains: list[tuple[int, float]] = []
+            dists_for_gains: list[float] = []
+            accum = 0.0
+
+            for pt_idx in sorted_candidates:
+                prev_in = self.membership[pt_idx, cluster_idx]
+                locked = self.locked_frac.get(pt_idx, 0.0)
+                # fraction of whole point still available to move into this cluster
+                avail_frac = (1 - prev_in) - locked
+                if avail_frac <= 0:
+                    continue
+                # x_mass available
+                avail_mass = self.x_feature[pt_idx] * avail_frac
+
+                if accum + avail_mass <= deficit:
+                    # lock entire available fraction
+                    frac_move = avail_frac
+                    gains.append((pt_idx, frac_move))
+                    dists_for_gains.append(
+                        np.linalg.norm(self.Y_features[pt_idx] - self.centroids[cluster_idx])
+                    )
+                    accum += avail_mass
+                    self.locked_frac[pt_idx] = locked + frac_move
+                else:
+                    # we only need part of the available mass
+                    needed = deficit - accum
+                    frac_of_mass = needed / self.x_feature[pt_idx]
+                    # but that is fraction of whole point:
+                    frac_move = min(frac_of_mass, avail_frac)
+                    if frac_move > 0:
+                        gains.append((pt_idx, frac_move))
+                        dists_for_gains.append(
+                            np.linalg.norm(self.Y_features[pt_idx] - self.centroids[cluster_idx])
+                        )
+                        self.locked_frac[pt_idx] = locked + frac_move
+                    break
+
+            if gains:
+                mean_dist = sum(dists_for_gains) / len(dists_for_gains)
+                score = mean_dist / deficit
+                cluster_gains[cluster_idx] = gains
+                cluster_scores[cluster_idx] = score
+
+        if not cluster_scores:
+            return {}
+
+        # choose the cluster with lowest score
+        best = min(cluster_scores, key=cluster_scores.get)
+        return {best: cluster_gains[best]}
+
+
+    def _assign_points_to_clusters(self, points_to_gain: dict[int, list[tuple[int, float]]]) -> None:
+        """
+        Assigns the computed gains for the single target cluster:
+        - `frac` is the additional fraction of the point to assign to the cluster.
+        - New cluster membership = previous membership + frac (capped at 1.0).
+        - The remaining mass (1 − new_membership) is redistributed among the other clusters
+        in proportion to their previous memberships.
+        """
+        if not points_to_gain:
+            return
+
+        # unpack the single entry
+        cluster_idx, gains = next(iter(points_to_gain.items()))
+        for pt_idx, add_frac in gains:
+            prev = self.membership[pt_idx].copy()
+            # print(f"pt_idx: {pt_idx}, prev: {prev}, add_frac: {add_frac}, cluster_idx: {cluster_idx}")
+            already_in = prev[cluster_idx]
+
+            # compute new membership for this cluster
+            new_cluster_mem = min(already_in + add_frac, 1.0)
+            rem_mass = 1.0 - new_cluster_mem
+
+            # prepare redistribution of remaining mass
+            prev_except = prev.copy()
+            prev_except[cluster_idx] = 0.0
+            total_prev_except = prev_except.sum()
+
+            # build the updated membership vector
+            new_mem = np.zeros_like(prev)
+            new_mem[cluster_idx] = new_cluster_mem
+
+            if total_prev_except > 0 and rem_mass > 0:
+                # scale other clusters proportionally
+                new_mem += (prev_except / total_prev_except) * rem_mass
+            elif rem_mass > 0:
+                # if no other membership existed, put remainder back into this cluster
+                new_mem[cluster_idx] += rem_mass
+
+            # apply the update
+            self.membership[pt_idx] = new_mem
+
+            # print(f"new_mem: {new_mem}")
+
+        # recompute cluster totals
+        self.Ti = self._generate_Ti()
+
+    def _update_centroids(self) -> None:
+        for i in range(self.k):
+            points = self.Y_features[self.membership[:, i] > 0]
+            self.centroids[i] = np.mean(points, axis=0)
+
+    def _stop_codition(self) -> bool:
+        return np.all(np.abs(self.Ti - self.goal) < 10**-self.tolerance)
+
+    def fit(self, Y_features: NDArray, x_feature: NDArray) -> None:
+        self.Y_features = Y_features
+        self.x_feature = x_feature
+        self.N = x_feature.size
+        self.goal = self._generate_goal()
+
+        kmeans = KMeans(
+            n_clusters=self.k,
+            init=self.centroids if self.centroids is not None else "k-means++",
+            n_init=10,
+            tol=10**-self.tolerance,
+        )
+        kmeans.fit(self.Y_features)
+
+        self.centroids = kmeans.cluster_centers_
+        self.labels = kmeans.labels_
+        self.membership = self._generate_membership()
+        self.Ti = self._generate_Ti()
+
+        iter_ = 0
+        self.locked_frac: dict[int, float] = {}
+
+        while not self._stop_codition() and iter_ < 5:
+            print("\nIteration:", iter_)
+            print("(BEF) Ti:", self.Ti)
+            points_to_gain = self._determine_points_to_gain()
+            if not points_to_gain:
+                break
+            self._assign_points_to_clusters(points_to_gain)
+            iter_ += 1
+
+            points = np.array(list(points_to_gain.values())[0])
+            points_id, points_frac = points[:, 0], points[:, 1]
+            points_id = points_id.astype(int)
+            points_frac = points_frac.astype(float)
+
+            mem = np.argmax(self.membership, axis=1)
+            # boarders = np.where(np.count_nonzero(self.membership, axis=1)>1)[0]
+            plt.figure(figsize=(5, 5), dpi=200)
+            plt.scatter(*Y_features.T, c=mem)
+            # plt.scatter(*Y_features[boarders].T, c='red', s=10)
+            plt.scatter(*Y_features[points_id].T, c='red', marker='*', s=10)
+            plt.scatter(*self.centroids.T, c='black', marker='x', s=100)
+
+            for i, point_id in enumerate(points_id):
+                x, y = Y_features[point_id]
+                plt.text(x, y - 0.05, f"{points_frac[i]:.2f}, {point_id}", color='black', fontsize=8, ha='center')  # Adjust y-offset as needed
+
+
+            print("(AFT) Ti:", self.Ti)
+            print("Diff Ti:", np.sum(np.abs(self.Ti - self.goal)))
+            print("Points to gain:", points_to_gain)
+
+            self._update_centroids()