HLP (Hyperlink prediction)¶

class HlpModule(L.LightningModule):
    """
    A LightningModule for HLP models with optional negative sampling.

    Attributes:
        encoder: Optional encoder module. Defaults to ``None`` as not
            all HLP model use an encoder.
        decoder: Decoder module to use to predict whether hyperedges are positive or negative.
        loss_fn: Loss function.
        metrics_log_kwargs: Additional keyword arguments to pass to all ``self.log`` calls
            for metrics. Useful for configuring distributed synchronization behavior of
            ``torchmetrics``. Defaults to ``None``.
        train_metrics: Optional metric collection for training.
        val_metrics: Optional metric collection for validation.
        test_metrics: Optional metric collection for testing.
    """

    def __init__(
        self,
        decoder: nn.Module,
        loss_fn: nn.Module,
        encoder: nn.Module | None = None,
        metrics: MetricCollection | None = None,
        metrics_log_kwargs: dict[str, Any] | None = None,
        negative_sampler: NegativeSampler | None = None,
        negative_sampling_schedule: NegativeSamplingSchedule = "every_epoch",
        negative_sampling_every_n: int = 1,
    ):
        """
        Initialize the HLP Lightning module.

        Args:
            decoder: Decoder module used to score hyperedges.
            loss_fn: Loss function.
            encoder: Optional encoder module. Defaults to ``None``.
            metrics: Optional metric collection cloned independently per stage.
                Defaults to ``None``.
            metrics_log_kwargs: Additional keyword arguments passed to metric log calls.
                Defaults to ``None``.
            negative_sampler: Optional negative sampler. Defaults to ``None``.
            negative_sampling_schedule: Schedule controlling when negatives are sampled.
                Defaults to ``"every_epoch"``.
            negative_sampling_every_n: Epoch interval for ``"every_n_epochs"`` scheduling.
                Defaults to ``1``.
        """
        super().__init__()
        self.encoder: nn.Module | None = encoder
        self.decoder: nn.Module = decoder
        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: MetricCollection | None = None
            self.val_metrics: MetricCollection | None = None
            self.test_metrics: MetricCollection | None = None

        self.__negative_sampling_scheduler: NegativeSamplingScheduler | None = None
        if negative_sampler is not None:
            self.__negative_sampling_scheduler = NegativeSamplingScheduler(
                negative_sampler,
                negative_sampling_schedule,
                negative_sampling_every_n,
            )

    @property
    def negative_sampling_config(self) -> dict[str, Any]:
        """
        Return the configured negative-sampling options.

        Returns:
            config: Scheduler configuration, or an empty dictionary when negative sampling is off.
        """
        if self.__negative_sampling_scheduler is None:
            return {}
        return self.__negative_sampling_scheduler.config

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

        Args:
            scores: The predicted scores from the model.
            labels: The true labels corresponding to the scores.
            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(scores, 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,
        scores: Tensor,
        labels: Tensor,
        batch_size: int,
        stage: Stage,
    ) -> None:
        """
        Compute and log metrics based on scores 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:
            scores: The predicted scores (logits) from the model.
            labels: The true labels corresponding to the scores.
            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
        self._configure_metric_distributed_available(stage_metrics)

        # Apply sigmoid to convert logits to probabilities as BinaryAUROC
        # and BinaryAveragePrecision expect probabilities in [0, 1]
        preds = torch.sigmoid(scores)
        targets = labels.long()

        # Accumulate predictions/targets for this batch
        stage_metrics.update(preds, targets)

        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, or ``None``.

        Args:
            stage: The current stage (train/val/test) for which to get metrics.

        Returns:
            metrics: The metric collection corresponding to the given stage, or ``None``
                if no metrics are configured.
        """
        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.

        TorchMetrics checks ``jit_distributed_available()`` by default. After DDP
        training, the process group can still be alive while a separate test trainer is
        intentionally single-device. In that case metric states must not all-gather.

        Args:
            metrics: The metric collection to configure.
        """
        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.

        This is needed as ``distributed_available_fn`` defaults to a check of
        ``torch.distributed.is_available()`` and ``torch.distributed.is_initialized()``.
        In our case, we can have a single-device test trainer after DDP training,
        so we want to disable metric synchronization if the trainer is not multi-process, not
        only if ``torch.distributed`` is not initialized.
        The issue is that without ``trainer.world_size > 1``, the single-device trainer
        tries to sync across DDP process group and hangs because it's the only one in that group.

        Returns:
            True when the attached trainer is multi-process and torch.distributed is initialized.
        """
        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_scores_and_labels(self, scores: Tensor, batch: HData) -> tuple[Tensor, Tensor]:
        """
        Select supervised hyperedge scores and labels from a full-context HLP batch.

        Args:
            scores: Model scores for all hyperedges in ``batch``.
            batch: Input HData batch containing ``target_hyperedge_mask``.

        Returns:
            target_scores: Scores for supervised target hyperedges.
            target_labels: Labels for supervised target hyperedges.

        Raises:
            ValueError: If the score tensor does not align with the batch hyperedge mask.
        """
        target_hyperedge_mask = batch.target_hyperedge_mask
        target_scores = scores[target_hyperedge_mask]
        target_labels = batch.y[target_hyperedge_mask]
        return target_scores, target_labels

    def _should_sample_negatives(self) -> bool:
        """
        Whether to resample negatives for the current epoch.
        """
        if self.__negative_sampling_scheduler is None:
            raise ValueError(
                "Asked to check negative sampling schedule but no negative sampler is configured."
            )
        return self.__negative_sampling_scheduler.should_sample(self.current_epoch)

    def _sample_negatives(self, batch: HData) -> HData:
        """
        Sample fresh negatives if the schedule requires it, otherwise return cache.

        Args:
            batch: The current batch of data for which to sample negatives.

        Returns:
            negatives: A batch of negative samples, either freshly sampled or from cache.
        """
        if self.__negative_sampling_scheduler is None:
            raise ValueError("Asked to sample negatives but no negative sampler is not configured.")
        return self.__negative_sampling_scheduler.sample(batch, self.current_epoch)

class CommonNeighborsHlpModule(HlpModule):
    """
    A LightningModule for the CommonNeighbors model with optional negative sampling.

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

    def __init__(
        self,
        train_hdata: HData,
        aggregation: Literal["mean", "min", "sum"] = "mean",
        decoder: nn.Module | None = None,
        loss_fn: nn.Module | None = None,
        metrics: MetricCollection | None = None,
        metrics_log_kwargs: dict[str, Any] | None = None,
    ):
        """
        Initialize the CommonNeighbors HLP module.

        Args:
            train_hdata: Training data used to precompute neighborhoods.
            aggregation: Common-neighbor aggregation method. Defaults to ``"mean"``.
            decoder: Optional decoder module. Defaults to ``CommonNeighbors``.
            loss_fn: Optional loss function. Defaults to ``BCEWithLogitsLoss``.
            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__(
            decoder=decoder if decoder is not None else CommonNeighbors(aggregation),
            loss_fn=loss_fn if loss_fn is not None else nn.BCEWithLogitsLoss(),
            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()

        # Disable automatic optimization since there is no training
        self.automatic_optimization: bool = False

    def forward(self, candidate_nodes: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Compute common neighbor scores for the given hyperedges.

        Args:
            candidate_nodes: Tensor containing node IDs to score of shape ``(num_nodes,)``.
            hyperedge_index: Tensor containing incidence information for the hyperedges to score.

        Returns:
            scores: A 1-D tensor of shape (num_hyperedges,) with CN scores.
        """
        return self.decoder(
            candidate_nodes=candidate_nodes,
            hyperedge_index=hyperedge_index,
            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 scores for a batch.

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

        Returns:
            scores: Predicted hyperedge scores.
        """
        return self.forward(
            # Use the global node IDs from the batch to predict scores for the hyperedges
            # as the model is trained on the node IDs from the training data
            candidate_nodes=batch.global_node_ids,
            hyperedge_index=batch.hyperedge_index,
        )

    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: The current stage of evaluation
                (e.g., ``Stage.TRAIN``, ``Stage.VAL``, ``Stage.TEST``).

        Returns:
            loss: The computed loss.
        """
        scores = self.forward(
            # Use the global node IDs from the batch to compute scores for the hyperedges
            # as the model is trained on the node IDs from the training data
            candidate_nodes=batch.global_node_ids,
            hyperedge_index=batch.hyperedge_index,
        )
        target_scores, target_labels = self._target_scores_and_labels(scores, batch)
        batch_size = target_labels.size(0)

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

        return loss

class GCNEncoderConfig(TypedDict):
    """
    Configuration for the GCN encoder in GCNHlpModule.

    Attributes:
        in_channels: Number of input features per node.
        out_channels: Number of output features (embedding size) per node.
        hidden_channels: Number of hidden units in the 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: Total number of nodes in the hypergraph. This is useful when setting is
            transductive but train dataset may not contain all hyperedges where some nodes appear,
            to ensure consistent encoding across splits.
        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 GCNHlpModule(HlpModule):
    """
    A LightningModule for GCN-based HLP.

    Uses a graph reduction of the input hypergraph to run GCN over nodes,
    aggregates node embeddings per hyperedge, and scores each hyperedge with a linear decoder.

    Attributes:
        encoder: GCN encoder module inherited from ``HlpModule``.
        decoder: SLP decoder module inherited from ``HlpModule``.
        loss_fn: Loss function inherited from ``HlpModule``.
        metrics_log_kwargs: Metric logging keyword arguments inherited from ``HlpModule``.
        train_metrics: Optional training metrics inherited from ``HlpModule``.
        val_metrics: Optional validation metrics inherited from ``HlpModule``.
        test_metrics: Optional test metrics inherited from ``HlpModule``.
        encoder_config: Configuration for the GCN encoder.
        aggregation: Method to aggregate node embeddings per hyperedge. Defaults to ``"mean"``.
        lr: Learning rate for the optimizer. Defaults to ``0.001``.
        weight_decay: L2 regularization. Defaults to ``0.0``.
    """

    def __init__(
        self,
        encoder_config: GCNEncoderConfig,
        aggregation: Literal["mean", "max", "min", "sum"] = "mean",
        loss_fn: nn.Module | None = None,
        lr: float = 0.001,
        weight_decay: float = 0.0,
        metrics: MetricCollection | None = None,
        metrics_log_kwargs: dict[str, Any] | None = None,
    ):
        """
        Initialize the GCN HLP module.

        Args:
            encoder_config: Configuration for the GCN encoder.
            aggregation: Method used to aggregate node embeddings per hyperedge.
                Defaults to ``"mean"``.
            loss_fn: Optional loss function. Defaults to ``BCEWithLogitsLoss``.
            lr: Learning rate for the optimizer. Defaults to ``0.001``.
            weight_decay: L2 regularization. Defaults to ``0.0``.
            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``.
        """
        encoder = GCN(
            in_channels=encoder_config["in_channels"],
            out_channels=encoder_config["out_channels"],
            hidden_channels=encoder_config.get("hidden_channels"),
            num_layers=encoder_config.get("num_layers", 2),
            drop_rate=encoder_config.get("drop_rate", 0.0),
            bias=encoder_config.get("bias", True),
            activation_fn=encoder_config.get("activation_fn"),
            activation_fn_kwargs=encoder_config.get("activation_fn_kwargs"),
            improved=encoder_config.get("improved", False),
            add_self_loops=encoder_config.get("add_self_loops", True),
            normalize=encoder_config.get("normalize", True),
            cached=encoder_config.get("cached", False),
        )
        decoder = SLP(in_channels=encoder_config["out_channels"], out_channels=1)

        super().__init__(
            encoder=encoder,
            decoder=decoder,
            loss_fn=loss_fn if loss_fn is not None else nn.BCEWithLogitsLoss(),
            metrics=metrics,
            metrics_log_kwargs=metrics_log_kwargs,
        )

        self.encoder_config: GCNEncoderConfig = encoder_config
        self.aggregation: Literal["mean", "max", "min", "sum"] = aggregation
        self.lr: float = lr
        self.weight_decay: float = weight_decay

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Reduce the hypergraph to a graph, encode nodes with GCN, aggregate per hyperedge, and score.

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

        Returns:
            scores: Logit scores of shape ``(num_hyperedges,)``.

        Raises:
            ValueError: If the encoder is not defined for this module.
        """
        if self.encoder is None:
            raise ValueError("Encoder is not defined for this HLP module.")

        # Reduce hypergraph to graph and remove self-loops
        reduced_edge_index = HyperedgeIndex(hyperedge_index).reduce(
            strategy=self.encoder_config.get(
                "graph_reduction_strategy",
                GraphReductionStrategyEnum.CLIQUE_EXPANSION,
            ),
            num_nodes=self.encoder_config.get("num_nodes"),
        )
        edge_index = EdgeIndex(reduced_edge_index).remove_selfloops().item

        # Encode nodes with GCN
        node_embeddings: Tensor = self.encoder(x, edge_index)

        # Aggregate node embeddings per hyperedge
        hyperedge_embeddings = HyperedgeAggregator(hyperedge_index, node_embeddings).pool(
            self.aggregation
        )

        return self.decoder(hyperedge_embeddings).squeeze(-1)

    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 hyperedge scores for a batch.

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

        Returns:
            scores: Predicted hyperedge scores.
        """
        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.
        """
        scores = self.forward(batch.x, batch.hyperedge_index)
        target_scores, target_labels = self._target_scores_and_labels(scores, batch)
        batch_size = target_labels.size(0)

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