Skip to content
Snippets Groups Projects
Commit 6b7167aa authored by Ben Webb's avatar Ben Webb
Browse files

Call ProtocolOutput hook & add dataset

If the PMI System object has ProtocolObject(s) attached,
notify them whenever we create a new EM restraint, and
provide a dataset pointing to the GMM file we used.
parent 7645c95d
No related branches found
No related tags found
No related merge requests found
......@@ -11,6 +11,7 @@ import IMP.container
import IMP.bayesianem
import IMP.isd
import IMP.pmi.tools
import IMP.pmi.mmcif
import IMP.isd.gmm_tools
import sys
from math import sqrt
......@@ -96,6 +97,7 @@ class GaussianEMRestraintWrapper(object):
print('will scale target mass by', target_mass_scale)
if target_fn != '':
self._set_dataset(target_fn)
self.target_ps = []
IMP.isd.gmm_tools.decorate_gmm_from_text(
target_fn,
......@@ -117,6 +119,9 @@ class GaussianEMRestraintWrapper(object):
ms=IMP.atom.Mass(p).get_mass()
IMP.atom.Mass(p).set_mass(ms*scale)
for p, state in IMP.pmi.tools._all_protocol_outputs([], densities[0]):
p.add_em3d_restraint(state, self.target_ps, self.densities, self)
# setup model GMM
self.model_ps = []
for h in self.densities:
......@@ -182,6 +187,11 @@ class GaussianEMRestraintWrapper(object):
self.rs.add_restraint(self.gaussianEM_restraint)
self.set_weight(weight)
def _set_dataset(self, target_fn):
"""Set the dataset to point to the input file"""
p = IMP.pmi.mmcif.GMMParser()
self.dataset = p.parse_file(target_fn)['dataset']
def center_target_density_on_model(self):
'''
aligns the center of mass of the target GMM on the center of mass of the model
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment