diff --git a/sanitize.py b/sanitize.py index 192e603754bad4fe54a4bba0310b387022c99441..c0e33c844fe7ef56c284760a96cc30fbe7f245df 100755 --- a/sanitize.py +++ b/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()