models/property_pred/prop_model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_scatter import scatter

from .prop_egnn import EnEquiEncoder
from ..common import compose_context_prop, ShiftedSoftplus


def get_encoder(config):
    if config.name == 'egnn' or config.name == 'egnn_enc':
        net = EnEquiEncoder(
            num_layers=config.num_layers,
            edge_feat_dim=config.edge_dim,
            hidden_dim=config.hidden_dim,
            num_r_gaussian=config.num_r_gaussian,
            act_fn=config.act_fn,
            norm=config.norm,
            update_x=False,
            k=config.knn,
            cutoff=config.cutoff,
        )
    else:
        raise ValueError(config.name)
    return net


class PropPredNet(nn.Module):
    def __init__(self, config, protein_atom_feature_dim, ligand_atom_feature_dim, output_dim=3):
        super(PropPredNet, self).__init__()
        self.config = config
        self.hidden_dim = config.hidden_channels
        self.output_dim = output_dim
        self.protein_atom_emb = nn.Linear(protein_atom_feature_dim, self.hidden_dim)
        self.ligand_atom_emb = nn.Linear(ligand_atom_feature_dim, self.hidden_dim)

        # self.mean = target_mean
        # self.std = target_std
        # self.register_buffer('target_mean', target_mean)
        # self.register_buffer('target_std', target_std)
        self.encoder = get_encoder(config.encoder)
        self.out_block = nn.Sequential(
            nn.Linear(self.hidden_dim, self.hidden_dim),
            ShiftedSoftplus(),
            nn.Linear(self.hidden_dim, output_dim),
        )

    def forward(self, protein_pos, protein_atom_feature, ligand_pos, ligand_atom_feature, batch_protein, batch_ligand,
                output_kind):
        h_protein = self.protein_atom_emb(protein_atom_feature)
        h_ligand = self.ligand_atom_emb(ligand_atom_feature)

        h_ctx, pos_ctx, batch_ctx = compose_context_prop(
            h_protein=h_protein,
            h_ligand=h_ligand,
            pos_protein=protein_pos,
            pos_ligand=ligand_pos,
            batch_protein=batch_protein,
            batch_ligand=batch_ligand,
        )

        h_ctx = self.encoder(
            node_attr=h_ctx,
            pos=pos_ctx,
            batch=batch_ctx,
        )  # (N_p+N_l, H)

        # Aggregate messages
        pre_out = scatter(h_ctx, index=batch_ctx, dim=0, reduce='sum')  # (N, F)
        output = self.out_block(pre_out)  # (N, C)
        if output_kind is not None:
            output_mask = F.one_hot(output_kind - 1, self.output_dim)
            output = torch.sum(output * output_mask, dim=-1, keepdim=True)
        return output

    def get_loss(self, batch, pos_noise_std, return_pred=False):
        protein_noise = torch.randn_like(batch.protein_pos) * pos_noise_std
        ligand_noise = torch.randn_like(batch.ligand_pos) * pos_noise_std
        pred = self(
            protein_pos=batch.protein_pos + protein_noise,
            protein_atom_feature=batch.protein_atom_feature.float(),
            ligand_pos=batch.ligand_pos + ligand_noise,
            ligand_atom_feature=batch.ligand_atom_feature_full.float(),
            batch_protein=batch.protein_element_batch,
            batch_ligand=batch.ligand_element_batch,
            output_kind=batch.kind,
            # output_kind=None
        )
        # pred = pred * y_std + y_mean
        loss_func = nn.MSELoss()
        loss = loss_func(pred.view(-1), batch.y)
        if return_pred:
            return loss, pred
        else:
            return loss


class PropPredNetEnc(nn.Module):
    def __init__(self, config, protein_atom_feature_dim, ligand_atom_feature_dim,
                 enc_ligand_dim, enc_node_dim, enc_graph_dim, enc_feature_type=None, output_dim=1):
        super(PropPredNetEnc, self).__init__()
        self.config = config
        self.hidden_dim = config.hidden_channels
        self.output_dim = output_dim
        self.enc_ligand_dim = enc_ligand_dim
        self.enc_node_dim = enc_node_dim
        self.enc_graph_dim = enc_graph_dim
        self.enc_feature_type = enc_feature_type

        self.protein_atom_emb = nn.Linear(protein_atom_feature_dim, self.hidden_dim)
        self.ligand_atom_emb = nn.Linear(ligand_atom_feature_dim + enc_ligand_dim, self.hidden_dim)
        # self.mean = target_mean
        # self.std = target_std
        # self.register_buffer('target_mean', target_mean)
        # self.register_buffer('target_std', target_std)
        self.encoder = get_encoder(config.encoder)
        if self.enc_node_dim > 0:
            self.enc_node_layer = nn.Sequential(
                nn.Linear(self.hidden_dim + self.enc_node_dim, self.hidden_dim),
                nn.ReLU(),
                nn.Linear(self.hidden_dim, self.hidden_dim),
            )

        self.out_block = nn.Sequential(
            nn.Linear(self.hidden_dim + self.enc_graph_dim, self.hidden_dim),
            ShiftedSoftplus(),
            nn.Linear(self.hidden_dim, output_dim),
        )

    def forward(self, protein_pos, protein_atom_feature, ligand_pos, ligand_atom_feature, batch_protein, batch_ligand,
                output_kind, enc_ligand_feature, enc_node_feature, enc_graph_feature):
        h_protein = self.protein_atom_emb(protein_atom_feature)
        if enc_ligand_feature is not None:
            ligand_atom_feature = torch.cat([ligand_atom_feature, enc_ligand_feature], dim=-1)
        h_ligand = self.ligand_atom_emb(ligand_atom_feature)

        h_ctx, pos_ctx, batch_ctx = compose_context_prop(
            h_protein=h_protein,
            h_ligand=h_ligand,
            pos_protein=protein_pos,
            pos_ligand=ligand_pos,
            batch_protein=batch_protein,
            batch_ligand=batch_ligand,
        )

        h_ctx = self.encoder(
            node_attr=h_ctx,
            pos=pos_ctx,
            batch=batch_ctx,
        )  # (N_p+N_l, H)

        if enc_node_feature is not None:
            h_ctx = torch.cat([h_ctx, enc_node_feature], dim=-1)
            h_ctx = self.enc_node_layer(h_ctx)

        # Aggregate messages
        pre_out = scatter(h_ctx, index=batch_ctx, dim=0, reduce='sum')  # (N, F)
        if enc_graph_feature is not None:
            pre_out = torch.cat([pre_out, enc_graph_feature], dim=-1)

        output = self.out_block(pre_out)  # (N, C)
        if output_kind is not None:
            output_mask = F.one_hot(output_kind - 1, self.output_dim)
            output = torch.sum(output * output_mask, dim=-1, keepdim=True)
        return output

    def get_loss(self, batch, pos_noise_std, return_pred=False):
        protein_noise = torch.randn_like(batch.protein_pos) * pos_noise_std
        ligand_noise = torch.randn_like(batch.ligand_pos) * pos_noise_std

        # add features
        enc_ligand_feature, enc_node_feature, enc_graph_feature = None, None, None
        if self.enc_feature_type == 'nll_all':
            enc_graph_feature = batch.nll_all  # [num_graphs, 22]
        elif self.enc_feature_type == 'nll':
            enc_graph_feature = batch.nll  # [num_graphs, 20]
        elif self.enc_feature_type == 'final_h':
            enc_node_feature = batch.final_h  # [num_pl_atoms, 128]
        elif self.enc_feature_type == 'pred_ligand_v':
            enc_ligand_feature = batch.pred_ligand_v  # [num_l_atoms, 13]
        elif self.enc_feature_type == 'pred_v_entropy_pre':
            enc_ligand_feature = batch.pred_v_entropy   # [num_l_atoms, 1]
        elif self.enc_feature_type == 'pred_v_entropy_post':
            enc_graph_feature = scatter(batch.pred_v_entropy, index=batch.ligand_element_batch, dim=0, reduce='sum')   # [num_graphs, 1]
        elif self.enc_feature_type == 'full':
            enc_graph_feature = torch.cat(
                [batch.nll_all, scatter(batch.pred_v_entropy, index=batch.ligand_element_batch, dim=0, reduce='sum')], dim=-1)
            enc_node_feature = batch.final_h
            enc_ligand_feature = torch.cat([batch.pred_ligand_v, batch.pred_v_entropy], -1)
        else:
            raise NotImplementedError

        pred = self(
            protein_pos=batch.protein_pos + protein_noise,
            protein_atom_feature=batch.protein_atom_feature.float(),
            ligand_pos=batch.ligand_pos + ligand_noise,
            ligand_atom_feature=batch.ligand_atom_feature_full.float(),
            batch_protein=batch.protein_element_batch,
            batch_ligand=batch.ligand_element_batch,
            output_kind=batch.kind,
            # output_kind=None,
            enc_ligand_feature=enc_ligand_feature,
            enc_node_feature=enc_node_feature,
            enc_graph_feature=enc_graph_feature
        )
        # pred = pred * y_std + y_mean
        loss_func = nn.MSELoss()
        loss = loss_func(pred.view(-1), batch.y)
        if return_pred:
            return loss, pred
        else:
            return loss