Predictors¶

@torch.no_grad()  # type: ignore
def fit(self, X: Float[torch.Tensor, "batch ambient_dim"], y: Real[torch.Tensor, "batch"]) -> ProductSpaceDT:
    """Reworked fit function for new version of ProductDT.

    Args:
        X: (batch, ambient_dim) tensor of trainind data (ambient coordinate representation)
        y: (batch,) tensor of labels (integer representation)

    Returns:
        None (fits tree in place)
    """
    # Pre-preprocessing step: aggregate special dimensions into a new Euclidean component
    if self.use_special_dims:
        X, self.pm = self._aggregate_special_dims(X)

    # Preprocess data
    angles, labels, comparisons_reshaped = self._preprocess(X=X, y=y)

    # Fit node
    self.tree = self._fit_node(angles=angles, labels=labels, comparisons=comparisons_reshaped, depth=self.max_depth)

    self.is_fitted_ = True  # Mark the model as fitted
    return self

@torch.no_grad()  # type: ignore
def predict_proba(self, X: Float[torch.Tensor, "batch intrinsic_dim"]) -> Float[torch.Tensor, "batch n_classes"]:
    """Predict class probabilities for samples in X."""
    if self.use_special_dims:
        X, _ = self._aggregate_special_dims(X)
    angles, _, _ = self._preprocess(X=X)
    if self.permutations is not None:
        angles = angles[:, self.permutations]
    return torch.vstack([self._traverse(angles_row, self.tree).probs for angles_row in angles])

@torch.no_grad()  # type: ignore
def fit(self, X: Float[torch.Tensor, "batch ambient_dim"], y: Real[torch.Tensor, "batch"]) -> ProductSpaceRF:
    """Preprocess and fit an ensemble of trees on subsampled data."""
    # Pre-preprocessing step: aggregate special dimensions
    if self.use_special_dims:
        X, self.pm = self.trees[0]._aggregate_special_dims(X)
        for tree in self.trees:
            tree.pm = self.pm

    # Can use any tree to preprocess X and y
    angles, labels, comparisons = self.trees[0]._preprocess(X=X, y=y)

    # Also update angle2man and special_first
    for tree in self.trees:
        tree.angle2man = self.trees[0].angle2man
        tree.special_first = self.trees[0].special_first
        tree.classes_ = self.trees[0].classes_
    self.classes_ = self.trees[0].classes_

    # Use seed here
    if self.random_state is not None:
        torch.manual_seed(self.random_state)

    # Subsample - just the indices
    n, d = angles.shape
    idx_sample_all, idx_dim_all = self._generate_subsample(n_rows=n, n_cols=d, n_trees=self.n_estimators)

    # Fit trees
    for tree, idx_sample, idx_dim in zip(self.trees, idx_sample_all, idx_dim_all, strict=False):
        tree.permutations = idx_dim
        if self.batched:
            comparisons_subsample = comparisons[idx_sample][:, idx_dim][:, :, idx_sample]
        else:
            comparisons_subsample = comparisons
        tree.tree = tree._fit_node(
            angles=angles[idx_sample][:, idx_dim],
            labels=labels[idx_sample],
            comparisons=comparisons_subsample,
            depth=self.max_depth,
        )

    self.is_fitted_ = True
    return self

@torch.no_grad()  # type: ignore
def predict_proba(self, X: Float[torch.Tensor, "batch intrinsic_dim"]) -> Float[torch.Tensor, "batch n_classes"]:
    """Predict class probabilities for samples in X."""
    return torch.stack([tree.predict_proba(X) for tree in self.trees]).mean(dim=0)

def forward(
    self,
    X: Float[torch.Tensor, "n_nodes dim"],
    A_hat: Float[torch.Tensor, "n_nodes n_nodes"] | None = None,
    aggregate_logits: bool = True,
    softmax: bool = False,
) -> (
    Float[torch.Tensor, "n_nodes n_classes"]
    | Float[torch.Tensor, "n_nodes"]
    | Float[torch.Tensor, "n_nodes n_nodes"]
):
    """Forward pass through the GCN layers and output head."""
    # Pass through main GCN layers
    H = self.gcn_layers(X, A_hat)

    # Task-specific output using the specialized layers
    if self.task == "link_prediction":
        return self.output_layer(H)  # Flattened for link prediction
    else:
        # For classification/regression, use stereographic logits
        logits = self.output_layer(H, A_hat, aggregate_logits=aggregate_logits)

        if softmax:
            logits = torch.softmax(logits, dim=-1)

        return logits.squeeze()

def fit(
    self,
    X: Float[torch.Tensor, "n_nodes dim"],
    y: Real[torch.Tensor, "n_nodes"],
    A: Float[torch.Tensor, "n_nodes n_nodes"] | None = None,
    epochs: int = 2_000,
    lr: float = 1e-2,
    use_tqdm: bool = True,
    tqdm_prefix: str | None = None,
) -> KappaGCN:
    """Fit the Kappa GCN model.

    Args:
        X (torch.Tensor): Feature matrix.
        y (torch.Tensor): Labels for training nodes.
        A (torch.Tensor): Adjacency or distance matrix.
        epochs: Number of training epochs (default=200).
        lr: Learning rate (default=1e-2).
        use_tqdm: Whether to use tqdm for progress bar.
        tqdm_prefix: Prefix for tqdm progress bar.
    """
    # Copy everything
    X = X.clone()
    y = y.clone()
    A = A.clone() if A is not None else None

    # Convert A to A_hat
    A_hat = get_A_hat(A, make_symmetric=True, add_self_loops=True) if A is not None else None

    # Collect all paramters
    euclidean_params = []
    riemannian_params = []
    for layer in self.gcn_layers.layers:
        euclidean_params.append(layer.W)
    if self.task == "link_prediction":
        euclidean_params += [self.output_layer.temperature, self.output_layer.bias]
    else:
        euclidean_params += [self.output_layer.W]
        riemannian_params += [self.output_layer.p_ks]

    # Optimizers
    opt = torch.optim.Adam(euclidean_params, lr=lr)
    ropt = geoopt.optim.RiemannianAdam(riemannian_params, lr=lr) if riemannian_params else None

    if self.task == "classification":
        loss_fn = torch.nn.CrossEntropyLoss()
        y = y.long()
    elif self.task == "regression":
        loss_fn = torch.nn.MSELoss()
        y = y.float()
    elif self.task == "link_prediction":
        loss_fn = torch.nn.BCEWithLogitsLoss()
        # y = y.flatten().float()
        y = y.float()
    else:
        raise ValueError("Invalid task!")

    self.train()
    if use_tqdm:
        my_tqdm = tqdm(total=epochs, desc=tqdm_prefix)

    losses = []
    for i in range(epochs):
        opt.zero_grad()
        if riemannian_params:
            ropt.zero_grad()  # type: ignore
        y_pred = self(X, A_hat)
        loss = loss_fn(y_pred, y)
        loss.backward()
        opt.step()
        if riemannian_params:
            ropt.step()  # type: ignore

        # Progress bar
        if use_tqdm:
            my_tqdm.update(1)
            my_tqdm.set_description(f"Epoch {i + 1}/{epochs}, Loss: {loss.item():.4f}")

        # Early termination for nan loss
        if torch.isnan(loss):
            print("Loss is NaN, stopping training.")
            break
        losses.append(loss.item())

    if use_tqdm:
        my_tqdm.close()

    self.is_fitted_ = True
    self.loss_history_["train"] = losses
    return self

def predict_proba(
    self, X: Float[torch.Tensor, "n_nodes dim"], A: Float[torch.Tensor, "n_nodes n_nodes"] | None = None
) -> Real[torch.Tensor, "n_nodes n_classes"] | Real[torch.Tensor, "n_nodes"]:
    """Predict class probabilities using the trained Kappa GCN.

    Args:
        X (torch.Tensor): Feature matrix (NxD).
        A (torch.Tensor): Adjacency or distance matrix (NxN).

    Returns:
        torch.Tensor: Predicted class probabilities / regression targets.
    """
    # Copy everything
    X = X.clone()
    A = A.clone() if A is not None else None
    A_hat = get_A_hat(A, make_symmetric=True, add_self_loops=True) if A is not None else None

    # Get edges for test set
    self.eval()
    y_pred = self(X, A_hat)
    return y_pred

def fit(
    self, X: Float[torch.Tensor, "n_samples n_manifolds"], y: Int[torch.Tensor, "n_samples"]
) -> ProductSpacePerceptron:
    """Trains the perceptron model using the provided data and labels.

    Args:
        X: Training data tensor.
        y: Class labels for the training data.

    Returns:
        self: Fitted perceptron model.
    """
    # Identify unique classes for multiclass classification
    self._store_classes(y)
    n_samples = X.shape[0]

    # Precompute kernel matrix
    Ks, _ = product_kernel(self.pm, X, None)
    K = torch.ones((n_samples, n_samples), dtype=X.dtype, device=X.device)
    for K_m, w in zip(Ks, self.weights, strict=False):
        K += w * K_m

    # Store training data and labels for prediction
    self.X_train_ = X
    self.y_train_ = y

    # Initialize dictionary to store alpha coefficients for each class
    self.alpha = {}

    # For patience checking
    best_epoch, least_errors = 0, n_samples + 1

    for class_label in self.classes_:
        class_label_item = class_label.item()

        # One-vs-rest labels
        y_binary = torch.where(y == class_label_item, 1, -1)  # Shape: (n_samples,)

        # Initialize alpha coefficients for this class
        alpha = torch.zeros(n_samples, dtype=X.dtype, device=X.device)

        for epoch in range(self.max_epochs):
            # Compute decision function: f = K @ (alpha * y_binary)
            f = K @ (alpha * y_binary)  # Shape: (n_samples,)

            # Compute predictions
            predictions = torch.sign(f)

            # Find misclassified samples
            misclassified = predictions != y_binary

            # If no misclassifications, break early
            if not misclassified.any():
                break

            # Test patience
            n_errors = misclassified.sum().item()
            if n_errors < least_errors:
                best_epoch, least_errors = epoch, n_errors
            if epoch - best_epoch >= self.patience:
                break

            # Update alpha coefficients for misclassified samples
            alpha[misclassified] += 1

        # Store the alpha coefficients for the current class
        self.alpha[class_label_item] = alpha

    self.is_fitted_ = True
    return self

def predict_proba(
    self,
    X: Float[torch.Tensor, "n_points n_features"],  # type: ignore[override]
) -> Float[torch.Tensor, "n_points n_classes"]:
    """Predicts the decision values for each class.

    Args:
        X: Test data tensor.

    Returns:
        decision_values: Decision values for each test sample and each class.
    """
    n_samples = X.shape[0]
    n_classes = len(self.classes_)
    decision_values = torch.zeros((n_samples, n_classes), dtype=X.dtype, device=X.device)

    # Compute kernel matrix between training data and test data
    Ks, _ = product_kernel(self.pm, self.X_train_, X)
    K_test = torch.ones((self.X_train_.shape[0], n_samples), dtype=X.dtype, device=X.device)
    for K_m, w in zip(Ks, self.weights, strict=False):
        K_test += w * K_m
    # K_test = self.X_train_ @ X.T

    for idx, class_label in enumerate(self.classes_):
        class_label_item = class_label.item()
        alpha = self.alpha[class_label_item]  # Shape: (n_samples_train,)
        y_binary = torch.where(self.y_train_ == class_label_item, 1, -1)  # Shape: (n_samples_train,)

        # Compute decision function for test samples
        f = (alpha * y_binary) @ K_test  # Shape: (n_samples_test,)
        decision_values[:, idx] = f

    return decision_values

def fit(
    self,
    X: Float[torch.Tensor, "n_samples n_manifolds"],
    y: Int[torch.Tensor, "n_samples"],
) -> ProductSpaceSVM:
    """Fit one-vs-rest SVMs on the product manifold data.

    Args:
        X: Training points tensor.
        y: Integer class labels tensor.

    Returns:
        self: Fitted ProductSpaceSVM instance.
    """
    # unique classes
    self._store_classes(y)
    n = X.shape[0]

    # aggregated kernel
    Ks, _ = product_kernel(self.pm, X, None)
    K_sum = torch.ones((n, n), dtype=X.dtype, device=X.device)
    for K_m, w in zip(Ks, self.weights, strict=False):
        K_sum += w * K_m

    X_np = X.detach().cpu().numpy()
    K_np = K_sum.detach().cpu().numpy()

    def sqrtm_psd(P: np.ndarray) -> Any:
        w, V = np.linalg.eigh(P)
        w_s = np.sqrt(np.clip(w, 0, None))
        B = V @ np.diag(w_s) @ V.T
        return (B + B.T) * 0.5

    # containers
    self.beta = {}
    self.zeta = {}
    self.epsilon = {}
    self.b = {}

    for cls in self.classes_:
        cls_item = cls.item() if isinstance(cls, torch.Tensor) else cls
        # one-vs-rest labels: +1 for cls, -1 for others
        y_bin = torch.where(y == cls_item, 1, -1)
        Y = torch.diagflat(y_bin).detach().cpu().numpy()

        # variables
        beta_var = cp.Variable(n)
        zeta = cp.Variable(n, nonneg=True)
        eps_var = cp.Variable(1)
        b_var = cp.Variable(1)

        # base constraints
        constraints = [eps_var >= 0]
        constraints.append(Y @ (K_np @ beta_var + b_var) >= eps_var - zeta)

        # per-manifold SOC
        for M, K_comp in zip(self.pm.P, Ks, strict=False):
            P_np = K_comp.detach().cpu().numpy()
            if M.type == "E" and self.e_constraints:
                B = sqrtm_psd(P_np)
                constraints.append(cp.norm(B @ beta_var, 2) <= 1.0)
            elif M.type == "S" and self.s_constraints:
                B = sqrtm_psd(P_np)
                constraints.append(cp.norm(B @ beta_var, 2) <= np.sqrt(np.pi / 2))
            elif M.type == "H" and self.h_constraints:
                # PSD split
                eigvals, eigvecs = np.linalg.eigh(P_np)
                plus = np.clip(eigvals, 0, None)
                minus = np.clip(-eigvals, 0, None)
                Kp = (eigvecs @ np.diag(plus) @ eigvecs.T + (eigvecs @ np.diag(plus) @ eigvecs.T).T) * 0.5
                Km = (eigvecs @ np.diag(minus) @ eigvecs.T + (eigvecs @ np.diag(minus) @ eigvecs.T).T) * 0.5
                Bp = sqrtm_psd(Kp)
                Bm = sqrtm_psd(Km)

                C_H = abs(M.curvature)
                R = -M.scale
                r_h = abs(np.arcsinh(-(R**2) * C_H))
                r = self.eps

                constraints.append(cp.norm(Bm @ beta_var, 2) <= np.sqrt(max(r, 0.0)))
                constraints.append(cp.norm(Bp @ beta_var, 2) <= np.sqrt(max(r + r_h, 0.0)))

        # solve
        prob = cp.Problem(cp.Minimize(-eps_var + cp.sum(zeta)), constraints)
        prob.solve(solver="SCS")

        # save results
        self.beta[cls_item] = np.ravel(beta_var.value)
        self.zeta[cls_item] = zeta.value
        self.epsilon[cls_item] = float(eps_var.value)
        self.b[cls_item] = float(b_var.value)

    # store training data
    self.X_train_ = torch.tensor(X_np, dtype=torch.float32)
    self.is_fitted_ = True
    return self

def predict_proba(
    self,
    X: Float[torch.Tensor, "n_samples n_manifolds"],
) -> Float[torch.Tensor, "n_samples n_classes"]:
    """Predict class probabilities using the fitted SVMs.

    Args:
        X: Test points tensor.

    Returns:
        class_probabilities: Class probabilities for each test sample.
    """
    X_tensor = torch.tensor(X, dtype=torch.float32) if not isinstance(X, torch.Tensor) else X
    X_tensor = X_tensor.to(self.X_train_.device)

    Ks_test, _ = product_kernel(self.pm, self.X_train_, X_tensor)
    Kt = torch.ones((self.X_train_.shape[0], X_tensor.shape[0]), device=X_tensor.device)
    for K_m, w in zip(Ks_test, self.weights, strict=False):
        Kt += w * K_m
    Kt_np = Kt.detach().cpu().numpy()

    n_test = X_tensor.shape[0]
    n_cls = len(self.classes_)
    dec = np.zeros((n_test, n_cls))
    for idx, cls in enumerate(self.classes_):
        cls_item = cls.item() if isinstance(cls, torch.Tensor) else cls
        beta_vec: np.ndarray = np.ravel(self.beta[cls_item])
        dec[:, idx] = Kt_np.T @ beta_vec + self.b[cls_item]

    exp_scores = np.exp(dec - dec.max(axis=1, keepdims=True))
    probs = exp_scores / exp_scores.sum(axis=1, keepdims=True)
    return torch.tensor(probs, dtype=torch.float32)