diff --git a/projects/novelfams/translate2modelcif.py b/projects/novelfams/translate2modelcif.py index 0d5ed81bf800c2fecf80b43ea49211d546458520..9188daca8267d9b96ce9eea92edcef58506bd919 100644 --- a/projects/novelfams/translate2modelcif.py +++ b/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):