Add global pLDDT from file

projects/novelfams/translate2modelcif.py

@@ -88,6 +88,14 @@ def _parse_args():
         metavar="<LIST FILE>",
         help="Path to a txt file with models build with AF2, 1 ID per line.",
     )
+    parser.add_argument(
+        "--global-plddt",
+        default=None,
+        type=str,
+        metavar="<LIST FILE>",
+        help="Path to a txt file with global pLDDT values, each line has a "
+        + "single ID followed by a global pLDDT value.",
+    )
     parser.add_argument(
         "--compress",
         default=False,
@@ -210,10 +218,11 @@ class _OST2ModelCIF(modelcif.model.AbInitioModel):
                 occupancy=atm.occupancy,
             )
 
-    def add_scores(self):
+    def add_scores(self, global_plddts):
         """Add QA metrics from AF2 scores."""
         # global scores
-        self.qa_metrics.append(_GlobalPLDDT(np.mean(self.plddts)))
+        if global_plddts is not None:
+            self.qa_metrics.append(_GlobalPLDDT(global_plddts[self.name]))
 
         # local scores
         i = 0
@@ -555,7 +564,7 @@ def _store_as_modelcif(
         scores_json=data_json,
         name=mdl_name,
     )
-    model.add_scores()
+    model.add_scores(data_json["global_plddts"])
     model_group = modelcif.model.ModelGroup([model])
     system.model_groups.append(model_group)
 
@@ -683,6 +692,7 @@ def _translate2modelcif_single(
     opts,
     mdl_details,
     af2_lst,
+    global_plddts,
 ):
     """Convert a single model with its accompanying data to ModelCIF."""
     # ToDo: re-enable Pylint
@@ -696,6 +706,7 @@ def _translate2modelcif_single(
     mdlcf_json["protocol"] = _get_protocol_steps_and_software(
         fam_name, af2_lst
     )
+    mdlcf_json["global_plddts"] = global_plddts
 
     # process coordinates
     target_entities, ost_ent = _get_entities(f_name, fam_name)
@@ -715,7 +726,7 @@ def _translate2modelcif_single(
     )
 
 
-def _translate2modelcif(f_name, af2_lst, opts):
+def _translate2modelcif(f_name, af2_lst, global_plddts, opts):
     """Convert a family of models with their accompanying data to ModelCIF."""
     # ToDo: re-enable Pylint
     # pylint: disable=too-many-locals
@@ -742,6 +753,7 @@ def _translate2modelcif(f_name, af2_lst, opts):
         opts,
         mdl_details,
         af2_lst,
+        global_plddts,
     )
 
 
@@ -763,6 +775,21 @@ def _read_af2_model_list(path):
     return af2_lst
 
 
+def _read_global_plddt_list(path):
+    """Read global pLDDT values from file."""
+    plddt = {}
+    if path is None:
+        return plddt
+
+    with open(path, encoding="ascii") as pfh:
+        for line in pfh:
+            line = line.strip()
+            line = line.split()
+            plddt[line[0]] = float(line[1])
+
+    return plddt
+
+
 def _main():
     """Run as script."""
     s_tmstmp = timer()
@@ -775,6 +802,9 @@ def _main():
     # read list of AF2 models
     af2_mdls = _read_af2_model_list(opts.af2_models)
 
+    # read global pLDDT values, if available
+    global_plddts = _read_global_plddt_list(opts.global_plddt)
+
     # iterate over models
     print(f"Processing {n_mdls} models.")
     tmstmp = s_tmstmp
@@ -784,6 +814,7 @@ def _main():
             _translate2modelcif(
                 f_name,
                 af2_mdls,
+                global_plddts,
                 opts,
             )
         except (_InvalidCoordinateError, _NoEntitiesError):