grover/model/models.py

"""
The GROVER models for pretraining, finetuning and fingerprint generating.
"""
from argparse import Namespace
from typing import List, Dict, Callable

import numpy as np
import torch
from torch import nn as nn

from grover.data import get_atom_fdim, get_bond_fdim
from grover.model.layers import Readout, GTransEncoder
from grover.util.nn_utils import get_activation_function


class GROVEREmbedding(nn.Module):
    """
    The GROVER Embedding class. It contains the GTransEncoder.
    This GTransEncoder can be replaced by any validate encoders.
    """

    def __init__(self, args: Namespace):
        """
        Initialize the GROVEREmbedding class.
        :param args:
        """
        super(GROVEREmbedding, self).__init__()
        self.embedding_output_type = args.embedding_output_type
        edge_dim = get_bond_fdim() + get_atom_fdim()
        node_dim = get_atom_fdim()
        if not hasattr(args, "backbone"):
            print("No backbone specified in args, use gtrans backbone.")
            args.backbone = "gtrans"
        if args.backbone == "gtrans" or args.backbone == "dualtrans":
            # dualtrans is the old name.
            self.encoders = GTransEncoder(args,
                                          hidden_size=args.hidden_size,
                                          edge_fdim=edge_dim,
                                          node_fdim=node_dim,
                                          dropout=args.dropout,
                                          activation=args.activation,
                                          num_mt_block=args.num_mt_block,
                                          num_attn_head=args.num_attn_head,
                                          atom_emb_output=self.embedding_output_type,
                                          bias=args.bias,
                                          cuda=args.cuda)

    def forward(self, graph_batch: List) -> Dict:
        """
        The forward function takes graph_batch as input and output a dict. The content of the dict is decided by
        self.embedding_output_type.

        :param graph_batch: the input graph batch generated by MolCollator.
        :return: a dict containing the embedding results.
        """
        output = self.encoders(graph_batch)
        if self.embedding_output_type == 'atom':
            return {"atom_from_atom": output[0], "atom_from_bond": output[1],
                    "bond_from_atom": None, "bond_from_bond": None}  # atom_from_atom, atom_from_bond
        elif self.embedding_output_type == 'bond':
            return {"atom_from_atom": None, "atom_from_bond": None,
                    "bond_from_atom": output[0], "bond_from_bond": output[1]}  # bond_from_atom, bond_from_bond
        elif self.embedding_output_type == "both":
            return {"atom_from_atom": output[0][0], "bond_from_atom": output[0][1],
                    "atom_from_bond": output[1][0], "bond_from_bond": output[1][1]}


class AtomVocabPrediction(nn.Module):
    """
    The atom-wise vocabulary prediction task. The atom vocabulary is constructed by the context.
    """
    def __init__(self, args, vocab_size, hidden_size=None):
        """
        :param args: the argument.
        :param vocab_size: the size of atom vocabulary.
        """
        super(AtomVocabPrediction, self).__init__()
        if not hidden_size:
            hidden_size = args.hidden_size
        self.linear = nn.Linear(hidden_size, vocab_size)
        self.logsoftmax = nn.LogSoftmax(dim=1)

    def forward(self, embeddings):
        """
        If embeddings is None: do not go through forward pass.
        :param embeddings: the atom embeddings, num_atom X fea_dim.
        :return: the prediction for each atom, num_atom X vocab_size.
        """
        if embeddings is None:
            return None
        return self.logsoftmax(self.linear(embeddings))


class BondVocabPrediction(nn.Module):
    """
    The bond-wise vocabulary prediction task. The bond vocabulary is constructed by the context.
    """
    def __init__(self, args, vocab_size, hidden_size=None):
        """
        Might need to use different architecture for bond vocab prediction.
        :param args:
        :param vocab_size: size of bond vocab.
        :param hidden_size: hidden size
        """
        super(BondVocabPrediction, self).__init__()
        if not hidden_size:
            hidden_size = args.hidden_size
        self.linear = nn.Linear(hidden_size, vocab_size)

        # ad-hoc here
        # If TWO_FC_4_BOND_VOCAB, we will use two distinct fc layer to deal with the bond and rev bond.
        self.TWO_FC_4_BOND_VOCAB = True
        if self.TWO_FC_4_BOND_VOCAB:
            self.linear_rev = nn.Linear(hidden_size, vocab_size)
        self.logsoftmax = nn.LogSoftmax(dim=1)

    def forward(self, embeddings):
        """
        If embeddings is None: do not go through forward pass.
        :param embeddings: the atom embeddings, num_bond X fea_dim.
        :return: the prediction for each atom, num_bond X vocab_size.
        """
        if embeddings is None:
            return None
        nm_bonds = embeddings.shape[0]  # must be an odd number
        # The bond and rev bond have odd and even ids respectively. See definition in molgraph.
        ids1 = [0] + list(range(1, nm_bonds, 2))
        ids2 = list(range(0, nm_bonds, 2))
        if self.TWO_FC_4_BOND_VOCAB:
            logits = self.linear(embeddings[ids1]) + self.linear_rev(embeddings[ids2])
        else:
            logits = self.linear(embeddings[ids1] + embeddings[ids2])

        return self.logsoftmax(logits)


class FunctionalGroupPrediction(nn.Module):
    """
    The functional group (semantic motifs) prediction task. This is a graph-level task.
    """
    def __init__(self, args, fg_size):
        """
        :param args: The arguments.
        :param fg_size: The size of semantic motifs.
        """
        super(FunctionalGroupPrediction, self).__init__()
        first_linear_dim = args.hidden_size
        hidden_size = args.hidden_size

        # In order to retain maximal information in the encoder, we use a simple readout function here.
        self.readout = Readout(rtype="mean", hidden_size=hidden_size)
        # We have four branches here. But the input with less than four branch is OK.
        # Since we use BCEWithLogitsLoss as the loss function, we only need to output logits here.
        self.linear_atom_from_atom = nn.Linear(first_linear_dim, fg_size)
        self.linear_atom_from_bond = nn.Linear(first_linear_dim, fg_size)
        self.linear_bond_from_atom = nn.Linear(first_linear_dim, fg_size)
        self.linear_bond_from_bond = nn.Linear(first_linear_dim, fg_size)

    def forward(self, embeddings: Dict, ascope: List, bscope: List) -> Dict:
        """
        The forward function of semantic motif prediction. It takes the node/bond embeddings, and the corresponding
        atom/bond scope as input and produce the prediction logits for different branches.
        :param embeddings: The input embeddings are organized as dict. The output of GROVEREmbedding.
        :param ascope: The scope for bonds. Please refer BatchMolGraph for more details.
        :param bscope: The scope for aotms. Please refer BatchMolGraph for more details.
        :return: a dict contains the predicted logits.
        """

        preds_atom_from_atom, preds_atom_from_bond, preds_bond_from_atom, preds_bond_from_bond = \
            None, None, None, None

        if embeddings["bond_from_atom"] is not None:
            preds_bond_from_atom = self.linear_bond_from_atom(self.readout(embeddings["bond_from_atom"], bscope))
        if embeddings["bond_from_bond"] is not None:
            preds_bond_from_bond = self.linear_bond_from_bond(self.readout(embeddings["bond_from_bond"], bscope))

        if embeddings["atom_from_atom"] is not None:
            preds_atom_from_atom = self.linear_atom_from_atom(self.readout(embeddings["atom_from_atom"], ascope))
        if embeddings["atom_from_bond"] is not None:
            preds_atom_from_bond = self.linear_atom_from_bond(self.readout(embeddings["atom_from_bond"], ascope))

        return {"atom_from_atom": preds_atom_from_atom, "atom_from_bond": preds_atom_from_bond,
                "bond_from_atom": preds_bond_from_atom, "bond_from_bond": preds_bond_from_bond}


class GroverTask(nn.Module):
    """
    The pretrain module.
    """
    def __init__(self, args, grover, atom_vocab_size, bond_vocab_size, fg_size):
        super(GroverTask, self).__init__()
        self.grover = grover
        self.av_task_atom = AtomVocabPrediction(args, atom_vocab_size)
        self.av_task_bond = AtomVocabPrediction(args, atom_vocab_size)
        self.bv_task_atom = BondVocabPrediction(args, bond_vocab_size)
        self.bv_task_bond = BondVocabPrediction(args, bond_vocab_size)

        self.fg_task_all = FunctionalGroupPrediction(args, fg_size)

        self.embedding_output_type = args.embedding_output_type

    @staticmethod
    def get_loss_func(args: Namespace) -> Callable:
        """
        The loss function generator.
        :param args: the arguments.
        :return: the loss fucntion for GroverTask.
        """
        def loss_func(preds, targets, dist_coff=args.dist_coff):
            """
            The loss function for GroverTask.
            :param preds: the predictions.
            :param targets: the targets.
            :param dist_coff: the default disagreement coefficient for the distances between different branches.
            :return:
            """
            av_task_loss = nn.NLLLoss(ignore_index=0, reduction="mean")  # same for av and bv

            fg_task_loss = nn.BCEWithLogitsLoss(reduction="mean")
            # av_task_dist_loss = nn.KLDivLoss(reduction="mean")
            av_task_dist_loss = nn.MSELoss(reduction="mean")
            fg_task_dist_loss = nn.MSELoss(reduction="mean")
            sigmoid = nn.Sigmoid()

            av_atom_loss, av_bond_loss, av_dist_loss = 0.0, 0.0, 0.0
            fg_atom_from_atom_loss, fg_atom_from_bond_loss, fg_atom_dist_loss = 0.0, 0.0, 0.0
            bv_atom_loss, bv_bond_loss, bv_dist_loss = 0.0, 0.0, 0.0
            fg_bond_from_atom_loss, fg_bond_from_bond_loss, fg_bond_dist_loss = 0.0, 0.0, 0.0

            if preds["av_task"][0] is not None:
                av_atom_loss = av_task_loss(preds['av_task'][0], targets["av_task"])
                fg_atom_from_atom_loss = fg_task_loss(preds["fg_task"]["atom_from_atom"], targets["fg_task"])

            if preds["av_task"][1] is not None:
                av_bond_loss = av_task_loss(preds['av_task'][1], targets["av_task"])
                fg_atom_from_bond_loss = fg_task_loss(preds["fg_task"]["atom_from_bond"], targets["fg_task"])

            if preds["bv_task"][0] is not None:
                bv_atom_loss = av_task_loss(preds['bv_task'][0], targets["bv_task"])
                fg_bond_from_atom_loss = fg_task_loss(preds["fg_task"]["bond_from_atom"], targets["fg_task"])

            if preds["bv_task"][1] is not None:
                bv_bond_loss = av_task_loss(preds['bv_task'][1], targets["bv_task"])
                fg_bond_from_bond_loss = fg_task_loss(preds["fg_task"]["bond_from_bond"], targets["fg_task"])

            if preds["av_task"][0] is not None and preds["av_task"][1] is not None:
                av_dist_loss = av_task_dist_loss(preds['av_task'][0], preds['av_task'][1])
                fg_atom_dist_loss = fg_task_dist_loss(sigmoid(preds["fg_task"]["atom_from_atom"]),
                                                      sigmoid(preds["fg_task"]["atom_from_bond"]))

            if preds["bv_task"][0] is not None and preds["bv_task"][1] is not None:
                bv_dist_loss = av_task_dist_loss(preds['bv_task'][0], preds['bv_task'][1])
                fg_bond_dist_loss = fg_task_dist_loss(sigmoid(preds["fg_task"]["bond_from_atom"]),
                                                      sigmoid(preds["fg_task"]["bond_from_bond"]))

            av_loss = av_atom_loss + av_bond_loss
            bv_loss = bv_atom_loss + bv_bond_loss
            fg_atom_loss = fg_atom_from_atom_loss + fg_atom_from_bond_loss
            fg_bond_loss = fg_bond_from_atom_loss + fg_bond_from_bond_loss

            fg_loss = fg_atom_loss + fg_bond_loss
            fg_dist_loss = fg_atom_dist_loss + fg_bond_dist_loss

            # dist_loss = av_dist_loss + bv_dist_loss + fg_dist_loss
            # print("%.4f %.4f %.4f %.4f %.4f %.4f"%(av_atom_loss,
            #                                       av_bond_loss,
            #                                       fg_atom_loss,
            #                                       fg_bond_loss,
            #                                       av_dist_loss,
            #                                       fg_dist_loss))
            # return av_loss + fg_loss + dist_coff * dist_loss
            overall_loss = av_loss + bv_loss + fg_loss + dist_coff * av_dist_loss + \
                           dist_coff * bv_dist_loss + fg_dist_loss

            return overall_loss, av_loss, bv_loss, fg_loss, av_dist_loss, bv_dist_loss, fg_dist_loss

        return loss_func

    def forward(self, graph_batch: List):
        """
        The forward function.
        :param graph_batch:
        :return:
        """
        _, _, _, _, _, a_scope, b_scope, _ = graph_batch
        a_scope = a_scope.data.cpu().numpy().tolist()

        embeddings = self.grover(graph_batch)

        av_task_pred_atom = self.av_task_atom(
            embeddings["atom_from_atom"])  # if None: means not go through this fowward
        av_task_pred_bond = self.av_task_bond(embeddings["atom_from_bond"])

        bv_task_pred_atom = self.bv_task_atom(embeddings["bond_from_atom"])
        bv_task_pred_bond = self.bv_task_bond(embeddings["bond_from_bond"])

        fg_task_pred_all = self.fg_task_all(embeddings, a_scope, b_scope)

        return {"av_task": (av_task_pred_atom, av_task_pred_bond),
                "bv_task": (bv_task_pred_atom, bv_task_pred_bond),
                "fg_task": fg_task_pred_all}


class GroverFpGeneration(nn.Module):
    """
    GroverFpGeneration class.
    It loads the pre-trained model and produce the fingerprints for input molecules.
    """
    def __init__(self, args):
        """
        Init function.
        :param args: the arguments.
        """
        super(GroverFpGeneration, self).__init__()

        self.fingerprint_source = args.fingerprint_source
        self.iscuda = args.cuda

        self.grover = GROVEREmbedding(args)
        self.readout = Readout(rtype="mean", hidden_size=args.hidden_size)

    def forward(self, batch, features_batch):
        """
        The forward function.
        It takes graph batch and molecular feature batch as input and produce the fingerprints of this molecules.
        :param batch:
        :param features_batch:
        :return:
        """
        _, _, _, _, _, a_scope, b_scope, _ = batch

        output = self.grover(batch)
        # Share readout
        mol_atom_from_bond_output = self.readout(output["atom_from_bond"], a_scope)
        mol_atom_from_atom_output = self.readout(output["atom_from_atom"], a_scope)

        if self.fingerprint_source == "bond" or self.fingerprint_source == "both":
            mol_bond_from_atom_output = self.readout(output["bond_from_atom"], b_scope)
            mol_bond_from_bodd_output = self.readout(output["bond_from_bond"], b_scope)

        if features_batch[0] is not None:
            features_batch = torch.from_numpy(np.stack(features_batch)).float()
            if self.iscuda:
                features_batch = features_batch.cuda()
            features_batch = features_batch.to(output["atom_from_atom"])
            if len(features_batch.shape) == 1:
                features_batch = features_batch.view([1, features_batch.shape[0]])
        else:
            features_batch = None

        if self.fingerprint_source == "atom":
            fp = torch.cat([mol_atom_from_atom_output, mol_atom_from_bond_output], 1)
        elif self.fingerprint_source == "bond":
            fp = torch.cat([mol_bond_from_atom_output, mol_bond_from_bodd_output], 1)
        else:
            # the both case.
            fp = torch.cat([mol_atom_from_atom_output, mol_atom_from_bond_output,
                            mol_bond_from_atom_output, mol_bond_from_bodd_output], 1)
        if features_batch is not None:
            fp = torch.cat([fp, features_batch], 1)
        return fp


class GroverFinetuneTask(nn.Module):
    """
    The finetune
    """
    def __init__(self, args):
        super(GroverFinetuneTask, self).__init__()

        self.hidden_size = args.hidden_size
        self.iscuda = args.cuda

        self.grover = GROVEREmbedding(args)

        if args.self_attention:
            self.readout = Readout(rtype="self_attention", hidden_size=self.hidden_size,
                                   attn_hidden=args.attn_hidden,
                                   attn_out=args.attn_out)
        else:
            self.readout = Readout(rtype="mean", hidden_size=self.hidden_size)

        self.mol_atom_from_atom_ffn = self.create_ffn(args)
        self.mol_atom_from_bond_ffn = self.create_ffn(args)
        #self.ffn = nn.ModuleList()
        #self.ffn.append(self.mol_atom_from_atom_ffn)
        #self.ffn.append(self.mol_atom_from_bond_ffn)

        self.classification = args.dataset_type == 'classification'
        if self.classification:
            self.sigmoid = nn.Sigmoid()

    def create_ffn(self, args: Namespace):
        """
        Creates the feed-forward network for the model.

        :param args: Arguments.
        """
        # Note: args.features_dim is set according the real loaded features data
        if args.features_only:
            first_linear_dim = args.features_size + args.features_dim
        else:
            if args.self_attention:
                first_linear_dim = args.hidden_size * args.attn_out
                # TODO: Ad-hoc!
                # if args.use_input_features:
                first_linear_dim += args.features_dim
            else:
                first_linear_dim = args.hidden_size + args.features_dim

        dropout = nn.Dropout(args.dropout)
        activation = get_activation_function(args.activation)
        # TODO: ffn_hidden_size
        # Create FFN layers
        if args.ffn_num_layers == 1:
            ffn = [
                dropout,
                nn.Linear(first_linear_dim, args.output_size)
            ]
        else:
            ffn = [
                dropout,
                nn.Linear(first_linear_dim, args.ffn_hidden_size)
            ]
            for _ in range(args.ffn_num_layers - 2):
                ffn.extend([
                    activation,
                    dropout,
                    nn.Linear(args.ffn_hidden_size, args.ffn_hidden_size),
                ])
            ffn.extend([
                activation,
                dropout,
                nn.Linear(args.ffn_hidden_size, args.output_size),
            ])

        # Create FFN model
        return nn.Sequential(*ffn)

    @staticmethod
    def get_loss_func(args):
        def loss_func(preds, targets,
                      dt=args.dataset_type,
                      dist_coff=args.dist_coff):

            if dt == 'classification':
                pred_loss = nn.BCEWithLogitsLoss(reduction='none')
            elif dt == 'regression':
                pred_loss = nn.MSELoss(reduction='none')
            else:
                raise ValueError(f'Dataset type "{args.dataset_type}" not supported.')

            # print(type(preds))
            # TODO: Here, should we need to involve the model status? Using len(preds) is just a hack.
            if type(preds) is not tuple:
                # in eval mode.
                return pred_loss(preds, targets)

            # in train mode.
            dist_loss = nn.MSELoss(reduction='none')
            # dist_loss = nn.CosineSimilarity(dim=0)
            # print(pred_loss)

            dist = dist_loss(preds[0], preds[1])
            pred_loss1 = pred_loss(preds[0], targets)
            pred_loss2 = pred_loss(preds[1], targets)
            return pred_loss1 + pred_loss2 + dist_coff * dist

        return loss_func

    def forward(self, batch, features_batch):
        _, _, _, _, _, a_scope, _, _ = batch

        output = self.grover(batch)
        # Share readout
        mol_atom_from_bond_output = self.readout(output["atom_from_bond"], a_scope)
        mol_atom_from_atom_output = self.readout(output["atom_from_atom"], a_scope)

        if features_batch[0] is not None:
            features_batch = torch.from_numpy(np.stack(features_batch)).float()
            if self.iscuda:
                features_batch = features_batch.cuda()
            features_batch = features_batch.to(output["atom_from_atom"])
            if len(features_batch.shape) == 1:
                features_batch = features_batch.view([1, features_batch.shape[0]])
        else:
            features_batch = None


        if features_batch is not None:
            mol_atom_from_atom_output = torch.cat([mol_atom_from_atom_output, features_batch], 1)
            mol_atom_from_bond_output = torch.cat([mol_atom_from_bond_output, features_batch], 1)

        if self.training:
            atom_ffn_output = self.mol_atom_from_atom_ffn(mol_atom_from_atom_output)
            bond_ffn_output = self.mol_atom_from_bond_ffn(mol_atom_from_bond_output)
            return atom_ffn_output, bond_ffn_output
        else:
            atom_ffn_output = self.mol_atom_from_atom_ffn(mol_atom_from_atom_output)
            bond_ffn_output = self.mol_atom_from_bond_ffn(mol_atom_from_bond_output)
            if self.classification:
                atom_ffn_output = self.sigmoid(atom_ffn_output)
                bond_ffn_output = self.sigmoid(bond_ffn_output)
            output = (atom_ffn_output + bond_ffn_output) / 2

        return output