vgae.py

# VGAE.py Module script 
#Code was taken from Pytorch Geometric VGAE user defined model where decoders and encoder can be defined 
#it was adapted to include property prediction functionalities and to have a graph level latent with pooling 

from typing import Optional, Tuple

import torch
from torch import Tensor
from torch.nn import Module
from torch_geometric.nn.inits import reset
from torch_geometric.utils import negative_sampling
import torch.nn.functional as F
import torch.nn as nn
from torch_geometric.data import Data
from torch_geometric.nn import AttentionalAggregation

EPS = 1e-15
MAX_LOGSTD = 10


# NodeDecoder used for encoder and matching algortihm testing the VGAE can accept a user specified decoder is module is None the NodeDecoder is used 
class NodeDecoder(torch.nn.Module):
    def __init__(self, dense_units=(512, 512, 1024, 1024), dropout=0.1):
        super().__init__()
        in_dim = LATENT_DIM
        blocks = []
        for u in dense_units:
            blocks += [nn.Linear(in_dim, u),
                       nn.GELU(),
                       nn.Dropout(dropout)]
            in_dim = u
        self.mlp = nn.Sequential(*blocks)
        self.feat_out = nn.Linear(in_dim, NUM_ATOMS * ATOM_DIM)

    def forward(self, z):
        x = self.mlp(z)
        feat = self.feat_out(x).view(-1, NUM_ATOMS, ATOM_DIM)
        return feat




class GAE(torch.nn.Module):
    r"""The Graph Auto-Encoder model from the
    `"Variational Graph Auto-Encoders" <https://arxiv.org/abs/1611.07308>`_
    paper based on user-defined encoder and decoder models.

    Args:
        encoder (torch.nn.Module): The encoder module.
        decoder (torch.nn.Module, optional): The decoder module. If set to
            :obj:`None`, will default to the
            :class:`torch_geometric.nn.models.InnerProductDecoder`.
            (default: :obj:`None`)
    """
    def __init__(self, encoder: Module, decoder: Optional[Module] = None):
        super().__init__()
        self.encoder = encoder
        self.decoder = NodeDecoder() if decoder is None else decoder
        GAE.reset_parameters(self)

    def reset_parameters(self):
        r"""Resets all learnable parameters of the module."""
        reset(self.encoder)
        reset(self.decoder)


    def forward(self, *args, **kwargs) -> Tensor:  
        r"""Alias for :meth:`encode`."""
        return self.encoder(*args, **kwargs)


    def encode(self, *args, **kwargs) -> Tensor:
        r"""Runs the encoder and computes node-wise latent variables."""
        return self.encoder(*args, **kwargs)


    def decode(self, *args, **kwargs) -> Tensor:
        r"""Runs the decoder and computes edge probabilities."""
        return self.decoder(*args, **kwargs)


    def recon_loss(self, z: Tensor, pos_edge_index: Tensor,
                   neg_edge_index: Optional[Tensor] = None) -> Tensor:
        r"""Given latent variables :obj:`z`, computes the binary cross
        entropy loss for positive edges :obj:`pos_edge_index` and negative
        sampled edges.

        Args:
            z (torch.Tensor): The latent space :math:`\mathbf{Z}`.
            pos_edge_index (torch.Tensor): The positive edges to train against.
            neg_edge_index (torch.Tensor, optional): The negative edges to
                train against. If not given, uses negative sampling to
                calculate negative edges. (default: :obj:`None`)
        """
        pos_loss = -torch.log(
            self.decoder(z, pos_edge_index, sigmoid=True) + EPS).mean()

        if neg_edge_index is None:
            neg_edge_index = negative_sampling(pos_edge_index, z.size(0))
        neg_loss = -torch.log(1 -
                              self.decoder(z, neg_edge_index, sigmoid=True) +
                              EPS).mean()

        return pos_loss + neg_loss


    def test(self, z: Tensor, pos_edge_index: Tensor,
             neg_edge_index: Tensor) -> Tuple[Tensor, Tensor]:
        r"""Given latent variables :obj:`z`, positive edges
        :obj:`pos_edge_index` and negative edges :obj:`neg_edge_index`,
        computes area under the ROC curve (AUC) and average precision (AP)
        scores.

        Args:
            z (torch.Tensor): The latent space :math:`\mathbf{Z}`.
            pos_edge_index (torch.Tensor): The positive edges to evaluate
                against.
            neg_edge_index (torch.Tensor): The negative edges to evaluate
                against.
        """
        from sklearn.metrics import average_precision_score, roc_auc_score

        pos_y = z.new_ones(pos_edge_index.size(1))
        neg_y = z.new_zeros(neg_edge_index.size(1))
        y = torch.cat([pos_y, neg_y], dim=0)

        pos_pred = self.decoder(z, pos_edge_index, sigmoid=True)
        neg_pred = self.decoder(z, neg_edge_index, sigmoid=True)
        pred = torch.cat([pos_pred, neg_pred], dim=0)

        y, pred = y.detach().cpu().numpy(), pred.detach().cpu().numpy()

        return roc_auc_score(y, pred), average_precision_score(y, pred)


class PropertyPredictor(nn.Module):
    r"""A property prediction head applied to the latent variable 
        :math:`\mathbf{z}`.

        Args:
        latent_dim (int): Dimensionality of the latent space 
            :math:`\mathbf{z}`.
        hidden_dim (int, optional): Dimensionality of hidden layers 
            in the predictor MLP. (default: :obj:`128`)
        out_dim (int, optional): Output dimensionality of the property 
            prediction. Use :obj:`1` for scalar regression or the number 
            of classes/tasks for classification. (default: :obj:`1`)
        dropout (float, optional): Dropout probability applied between 
            hidden layers. (default: :obj:`0.1`)

        The predictor is a feed-forward MLP trained jointly with the VGAE. 
        It enables supervised learning of molecular or graph-level properties 
        while retaining generative latent space structure, allowing latent space shaping.
        """
    def __init__(self, latent_dim: int, hidden_dim: int = 128, out_dim: int = 1, dropout: float = 0.1):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, out_dim)
        )

    def forward(self, z: Tensor) -> Tensor:
        return self.mlp(z)



class VGAE(GAE):
    r"""VGAE with optional graph-level latent, attention pooling and property predictor for latent space shaping.
    
        Args:
        encoder (torch.nn.Module): The encoder module.
        decoder (torch.nn.Module, optional): The decoder module. Defaults
            to :class:`NodeDecoder` if not specified.
        latent_level (str, optional): Whether to use a graph-level or
            node-level latent representation. (default: :obj:`'graph'`)
        latent_dim (int, optional): Dimensionality of the latent space. 
            (default: :obj:`64`)
        property_out_dim (int, optional): Dimensionality of the property 
            prediction output. Use :obj:`1` for scalar regression or the 
            number of classes/tasks for classification/multitask 
            learning. (default: :obj:`1`)
    
    """
    def __init__(self, encoder: Module, decoder: Optional[Module] = None,
                 latent_level: str = 'graph', latent_dim: int = 64, property_out_dim: int = 1):
        super().__init__(encoder, decoder)
        self.latent_level = latent_level
        self.latent_dim   = latent_dim
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.encoder.to(self.device)
        self.decoder.to(self.device)

        H = getattr(self.encoder, "out_channels", None)
        if H is None:
            raise ValueError("Encoder must expose .out_channels")

        mlp_gate = nn.Sequential(
            nn.Linear(H, 128), nn.LeakyReLU(),
            nn.Linear(128, 128), nn.LeakyReLU(),
            nn.Linear(128, 1)
        ).to(self.device)
        mlp_end  = nn.Sequential(
            nn.Linear(H, 128), nn.LeakyReLU(),
            nn.Linear(128, 128), nn.LeakyReLU(),
            nn.Linear(128, H)
        ).to(self.device)
        self.pool = AttentionalAggregation(mlp_gate, mlp_end.to(self.device))

        # Gaussian heads
        self.mu    = nn.Linear(H, latent_dim).to(self.device)
        self.sigma = nn.Linear(H, latent_dim).to(self.device)   
        
        #Property prediction head
        self.property_predictor = PropertyPredictor(latent_dim, hidden_dim=128, out_dim=property_out_dim).to(self.device)

    def reparametrize(self, mu: Tensor, logstd: Tensor) -> Tensor:
        return mu + torch.randn_like(logstd) * torch.exp(logstd) if self.training else mu

    def encode(self, data: Data, **kwargs):
        """Return (z, mu, logstd, node_level)."""
        node_level = self.encoder(data, **kwargs) 

        if self.latent_level == 'graph':
            x = self.pool(node_level, data.batch) 
        else: 
            x = node_level                        

        mu     = self.mu(x)                      
        logstd = self.sigma(x).clamp(max=MAX_LOGSTD)
        z      = self.reparametrize(mu, logstd)
        return z, mu, logstd, node_level

    @staticmethod
    def kl_loss(mu: Tensor, logstd: Tensor) -> Tensor:
        kl = 1 + 2 * logstd - mu**2 - logstd.exp()**2
        kl = torch.sum(kl, dim=1)
        return -0.5 * torch.mean(kl)

    def forward(self, data: Data, **kwargs):
        z, mu, sigma, node_level = self.encode(data, **kwargs)

        # Reconstruction
        recon, pred_num = self.decoder(z, data)
        prop_pred = self.property_predictor(mu)

        return recon, mu, sigma, pred_num, prop_pred