refactor reconstruct_multi_models

Author: brianhie

bin/amis.py

@@ -203,9 +203,22 @@ def diff(seq_old, seq_new, start=0, end=None):
             different_muts.append((idx, ch_old, ch_new))
     return different_muts
 
-def reconstruct_multi_models(wt_seq, models, alpha=None, return_names=False):
+def reconstruct_multi_models(
+        wt_seq,
+        model_names=[
+            'esm1b',
+            'esm1v1',
+            'esm1v2',
+            'esm1v3',
+            'esm1v4',
+            'esm1v5',
+        ],
+        alpha=None,
+        return_names=False,
+):
     mutations_models, mutations_model_names = {}, {}
-    for model in models:
+    for model_name in model_names:
+        model = get_model_name(model_name)
         if alpha is None:
             wt_new = reconstruct(
                 wt_seq, model, decode_kwargs={ 'exclude': 'unnatural' }
@@ -221,12 +234,13 @@ def reconstruct_multi_models(wt_seq, models, alpha=None, return_names=False):
                 mutations_model_names[mutation] = []
             mutations_models[mutation] += 1
             mutations_model_names[mutation].append(model.name_)
+        del model
 
     if return_names:
         return mutations_models, mutations_model_names
 
     return mutations_models
-    
+
 def evolve(seq_uca, model, n_generations=1):
     print(f'Gen 0: {seq_uca}')
     seq_curr = seq_uca

bin/dms.py

@@ -2,10 +2,7 @@
 import pandas as pd
 import scipy.stats as ss
 
-from amis import (
-    get_model_name,
-    reconstruct_multi_models,
-)
+from amis import reconstruct_multi_models
 
 def parse_args():
     import argparse
@@ -26,7 +23,7 @@ def parse_args():
     args = parser.parse_args()
     return args
 
-def dms_results(fname, models, alpha=None,):
+def dms_results(fname, model_names, alpha=None,):
     namespace = fname.split('/')[-1].split('.')[0]
     with open(fname):
         df = pd.read_csv(fname, delimiter=',')
@@ -39,7 +36,7 @@ def dms_results(fname, models, alpha=None,):
         else:
             wt_seq += aa_orig
 
-    mutations_models = reconstruct_multi_models(wt_seq, models, alpha=alpha)
+    mutations_models = reconstruct_multi_models(wt_seq, model_names, alpha=alpha)
 
     mutation_scores, studies = {}, []
     for column in df.columns:
@@ -70,10 +67,6 @@ def dms_results(fname, models, alpha=None,):
 if __name__ == '__main__':
     args = parse_args()
 
-    models = [
-        get_model_name(model_name) for model_name in args.model_names
-    ]
-    
     fnames = [
         'data/dms/dms_adrb2.csv',
         'data/dms/dms_bla.csv',
@@ -87,4 +80,4 @@ def dms_results(fname, models, alpha=None,):
     ]
 
     for fname in fnames:
-        dms_results(fname, models, alpha=args.alpha)
+        dms_results(fname, args.model_names, alpha=args.alpha)

bin/dms_esm.py

@@ -2,10 +2,7 @@
 import pandas as pd
 import scipy.stats as ss
 
-from amis import (
-    get_model_name,
-    reconstruct_multi_models,
-)
+from amis import reconstruct_multi_models
 from predict_esm import predict_esm
 
 def parse_args():
@@ -50,17 +47,14 @@ def dms_esm(fname, model_locations):
         ]))
     df_esm['ESM_summed'] = summed_log_probs
 
-    models = [
-        get_model_name(model_name) for model_name in
-        [ 'esm1b', 'esm1v1', 'esm1v2', 'esm1v3', 'esm1v4', 'esm1v5', ] 
-    ]
+    model_names = [ 'esm1b', 'esm1v1', 'esm1v2', 'esm1v3', 'esm1v4', 'esm1v5', ] 
     if '_infa' in fname:
         alpha = 0.5
     else:
         alpha = None
     mutations_models, mutations_model_names = reconstruct_multi_models(
         wt_seq,
-        models,
+        models_names,
         alpha=alpha,
         return_names=True,
     )

bin/originality.py

@@ -1,8 +1,5 @@
 from utils import *
-from amis import (
-    get_model_name,
-    reconstruct_multi_models,
-)
+from amis import reconstruct_multi_models
 
 seqs_abs = {
     'medi_vh': 'QVQLQQSGPGLVKPSQTLSLTCAISGDSVSSYNAVWNWIRQSPSRGLEWLGRTYYRSGWYNDYAESVKSRITINPDTSKNQFSLQLNSVTPEDTAVYYCARSGHITVFGVNVDAFDMWGQGTMVTVSS',
@@ -62,14 +59,10 @@ def load_ratios(name, database='abysis'):
 if __name__ == '__main__':
     model_names = [ 'esm1b', 'esm1v1', 'esm1v2', 'esm1v3', 'esm1v4', 'esm1v5', ]
 
-    models = [
-        get_model_name(model_name) for model_name in model_names
-    ]
-
     for name in seqs_abs:
         seq = seqs_abs[name]
         print(f'\n{name}')
-        mutations_models = reconstruct_multi_models(seq, models, alpha=0.5)
+        mutations_models = reconstruct_multi_models(seq, model_names, alpha=0.5)
 
         abysis_likelihood_ratios = load_ratios(name, 'abysis')
         uniref_likelihood_ratios = load_ratios(name, 'uniref')

bin/recommend.py

@@ -1,7 +1,4 @@
-from amis import (
-    get_model_name,
-    reconstruct_multi_models,
-)
+from amis import reconstruct_multi_models
 
 def parse_args():
     import argparse
@@ -29,11 +26,11 @@ def parse_args():
 if __name__ == '__main__':
     args = parse_args()
 
-    models = [
-        get_model_name(model_name) for model_name in args.model_names
-    ]
-
-    mutations_models = reconstruct_multi_models(args.sequence, models, alpha=args.alpha)
+    mutations_models = reconstruct_multi_models(
+        args.sequence,
+        args.model_names,
+        alpha=args.alpha,
+    )
     for k, v in sorted(mutations_models.items(), key=lambda item: -item[1]):
         mut_str = f'{k[1]}{k[0] + 1}{k[2]}'
         print(f'{mut_str}\t{v}')

bin/therapeutic_antibodies.py

@@ -1,10 +1,7 @@
 import pandas as pd
 import time
 
-from amis import (
-    get_model_name,
-    reconstruct_multi_models,
-)
+from amis import reconstruct_multi_models
 from utils import tprint
 
 
@@ -37,10 +34,6 @@ def parse_args():
 if __name__ == '__main__':
     args = parse_args()
 
-    models = [
-        get_model_name(model_name) for model_name in args.model_names
-    ]
-    
     df = pd.read_csv(args.seq_fname, sep=',')
 
     data = []
@@ -60,7 +53,7 @@ def parse_args():
         for wt_seq in [ vh, vl, vh_bi, vl_bi ]:
             if wt_seq == 'na':
                 continue
-            mutations_models = reconstruct_multi_models(wt_seq, models)
+            mutations_models = reconstruct_multi_models(wt_seq, args.model_names)
             for mutation, n_models in sorted(
                     mutations_models.items(), key=lambda item: item[1]
             ):