Add optional constraint conversion to sanitize.py

sanitize.py

@@ -10,24 +10,55 @@ def main(args):
                          "dihedrals":   8,
                          "constraints": 4,
                          "pairs":      -1}
+        has_constr = False
+        constr_lines = []
+
         for line in infile:
             line = line.strip("\n")
             if line.startswith("["):
                 section = line.strip("[ ]")
+                if args.fc is not None and section == "constraints":
+                    print(line, file=outfile)
+                    has_constr = True
+                    for c_line in constr_lines:
+                        print(c_line, file=outfile)
+                    continue
+
+                if section == "system" and not has_constr and constr_lines:
+                    print("[ constraints ]", file=outfile)
+                    for c_line in constr_lines:
+                        print(c_line, file=outfile)
+                    print(file=outfile)
+
+                print(line, file=outfile)
+                continue
+
             elif line.startswith("#") or line.startswith(";") or line == "":
-                pass
+                print(line, file=outfile)
+                continue
+
             elif len(line.split()) < expected_toks.get(section, 0):
                 continue
+
             elif expected_toks.get(section, 0) < 0:
                 continue
 
+            # Convert high force constants to constraints
+            if args.fc is not None and section == "bonds" and float(line.split()[4]) >= args.fc:
+                constr_lines.append(line[:line.rindex(line.split()[4]) - 1])
+                continue
+
+            if args.fc is not None and section == "constraints":
+
             print(line, file=outfile)
 
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Sanitize PDB2GMX output topology.")
     parser.add_argument('-i', type=str, required=True, help="Input topology file")
     parser.add_argument('-o', type=str, required=True, help="Output topology file")
+    parser.add_argument('-fc', type=float, default=None, help="Convert high force constants to constraints")
 
     args = parser.parse_args()