Node classification (NC)¶

class NCClassifier(L.LightningModule):
    """
    A LightningModule base class for multiclass node-classification models.

    Attributes:
        encoder: Optional encoder module. Defaults to ``None``.
        classifier: Module that maps node features to class logits.
        loss_fn: Multiclass loss function.
        metrics_log_kwargs: Keyword arguments passed to metric log calls.
        train_metrics: Optional training metrics.
        val_metrics: Optional validation metrics.
        test_metrics: Optional test metrics.
    """

    def __init__(
        self,
        classifier: nn.Module,
        loss_fn: nn.Module,
        encoder: nn.Module | None = None,
        metrics: MetricCollection | None = None,
        metrics_log_kwargs: dict[str, Any] | None = None,
    ) -> None:
        """
        Initialize the NC classifier.

        Args:
            classifier: Module producing one class-logit row per node.
            loss_fn: Loss function used on supervised target nodes.
            encoder: Optional encoder module. Defaults to ``None``.
            metrics: Optional metric collection cloned independently per stage.
            metrics_log_kwargs: Additional keyword arguments passed to metric log calls.
        """
        super().__init__()
        self.encoder: nn.Module | None = encoder
        self.classifier: nn.Module = classifier
        self.loss_fn: nn.Module = loss_fn
        self.metrics_log_kwargs: dict[str, Any] = metrics_log_kwargs or {}

        if metrics is not None:
            self.train_metrics: MetricCollection | None = metrics.clone(
                prefix=stage_metric_prefix(Stage.TRAIN)
            )
            self.val_metrics: MetricCollection | None = metrics.clone(
                prefix=stage_metric_prefix(Stage.VAL)
            )
            self.test_metrics: MetricCollection | None = metrics.clone(
                prefix=stage_metric_prefix(Stage.TEST)
            )
        else:
            self.train_metrics = None
            self.val_metrics = None
            self.test_metrics = None

    def _compute_loss(
        self,
        logits: Tensor,
        labels: Tensor,
        batch_size: int,
        stage: Stage,
    ) -> Tensor:
        """
        Compute and log loss based on logits and labels.

        Args:
            logits: The predicted logits from the model.
            labels: The true labels corresponding to the logits.
            batch_size: The size of the current batch, used for logging.
            stage: The current stage (train/val/test) for logging purposes.

        Returns:
            loss: The computed loss tensor.
        """
        loss = self.loss_fn(logits, labels)
        self.log(
            name=stage_metric_name(stage, "loss"),
            value=loss,
            prog_bar=True,
            batch_size=batch_size,
            **self.metrics_log_kwargs,
        )
        return loss

    def _compute_metrics(
        self,
        logits: Tensor,
        labels: Tensor,
        batch_size: int,
        stage: Stage,
    ) -> None:
        """
        Compute and log metrics based on logits and labels.

        Uses class-based torchmetrics with proper multi-batch accumulation:
        1. ``update()`` accumulates predictions/targets across batches.
        2. Passing the MetricCollection to ``self.log_dict()`` tells Lightning to call
            ``compute()`` at epoch end and ``reset()`` automatically.

        Args:
            logits: The predicted logits from the model.
            labels: The true labels corresponding to the logits.
            batch_size: The size of the current batch, used for logging.
            stage: The current stage (train/val/test) for logging purposes.
        """
        stage_metrics = self._get_stage_metrics(stage)
        if stage_metrics is None:
            return  # No metrics to compute for this stage
        self._configure_metric_distributed_available(stage_metrics)

        # Accumulate predictions/targets for this batch
        stage_metrics.update(logits, labels)
        self.log_dict(
            stage_metrics,
            prog_bar=True,
            on_step=False,
            on_epoch=True,  # Compute and log metrics at epoch end for proper accumulation
            batch_size=batch_size,
            **self.metrics_log_kwargs,
        )

    def _get_stage_metrics(self, stage: Stage) -> MetricCollection | None:
        """
        Return the metric collection for the given stage.
        """
        match stage:
            case Stage.TRAIN:
                return self.train_metrics
            case Stage.VAL:
                return self.val_metrics
            case Stage.TEST:
                return self.test_metrics
            case _:
                raise ValueError(f"Unrecognized stage: {stage}")

    def _configure_metric_distributed_available(self, metrics: MetricCollection) -> None:
        """
        Make torchmetrics sync decisions follow the active Lightning trainer.
        """
        for metric in metrics.values(copy_state=False):
            metric.distributed_available_fn = self._distributed_available_fn

    def _distributed_available_fn(self) -> bool:
        """
        Return whether metrics should synchronize for the current trainer.
        """
        if self._trainer is None:
            return False
        return (
            self._trainer.world_size > 1
            and torch.distributed.is_available()
            and torch.distributed.is_initialized()
        )

    def _target_logits_and_labels(self, logits: Tensor, batch: HData) -> tuple[Tensor, Tensor]:
        """
        Select supervised node logits and labels from a batch.

        Args:
            logits: Node logits of shape ``[num_nodes, num_classes]``.
            batch: Batch containing node labels and an optional target mask.

        Returns:
            target_logits: Logits for supervised nodes.
            target_labels: Labels for supervised nodes.
        """
        target_node_mask = batch.target_node_mask
        target_logits = logits[target_node_mask]
        target_labels = batch.y[target_node_mask]
        return target_logits, target_labels

class CommonNeighborsClassifier(NCClassifier):
    """
    A LightningModule for the CommonNeighbors-based NC classifier.

    Attributes:
        encoder: Optional encoder module inherited from ``NCClassifier``. Defaults to ``None``.
        classifier: Identity placeholder inherited from ``NCClassifier``.
        loss_fn: Loss function inherited from ``NCClassifier``.
        metrics_log_kwargs: Metric logging keyword arguments inherited from ``NCClassifier``.
        train_metrics: Optional training metrics inherited from ``NCClassifier``.
        val_metrics: Optional validation metrics inherited from ``NCClassifier``.
        test_metrics: Optional test metrics inherited from ``NCClassifier``.
        node_to_neighbors: Precomputed training-world node neighborhoods.
        automatic_optimization: Disabled because this module has no trainable optimization.
    """

    def __init__(
        self,
        train_hdata: HData,
        num_classes: int,
        aggregation: Literal["mean", "min", "sum"] = "sum",
        exclude_self_reference: bool = True,
        classifier: nn.Module | None = None,
        loss_fn: nn.Module | None = None,
        metrics: MetricCollection | None = None,
        metrics_log_kwargs: dict[str, Any] | None = None,
    ) -> None:
        """
        Initialize the CommonNeighbors-based NC classifier.

        Args:
            train_hdata: Training data used to precompute neighborhoods and labeled
                reference nodes.
            num_classes: Number of node classes.
            aggregation: Method used to aggregate reference-node scores per class.
                Defaults to ``"sum"``.
            exclude_self_reference: Whether a node should ignore itself when it also
                appears among labeled reference nodes. Defaults to ``True``.
            classifier: Optional classifier module.
            loss_fn: Optional loss function. Defaults to ``CrossEntropyLoss``.
            metrics: Optional metric collection for evaluation. Defaults to ``None``.
            metrics_log_kwargs: Additional keyword arguments passed to metric log calls.
                Defaults to ``None``.
        """
        super().__init__(
            classifier=classifier
            if classifier is not None
            else self.__default_classifier(
                train_hdata=train_hdata,
                num_classes=num_classes,
                aggregation=aggregation,
                exclude_self_reference=exclude_self_reference,
            ),
            loss_fn=loss_fn if loss_fn is not None else nn.CrossEntropyLoss(),
            metrics=metrics,
            metrics_log_kwargs=metrics_log_kwargs,
        )

        # Pre-compute neighbors of training nodes based on training edges only
        # to create a "known world" for the model to make predictions from
        self.node_to_neighbors: dict[int, Neighborhood] = Hypergraph.from_hyperedge_index(
            hyperedge_index=HyperedgeIndex(
                hyperedge_index=train_hdata.hyperedge_index,
            )
            .to_global(global_node_ids=train_hdata.global_node_ids)
            .item
        ).neighbors_of_all()

        self.automatic_optimization: bool = False

    def forward(self, global_node_ids: Tensor) -> Tensor:
        """
        Compute common-neighbor class scores for the given nodes.

        Args:
            global_node_ids: Stable node IDs matching the training-world node IDs.

        Returns:
            logits: Node-class scores of shape ``(num_nodes, num_classes)``.
        """
        return self.classifier(
            candidate_nodes=global_node_ids,
            node_to_neighbors=self.node_to_neighbors,
        )

    def on_fit_start(self) -> None:
        """
        Warn users if they are running unnecessary training epochs.
        """
        if self.trainer.max_epochs is None or self.trainer.max_epochs > 0:
            warnings.warn(
                f"{self.__class__.__name__} is a non-trainable heuristic model. "
                "No optimization occurs. Set max_epochs=0 in your trainer for instant evaluation.",
                UserWarning,
                stacklevel=2,
            )

    def training_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Return a zero loss for the non-trainable heuristic.

        Args:
            batch: Input batch, unused.
            batch_idx: Batch index, unused.

        Returns:
            loss: Zero scalar tensor on the module device.
        """
        return torch.tensor(0.0, dtype=torch.float, device=self.device)

    def validation_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Return a zero validation loss for the non-trainable heuristic.

        Args:
            batch: Input batch, unused.
            batch_idx: Batch index, unused.

        Returns:
            loss: Zero scalar tensor on the module device.
        """
        return torch.tensor(0.0, dtype=torch.float, device=self.device)

    def test_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Evaluate a test batch.

        Args:
            batch: Test batch.
            batch_idx: Batch index, unused.

        Returns:
            loss: Test loss.
        """
        return self.__step(batch, stage=Stage.TEST)

    def predict_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Predict common-neighbor class scores for a batch.

        Args:
            batch: Prediction batch.
            batch_idx: Batch index, unused.

        Returns:
            logits: Predicted node-class scores.
        """
        return self.forward(batch.global_node_ids)

    def configure_optimizers(self) -> None:
        """
        Return no optimizer for the non-trainable heuristic.

        Returns:
            optimizer: No optimizer is configured, so always returns ``None``.
        """
        # No training, so no optimizers needed

    def __step(self, batch: HData, stage: Stage) -> Tensor:
        """
        Shared evaluation logic for all stages.

        Args:
            batch: `HData` object containing the hypergraph.
            stage: Current evaluation stage.

        Returns:
            loss: Computed loss.
        """
        scores = self.forward(batch.global_node_ids)
        target_logits, target_labels = self._target_logits_and_labels(scores, batch)
        batch_size = int(target_labels.size(0))

        loss = self._compute_loss(target_logits, target_labels, batch_size, stage)
        self._compute_metrics(target_logits, target_labels, batch_size, stage)

        return loss

    def __class_to_node_ids(self, train_hdata: HData, num_classes: int) -> dict[int, list[int]]:
        """
        Group labeled training nodes by class.

        Args:
            train_hdata: Training data containing node labels and train node mask.
            num_classes: Number of node classes.

        Returns:
            class_to_node_ids: Dictionary mapping class IDs to lists of node IDs.
        """
        class_to_node_ids: dict[int, list[int]] = {class_id: [] for class_id in range(num_classes)}
        train_node_ids = train_hdata.global_node_ids[train_hdata.target_node_mask].tolist()
        train_class_labels = train_hdata.y[train_hdata.target_node_mask].tolist()

        for node_id, class_label in zip(train_node_ids, train_class_labels, strict=True):
            class_to_node_ids[int(class_label)].append(int(node_id))

        return class_to_node_ids

    def __default_classifier(
        self,
        train_hdata: HData,
        num_classes: int,
        aggregation: Literal["mean", "min", "sum"],
        exclude_self_reference: bool = True,
    ) -> CommonNeighbors:
        """
        Create a default CommonNeighbors classifier with a CommonNeighborsNodeScorer.

        Args:
            train_hdata: Training data containing node labels and train node mask.
            num_classes: Number of node classes.
            aggregation: Method used to aggregate reference-node scores per class.
            exclude_self_reference: Whether a node should ignore itself when it also
                appears among labeled reference nodes.

        Returns:
            classifier: A CommonNeighbors classifier with a CommonNeighborsNodeScorer.
        """
        class_to_node_ids = self.__class_to_node_ids(
            train_hdata=train_hdata,
            num_classes=num_classes,
        )

        scorer = CommonNeighborsNodeScorer(
            num_classes=num_classes,
            class_to_node_ids=class_to_node_ids,
            aggregation=aggregation,
            exclude_self_reference=exclude_self_reference,
        )

        return CommonNeighbors(aggregation=aggregation, scorer=scorer)

class GCNClassifierConfig(TypedDict):
    """
    Configuration for the GCN classifier in ``GCNClassifier``.

    Attributes:
        in_channels: Number of input features per node.
        out_channels: Number of node classes.
        hidden_channels: Number of hidden units in intermediate GCN layers.
        num_layers: Number of GCN layers. Defaults to ``2``.
        drop_rate: Dropout rate applied after each hidden GCN layer. Defaults to ``0.0``.
        bias: Whether to include bias terms. Defaults to ``True``.
        improved: Whether to use the improved GCN normalization. Defaults to ``False``.
        add_self_loops: Whether to add self-loops before convolution. Defaults to ``True``.
        normalize: Whether to normalize the adjacency matrix in ``GCNConv``. Defaults to ``True``.
        cached: Whether to cache the normalized graph in ``GCNConv``. Defaults to ``False``.
        graph_reduction_strategy: Strategy for reducing the hypergraph to a graph.
            Defaults to ``"clique_expansion"``.
        num_nodes: Optional total number of nodes to use during graph reduction.
        activation_fn: Activation function to use after each hidden layer.
            Defaults to ``nn.ReLU``.
        activation_fn_kwargs: Keyword arguments for the activation function.
            Defaults to empty dict.
    """

    in_channels: int
    out_channels: int
    hidden_channels: NotRequired[int]
    num_layers: NotRequired[int]
    drop_rate: NotRequired[float]
    bias: NotRequired[bool]
    improved: NotRequired[bool]
    add_self_loops: NotRequired[bool]
    normalize: NotRequired[bool]
    cached: NotRequired[bool]
    graph_reduction_strategy: NotRequired[GraphReductionStrategy]
    num_nodes: NotRequired[int]
    activation_fn: NotRequired[ActivationFn]
    activation_fn_kwargs: NotRequired[dict]

class GCNClassifier(NCClassifier):
    """
    A LightningModule for GCN-based NC classifier.

    Reduces the hypergraph to a graph, then applies GCN
    layers to produce per-node class logits.

    Attributes:
        encoder: Optional encoder module inherited from ``NCClassifier``. Defaults to ``None``.
        classifier: GCN classifier module inherited from ``NCClassifier``.
        loss_fn: Loss function inherited from ``NCClassifier``.
        metrics_log_kwargs: Metric logging keyword arguments inherited from ``NCClassifier``.
        train_metrics: Optional training metrics inherited from ``NCClassifier``.
        val_metrics: Optional validation metrics inherited from ``NCClassifier``.
        test_metrics: Optional test metrics inherited from ``NCClassifier``.
        lr: Learning rate for the optimizer. Defaults to ``0.01``.
        weight_decay: L2 regularization. Defaults to ``5e-4``.
    """

    def __init__(
        self,
        classifier_config: GCNClassifierConfig,
        loss_fn: nn.Module | None = None,
        lr: float = 0.01,
        weight_decay: float = 5e-4,
        metrics: MetricCollection | None = None,
        metrics_log_kwargs: dict[str, Any] | None = None,
    ) -> None:
        """
        Initialize the GCN-based NC classifier.

        Args:
            classifier_config: Configuration for the GCN classifier.
            loss_fn: Optional loss function. Defaults to ``CrossEntropyLoss``.
            lr: Learning rate for the optimizer. Defaults to ``0.01``.
            weight_decay: L2 regularization. Defaults to ``5e-4``.
            metrics: Optional metric collection for evaluation. Defaults to ``None``.
            metrics_log_kwargs: Additional keyword arguments passed to metric log calls.
                Useful for configuring distributed synchronization behavior
                of ``torchmetrics``. Defaults to ``None``.
        """
        classifier = GCN(
            in_channels=classifier_config["in_channels"],
            out_channels=classifier_config["out_channels"],
            hidden_channels=classifier_config.get("hidden_channels"),
            num_layers=classifier_config.get("num_layers", 2),
            drop_rate=classifier_config.get("drop_rate", 0.0),
            bias=classifier_config.get("bias", True),
            activation_fn=classifier_config.get("activation_fn"),
            activation_fn_kwargs=classifier_config.get("activation_fn_kwargs"),
            improved=classifier_config.get("improved", False),
            add_self_loops=classifier_config.get("add_self_loops", True),
            normalize=classifier_config.get("normalize", True),
            cached=classifier_config.get("cached", False),
        )

        super().__init__(
            classifier=classifier,
            loss_fn=loss_fn if loss_fn is not None else nn.CrossEntropyLoss(),
            metrics=metrics,
            metrics_log_kwargs=metrics_log_kwargs,
        )

        self.classifier_config: GCNClassifierConfig = classifier_config
        self.lr: float = lr
        self.weight_decay: float = weight_decay

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Predict node-class logits from node features and hypergraph structure.

        Examples:
            Given 4 nodes with 3 features each and 2 output classes:
                >>> x.shape
                ... torch.Size([4, 3])
                >>> hyperedge_index.shape
                ... torch.Size([2, 6])

            The forward pass maps each node to one row of class logits:
                >>> logits = model.forward(x, hyperedge_index)
                >>> logits.shape
                ... torch.Size([4, 2])

        Args:
            x: Node feature matrix of shape ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge connectivity of shape ``(2, num_incidences)``.

        Returns:
            logits: Node-class logits of shape ``(num_nodes, num_classes)``.
        """
        # Reduce hypergraph to graph and remove self-loops
        reduced_edge_index = HyperedgeIndex(hyperedge_index).reduce(
            strategy=self.classifier_config.get(
                "graph_reduction_strategy",
                GraphReductionStrategyEnum.CLIQUE_EXPANSION,
            ),
            num_nodes=self.classifier_config.get("num_nodes"),
        )
        edge_index = EdgeIndex(reduced_edge_index).remove_selfloops().item

        return self.classifier(x, edge_index)

    def training_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Run a training step.

        Args:
            batch: Training batch.
            batch_idx: Batch index, unused.

        Returns:
            loss: Training loss.
        """
        return self.__eval_step(batch, Stage.TRAIN)

    def validation_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Run a validation step.

        Args:
            batch: Validation batch.
            batch_idx: Batch index, unused.

        Returns:
            loss: Validation loss.
        """
        return self.__eval_step(batch, Stage.VAL)

    def test_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Run a test step.

        Args:
            batch: Test batch.
            batch_idx: Batch index, unused.

        Returns:
            loss: Test loss.
        """
        return self.__eval_step(batch, Stage.TEST)

    def predict_step(self, batch: HData, batch_idx: int) -> Tensor:
        """
        Predict node-class logits for a batch.

        Args:
            batch: Prediction batch.
            batch_idx: Batch index, unused.

        Returns:
            logits: Predicted node-class logits.
        """
        return self.forward(batch.x, batch.hyperedge_index)

    def configure_optimizers(self) -> optim.Adam:
        """
        Configure the optimizer.

        Returns:
            optimizer: Adam optimizer.
        """
        return optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)

    def __eval_step(self, batch: HData, stage: Stage) -> Tensor:
        """
        Run shared evaluation logic for a stage.

        Args:
            batch: Input batch.
            stage: Current evaluation stage.

        Returns:
            loss: Computed loss.
        """
        logits = self.forward(batch.x, batch.hyperedge_index)
        target_logits, target_labels = self._target_logits_and_labels(logits, batch)
        batch_size = int(target_labels.size(0))

        loss = self._compute_loss(target_logits, target_labels, batch_size, stage)
        self._compute_metrics(target_logits, target_labels, batch_size, stage)
        return loss