Kappa Gcn¶

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