Add local pLDDT from b-factors

projects/novelfams/translate2modelcif.py

@@ -165,7 +165,6 @@ def _get_ch_name(ch, use_auth=False):
     return ch.name
 
 
-# ToDo: are those Ab Initio Models?
 class _OST2ModelCIF(modelcif.model.AbInitioModel):
     """Map OST entity elements to ihm.model"""
 
@@ -210,6 +209,23 @@ class _OST2ModelCIF(modelcif.model.AbInitioModel):
                 occupancy=atm.occupancy,
             )
 
+    def add_scores(self):
+        """Add QA metrics from AF2 scores."""
+        # local scores
+        i = 0
+        for chn_i in self.ost_entity.chains:
+            ch_name = _get_ch_name(chn_i, self.use_auth)
+            for res_i in chn_i.residues:
+                # local pLDDT
+                res_num = _get_res_num(res_i, self.use_auth)
+                self.qa_metrics.append(
+                    _LocalPLDDT(
+                        self.asym[ch_name].residue(res_num),
+                        self.plddts[i],
+                    )
+                )
+                i += 1
+
 
 def _get_metadata(metadata_file):
     """Read csv file with metedata and prepare for next steps."""
@@ -535,7 +551,7 @@ def _store_as_modelcif(
         scores_json=data_json,
         name=mdl_name,
     )
-
+    model.add_scores()
     model_group = modelcif.model.ModelGroup([model])
     system.model_groups.append(model_group)
 
@@ -700,14 +716,14 @@ def _translate2modelcif(f_name, af2_lst, opts):
     # ToDo: re-enable Pylint
     # pylint: disable=too-many-locals
 
-    # skip if done already
-    if opts.compress:
-        cifext = "cif.gz"
-    else:
-        cifext = "cif"
-    mdl_path = os.path.join(opts.out_dir, f"{f_name}.{cifext}")
     # ToDo: re-enable for final version
-    #if os.path.exists(mdl_path):
+    # skip if done already
+    # if opts.compress:
+    #     cifext = "cif.gz"
+    # else:
+    #    cifext = "cif"
+    # mdl_path = os.path.join(opts.out_dir, f"{f_name}.{cifext}")
+    # if os.path.exists(mdl_path):
     #    print(f"  {f_name} already done")
     #    return