Fix bugs in build_no_seq with entropy score

3dem · Nov 28, 2023 · 67a9b96 · 67a9b96
1 parent 2d13e5a
commit 67a9b96
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/model_angelo/apps/build_no_seq.py b/model_angelo/apps/build_no_seq.py
@@ -222,7 +222,7 @@ def main(parsed_args):
         os.replace(raw_file_src, raw_file_dst)
 
         # Entropy file
-        raw_es_file_src = gnn_output.replace("output.cif", "output_fixed_aa_entropy_score.cif")
+        raw_es_file_src = gnn_output.replace("output.cif", "output_entropy_score.cif")
         raw_es_file_dst = os.path.join(parsed_args.output_dir, f"{name}_entropy_score.cif")
 
         os.replace(raw_es_file_src, raw_es_file_dst)

diff --git a/model_angelo/gnn/flood_fill.py b/model_angelo/gnn/flood_fill.py
@@ -208,6 +208,7 @@ def final_results_to_cif(
         )
         * 100
     )
+    entropy_scores = final_results["entropy_score"][existence_mask] * 100
 
     if refine:
         protein.atomc_positions = all_atoms
@@ -251,6 +252,13 @@ def final_results_to_cif(
             cif_path,
             bfactors=[bfactors[c] for c in chains],
         )
+        chain_atom14_to_cif(
+            [aatype[c] for c in chains],
+            [all_atoms[c] for c in chains],
+            [atom_mask[c] for c in chains],
+            cif_path.replace(".cif", "_entropy_score.cif"),
+            bfactors=[entropy_scores[c] for c in chains],
+        )
 
     chain_aa_logits = [final_results["aa_logits"][existence_mask][c] for c in chains]
     chain_prot_mask = [prot_mask[c] for c in chains]