clean_dof_dfr.py

#!/usr/bin/env python3
"""
Receives a complete .dfr (generated by either fragGen or gromacs2dice)
and based on the 'atoms fragments' section, remove unecessary degrees
of freedom from the input.

Author: Henrique Musseli Cezar
Date: SEP/2016
"""

import os
import sys
import argparse
import io

def clean_dofs(fname):
  with open(fname,'r') as f:
    # read lines until it gets to 'atoms fragments'
    line = f.readline()
    while line.strip() != '$atoms fragments':
      line = f.readline()

    # read fragments and store the rigid ones
    frags = {}
    rigidFrags = []
    line = f.readline()
    fullline = True
    while '$end' not in line.strip():
      if not line.strip():
        line = f.readline()
        continue
      if fullline:
        fragnum = int(line.split()[0])
      if "[" in line and "]" in line:
        frags[fragnum] = [int(x) for x in line.split()[2:-2]]
        frags[fragnum].append(line.split()[-1])
      elif "[" in line and "]" not in line:
        fullline = False
        frags[fragnum] = [int(x) for x in line.split()[2:]]
      elif "[" not in line and "]" not in line:
        if fullline:
          print("You're supposed to use [] to delimit your fragments.")
          sys.exit(0)
        for frag in [int(x) for x in line.split()]:
          frags[fragnum].append(frag)
      else:
        fullline = True
        for frag in [int(x) for x in line.split()[:-2]]:
          frags[fragnum].append(frag)
        frags[fragnum].append(line.split()[-1])

      if fullline:
        if frags[fragnum][-1].upper() == 'R': 
          rigidFrags.append(fragnum)
      line = f.readline()
    nfrags = len(frags.keys())

    # get the fragment connections
    fConnLines = []
    line = f.readline()
    while line.strip() != '$fragment connection':
      line = f.readline()
    line = f.readline()
    while '$end' not in line.strip():
      if not line.strip():
        line = f.readline()
        continue
      fConnLines.append(line)
      line = f.readline()

    # start removing the dofs
    # just append all bonds, since they are used to determine the fnb in DICE
    bonds = []
    line = f.readline()
    while line.strip() != '$bond':
      line = f.readline()
    line = f.readline()
    while '$end' not in line.strip():
      if not line.strip():
        line = f.readline()
        continue    
      bonds.append(line)
      line = f.readline()

    # remove every unused angle
    angles = []
    line = f.readline()
    while line.strip() != '$angle':
      line = f.readline()
    line = f.readline()
    while '$end' not in line.strip():
      if not line.strip():
        line = f.readline()
        continue
      save = True
      a1, a2, a3 = [int(x) for x in line.split()[:3]]
      for rig in rigidFrags:
        if (a1 in frags[rig]) and (a2 in frags[rig]) and (a3 in frags[rig]):
          save = False
      if save:
        angles.append(line)
      line = f.readline()

    # remove every unused dihedral
    dihedrals = []
    line = f.readline()
    while line.strip() != '$dihedral':
      line = f.readline()
    line = f.readline()
    while '$end' not in line.strip():
      if not line.strip():
        line = f.readline()
        continue
      save = True
      a1, a2, a3, a4 = [int(x) for x in line.split()[:4]]
      for rig in rigidFrags:
        if (a1 in frags[rig]) and (a2 in frags[rig]) and (a3 in frags[rig]) and (a4 in frags[rig]):
          save = False
      if save:
        dihedrals.append(line)
      line = f.readline()

    # remove every unused improper dihedral
    imp_dihedrals = []
    line = f.readline()
    impd = True
    while line.strip() != '$improper dihedral':
      line = f.readline()
      if line == '':
        impd = False
        break
    if impd:
      line = f.readline()
      while '$end' not in line.strip():
        if not line.strip():
          line = f.readline()
          continue
        save = True
        a1, a2, a3, a4 = [int(x) for x in line.split()[:4]]
        for rig in rigidFrags:
          if (a1 in frags[rig]) and (a2 in frags[rig]) and (a3 in frags[rig]) and (a4 in frags[rig]):
            save = False
        if save:
          imp_dihedrals.append(line)
        line = f.readline()

    # print the simplified dfr to string
    output = io.StringIO()
    print('$atoms fragments', file=output)
    for frag in frags.keys():
      string = ""
      string += str(frag)+"\t[ "
      for i, el in enumerate(frags[frag]):
        if i == len(frags[frag])-1:
          string += "] "+str(el)
        else:
          string += str(el)+"\t"
      print(string, file=output)
    print('$end atoms fragments\n', file=output)

    print('$fragment connection', file=output)
    for line in fConnLines:
      print("%s" % line.strip(), file=output)
    print('$end fragment connection\n', file=output)

    print('$bond', file=output)
    for line in bonds:
      print("%s" % line.strip(), file=output)
    print('$end bond\n', file=output)

    if len(angles) > 0:
      print('$angle', file=output)
      for line in angles:
        print("%s" % line.strip(), file=output)
      print('$end angle\n', file=output)

    print('$dihedral', file=output)
    for line in dihedrals:
      print("%s" % line.strip(), file=output)
    print('$end dihedral', file=output)

    if (len(imp_dihedrals) > 0) and impd:
      print('\n$improper dihedral', file=output)
      for line in imp_dihedrals:
        print("%s" % line.strip(), file=output)
      print('$end improper dihedral', file=output)

    contents = output.getvalue()
    output.close()
    return contents


if __name__ == '__main__':
  parser = argparse.ArgumentParser(description="Receives a complete .dfr (generated by either fragGen or gromacs2dice) and based on the 'atoms fragments' section, remove unecessary degrees of freedom from the input.")
  parser.add_argument("filename", help="the dfr containing all the degrees of freedom information")
  args = parser.parse_args()

  filename = os.path.realpath(args.filename)
  base, ext = os.path.splitext(filename)

  print(clean_dofs(filename))