Made able to map solvent without bonds

Progress bar takes while pre/post tests

pycgtool.py

@@ -28,18 +28,15 @@ def main(args, config):
 
     # Main loop - perform mapping and measurement on every frame in XTC
     numframes = frame.numframes if args.frames == -1 else args.frames
-    for _ in Progress(numframes):
+    for _ in Progress(numframes, postwhile=frame.next_frame):
         if args.map:
-            cgframe = mapping.apply(frame, exclude={"SOL"})
+            cgframe = mapping.apply(frame)
         else:
             cgframe = frame
 
         if args.bnd:
             bonds.apply(cgframe)
 
-        if not frame.next_frame():
-            break
-
     if args.map:
         cgframe.output("out.gro", format=config["output"])
 

pycgtool/bondset.py

@@ -248,6 +248,7 @@ class BondSet:
             return ang if direction > 0 else -ang
 
         for prev_res, res, next_res in sliding(frame):
+            try:
                 mol_meas = self._molecules[res.name]
                 for bond in mol_meas:
                     try:
@@ -264,6 +265,9 @@ class BondSet:
                         # NotImplementedError is raised if form is not implemented
                         # TypeError is raised when residues on end of chain calc bond to next
                         pass
+            except KeyError:
+                # Bonds have not been specified for molecule - probably water
+                pass
 
     def boltzmann_invert(self):
         """

pycgtool/forcefield.py

@@ -29,6 +29,11 @@ class ForceField:
         # Copy main MARTINI itp
         shutil.copyfile(os.path.join(dir_up(os.path.realpath(__file__), 2), "data", "martini_v2.2.itp"),
                         os.path.join(self.dirname, "martini_v2.2.itp"))
+        # Copy water models
+        shutil.copyfile(os.path.join(dir_up(os.path.realpath(__file__), 2), "data", "watermodels.dat"),
+                        os.path.join(self.dirname, "watermodels.dat"))
+        shutil.copyfile(os.path.join(dir_up(os.path.realpath(__file__), 2), "data", "w.itp"),
+                        os.path.join(self.dirname, "w.itp"))
 
         # Create atomtypes.atp required for correct masses with pdb2gmx
         with CFG(os.path.join(dir_up(os.path.realpath(__file__), 2), "data", "martini_v2.2.itp"), allow_duplicate=True) as cfg,\
@@ -52,9 +57,15 @@ class ForceField:
         # TODO print everything to file
         with open(os.path.join(self.dirname, name), "w") as rtp:
             print("[ bondedtypes ]", file=rtp)
-            print(("{:4d}"*4).format(1, 1, 1, 1), file=rtp)
+            print(("{:4d}"*8).format(1, 1, 1, 1, 1, 1, 0, 0), file=rtp)
 
             for mol in mapping:
+                try:
+                    bonds[mol]
+                except KeyError:
+                    # Skip molecules without bonds
+                    continue
+
                 print("[ {0} ]".format(mol), file=rtp)
 
                 print(" [ atoms ]", file=rtp)

pycgtool/util.py

@@ -9,7 +9,7 @@ np.seterr(all="raise")
 
 
 class Progress:
-    def __init__(self, maxits, length=20):
+    def __init__(self, maxits, length=20, prewhile=None, postwhile=None):
         """
         Class to handle printing of a progress bar within loops.
 
@@ -18,6 +18,8 @@ class Progress:
         """
         self._maxits = maxits
         self._length = length
+        self._prewhile = prewhile
+        self._postwhile = postwhile
         self._its = 0
 
     def __enter__(self):
@@ -30,23 +32,25 @@ class Progress:
         return self
 
     def __next__(self):
+        if self._postwhile is not None and self._its > 0 and not self._postwhile():
+            self._stop()
+
+        if self._prewhile is not None and not self._prewhile():
+            self._stop()
+
         self._its += 1
         if self._its % 10 == 0:
             self._display()
 
-        if self._its >= self._maxits:
-            self._display()
-            print()
-            raise StopIteration
+        if self._its > self._maxits:
+            self._stop()
 
         return self._its
 
-    def iteration(self):
-        self._its += 1
-        if self._its % 10 == 0 or self._its == self._maxits:
+    def _stop(self):
         self._display()
-            if self._its == self._maxits:
         print()
+        raise StopIteration
 
     def _display(self):
         done = int(self._length * (self._its / self._maxits))