Clustering

@torch.no_grad()
def fit(self, x: Any) -> FuzzyCMeans:
    """Fit the Fuzzy C-Means model to the input data.

    Args:
        x: Input samples of shape (n_samples, n_features).

    Returns:
        The fitted estimator.
    """
    x_t = _as_tensor(x)
    generator = _build_generator(x_t.device, self.random_state)
    n_samples = x_t.shape[0]

    u = torch.rand((n_samples, self.n_clusters), generator=generator, device=x_t.device)
    u /= u.sum(dim=1, keepdim=True)

    for _ in range(self.max_iter):
        u_pow = u.pow(self.m)
        denominator = u_pow.sum(dim=0).view(self.n_clusters, 1)
        centers = (u_pow.T @ x_t) / denominator.clamp_min(self.eps)

        distances = torch.cdist(x_t, centers, p=2)
        zero_mask = distances == 0.0
        if zero_mask.any():
            u = torch.zeros_like(u)
            chosen = zero_mask.float().argmax(dim=1)
            u[torch.arange(n_samples, device=x_t.device), chosen] = 1.0
        else:
            distances = distances.clamp_min(self.eps)
            power = -2.0 / (self.m - 1.0)
            inv = distances.pow(power)
            u = inv / inv.sum(dim=1, keepdim=True)

        if self.cluster_centers_ is not None:
            center_shift = torch.norm(centers - self.cluster_centers_, dim=1).max().item()
        else:
            center_shift = float("inf")
        self.cluster_centers_ = centers
        self.membership_ = u
        if center_shift <= self.tol:
            break

    self.cluster_centers_ = centers
    self.membership_ = u
    return self

@torch.no_grad()
def fit_predict(self, x: Any) -> Tensor:
    """Fit the model and return the predicted labels.

    Args:
        x: Input samples of shape (n_samples, n_features).

    Returns:
        Hard cluster labels for each sample.
    """
    self.fit(x)
    if self.membership_ is None:
        raise RuntimeError("FuzzyCMeans did not produce membership values")
    return self.membership_.argmax(dim=1)

@torch.no_grad()
def predict(self, x: Any) -> Tensor:
    """Predict cluster labels for new samples.

    Args:
        x: New samples of shape (n_samples, n_features).

    Returns:
        Hard cluster labels derived from the current membership matrix.
    """
    if self.membership_ is None:
        raise RuntimeError("FuzzyCMeans instance is not fitted yet")
    x_t = _as_tensor(x)
    distances = torch.cdist(x_t, self.cluster_centers_, p=2).clamp_min(self.eps)
    power = -2.0 / (self.m - 1.0)
    inv = distances.pow(power)
    u = inv / inv.sum(dim=1, keepdim=True)
    return u.argmax(dim=1)

@torch.no_grad()
def fit(self, x: Any) -> KMeans:
    """Fit the KMeans model to the input data.

    Args:
        x: Input samples of shape (n_samples, n_features).

    Returns:
        The fitted estimator.
    """
    x_t = _as_tensor(x)
    best_inertia = float("inf")
    best_centers: Tensor | None = None
    best_labels: Tensor | None = None

    for init in range(max(1, self.n_init)):
        generator = _build_generator(x_t.device, None if self.random_state is None else self.random_state + init)
        centers = _initialize_centroids(x_t, self.n_clusters, generator)

        for _ in range(self.max_iter):
            distances = torch.cdist(x_t, centers, p=2)
            labels = distances.argmin(dim=1)

            new_centers = []
            for idx in range(self.n_clusters):
                mask = labels == idx
                if mask.any():
                    new_centers.append(x_t[mask].mean(dim=0))
                else:
                    fallback_idx = torch.randint(x_t.shape[0], (), device=x_t.device)
                    new_centers.append(x_t[fallback_idx])
            new_centers = torch.stack(new_centers, dim=0)

            center_shift = torch.norm(centers - new_centers, dim=1).max().item()
            centers = new_centers
            if center_shift <= self.tol:
                break

        distances = torch.cdist(x_t, centers, p=2)
        labels = distances.argmin(dim=1)
        inertia = float((distances[torch.arange(x_t.shape[0]), labels] ** 2).sum().item())

        if inertia < best_inertia:
            best_inertia = inertia
            best_centers = centers.clone()
            best_labels = labels.clone()

    self.cluster_centers_ = best_centers
    self.labels_ = best_labels
    self.inertia_ = best_inertia
    return self

@torch.no_grad()
def fit_predict(self, x: Any) -> Tensor:
    """Fit the model and return the predicted labels.

    Args:
        x: Input samples of shape (n_samples, n_features).

    Returns:
        Cluster labels for each sample.
    """
    self.fit(x)
    if self.labels_ is None:
        raise RuntimeError("KMeans did not produce labels")
    return self.labels_

@torch.no_grad()
def predict(self, x: Any) -> Tensor:
    """Predict cluster labels for new samples.

    Args:
        x: New samples of shape (n_samples, n_features).

    Returns:
        Cluster labels for each sample.
    """
    if self.cluster_centers_ is None:
        raise RuntimeError("KMeans instance is not fitted yet")
    x_t = _as_tensor(x)
    distances = torch.cdist(x_t, self.cluster_centers_, p=2)
    return distances.argmin(dim=1)

@torch.no_grad()
def fit(self, x: Any) -> MiniBatchKMeans:
    """Fit the Mini-Batch K-Means model to the input data."""
    x_t = _as_tensor(x)
    n_samples = x_t.shape[0]
    generator = _build_generator(x_t.device, self.random_state)
    centers = _initialize_centroids(x_t, self.n_clusters, generator)

    # Running counts for online centroid updates
    counts = torch.zeros(self.n_clusters, device=x_t.device, dtype=x_t.dtype)
    batch_size = min(self.batch_size, n_samples)

    for _ in range(self.max_iter):
        idx = torch.randperm(n_samples, generator=generator, device=x_t.device)[:batch_size]
        batch = x_t[idx]

        dists = torch.cdist(batch, centers, p=2)
        labels = dists.argmin(dim=1)

        for k in range(self.n_clusters):
            mask = labels == k
            if not mask.any():
                continue
            n_k = mask.sum()
            counts[k] += n_k
            lr = n_k.float() / counts[k]
            centers[k] += lr * (batch[mask].mean(dim=0) - centers[k])

    # Assign all points to final centers
    dists = torch.cdist(x_t, centers, p=2)
    self.labels_ = dists.argmin(dim=1)
    self.cluster_centers_ = centers
    return self

@torch.no_grad()
def predict(self, x: Any) -> Tensor:
    """Predict cluster labels for new samples."""
    if self.cluster_centers_ is None:
        raise RuntimeError("MiniBatchKMeans instance is not fitted yet")
    x_t = _as_tensor(x)
    return torch.cdist(x_t, self.cluster_centers_, p=2).argmin(dim=1)