Skip to content
Snippets Groups Projects
Commit dbb0db78 authored by Bienchen's avatar Bienchen
Browse files

Add model, produce coordinates

parent 2a2b3b50
No related branches found
No related tags found
No related merge requests found
......@@ -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}.")
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment