add masking chains by CDRs, fix naming of runction chain masking

training/model_utils.py

@@ -1,6 +1,7 @@
 from __future__ import print_function
 import json, time, os, sys, glob
 import shutil
+from typing import Sequence
 import numpy as np
 import torch
 from torch import optim
@@ -15,7 +16,10 @@
 import random
 import itertools
 
-def get_index_of_masks(len_seq:int,max_parts:int,max_length:int)->list[int]:
+from pathlib import Path
+from proteinlib.structure.antibody_antigen_complex import AntibodyAntigenComplex, NumberingScheme
+
+def get_mask_random(len_seq:int,max_parts:int,max_length:int)->list[int]:
     """
     Randomly generates mask with length equal to len_seq 
     Args:
@@ -27,7 +31,7 @@ def get_index_of_masks(len_seq:int,max_parts:int,max_length:int)->list[int]:
     """
     parts=random.randint(1,max_parts)
     tuples=[]
-    mask=[0]*len_seq
+    mask=np.zeros(len_seq)
     for _ in range(parts):
         s=random.randint(0,len_seq-max_length-1)
         e=random.randint(s,s+max_length-1)
@@ -36,13 +40,35 @@ def get_index_of_masks(len_seq:int,max_parts:int,max_length:int)->list[int]:
             if s>=t[0] and s<=t[1] or e>=t[0] and e<=t[1]:
                 f=True
                 break
-        if not f:
+        is_previous_masked=(s==0 or not mask[s-1])
+        is_next_masked=(e==len_seq-1 or not mask[e+1])
+        if not f and is_previous_masked and is_next_masked:
             tuples.append((s,e))
-            for i in range(s,e+1):
-                mask[i]=1
+            mask[s:(e+1)]=1
+    return mask.tolist()
+
+def get_mask_cdrs_one_chain(chain,indexes_of_cdrs:list[int]):
+    r=chain.region_boundaries
+    cdrs_index=[(r[i],r[i+1]) for i in range(1,len(r)-1,2)]
+    mask=np.zeros(len(chain.sequence))
+    for i in indexes_of_cdrs:
+        t=cdrs_index[i]
+        for i in range(t[0],t[1]):
+            mask[i]=1
     return mask
 
-def featurize(batch, device,max_parts=6,max_length=6):
+def get_mask_cdrs(pdb:Path,heavy_chain_id:str,light_chain_id:str,antigen_chain_ids:Sequence[str],indexes_of_cdrs:list[int],numbering=NumberingScheme.CHOTHIA,):
+    ab_complex = AntibodyAntigenComplex.from_pdb(
+        pdb=pdb,
+        heavy_chain_id=heavy_chain_id,
+        light_chain_id=light_chain_id,
+        antigen_chain_ids=antigen_chain_ids,
+        numbering=numbering,
+    )
+    return get_mask_cdrs_one_chain(ab_complex.antibody.heavy_chain,indexes_of_cdrs),get_mask_cdrs_one_chain(ab_complex.antibody.light_chain,indexes_of_cdrs)
+
+def featurize(batch, device,max_parts=8,max_length=20,mode='train',pdb_dir=Path('/mnt/sabdab/chothia'),cdr_indexes=[0,0]):
+    ind,jnd=cdr_indexes
     alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
     B = len(batch)
     lengths = np.array([len(b['seq']) for b in batch], dtype=np.int32) #sum of chain seq lengths
@@ -106,7 +132,22 @@ def featurize(batch, device,max_parts=6,max_length=6):
                 chain_seq = b[f'seq_chain_{letter}']
                 chain_length = len(chain_seq)
                 chain_coords = b[f'coords_chain_{letter}'] #this is a dictionary
-                chain_mask=np.array(get_index_of_masks(len(chain_seq),max_parts,max_length))
+                if mode=='test':
+                    chain_mask=np.array(get_mask_random(chain_length,max_parts,max_length)) 
+                elif mode=='valid':
+                    if masked_chains.index(letter)==ind:
+                        chain_mask= np.array(get_mask_cdrs(pdb_dir/f"{b['name']}.pdb",b['masked_list'][0],b['masked_list'][1],b['visible_list'],[jnd])[ind])
+                    else:
+                        chain_mask = np.zeros(chain_length)
+                    if chain_length!=chain_mask.shape[0]:
+                        if masked_chains.index(letter)==ind:
+                            chain_mask= np.array(get_mask_cdrs(pdb_dir/f"{b['name']}.pdb",b['masked_list'][0],b['masked_list'][1],b['visible_list'],max_parts,max_length,[jnd])[ind])
+                        else:
+                            chain_mask = np.zeros(chain_length)
+
+                    assert chain_length==chain_mask.shape[0]
+                else:
+                    chain_mask=np.array(get_mask_cdrs(pdb_dir/f"{b['name']}.pdb",b['masked_list'][0],b['masked_list'][1],b['visible_list'],[0,1,2])[masked_chains.index(letter)])
                 x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_lenght,4,3]
                 x_chain_list.append(x_chain)
                 chain_mask_list.append(chain_mask)

training/test_get_index_of_masks.py

@@ -1,15 +1,16 @@
 from itertools import groupby
+from pathlib import Path
 import random
 import string
-from model_utils import get_index_of_masks
-import numpy as np
+from model_utils import get_mask_random,get_mask_cdrs
+from proteinlib.structure.antibody_antigen_complex import AntibodyAntigenComplex, NumberingScheme
 
-def test_mask_indexing():
+def test_mask_random():
     random.seed(42)
     len_seq=100
     max_parts=10
     max_length=5
-    mask=get_index_of_masks(len_seq,max_parts,max_length)
+    mask=get_mask_random(len_seq,max_parts,max_length)
     assert sum([k for k,_ in groupby(mask)])<=max_parts
     assert max(len(list(g)) for k, g in groupby(mask) if k==1) <=max_length
     assert mask== [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0,
@@ -18,5 +19,29 @@ def test_mask_indexing():
     s=''.join((1+len_seq//len(string.ascii_uppercase))*string.ascii_uppercase)[:len_seq]
     s_masked=''.join(['_' if mask[i] else s[i] for i in range(len_seq)])
     assert s_masked=='ABC___GHIJKLMNOPQRSTUVWXYZABCDE__HIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUV'
+
+def test_mask_cdrs():
+    max_parts=10
+    max_length=5
+
+    pdb=Path('data/7pa6.pdb')
+    heavy_chain_id="K"
+    light_chain_id="k"
+    antigen_chain_ids=["A"]
+    numbering=NumberingScheme.CHOTHIA
+    mask_vh,mask_vl=get_mask_cdrs(pdb=pdb,heavy_chain_id=heavy_chain_id,light_chain_id=light_chain_id,antigen_chain_ids=antigen_chain_ids,max_parts=max_parts,max_length=max_length,numbering=numbering)
+    assert sum([k for k,_ in groupby(mask_vh)])<=max_parts
+    assert sum([k for k,_ in groupby(mask_vl)])<=max_parts
+    assert max(len(list(g)) for k, g in groupby(mask_vh) if k==1)<=max_length
+    assert max(len(list(g)) for k, g in groupby(mask_vl) if k==1)<=max_length
+    ab_complex = AntibodyAntigenComplex.from_pdb(pdb=pdb,heavy_chain_id=heavy_chain_id,light_chain_id=light_chain_id,antigen_chain_ids=antigen_chain_ids,numbering=numbering)
+    r_h=ab_complex.antibody.heavy_chain.region_boundaries
+    r_l=ab_complex.antibody.light_chain.region_boundaries
+    assert len(r_h)==len(r_l)
+    for i in range(0,len(r_h)-1,2):
+        assert not any(mask_vh[r_h[i]:r_h[i+1]])
+        assert not any(mask_vl[r_l[i]:r_l[i+1]])
+
 
-test_mask_indexing()
+test_mask_random()
+test_mask_cdrs()

training/training.py

@@ -2,8 +2,31 @@
 import os.path
 from pathlib import Path
 import pickle
+from tqdm import tqdm
 
-def main(args, preprocessed_path):
+def dict_train_test_split(complex_list:list[str],train_ids:list[str],val_ids:list[str],test_ids:list[str],del_train_ids:list[str],del_test_ids:list[str]):
+    d_train=[]
+    d_val=[]
+    complex_ids=set()
+    for c in complex_list:
+        if c['name'] in complex_ids:
+            continue
+        complex_ids.add(c['name'])
+        if c['name_long'] in train_ids:
+            d_train.append(c)
+        elif c['name_long'] in val_ids:
+            d_val.append(c)
+        elif c['name_long'] in test_ids:
+            pass
+        elif c['name_long'] in del_train_ids:
+            pass
+        elif c['name_long'] in del_test_ids:
+            pass
+        else:
+            print(f'Complex id {c["name_long"]} is not presented in train or test IDs')
+    return d_train,d_val
+
+def main(args, preprocessed_path:Path, train_ids:list[str], val_ids:list[str],test_ids:list[str],del_train_ids:list[str],del_test_ids:list[str],indexes_of_cdrs:list[tuple[int,int]]):
     import json, time, os, sys, glob
     import shutil
     import warnings
@@ -29,7 +52,10 @@ def main(args, preprocessed_path):
     device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
 
     base_folder = time.strftime(args.path_for_outputs, time.localtime())
-
+
+    df={'epoch': [], 'step':[], 'time': [], 'train': [], 'valid':[], 'train_acc': [], 'valid_acc': [],
+            (0,0):[],(0,1):[],(0,2):[],
+            (1,0):[],(1,1):[],(1,2):[],}
     if base_folder[-1] != '/':
         base_folder += '/'
     if not os.path.exists(base_folder):
@@ -78,30 +104,26 @@ def main(args, preprocessed_path):
                         dropout=args.dropout, 
                         augment_eps=args.backbone_noise)
     model.to(device)
-
-
     if PATH:
         checkpoint = torch.load(PATH)
-        total_step = checkpoint['step'] #write total_step from the checkpoint
-        epoch = checkpoint['epoch'] #write epoch from the checkpoint
+        total_step = 0
+        epoch = 0
         model.load_state_dict(checkpoint['model_state_dict'])
     else:
         total_step = 0
         epoch = 0
 
     optimizer = get_std_opt(model.parameters(), args.hidden_dim, total_step)
 
-
-    if PATH:
-        optimizer.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-
-    dataset_train = StructureDataset(d, truncate=None, max_length=args.max_protein_length) 
-    dataset_valid = StructureDataset(d, truncate=None, max_length=args.max_protein_length)
+    d_train,d_test=dict_train_test_split(d,train_ids,val_ids,test_ids,del_train_ids,del_test_ids)
+    dataset_train = StructureDataset(d_train, truncate=None, max_length=args.max_protein_length) 
+    dataset_valid = StructureDataset(d_test, truncate=None, max_length=args.max_protein_length)
 
     loader_train = StructureLoader(dataset_train, batch_size=args.batch_size)
     loader_valid = StructureLoader(dataset_valid, batch_size=args.batch_size)
 
     reload_c = 0 
+    max_val_acc=0
     for e in range(args.num_epochs):
         t0 = time.time()
         e = epoch + e
@@ -110,7 +132,7 @@ def main(args, preprocessed_path):
         train_acc = 0.
         for _, batch in enumerate(loader_train):
             start_batch = time.time()
-            X, S, mask, lengths, chain_M, residue_idx, mask_self, chain_encoding_all = featurize(batch, device)
+            X, S, mask, lengths, chain_M, residue_idx, mask_self, chain_encoding_all = featurize(batch, device,mode='valid',cdr_indexes=indexes_of_cdrs[0])
             elapsed_featurize = time.time() - start_batch
             optimizer.zero_grad()
             mask_for_loss = mask*chain_M
@@ -144,38 +166,84 @@ def main(args, preprocessed_path):
             train_weights += torch.sum(mask_for_loss).cpu().data.numpy()
 
             total_step += 1
-
-        model.eval()
+        dict_val={}
+
+
+        for ind,jnd in indexes_of_cdrs:
+            model.eval()
+            with torch.no_grad():
+                validation_sum, validation_weights = 0., 0.
+                validation_acc = 0.
+                for _, batch in enumerate(loader_valid):
+                    X, S, mask, lengths, chain_M, residue_idx, mask_self, chain_encoding_all = featurize(batch, device,mode='valid',cdr_indexes=[ind,jnd])
+                    log_probs = model(X, S, mask, chain_M, residue_idx, chain_encoding_all)
+                    mask_for_loss = mask*chain_M
+                    loss, loss_av, true_false = loss_nll(S, log_probs, mask_for_loss)
+
+                    validation_sum += torch.sum(loss * mask_for_loss).cpu().data.numpy()
+                    validation_acc += torch.sum(true_false * mask_for_loss).cpu().data.numpy()
+                    validation_weights += torch.sum(mask_for_loss).cpu().data.numpy()
+
+            validation_loss = validation_sum / validation_weights
+            validation_accuracy = validation_acc / validation_weights
+            validation_perplexity = np.exp(validation_loss)
+
+            validation_perplexity_ = np.format_float_positional(np.float32(validation_perplexity), unique=False, precision=3)
+            validation_accuracy_ = np.format_float_positional(np.float32(validation_accuracy), unique=False, precision=3)
+            df[(ind,jnd)].append(validation_accuracy_)
         with torch.no_grad():
             validation_sum, validation_weights = 0., 0.
             validation_acc = 0.
             for _, batch in enumerate(loader_valid):
-                X, S, mask, lengths, chain_M, residue_idx, mask_self, chain_encoding_all = featurize(batch, device)
+                X, S, mask, lengths, chain_M, residue_idx, mask_self, chain_encoding_all = featurize(batch, device,mode='valid',cdr_indexes=[ind,jnd])
                 log_probs = model(X, S, mask, chain_M, residue_idx, chain_encoding_all)
                 mask_for_loss = mask*chain_M
                 loss, loss_av, true_false = loss_nll(S, log_probs, mask_for_loss)
 
                 validation_sum += torch.sum(loss * mask_for_loss).cpu().data.numpy()
                 validation_acc += torch.sum(true_false * mask_for_loss).cpu().data.numpy()
                 validation_weights += torch.sum(mask_for_loss).cpu().data.numpy()
-
-        train_loss = train_sum / train_weights
-        train_accuracy = train_acc / train_weights
-        train_perplexity = np.exp(train_loss)
+
         validation_loss = validation_sum / validation_weights
         validation_accuracy = validation_acc / validation_weights
         validation_perplexity = np.exp(validation_loss)
+
+        train_loss = train_sum / train_weights
+        train_accuracy = train_acc / train_weights
+        train_perplexity = np.exp(train_loss)
 
         train_perplexity_ = np.format_float_positional(np.float32(train_perplexity), unique=False, precision=3)     
         validation_perplexity_ = np.format_float_positional(np.float32(validation_perplexity), unique=False, precision=3)
         train_accuracy_ = np.format_float_positional(np.float32(train_accuracy), unique=False, precision=3)
         validation_accuracy_ = np.format_float_positional(np.float32(validation_accuracy), unique=False, precision=3)
-
+        if validation_accuracy>max_val_acc:
+            checkpoint_filename_best = base_folder+'model_weights/epoch_best.pt'.format(e+1, total_step)
+            torch.save({
+                        'epoch': e+1,
+                        'step': total_step,
+                        'num_edges' : args.num_neighbors,
+                        'noise_level': args.backbone_noise,
+                        'model_state_dict': model.state_dict(),
+                        'optimizer_state_dict': optimizer.optimizer.state_dict(),
+                        }, checkpoint_filename_best)
+            max_val_acc=validation_accuracy
         t1 = time.time()
         dt = np.format_float_positional(np.float32(t1-t0), unique=False, precision=1) 
         with open(logfile, 'a') as f:
             f.write(f'epoch: {e+1}, step: {total_step}, time: {dt}, train: {train_perplexity_}, valid: {validation_perplexity_}, train_acc: {train_accuracy_}, valid_acc: {validation_accuracy_}\n')
         print(f'epoch: {e+1}, step: {total_step}, time: {dt}, train: {train_perplexity_}, valid: {validation_perplexity_}, train_acc: {train_accuracy_}, valid_acc: {validation_accuracy_}')
+        for ind,jnd in indexes_of_cdrs:
+            print(f'({ind},{jnd}): {df[(ind,jnd)][-1]}',end=' ')
+
+        print()
+        df['epoch'].append(e+1)
+        df['step'].append(total_step) 
+        df['time'].append(dt)
+        df['train'].append(train_perplexity_)
+        df['valid'].append(validation_perplexity_)
+        df['train_acc'].append(train_accuracy_)
+        df['valid_acc'].append(validation_accuracy_)
+
 
         checkpoint_filename_last = base_folder+'model_weights/epoch_last.pt'.format(e+1, total_step)
         torch.save({
@@ -197,8 +265,9 @@ def main(args, preprocessed_path):
                     'model_state_dict': model.state_dict(),
                     'optimizer_state_dict': optimizer.optimizer.state_dict(),
                     }, checkpoint_filename)
-
-
+
+    with open(base_folder+'val_accuracy.pkl', 'wb') as fp:
+        pickle.dump(df,fp)
 if __name__ == "__main__":
     argparser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
@@ -221,7 +290,27 @@ def main(args, preprocessed_path):
     argparser.add_argument("--debug", type=bool, default=False, help="minimal data loading for debugging")
     argparser.add_argument("--gradient_norm", type=float, default=-1.0, help="clip gradient norm, set to negative to omit clipping")
     argparser.add_argument("--mixed_precision", type=bool, default=True, help="train with mixed precision")
-
+
+    argparser.add_argument("--preprocessed_path", type=Path, default="/mnt/proteinmpnn/ProteinMPNN_preprocessed_chothia_proteinlib_logging.pickle")
+    argparser.add_argument("--regions", type=str, default="H3")
+    argparser.add_argument("--comment", type=str, default="")
+
     args = argparser.parse_args() 
-    preprocessed_path = Path("/mnt/proteinmpnn/ProteinMPNN_preprocessed_chothia.pickle")
-    main(args, preprocessed_path)   
+    preprocessed_path = Path(args.preprocessed_path)
+    regions=args.regions
+    args.path_for_outputs=f'exp_{regions}{args.comment}'
+    train_file=Path(f'train_val_test_{regions}/train_renamed_clusterRes_0.5_DB_CDR_{regions}.fasta_cluster.txt')
+    val_file=Path(f'train_val_test_{regions}/val_renamed_clusterRes_0.5_DB_CDR_{regions}.fasta_cluster.txt')
+    test_file=Path(f'train_val_test_{regions}/test_renamed_clusterRes_0.5_DB_CDR_{regions}.fasta_cluster.tsv')
+    del_train_file=Path(f'train_val_test_{regions}/deleted_train_and_val_renamed_clusterRes_0.5_DB_CDR_{regions}.fasta_cluster.tsv')
+    del_test_file=Path(f'train_val_test_{regions}/deleted_train_and_val_renamed_clusterRes_0.5_DB_CDR_{regions}.fasta_cluster.tsv')
+
+
+    train_ids=train_file.read_text().splitlines()
+    val_ids=val_file.read_text().splitlines()
+    test_ids=test_file.read_text().splitlines()
+    del_train_ids=del_train_file.read_text().splitlines()
+    del_test_ids=del_test_file.read_text().splitlines()
+    list_of_cdrs=['H1','H2','H3']
+    indexes_of_cdrs=[(0,list_of_cdrs.index(regions))]
+    main(args, preprocessed_path,train_ids,val_ids,test_ids,del_train_ids,del_test_ids,indexes_of_cdrs)