Added support for harmonic restraints to arbdmodel.py

mrdna/model/arbdmodel.py

@@ -127,6 +127,12 @@ class Parent():
         # for b in (i,j): assert(b in beads)
         self.exclusions.append( (i,j) )
 
+    def get_restraints(self):
+        ret = []
+        for c in self.children:
+            ret.extend( c.get_restraints() )
+        return ret
+
     def get_bonds(self):
         ret = self.bonds
         for c in self.children:
@@ -361,10 +367,18 @@ class PointParticle(Transformable, Child):
         self.idx     = None
         self.name = name
         self.counter = 0
+        self.restraints = []
 
         for key,val in kwargs.items():
             self.__dict__[key] = val
         
+    def add_restraint(self, restraint):
+        ## TODO: how to handle duplicating and cloning bonds
+        self.restraints.append( restraint )
+
+    def get_restraints(self):
+        return [(self,r) for r in self.restraints]
+
 
     def __getattr__(self, name):
         """
@@ -760,6 +774,7 @@ class ArbdModel(PdbModel):
         d = self.potential_directory = "potentials"
         if not os.path.exists(d):
             os.makedirs(d)
+        self._restraint_filename = "%s/%s.restraint.txt" % (d, prefix)
         self._bond_filename = "%s/%s.bonds.txt" % (d, prefix)
         self._angle_filename = "%s/%s.angles.txt" % (d, prefix)
         self._dihedral_filename = "%s/%s.dihedrals.txt" % (d, prefix)
@@ -767,6 +782,7 @@ class ArbdModel(PdbModel):
         
         # self._writeArbdCoordFile( prefix + ".coord.txt" )
         self._writeArbdParticleFile( prefix + ".particles.txt" )
+        self._writeArbdRestraintFile()
         self._writeArbdBondFile()
         self._writeArbdAngleFile()
         self._writeArbdDihedralFile()
@@ -877,6 +893,7 @@ tabulatedPotential  1
                 fh.write("tabulatedFile %d@%d@%s\n" % (i,j,f))
 
             ## Bonded interactions
+            restraints = self.get_restraints()
             bonds = self.get_bonds()
             angles = self.get_angles()
             dihedrals = self.get_dihedrals()
@@ -894,10 +911,16 @@ tabulatedPotential  1
                 for b in list(set([b for i,j,k,l,b in dihedrals])):
                     fh.write("tabulatedDihedralFile %s\n" % b)
 
-            fh.write("inputBonds %s\n" % self._bond_filename)
-            fh.write("inputAngles %s\n" % self._angle_filename)
-            fh.write("inputDihedrals %s\n" % self._dihedral_filename)
-            fh.write("inputExcludes %s\n" % self._exclusion_filename)
+            if len(restraints) > 0:
+                fh.write("inputRestraints %s\n" % self._restraint_filename)
+            if len(bonds) > 0:
+                fh.write("inputBonds %s\n" % self._bond_filename)
+            if len(angles) > 0:
+                fh.write("inputAngles %s\n" % self._angle_filename)
+            if len(dihedrals) > 0:
+                fh.write("inputDihedrals %s\n" % self._dihedral_filename)
+            if len(exclusions) > 0:
+                fh.write("inputExcludes %s\n" % self._exclusion_filename)
      
         write_null_dx = False
         for pt,num in self.getParticleTypesAndCounts():
@@ -958,6 +981,20 @@ component "data" value 3
     def _getNonbondedPotential(self,x,a,b):
         return a*(np.exp(-x/b))    
 
+    def _writeArbdRestraintFile( self ):
+        with open(self._restraint_filename,'w') as fh:
+            for i,restraint in self.get_restraints():
+                item = [i.idx]
+                if len(restraint) == 1:
+                    item.append(restraint[0])
+                    item.extend(i.get_collapsed_position())
+                elif len(restraint) == 2:
+                    item.append(restraint[0])
+                    item.extend(restraint[1])
+                elif len(restraint) == 5:
+                    item.extend(restraint)
+                fh.write("RESTRAINT %d %f %f %f %f\n" % tuple(item))
+
     def _writeArbdBondFile( self ):
         for b in list( set( [b for i,j,b,ex in self.get_bonds()] ) ):
             if type(b) is not str: