before merging

Author: tyfei0216

VirusDataset.py

@@ -0,0 +1,172 @@
+import os
+
+from pytorch_lightning.utilities.types import TRAIN_DATALOADERS 
+import ioutils 
+import random
+
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+import torch
+import pytorch_lightning as L
+
+def readVirusSequences(pos=None, trunc=1498, sample = 300, seed=1509):
+    random.seed(seed)
+    print("read positive samples")
+    seqs = {}
+    if pos is None:
+        pos = os.listdir("/home/tyfei/datasets/ion_channel/Interprot/ion_channel/0.99")
+    for i in pos:
+        # print(i)
+        try:
+            if i.endswith(".fas"): 
+                # print(i, i[:i.find(".")] in df["Accession"].values)
+                gen = ioutils.readFasta("/home/tyfei/datasets/ion_channel/Interprot/ion_channel/0.99/"+i, truclength=trunc)
+                seqs[i[:i.find(".")]] = [i for i in gen]
+                # print(i, "success") 
+        except:
+            pass
+            # print(i, "failed")
+
+    sequences = [] 
+    labels = [] 
+    for i in seqs:
+        sampled = random.sample(seqs[i], min(sample, len(seqs[i])))
+        sequences.extend(sampled)
+        labels.extend([1]*len(sampled))
+
+    print("read negative samples")
+    gen = ioutils.readFasta("/home/tyfei/datasets/ion_channel/Interprot/Negative_sample/decoy_1m_new.fasta", truclength=trunc)
+    seqs["neg"] = [i for i in gen]
+    sampled = random.sample(seqs["neg"], min(len(labels), len(seqs["neg"])))
+    sequences.extend(sampled)
+    labels.extend([0]*len(sampled))
+
+    print("read virus sequences")
+    allvirus = []
+    for i in os.listdir("/home/tyfei/datasets/NCBI_virus/genbank_csv/"):
+        allvirus.extend(ioutils.readNCBICsv("/home/tyfei/datasets/NCBI_virus/genbank_csv/"+i, truclength=trunc))
+
+    return sequences, labels, allvirus
+
+
+class SeqDataset2(Dataset):
+    def __init__(self, seq, label, seqtest):
+        
+        if not isinstance(seq, torch.Tensor):
+            seq = torch.tensor(seq).long()
+        self.seq = seq
+
+        if not isinstance(label, torch.Tensor):
+            label = torch.tensor(label).long()
+        self.label = label
+        
+        if not isinstance(seqtest, torch.Tensor):
+            seqtest = torch.tensor(seqtest).long()
+        self.seqtest = seqtest
+
+        self.seqlen = seq.shape[0]
+        self.seqtestlen = seqtest.shape[0]
+
+    def __len__(self):
+        return max(self.seqlen, self.seqtestlen)
+
+    def __getitem__(self, idx):
+        return self.seq[idx%self.seqlen], self.label[idx%self.seqlen], self.seqtest[idx%self.seqtestlen]
+
+class TestDataset(Dataset):
+    def __init__(self, seq):
+        if not isinstance(seq, torch.Tensor):
+            seq = torch.tensor(seq).long()
+        self.seq = seq
+
+    def __len__(self):
+        return self.seq.shape[0]
+
+    def __getitem__(self, idx):
+        
+        return self.seq[idx]
+    
+
+class SeqDataset(Dataset):
+    def __init__(self, seq, label):
+        if not isinstance(seq, torch.Tensor):
+            seq = torch.tensor(seq).long()
+        self.seq = seq
+
+        if not isinstance(label, torch.Tensor):
+            label = torch.tensor(label).long()
+        self.label = label
+
+    def __len__(self):
+        return self.seq.shape[0]
+
+    def __getitem__(self, idx):
+        
+        return self.seq[idx], self.label[idx]
+    
+class SeqdataModule(L.LightningDataModule):
+    def __init__(self, trainset=None, testset=None, path="/home/tyfei/datasets/pts/virus", batch_size = 12, train_test_split=[0.8, 0.2], seed = 1509) -> None:
+        super().__init__()
+        
+        self.train_test_split = train_test_split
+        self.batch_size = batch_size
+        self.path = path 
+        self.seed = seed
+        
+        if isinstance(testset, str):
+            self.test_set = torch.load(testset)
+        else:
+            self.test_set = testset 
+        
+        if isinstance(trainset, str):
+            self.trainset = torch.load(trainset)
+        else:
+            self.trainset = trainset
+
+        # if self.trainset is not None:
+        #     train_set, val_set = torch.utils.data.random_split(trainset, train_test_split)
+        
+        #     self.train_set = train_set  
+        #     self.val_set = val_set 
+
+    def saveDataset(self):
+        torch.save(self.trainset, os.path.join(self.path, "train.pt"))
+        torch.save(self.test_set, os.path.join(self.path, "test.pt"), self.test_set)
+        
+    
+    def setup(self, stage):
+        if stage == "fit" or stage == "validate":
+            if self.trainset is None:
+                if os.path.exists(os.path.join(self.path, "train.pt")):
+                    self.trainset = torch.load(os.path.join(self.path, "train.pt"))
+                else:
+                    raise FileExistsError
+            
+            if not hasattr(self, "train_set"):
+                torch.manual_seed(self.seed)
+                train_set, val_set = torch.utils.data.random_split(self.trainset, self.train_test_split)
+                self.train_set = train_set  
+                self.val_set = val_set 
+        
+        if stage == "predict":
+            if self.test_set is None:
+                if os.path.exists(os.path.join(self.path, "test.pt")):
+                    self.test_set = torch.load(os.path.join(self.path, "test.pt"))
+                else:
+                    raise FileExistsError
+                
+        if stage == "test":
+            raise NotImplementedError
+
+
+        return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True, num_workers=4)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size, shuffle=False)
+
+    def predict_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size, shuffle=False)
+    
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True, num_workers=4)
+    

datautils.py

@@ -12,6 +12,8 @@
 import pytorch_lightning as L 
 tqdm.pandas()
 
+
+
 def parseline(line):
     try:
         line = line.decode("utf-8").strip()

ionchannel.py

@@ -0,0 +1,211 @@
+import torch
+import esm 
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as L
+import numpy as np
+
+# from torchmetrics import Metric
+
+import torchmetrics
+
+from torch.autograd import Function
+
+class GradientR(Function):
+    @staticmethod
+    def forward(ctx, x, alpha):
+        ctx.save_for_backward(x, alpha)
+        return x
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = None
+        _, alpha = ctx.saved_tensors
+        if ctx.needs_input_grad[0]:
+            grad_input = - alpha*grad_output
+        return grad_input, None
+
+
+class GradientReversal(nn.Module):
+    def __init__(self, alpha):
+        super().__init__()
+        self.alpha = torch.tensor(alpha, requires_grad=False)
+
+    def forward(self, x):
+        return GradientR.apply(x, self.alpha)
+
+class ionclf(L.LightningModule):
+    def __init__(self, esm_model, unfix = None, addadversial=True, lamb=0.1, lr=5e-4) -> None:
+        super().__init__()
+        self.num_layers = esm_model.num_layers 
+        self.embed_dim = esm_model.embed_dim 
+        self.attention_heads = esm_model.attention_heads 
+        self.alphabet = esm.data.Alphabet.from_architecture("ESM-1b") 
+        self.alphabet_size = len(self.alphabet)
+        self.addadversial = addadversial
+        self.lamb =  lamb
+        self.lr = lr
+
+        self.esm_model = esm_model 
+
+        self.cls = nn.Sequential(nn.Linear(self.embed_dim, self.embed_dim // 2), 
+                                  nn.LayerNorm(self.embed_dim // 2), 
+                                  nn.GELU(), 
+                                  nn.Linear(self.embed_dim // 2, self.embed_dim // 4), 
+                                  nn.LayerNorm(self.embed_dim // 4), 
+                                  nn.GELU(), 
+                                  nn.Linear(self.embed_dim // 4, 1)
+                                  )
+        
+        self.dis = nn.Sequential(nn.Linear(self.embed_dim, self.embed_dim // 2), 
+                                  nn.LayerNorm(self.embed_dim // 2), 
+                                  nn.GELU(), 
+                                  nn.Linear(self.embed_dim // 2, self.embed_dim // 4), 
+                                  nn.LayerNorm(self.embed_dim // 4), 
+                                  nn.GELU(), 
+                                  nn.Linear(self.embed_dim // 4, 1)
+                                  )
+
+        self.reverse = GradientReversal(1)
+
+        if unfix is None:
+            self.fixParameters() 
+        else:
+            self.fixParameters(unfix)
+
+        self.acc = torchmetrics.Accuracy(task="binary")
+
+        self.training_step_outputs = []
+        self.validation_step_outputs = []
+
+
+
+    def fixParameters(self, unfix=["9", "10", "11"]):
+        for i, j in self.named_parameters():
+            flag = 1
+            if "esm_model" not in i:
+                    flag = 0
+            for k in unfix:
+                if k in i:
+                    flag = 0
+    
+            if flag == 1:
+                j.requires_grad = False 
+            else:
+                j.requires_grad = True
+     
+    def forward(self, x):
+        representations = self.esm_model(x, repr_layers=[self.num_layers])
+
+        x = representations["representations"][self.num_layers][:, 0] 
+        x1 = self.reverse(x)
+        pre = self.cls(x) 
+        pre = F.sigmoid(pre)
+
+        y = self.dis(x1)
+        y = F.sigmoid(y)
+
+        return  pre, y
+    
+    def _common_training_step(self, batch):
+        X1, y, X2 = batch 
+        y_pre, dis_pre_x1 = self(X1) 
+        _y, dis_pre_x2 = self(X2) 
+
+        loss1 = F.binary_cross_entropy(y_pre.squeeze(), y.float()) 
+        loss2 = F.binary_cross_entropy(dis_pre_x1, torch.zeros_like(dis_pre_x1)) + \
+            F.binary_cross_entropy(dis_pre_x2, torch.ones_like(dis_pre_x1))
+        
+        if self.addadversial:
+            loss = loss1+loss2*self.lamb
+        else:
+            loss = loss1
+        
+        return loss, loss1, loss2, y_pre, y
+
+
+    def training_step(self, batch, batch_idx):
+
+        loss, loss1, loss2, y_pre, y = self._common_training_step(batch)
+
+        acc = self.acc(y_pre.squeeze(), y)
+
+        self.log_dict({"predict loss":loss1.item(), "adversial loss":loss2.item(), "acc":acc}, prog_bar=True, on_step=True)
+        self.training_step_outputs.append({"loss":loss.detach().cpu(), "y":y_pre.detach().squeeze().cpu(), "true_label":y.cpu()})
+        
+        return loss
+
+    def _common_epoch_end(self, outputs):
+
+        loss = torch.stack([x["loss"] for x in outputs]).mean()
+        scores = torch.concatenate([x["y"] for x in outputs])
+        y = torch.concatenate([x["true_label"] for x in outputs]) 
+
+        outputs.clear() 
+        return loss, self.acc(scores, y)
+
+    def on_training_epoch_end(self):
+
+        loss, acc = self._common_epoch_end(self.training_step_outputs)
+
+        # print("finish training epoch, loss %f, acc %f"%(loss, acc))
+        self.log_dict(
+            {
+                "mean_loss":loss, 
+                "train_acc": acc,
+            },
+            on_step=False,
+            on_epoch=True,
+            prog_bar=False,
+        )
+    
+    def validation_step(self, batch, batch_idx):
+
+        loss, loss1, loss2, y_pre, y = self._common_training_step(batch)
+
+        acc = self.acc(y_pre.squeeze(), y)
+
+        self.log_dict({"predict loss":loss1.item(), "adversial loss":loss2.item(), "acc":acc}, prog_bar=True, on_step=True)
+        
+        self.validation_step_outputs.append({"loss":loss.cpu(), "y":y_pre.squeeze().cpu(), "true_label":y.cpu()})
+        
+        return loss
+    
+    def on_validation_epoch_end(self):
+        loss, acc = self._common_epoch_end(self.validation_step_outputs)
+        # print("finish validating, loss %f, acc %f"%(loss, acc))
+        self.log_dict(
+            {
+                "loss":loss, 
+                "validate_acc": acc,
+            },
+            on_step=False,
+            on_epoch=True,
+            prog_bar=False,
+        )
+    
+    def test_step(self, batch, batch_idx):
+        x = batch
+        y_pre, _ = self(x) 
+        return y_pre 
+    
+    def predict_step(self, batch, batch_idx):
+        if isinstance(batch, tuple):
+            if len(batch) == 3:
+                X1, y, X2 = batch 
+            elif len(batch) == 2:
+                X1, y  = batch  
+            else:
+                raise ValueError   
+        else:
+            X1 = batch 
+        pre, _ = self(X1)
+
+        pre = pre.squeeze() 
+        return pre
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(filter(lambda p:p.requires_grad, self.parameters()), lr=self.lr)
+
+        return optimizer    
+