From 6b7167aac4aff58455abfe4edad0631724dd62e1 Mon Sep 17 00:00:00 2001
From: Ben Webb <ben@salilab.org>
Date: Sun, 20 Oct 2019 23:57:40 -0700
Subject: [PATCH] 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.
---
pyext/src/restraint.py | 10 ++++++++++
1 file changed, 10 insertions(+)
diff --git a/pyext/src/restraint.py b/pyext/src/restraint.py
index d44b46e..6c8a201 100755
--- a/pyext/src/restraint.py
+++ b/pyext/src/restraint.py
@@ -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
--
GitLab