Merge pull request #263 from isblab/develop

Additions to EM and crosslinking restraints. Minor modifications in ReplicaExchange macro

pyext/src/macros.py

@@ -26,6 +26,9 @@
 import warnings
 import math
 
+import pickle
+import RMF
+
 
 class _MockMPIValues(object):
     """Replace samplers.MPI_values when in test mode"""
@@ -109,7 +112,10 @@ def __init__(self, model, root_hier,
                  atomistic=False,
                  replica_exchange_object=None,
                  test_mode=False,
-                 score_moved=False):
+                 score_moved=False,
+                 use_nestor=False,
+                 nestor_restraints=None,
+                 nestor_rmf_fname_prefix="nested",):
         """Constructor.
            @param model The IMP model
            @param root_hier Top-level (System)hierarchy
@@ -174,6 +180,13 @@ def __init__(self, model, root_hier,
            @param score_moved   If True, attempt to speed up Monte Carlo
                   sampling by caching scoring function terms on particles
                   that didn't move.
+           @param use_nestor   If True, follows the Nested Sampling workflow
+                  of the NestOR module and skips writing stat files and
+                  replica stat files.
+           @param nestor_restraints   A list of restraints for which
+                  likelihoods are to be computed for use by NestOR module.
+           @param nestor_rmf_fname_prefix   Prefix to be used for storing .rmf3
+                  files generated by NestOR .
         """
         self.model = model
         self.vars = {}
@@ -258,6 +271,9 @@ def __init__(self, model, root_hier,
         self.vars["geometries"] = None
         self.test_mode = test_mode
         self.score_moved = score_moved
+        self.nest = use_nestor
+        self.nestor_restraints = nestor_restraints
+        self.nestor_rmf_fname = nestor_rmf_fname_prefix
 
     def add_geometries(self, geometries):
         if self.vars["geometries"] is None:
@@ -276,6 +292,7 @@ def show_info(self):
         keys.sort()
         for v in keys:
             print("------", v.ljust(30), self.vars[v])
+        print("Use nestor: ", self.nest)
 
     def get_replica_exchange_object(self):
         return self.replica_exchange_object
@@ -375,7 +392,7 @@ def execute_macro(self):
         rmf_dir = globaldir + self.vars["rmf_dir"]
         pdb_dir = globaldir + self.vars["best_pdb_dir"]
 
-        if not self.test_mode:
+        if not self.test_mode and not self.nest:
             if self.vars["do_clean_first"]:
                 pass
 
@@ -410,33 +427,37 @@ def execute_macro(self):
         if self.rmf_output_objects is not None:
             self.rmf_output_objects.append(sw)
 
-        print("Setting up stat file")
-        output = IMP.pmi.output.Output(atomistic=self.vars["atomistic"])
-        low_temp_stat_file = globaldir + \
-            self.vars["stat_file_name_suffix"] + "." + str(myindex) + ".out"
+        if not self.nest:
+            print("Setting up stat file")
+            output = IMP.pmi.output.Output(atomistic=self.vars["atomistic"])
+            low_temp_stat_file = globaldir + \
+                self.vars["stat_file_name_suffix"] + "." + \
+                str(myindex) + ".out"
 
         # Ensure model is updated before saving init files
         if not self.test_mode:
             self.model.update()
 
-        if not self.test_mode:
+        if not self.test_mode and not self.nest:
             if self.output_objects is not None:
                 output.init_stat2(low_temp_stat_file,
                                   self.output_objects,
                                   extralabels=["rmf_file", "rmf_frame_index"])
         else:
             print("Stat file writing is disabled")
 
-        if self.rmf_output_objects is not None:
+        if self.rmf_output_objects is not None and not self.nest:
             print("Stat info being written in the rmf file")
 
-        print("Setting up replica stat file")
-        replica_stat_file = globaldir + \
-            self.vars["replica_stat_file_suffix"] + "." + str(myindex) + ".out"
-        if not self.test_mode:
-            output.init_stat2(replica_stat_file, [rex], extralabels=["score"])
+        if not self.test_mode and not self.nest:
+            print("Setting up replica stat file")
+            replica_stat_file = globaldir + \
+                self.vars["replica_stat_file_suffix"] + "." + \
+                str(myindex) + ".out"
+            if not self.test_mode:
+                output.init_stat2(replica_stat_file, [rex],
+                                extralabels=["score"])
 
-        if not self.test_mode:
             print("Setting up best pdb files")
             if not self.is_multi_state:
                 if self.vars["number_of_best_scoring_models"] > 0:
@@ -478,7 +499,7 @@ def execute_macro(self):
         else:
             output_hierarchies = [self.root_hier]
 
-        if not self.test_mode:
+        if not self.test_mode and not self.nest:
             print("Setting up and writing initial rmf coordinate file")
             init_suffix = globaldir + self.vars["initial_rmf_name_suffix"]
             output.init_rmf(init_suffix + "." + str(myindex) + ".rmf3",
@@ -498,7 +519,7 @@ def execute_macro(self):
 
         self._add_provenance(sampler_md, sampler_mc)
 
-        if not self.test_mode:
+        if not self.test_mode and not self.nest:
             print("Setting up production rmf files")
             rmfname = rmf_dir + "/" + str(myindex) + ".rmf3"
             output.init_rmf(rmfname, output_hierarchies,
@@ -510,12 +531,14 @@ def execute_macro(self):
 
         ntimes_at_low_temp = 0
 
-        if myindex == 0:
+        if myindex == 0 and not self.nest:
             self.show_info()
         self.replica_exchange_object.set_was_used(True)
         nframes = self.vars["number_of_frames"]
         if self.test_mode:
             nframes = 1
+
+        sampled_likelihoods = []
         for i in range(nframes):
             if self.test_mode:
                 score = 0.
@@ -529,7 +552,8 @@ def execute_macro(self):
                 score = IMP.pmi.tools.get_restraint_set(
                     self.model).evaluate(False)
                 mpivs.set_value("score", score)
-            output.set_output_entry("score", score)
+            if not self.nest:
+                output.set_output_entry("score", score)
 
             my_temp_index = int(rex.get_my_temp() * temp_index_factor)
 
@@ -552,7 +576,16 @@ def execute_macro(self):
             if save_frame:
                 print("--- frame %s score %s " % (str(i), str(score)))
 
-                if not self.test_mode:
+                if self.nest:
+                    if math.isnan(score):
+                        sampled_likelihoods.append(math.nan)
+                    else:
+                        likelihood_for_sample = 1
+                        for rstrnt in self.nestor_restraints:
+                            likelihood_for_sample *= rstrnt.get_likelihood()
+                        sampled_likelihoods.append(likelihood_for_sample)
+
+                if not self.test_mode and not self.nest:
                     if i % self.vars["nframes_write_coordinates"] == 0:
                         print('--- writing coordinates')
                         if self.vars["number_of_best_scoring_models"] > 0:
@@ -568,14 +601,32 @@ def execute_macro(self):
                         output.write_stat2(low_temp_stat_file)
                 ntimes_at_low_temp += 1
 
-            if not self.test_mode:
+            if not self.test_mode and not self.nest:
                 output.write_stat2(replica_stat_file)
             if self.vars["replica_exchange_swap"]:
                 rex.swap_temp(i, score)
+
+        if self.nest and len(sampled_likelihoods) > 0:
+            with open(
+                "likelihoods_" + \
+                str(self.replica_exchange_object.get_my_index()),"wb") as lif:
+                pickle.dump(sampled_likelihoods, lif)
+
+            nestor_rmf_fname = str(self.nestor_rmf_fname) + '_' + \
+                str(self.replica_exchange_object.get_my_index()) +'.rmf3'
+
+            self.rmf_h = RMF.create_rmf_file(nestor_rmf_fname)
+            IMP.rmf.add_hierarchy(self.rmf_h, self.root_hier)
+            try:
+                IMP.rmf.save_frame(self.rmf_h)
+            except Exception as e:
+                print(e)
+                pass
+
         for p, state in IMP.pmi.tools._all_protocol_outputs(self.root_hier):
             p.add_replica_exchange(state, self)
 
-        if not self.test_mode:
+        if not self.test_mode and not self.nest:
             print("closing production rmf files")
             output.close_rmf(rmfname)
 

pyext/src/restraints/crosslinking.py

@@ -455,6 +455,14 @@ def get_output(self):
 
         return output
 
+    def get_likelihood(self):
+        """Get the unweighted likelihood of the restraint"""
+        likelihood = 1
+        for restraint in self.xl_restraints:
+            likelihood *= restraint.get_probability()
+
+        return likelihood
+
     def get_movers(self):
         """ Get all need data to construct a mover in IMP.pmi.dof class"""
         movers = []

pyext/src/restraints/em.py

@@ -358,6 +358,11 @@ def get_output(self):
                self.label] = str(self.sigmaglobal.get_scale())
         return output
 
+    def get_likelihood(self):
+        """Get the unweighted likelihood of the restraint"""
+        likelihood = self.gaussianEM_restraint.get_probability()
+        return likelihood
+
     def evaluate(self):
         return self.weight * self.rs.unprotected_evaluate(None)