added compile_metrics.py

scripts/compile_metrics.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python
+#
+# Parses and compiles metrics previously computed by calc_metrics.sh.
+#
+# Usage: ./compile_metrics.py FOLDER
+# 
+# FOLDER should contain subfolders like lddt, trr_score, etc. This will output
+# a file, combined_metrics.csv, in FOLDER.
+#
+import pandas as pd
+import numpy as np
+import os, glob, argparse, sys
+from collections import OrderedDict
+
+p = argparse.ArgumentParser()
+p.add_argument('folder', help='Folder of outputs to process')
+p.add_argument('--out',  help='Output file name.')
+args = p.parse_args()
+
+if args.out is None:
+    args.out = os.path.join(args.folder,'combined_metrics.csv')
+
+if not os.path.isdir(args.folder):
+    sys.exit(f'ERROR: Input path {args.folder} not a folder.')
+
+def parse_fastdesign_filters(folder):
+    files = glob.glob(os.path.join(folder,'*.pdb'))
+    records = []
+    for f in files:
+        row = OrderedDict()
+        row['name'] = os.path.basename(f)[:-4]
+        recording = False
+        with open(f) as inf:
+            for line in inf:
+                if recording and len(line)>1:
+                    tokens = line.split()
+                    if len(tokens) == 2:
+                        row[tokens[0]] = float(tokens[1])
+                if '#END_POSE_ENERGIES_TABLE' in line: 
+                    recording=True
+                if line.startswith('pose'):
+                    row['rosetta_energy'] = float(line.split()[-1])
+        records.append(row)
+    if len(records)>0: return pd.DataFrame.from_records(records)
+    return pd.DataFrame({'name':[]})
+
+def parse_lddt(folder):
+    data = {'name':[], 'lddt':[]}
+    files = glob.glob(os.path.join(folder,'*.npz'))
+    if len(files)==0:
+        return pd.DataFrame({'name':[]})
+    for f in files:
+        prefix = os.path.basename(f).replace('.npz','')
+        lddt_data = np.load(f)
+        data['lddt'].append(lddt_data['lddt'].mean())
+        data['name'].append(prefix)
+    return pd.DataFrame.from_dict(data)
+
+def parse_rosetta_energy_from_pdb(folder):
+    files = glob.glob(os.path.join(folder,'*.pdb'))
+    records = []
+    for pdbfile in files:
+        with open(pdbfile) as inf:
+            name = os.path.basename(pdbfile).replace('.pdb','')
+            rosetta_energy = np.nan
+            for line in inf.readlines():
+                if line.startswith('pose'):
+                    rosetta_energy = float(line.split()[-1])
+            row = OrderedDict()
+            row['name'] = name
+            row['rosetta_energy'] = rosetta_energy
+            records.append(row)
+    if len(records)==0: return pd.DataFrame({'name':[]})
+    return pd.DataFrame.from_records(records)
+
+def parse_frag_qual(folder):
+    records = []
+    for frag_folder in glob.glob(os.path.join(folder,'*_fragments')):
+        fn = os.path.join(frag_folder,'frag_qual.dat')
+        if not os.path.exists(fn): continue
+        with open(fn) as inf:
+            lines = inf.readlines()
+            index=1
+            y_index=[]
+            y_avg=[]
+            y_bestmer=[]
+            for line in lines:
+                if int(line.split()[1]) == index:
+                    y_index.append(float(line.split()[3]))
+                else:
+                    y_avg.append(np.average(np.array(y_index)))
+                    y_bestmer.append(np.amin(np.array(y_index)))
+                    y_index=[]
+                index=int(line.split()[1])
+            avg_all_frags=np.average(y_avg)
+            avg_best_frags=np.average(y_bestmer)
+        row = OrderedDict()
+        row['name'] = os.path.basename(frag_folder).replace('_fragments','')
+        row['avg_all_frags'] = avg_all_frags
+        row['avg_best_frags'] = avg_best_frags
+        records.append(row)
+    if len(records)==0: return pd.DataFrame({'name':[]})
+    return pd.DataFrame.from_records(records)
+
+def parse_cce(folder):
+    df = pd.DataFrame()
+    for fn in glob.glob(os.path.join(folder,'*.trR_scored.txt')):
+        row = pd.read_csv(fn)
+        df = df.append(row)
+    if df.shape[0]>0:
+        df.columns = ['name','cce10','cce_1d','acc']
+        return df[['name','cce10']]
+    return pd.DataFrame({'name':[]})
+
+def csv2df(fn,**kwargs):
+    if os.path.exists(fn): return pd.read_csv(fn,**kwargs)
+    return pd.DataFrame({'name':[]})
+
+def parse_all_metrics(folder):
+    df = pd.DataFrame({'name':[]})
+
+    print(f'Parsing metrics in {folder}: ',end='')
+    tmp = parse_lddt(os.path.join(folder,'lddt'))
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'lddt ({tmp.shape[0]}), ',end='',flush=True)
+
+    fn = os.path.join(folder,'pymol_metrics.csv')
+    tmp = parse_fastdesign_filters(os.path.join(folder))
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'rosetta metrics from PDB file ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = parse_cce(os.path.join(folder,'trr_score'))
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'cce ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = parse_frag_qual(os.path.join(folder,'frags'))
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'fragment quality ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = csv2df(os.path.join(folder,'tmscores.csv'),index_col=0)
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'TM-scores ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = csv2df(os.path.join(folder,'pymol_metrics.csv'),index_col=0)
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'Pymol metrics ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = csv2df(os.path.join(folder,'pyrosetta_metrics.csv'),index_col=0)
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'Pyrosetta metrics ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = csv2df(os.path.join(folder,'af2_metrics.csv'),index_col=0)
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'AlphaFold2 metrics ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = csv2df(os.path.join(folder,'../complex/interface_metrics.csv'),index_col=0)
+    df = df.merge(tmp,on='name',how='outer')
+    print(f'interface metrics ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = csv2df(os.path.join(folder,'rmsd_trr.csv'))
+    if len(tmp)>0:
+        df = df.merge(tmp[['name','rmsd']].rename(columns={'rmsd':'rmsd_trr'}),on='name',how='outer')
+    print(f'TrR RMSD ({tmp.shape[0]}), ',end='',flush=True)
+
+    tmp = csv2df(os.path.join(folder,'rmsd_trunk.csv'))
+    if len(tmp)>0:
+        df = df.merge(tmp[['name','rmsd']].rename(columns={'rmsd':'rmsd_trunk'}),on='name',how='outer')
+    print(f'Trunk RMSD ({tmp.shape[0]})',flush=True)
+
+    print(f'final dataframe shape: {df.shape}')
+    print(f'final dataframe columns: {df.columns.values}')
+    return df
+
+df = parse_all_metrics(args.folder)
+df.to_csv(args.out)