Added support for ARBD group sites

mrdna/arbdmodel/__init__.py

@@ -730,6 +730,22 @@ class PdbModel(Transformable, Parent):
 
 
 class ArbdModel(PdbModel):
+
+    class _GroupSite():
+        """ Class to represent a collection of particles that can be used by bond potentials """
+        def __init__(self, particles, weights=None):
+            if weights is not None:
+                raise NotImplementedError
+            self.particles = particles
+            self.idx = None
+
+        def get_center(self):
+            c = np.array((0,0,0))
+            for p in self.particles:
+                c = c + p.collapsedPosition()
+            c = c / len(self.particles)
+            return c
+
     def __init__(self, children, origin=None, dimensions=(1000,1000,1000), temperature=291, timestep=50e-6,
                  particle_integrator = 'Brown',
                  cutoff=50, decomp_period=1000, pairlist_distance=None, nonbonded_resolution=0.1,
@@ -764,6 +780,13 @@ class ArbdModel(PdbModel):
 
         self.cacheUpToDate = False
 
+        self.group_sites = []
+
+    def add_group_site(self, particles, weights=None):
+        g = ArbdModel._GroupSite(particles, weights)
+        self.group_sites.append(g)
+        return g
+
     def clear_all(self, keep_children=False):
         Parent.clear_all(self, keep_children=keep_children)
         self.particles = []
@@ -771,6 +794,7 @@ class ArbdModel(PdbModel):
         self.type_counts = None
         self._nbParamFiles = []
         self._written_bond_files = dict()
+        self.group_sites = []
 
     def _getNbScheme(self, typeA, typeB):
         scheme = None
@@ -804,8 +828,12 @@ class ArbdModel(PdbModel):
         # self.particles = sorted(particles, key=lambda p: (p.type_, p.idx))
         
         ## Update particle indices
-        for p,i in zip(self.particles,range(len(self.particles))):
+        for i,p in enumerate(self.particles):
             p.idx = i
+
+        ## TODO recurse through childrens' group_sites
+        for i,g in enumerate(self.group_sites):
+            g.idx = len(self.particles)+i
             
         # self.initialCoords = np.array([p.initialPosition for p in self.particles])
 
@@ -916,6 +944,7 @@ class ArbdModel(PdbModel):
         self._exclusion_filename = "%s/%s.exculsions.txt" % (d, prefix)
         self._bond_angle_filename = "%s/%s.bondangles.txt" % (d, prefix)
         self._product_potential_filename = "%s/%s.prodpot.txt" % (d, prefix)
+        self._group_sites_filename = "%s/%s.groups.txt" % (d, prefix)
         
         # self._writeArbdCoordFile( prefix + ".coord.txt" )
         self._writeArbdParticleFile( prefix + ".particles.txt" )
@@ -926,6 +955,7 @@ class ArbdModel(PdbModel):
         self._writeArbdExclusionFile()
         self._writeArbdBondAngleFile()
         self._writeArbdProductPotentialFile()
+        self._writeArbdGroupSitesFile()
         self._writeArbdPotentialFiles( prefix, directory = d )
         self._writeArbdConf( prefix, numSteps=numSteps, outputPeriod=outputPeriod, restart_file=restart_file )
         
@@ -1123,6 +1153,9 @@ tabulatedPotential  1
             exclusions = self.get_exclusions()
             bond_angles = self.get_bond_angles()
             prod_pots = self.get_product_potentials()
+            # group_sites = self.get_group_sites()
+            group_sites = self.group_sites
+
             if len(bonds) > 0:
                 for b in self._get_bond_potentials():
                     fh.write("tabulatedBondFile %s\n" % b)
@@ -1149,6 +1182,8 @@ tabulatedPotential  1
                 fh.write("inputBondAngles %s\n" % self._bond_angle_filename)
             if len(prod_pots) > 0:
                 fh.write("inputProductPotentials %s\n" % self._product_potential_filename)
+            if len(group_sites) > 0:
+                fh.write("inputGroups %s\n" % self._group_sites_filename)
      
         write_null_dx = False
         for pt,num in self.getParticleTypesAndCounts():
@@ -1290,6 +1325,13 @@ component "data" value 3
                         line = line+" ".join([str(x) for x in [type_,b] if x != ""])+" "
                     fh.write(line)
 
+    def _writeArbdGroupSitesFile( self ):
+        if len(self.group_sites) > 0:
+            with open(self._group_sites_filename,'w') as fh:
+                for i,g in enumerate(self.group_sites):
+                    assert( i+len(self.particles) == g.idx )
+                    ids = " ".join([str(int(p.idx)) for p in g.particles])
+                    fh.write("GROUP %s\n" % ids)
 
     def dimensions_from_structure( self, padding_factor=1.5, isotropic=False ):
         ## TODO: cache coordinates using numpy arrays for quick min/max