diff --git a/translate2modelcif.py b/translate2modelcif.py index 866dd20f3c50906c6fa76aa5830570d0e2850215..a5572989bfa51aded4582e46a055180d351d9540 100644 --- a/translate2modelcif.py +++ b/translate2modelcif.py @@ -10,6 +10,7 @@ import requests import ihm import modelcif import modelcif.dumper +import modelcif.model import modelcif.reference from ost import io @@ -41,6 +42,34 @@ def _parse_args(): return opts +class _OST2ModelCIF(modelcif.model.AbInitioModel): + """Map OST entity elements to ihm.model""" + + def __init__(self, *args, **kwargs): + """Initialise a model""" + # check that ost_entity exists, TypeError otherwise + self.ost_entity = kwargs.pop("ost_entity") + self.asym = kwargs.pop("asym") + + super().__init__(*args, **kwargs) + + def get_atoms(self): + # ToDo [internal]: Take B-factor out since its not a B-factor? + for atm in self.ost_entity.atoms: + yield modelcif.model.Atom( + asym_unit=self.asym[atm.chain.name], + seq_id=atm.residue.number.num, + atom_id=atm.name, + type_symbol=atm.element, + x=atm.pos[0], + y=atm.pos[1], + z=atm.pos[2], + het=atm.is_hetatom, + biso=atm.b_factor, + occupancy=atm.occupancy, + ) + + def _abort_msg(msg, exit_code=1): """Write error message and exit with exit_code.""" print(f"{msg}\nAborting.", file=sys.stderr) @@ -227,10 +256,10 @@ def _get_entities(pdb_file, up_acs): cif_ent.update(upkb) entities.append(cif_ent) - return entities + return entities, ost_ent -def _store_as_modelcif(interaction_name, data_json, file_prfx): +def _store_as_modelcif(interaction_name, data_json, ost_ent, file_prfx): """Mix all the data into a ModelCIF file.""" # create system to gather all the data system = modelcif.System( @@ -280,11 +309,24 @@ def _store_as_modelcif(interaction_name, data_json, file_prfx): assembly = modelcif.Assembly( asym_units.values(), name="ma_struct_assembly_details.assembly_name" ) - system.assemblies.append(assembly) # audit_authors system.authors.extend(data_json["audit_authors"]) + # set up the model to produce coordinates + # ToDo [input]: Get ma_model_list.model_name + model = _OST2ModelCIF( + assembly=assembly, + asym=asym_units, + ost_entity=ost_ent, + name="ma_model_list.model_name", + ) + # ToDo [input]: Get name + model_group = modelcif.model.ModelGroup( + [model], name="ma_model_list.model_group_name" + ) + system.model_groups.append(model_group) + # write modelcif System to file with open(f"{file_prfx}.cif", "w", encoding="utf8") as mmcif_fh: modelcif.dumper.write(mmcif_fh, [system]) @@ -305,8 +347,8 @@ def _create_interaction_json(): def _create_model_json(data, pdb_file, up_acs): """Create a dictionary (mimicking JSON) that contains all the data.""" - data["target_entities"] = _get_entities(pdb_file, up_acs) - return data + data["target_entities"], ost_ent = _get_entities(pdb_file, up_acs) + return ost_ent def _main(): @@ -332,9 +374,9 @@ def _main(): fle = os.path.join(opts.model_dir, fle) # gather data into JSON-like structure - _create_model_json(mdlcf_json, fle, up_acs) + ost_ent = _create_model_json(mdlcf_json, fle, up_acs) - _store_as_modelcif(uid, mdlcf_json, file_prfx) + _store_as_modelcif(uid, mdlcf_json, ost_ent, file_prfx) # ToDo [internal]: wipe data or is it overwritten in mdlcf_json? print(f"... done with {opts.model_dir}.")