add index correction

Author: chrischen1

EnQA.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from data.loader import expand_sh
+from data.process_label import parse_pdbfile
 from feature import create_basic_features, get_base2d_feature
 from data.process_alphafold import process_alphafold_target_ensemble, process_alphafold_model
 from network.resEGNN import resEGNN, resEGNN_with_mask, resEGNN_with_ne
@@ -22,7 +23,7 @@
                         help='Path to alphafold prediction results.')
     parser.add_argument('--alphafold_feature_cache', type=str, required=False, default='')
     parser.add_argument('--af2_pdb', type=str, required=False, default='',
-                        help='Optional. PDBs from AlphaFold2 predcition for index correction with input pdb')
+                        help='Optional. PDBs from AlphaFold2 predcition for index correction with input pdb. Must contain all residues in input pdb.')
 
     args = parser.parse_args()
     if args.alphafold_feature_cache == '':
@@ -65,6 +66,17 @@
             if args.alphafold_feature_cache is not None:
                 pickle.dump({'plddt': plddt, 'cmap': cmap, 'dict_2d': dict_2d},
                             open(args.alphafold_prediction_cache, 'wb'))
+        if args.af2_pdb != '':
+            pose_input = parse_pdbfile(args.input)
+            input_idx = np.array([i['rindex'] for i in pose_input])
+            pose_af2 = parse_pdbfile(args.af2_pdb)
+            af2_idx = np.array([i['rindex'] for i in pose_af2])
+            mask = af2_idx in input_idx
+            af2_qa = af2_qa[:, mask]
+            plddt = plddt[:, mask]
+            cmap = cmap[:, mask][mask, :]
+            for f2d_type in dict_2d.keys():
+                dict_2d[f2d_type] = dict_2d[f2d_type][:, :, mask][:, mask, :]
     else:
         dict_2d['f2d_dan'] = get_base2d_feature(args.input, args.output)
     with torch.no_grad():

README.md

@@ -65,3 +65,15 @@ Also, there are [binaries](https://gitlab.inria.fr/GruLab/s-gcn/-/tree/master/#s
 ## Generating AlphaFold2 models for assisted quality assessment
 
 For generating models using AlphaFold2, an installation of AlphaFold2 following its [Official Repo](https://github.com/deepmind/alphafold) is required. For our experiments, we use its original model used at CASP14 with no ensembling (--model_preset=monomer), with all genetic databases used at CASP14 (--db_preset=full_dbs), and restricts templates only to structures that were available at the start of CASP14 (--max_template_date=2020-05-14).
+
+## Quick configuration
+
+```
+pip install biopandas biopython scipy torch==1.10.1
+pip install dgl-cu113 -f https://data.dgl.ai/wheels/repo.html
+pip install packaging wandb line_profiler
+pip install git+https://github.com/AMLab-Amsterdam/lie_learn
+git clone https://github.com/FabianFuchsML/se3-transformer-public.git
+cd se3-transformer-public
+pip install -e .
+```