Drug-combination-synergy/predictor/model/models.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import sys 

PROJ_DIR = os.path.dirname(os.path.abspath(os.path.join(os.path.dirname( __file__ ), '..')))

sys.path.insert(0, PROJ_DIR)
from drug.models import GCN
from drug.datasets import DDInteractionDataset
from model.utils import get_FP_by_negative_index



class Connector(nn.Module):
    def __init__(self, gpu_id=None):
        super(Connector, self).__init__()
        # self.ddiDataset = DDInteractionDataset(gpu_id = gpu_id)
        self.gcn = None
        
        #Cell line features
        # np.load('cell_feat.npy')

    def forward(self, drug1_idx, drug2_idx, cell_feat, subgraph):
        if self.gcn == None:
            self.gcn = GCN(subgraph.num_features, subgraph.num_features // 2)
        x = subgraph.get().x
        edge_index = subgraph.edge_index
        x = self.gcn(x, edge_index)
        drug1_idx = torch.flatten(drug1_idx)
        drug2_idx = torch.flatten(drug2_idx)
        drug1_feat = x[drug1_idx]
        drug2_feat = x[drug2_idx]
        for i, idx in enumerate(drug1_idx):
            if idx < 0:
                drug1_feat[i] = get_FP_by_negative_index(idx)
        for i, idx in enumerate(drug2_idx):
            if idx < 0:
                drug2_feat[i] = get_FP_by_negative_index(idx)
        feat = torch.cat([drug1_feat, drug2_feat, cell_feat], 1)
        return feat


class MLP(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, gpu_id=None):
        super(MLP, self).__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.BatchNorm1d(hidden_size),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.BatchNorm1d(hidden_size // 2),
            nn.Linear(hidden_size // 2, 1)
        )

        self.connector = Connector(gpu_id)
    
    # prev input: self, drug1_feat: torch.Tensor, drug2_feat: torch.Tensor, cell_feat: torch.Tensor, subgraph: related subgraph for the batch
    def forward(self, drug1_idx, drug2_idx, cell_feat, subgraph):
        feat = self.connector(drug1_idx, drug2_idx, cell_feat, subgraph)
        out = self.layers(feat)
        return out


# other PRODeepSyn models have been deleted for now