Clustering¶

def fit(self, X: Float[torch.Tensor, "n_points n_features"], y: None = None) -> "RiemannianFuzzyKMeans":
    """Fit the Riemannian Fuzzy K-Means model to the data X.

    Args:
        X: Input data. Features should match the manifold's geometry.
        y: Ignored, present for compatibility with scikit-learn's API.

    Returns:
        self: Fitted `RiemannianFuzzyKMeans` instance.

    Raises:
        ValueError: If the input data's dimension does not match the manifold's ambient dimension.
        RuntimeError: If the optimizer is not set correctly or if the model has not been initialized properly.
    """
    if isinstance(X, np.ndarray):
        X = torch.from_numpy(X).type(torch.get_default_dtype())
    elif not isinstance(X, torch.Tensor):
        X = torch.tensor(X, dtype=torch.get_default_dtype())

    # Ensure X is on the same device as the manifold
    X = X.to(self.pm.device)

    if X.shape[1] != self.pm.ambient_dim:
        raise ValueError(
            f"Input data X's dimension ({X.shape[1]}) in fit() does not match "
            f"the manifold's ambient dimension ({self.pm.ambient_dim})."
        )

    self._init_centers(X)
    m, tol = self.m, self.tol
    losses = []
    for i in range(self.max_iter):
        self.opt_.zero_grad()
        # self.pm.dist is implemented in manifolds.py and handles broadcasting
        d = self.pm.dist(X, self.mu_)  # X is (N,D), mu_ is (K,D) -> d is (N,K)
        # Original RFK: d = self.pm.dist(X.unsqueeze(1), self.mu_.unsqueeze(0))
        # The .dist in manifolds.py uses X[:, None] and Y[None, :], so direct call should work if mu_ is (K,D)

        S = torch.sum(d.pow(-2 / (m - 1)) + 1e-8, dim=1)  # Add epsilon for stability
        loss = torch.sum(S.pow(1 - m))
        loss.backward()
        losses.append(loss.item())
        self.opt_.step()
        if self.verbose:
            print(f"RFK iter {i + 1}, loss={loss.item():.4f}")
        if i > 0 and abs(losses[-1] - losses[-2]) < tol:
            break

    # save the result
    self.losses_ = np.array(losses)
    with torch.no_grad():  # Ensure no gradients are computed for final calculations
        dfin = self.pm.dist(X, self.mu_)  # Re-calculate dist to final centers
        inv = dfin.pow(-2 / (m - 1)) + 1e-8  # Add epsilon
        u_final = inv / (inv.sum(dim=1, keepdim=True) + 1e-8)  # Add epsilon
    self.u_ = u_final.detach().cpu().numpy()
    self.labels_ = np.argmax(self.u_, axis=1)
    self.cluster_centers_ = self.mu_.data.clone().detach().cpu().numpy()
    return self

def predict(self, X: Float[torch.Tensor, "n_points n_features"]) -> Int[torch.Tensor, "n_points"]:
    """Predict the closest cluster each sample in X belongs to.

    Args:
        X: Input data. Features should match the manifold's geometry.

    Returns:
        labels: Cluster labels for each sample in X.

    Raises:
        ValueError: If the input data's dimension does not match the manifold's ambient dimension.
        RuntimeError: If the model has not been fitted yet.
    """
    if isinstance(X, np.ndarray):
        X = torch.from_numpy(X).type(torch.get_default_dtype())
    elif not isinstance(X, torch.Tensor):
        X = torch.tensor(X, dtype=torch.get_default_dtype())

    # Ensure X is on the same device as the manifold
    X = X.to(self.pm.device)

    if X.shape[1] != self.pm.ambient_dim:
        raise ValueError(
            f"Input data X's dimension ({X.shape[1]}) in predict() does not match "
            f"the manifold's ambient dimension ({self.pm.ambient_dim})."
        )

    if not hasattr(self, "mu_") or self.mu_ is None:
        raise RuntimeError("The RFK model has not been fitted yet. Call 'fit' before 'predict'.")

    with torch.no_grad():
        dmat = self.pm.dist(X, self.mu_)  # X is (N,D), mu_ is (K,D) -> dmat is (N,K)
        inv = dmat.pow(-2 / (self.m - 1)) + 1e-8  # Add epsilon
        u = inv / (inv.sum(dim=1, keepdim=True) + 1e-8)  # Add epsilon
        labels = torch.argmax(u, dim=1).cpu().numpy()
    return labels