model_lightning_module.py

"""
This file contains the ONNX compatible LightningModule. 

Author: Jonathan Renusch
Date: 2025-15-04
"""

import torch
import os
import pytorch_lightning as L

from sklearn.metrics import f1_score, roc_auc_score

from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping

from torch.optim.lr_scheduler import (
    ReduceLROnPlateau,
    CosineAnnealingLR,
    CosineAnnealingWarmRestarts,
    LRScheduler,
)

from Classifier.utils.plot_functions import *
from Classifier.utils.utility_classes import FocalLoss


class ModelLightningModule(L.LightningModule):
    def __init__(
        self,
        model: torch.nn.Module,
        # train_dataset: Dataset = None,
        comet_ml: bool = False,
        learning_rate: float = 0.001,
        batch_size: int = 512,
        num_workers: int = 12,
        lr_patience: int = 3,
        metric_callbacks: str = "val_auc_epoch",
        early_stopping_patience: int = 15,
        weight_decay: float = 0.1,
        momentum: float = 0.1,
        optimizer_name: str = "AdamW",
        lr_scheduler_name: str = "ReduceLROnPlateau",
        loss_name: str = "BCEWithLogitsLoss",
        efficiency: float = 0.99,
    ) -> None:

        super().__init__()
        self.model = model

        self.criterion = torch.nn.BCELoss()
        # self.train_dataset = train_dataset
        self.lr = learning_rate
        self.lr_patience = lr_patience
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.metric_callbacks = metric_callbacks
        self.early_stopping_patience = early_stopping_patience
        self.efficiency = efficiency
        self.comet_ml = comet_ml
        self.test_loss: list[torch.Tensor] = []
        self.test_predictions: list[torch.Tensor] = []
        self.test_labels: list[torch.Tensor] = []
        self.val_loss: list[torch.Tensor] = []
        self.val_predictions: list[torch.Tensor] = []
        self.val_labels: list[torch.Tensor] = []

        if "acc" in self.metric_callbacks:
            self.callback_mode = "max"
        elif "auc" in self.metric_callbacks:
            self.callback_mode = "max"
        elif "rejection_rate" in self.metric_callbacks:
            self.callback_mode = "max"
        elif "loss" in self.metric_callbacks:
            self.callback_mode = "min"
        else:
            raise Exception(
                "Having trouble finding callback mode 'min' or 'max'. \n \
                Please give metric_calbacks argument containing string 'acc' or 'loss'."
            )

        self.optimizer_name = optimizer_name
        self.loss_name = loss_name
        self.lr_scheduler_name = lr_scheduler_name
        self.momentum = momentum
        self.weight_decay = weight_decay

    def training_step(
        self, batch: dict[str, torch.Tensor], batch_idx: torch.Tensor
    ) -> torch.Tensor:
        out, loss, y = self.shared_step(batch)
        acc = self.accurarcy(out, y)
        auc = roc_auc_score(y.detach().cpu().numpy(), out.detach().cpu().numpy())
        logs = {"train_loss": loss, "train_acc": acc, "train_auc": auc}
        self.log_dict(
            logs,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            batch_size=self.batch_size,
            sync_dist=True,
        )
        return loss

    def validation_step(self, batch: dict[str, torch.Tensor], batch_idx: torch.Tensor) -> dict:
        out, loss, y = self.shared_step(batch)
        self.val_predictions.append(out)
        self.val_labels.append(y)
        self.val_loss.append(loss.cpu().item())
        return loss

    def on_validation_epoch_end(self):
        val_predictions = torch.concatenate(self.val_predictions).cpu().numpy().ravel()
        val_labels = torch.concatenate(self.val_labels).cpu().numpy().ravel()

        # logging of final test loss, accuracy and auc
        acc = self.accurarcy(torch.tensor(val_predictions), torch.tensor(val_labels))
        auc = roc_auc_score(np.array(val_labels), np.array(val_predictions))
        logs = {
            "val_loss_epoch": np.mean(self.val_loss),
            "val_acc_epoch": acc,
            "val_auc_epoch": auc,
        }
        self.log_dict(
            logs,
            sync_dist=True,
        )
        val_predictions = torch.concatenate(self.val_predictions).cpu().numpy().ravel()
        val_labels = torch.concatenate(self.val_labels).cpu().numpy().ravel()
        rejection_rate = calculate_rejection_rate(
            val_labels, val_predictions, efficiency=self.efficiency
        )
        logs = {f"val_epoch_rejection_rate": rejection_rate}
        self.log_dict(
            logs,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            # batch_size=self.batch_size,
            sync_dist=True,
        )
        self.val_predictions = []
        self.val_labels = []

    def shared_step(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        out = self.model(batch["x"], batch["edge_index"])
        loss = self.criterion(out, batch["y"])
        return out, loss, batch["y"]

    def test_step(self, batch, batch_idx) -> torch.Tensor:
        out, loss, y = self.shared_step(batch)
        self.test_predictions.append(out)
        self.test_labels.append(y)
        self.test_loss.append(loss.cpu().item())
        return out

    def on_test_epoch_end(self, cuts=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]) -> None:
        test_predictions = torch.concatenate(self.test_predictions).cpu().numpy().ravel()
        test_labels = torch.concatenate(self.test_labels).cpu().numpy().ravel()

        # logging of final test loss, accuracy and auc
        acc = self.accurarcy(torch.tensor(test_predictions), torch.tensor(test_labels))
        auc = roc_auc_score(np.array(test_labels), np.array(test_predictions))
        logs = {"test_loss": np.mean(self.test_loss), "test_acc": acc, "test_auc": auc}
        self.log_dict(
            logs,
            sync_dist=True,
        )

        for cut in cuts:
            figure, rejection_rate = plot_confusion_metrics_matrix(
                test_labels, test_predictions, cut=cut
            )
            if self.comet_ml:
                self.logger.experiment.log_figure(
                    figure_name=f"confusion matrix {cut}",
                    figure=figure,
                )
            else:
                figure.savefig(os.path.join(self.logger.log_dir, f"confusion_matrix_{cut}.png"))

        if self.comet_ml:

            self.logger.experiment.log_figure(
                figure_name="ROC curve", figure=plot_roc_curve(test_labels, test_predictions)
            )
        else:
            figure = plot_roc_curve(test_labels, test_predictions)
            figure.savefig(os.path.join(self.logger.log_dir, "roc_curve.png"))

        # rejection active
        figure, rejection_rate = plot_confusion(
            test_labels, test_predictions, efficiency=self.efficiency
        )

        self.logger.log_metrics({f"rejection_rate {self.efficiency}": rejection_rate})

        if self.comet_ml:
            self.logger.experiment.log_figure(
                figure_name=f"confusion {self.efficiency}",
                figure=figure,
            )
        else:
            figure.savefig(os.path.join(self.logger.log_dir, f"confusion_{self.efficiency}.png"))
        # rejection legacy
        figure, rejection_rate = plot_confusion(test_labels, test_predictions, efficiency=0.995)

        self.logger.log_metrics({f"rejection_rate 0.995": rejection_rate})

        if self.comet_ml:
            self.logger.experiment.log_figure(
                figure_name=f"confusion 0.995",
                figure=figure,
            )
        else:
            figure.savefig(os.path.join(self.logger.log_dir, f"confusion_0.995.png"))

        test_predictions = (test_predictions >= 0.5).astype(int)
        self.logger.log_metrics({"test_f1_score": f1_score(test_labels, test_predictions)})

    def configure_optimizers(self) -> dict:

        # Select optimizer

        if self.optimizer_name == "Adam":
            self.optimizer = torch.optim.Adam(
                self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay
            )
        elif self.optimizer_name == "AdamW":
            self.optimizer = torch.optim.AdamW(
                self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay
            )
        elif self.optimizer_name == "RMSprop":
            self.optimizer = torch.optim.RMSprop(
                self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay
            )
        elif self.optimizer_name == "SGD":
            self.optimizer = torch.optim.SGD(
                self.model.parameters(),
                lr=self.lr,
                weight_decay=self.weight_decay,
                momentum=self.momentum,
            )
        elif self.optimizer_name == "Adagrad":
            self.optimizer = torch.optim.Adagrad(
                self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay
            )
        else:
            raise Exception("The requested optimizer is not implemented!")

        # Select loss function

        if self.loss_name == "BCEWithLogitsLoss":
            self.criterion = torch.nn.BCEWithLogitsLoss()
        elif self.loss_name == "BCELoss":
            self.criterion = torch.nn.BCELoss()
        elif self.loss_name == "FocalLoss":
            self.criterion = FocalLoss(alpha=0.8, gamma=2)
        elif self.loss_name == "MSELoss":
            self.criterion = torch.nn.MSELoss()
        elif self.loss_name == "SoftMarginLoss":
            self.criterion = torch.nn.SoftMarginLoss()
        else:
            raise Exception("The requested loss is not implemented!")

        # select lr scheduler:

        if self.lr_scheduler_name == "ReduceLROnPlateau":
            self.scheduler = ReduceLROnPlateau(
                self.optimizer, mode=self.callback_mode, factor=0.1, patience=self.lr_patience
            )
        elif self.lr_scheduler_name == "CosineAnnealingLR":
            self.scheduler = CosineAnnealingLR(self.optimizer, eta_min=1e-9, T_max=10)
        elif self.lr_scheduler_name == "CosineAnnealingWarmRestarts":
            self.scheduler = CosineAnnealingWarmRestarts(
                self.optimizer, eta_min=1e-9, T_0=10, T_mult=2
            )
        elif self.lr_scheduler_name == "LRScheduler":
            self.scheduler = LRScheduler(
                self.optimizer,
            )
        else:
            raise Exception("The requested lr scheduler is not implemented!")

        return {
            "optimizer": self.optimizer,
            "lr_scheduler": {
                "scheduler": self.scheduler,
                "monitor": self.metric_callbacks,
                "frequency": 1,
                "interval": "epoch",
                "strict": True,
            },
        }

    def accurarcy(self, out: torch.Tensor, y: torch.Tensor) -> float:
        predictions = (out > 0.5).float()
        correct = (predictions == y).sum().item()
        total = y.size(0)
        return correct / total

    def configure_callbacks(self) -> list:

        checkpoints = ModelCheckpoint(
            monitor=self.metric_callbacks,
            mode=self.callback_mode,
            filename="{epoch}-{val_loss_epoch:.3f}-{val_acc_epoch:.3f}",
            save_top_k=1,
        )

        early_stop = EarlyStopping(
            monitor=self.metric_callbacks,
            mode=self.callback_mode,
            patience=self.early_stopping_patience,
        )

        return [checkpoints, early_stop]