reorder_ligpargen.py

#!/usr/bin/env python3
"""
Receives original .pdb submitted to LigParGen and the output of LigParGen (.gro and .itp)
to reorder the output to be in the same order of the original input

Author: Henrique Musseli Cezar
Date: MAY/2019
"""

import os
import argparse
try: 
  import pybel
  import openbabel
  ob3 = False
except:
  from openbabel import pybel
  from openbabel import openbabel
  ob3 = True
import rmsd
import numpy as np
import re

def get_mol_info(mol):
  # table to convert atomic number to symbols
  if not ob3:
    etab = openbabel.OBElementTable()

  q_atoms = []
  q_all = []
  for atom in mol:
    if ob3:
      q_atoms.append(openbabel.GetSymbol(atom.atomicnum))
    else:
      q_atoms.append(etab.GetSymbol(atom.atomicnum))
    q_all.append(atom.coords)

  return np.asarray(q_atoms), np.asarray(q_all)  

def get_atom_correspondence(pdb, gro):
  # read pdb
  pdbmol = pybel.readfile("pdb",pdb).__next__()
  p_atoms, p_all = get_mol_info(pdbmol)

  # read gro
  gromol = pybel.readfile("gro",gro).__next__()
  q_atoms, q_all = get_mol_info(gromol)

  return rmsd.reorder_hungarian(p_atoms, q_atoms, p_all, q_all), rmsd.reorder_hungarian(q_atoms, p_atoms, q_all, p_all)

def reorder_gro(gro, mapp):
  fout = open("reordered_"+os.path.basename(gro), "w")
  with open(gro, "r") as f:
    fout.write(f.readline())
    natoms = f.readline()
    fout.write(natoms)
    natoms = int(natoms)
    
    # read the lines
    lines = []
    for i in range(natoms):
      line = f.readline()
      lines.append(line)

    # write reordering
    for i in range(len(lines)):
      fout.write(lines[mapp[i]])

    fout.write(f.readline())
  fout.close()

def reorder_itp(itp, mapp, imapp):
  fout = open("reordered_"+os.path.basename(itp), "w")
  with open(itp, "r") as f:
    line = f.readline()
    while ";   nr" not in line:
      if "GENERATED BY LigParGen Server" in line:
        line = "; GENERATED BY LigParGen Server and reordered by reordered_ligpargen\n"
      fout.write(line)
      line = f.readline()

    fout.write(line)
    line = f.readline()

    # read atoms
    lines = []
    while ("opls" in line) or line.strip().startswith(";"):
      if line.strip().startswith(";"):
        fout.write(line)
        f.readline()
        continue
      lines.append(line)
      line = f.readline()

    # write atoms in the right order
    for i in range(len(lines)):
      replin = re.sub(r'\d+',"%d"%(i+1),lines[mapp[i]],1)
      diffslen = len(replin) - len(lines[i])
      if diffslen < 0 :
        replin = -diffslen*" " + replin
      elif diffslen > 0:
        replin = replin[diffslen:]
      fout.write(replin)

    while "[ bonds ]" not in line:
      fout.write(line)
      line = f.readline()

    fout.write(line)
    line = f.readline()

    # write the bonds using the right labels
    while "[ angles ]" not in line:
      if not line.strip() or line.strip().startswith(";"):
        fout.write(line)
        line = f.readline()
        continue

      a1 = int(line.split()[0])-1
      a2 = int(line.split()[1])-1

      constants = ""
      for val in line.split()[2:]:
        constants += val+"\t"

      fout.write("%5d %5d     %s\n" % (imapp[a1]+1,imapp[a2]+1,constants))
      line = f.readline()

    # write the angles
    fout.write(line)
    line = f.readline()
    while "[ dihedrals ]" not in line:
      if not line.strip() or line.strip().startswith(";"):
        fout.write(line)
        line = f.readline()
        continue

      a1 = int(line.split()[0])-1
      a2 = int(line.split()[1])-1
      a3 = int(line.split()[2])-1

      constants = ""
      for val in line.split()[3:]:
        constants += val+"\t"

      fout.write("%5d %5d %5d    %s\n" % (imapp[a1]+1,imapp[a2]+1,imapp[a3]+1,constants))
      line = f.readline()

    # write the dihedrals
    fout.write(line)
    line = f.readline()
    while "[ dihedrals ]" not in line and "[ pairs ]" not in line:
      if not line.strip() or line.strip().startswith(";"):
        fout.write(line)
        line = f.readline()
        continue

      a1 = int(line.split()[0])-1
      a2 = int(line.split()[1])-1
      a3 = int(line.split()[2])-1
      a4 = int(line.split()[3])-1

      constants = ""
      for val in line.split()[4:]:
        constants += val+"\t"

      fout.write("%5d %5d %5d %5d   %s\n" % (imapp[a1]+1,imapp[a2]+1,imapp[a3]+1,imapp[a4]+1,constants))
      line = f.readline()

    # write the dihedrals if necessary
    if "[ dihedrals ]" in line:
      fout.write(line)
      line = f.readline()
      while "[ pairs ]" not in line:
        if not line.strip() or line.strip().startswith(";"):
          fout.write(line)
          line = f.readline()
          continue

        a1 = int(line.split()[0])-1
        a2 = int(line.split()[1])-1
        a3 = int(line.split()[2])-1
        a4 = int(line.split()[3])-1

        constants = ""
        for val in line.split()[4:]:
          constants += val+"\t"

        fout.write("%5d %5d %5d %5d   %s\n" % (imapp[a1]+1,imapp[a2]+1,imapp[a3]+1,imapp[a4]+1,constants))
        line = f.readline()

    # write the pairs
    fout.write(line)
    line = f.readline()
    while line:
      if not line.strip() or line.strip().startswith(";"):
        fout.write(line)
        line = f.readline()
        continue

      a1 = int(line.split()[0])-1
      a2 = int(line.split()[1])-1

      constants = ""
      for val in line.split()[2:]:
        constants += val+"\t"

      fout.write("%5d %5d     %s\n" % (imapp[a1]+1,imapp[a2]+1,constants))
      line = f.readline()

  fout.close()


if __name__ == '__main__':
  parser = argparse.ArgumentParser(description="Receives the original .pdb sent to LigParGen to reorder the output of LigParGen to have the atoms in the same order.")
  parser.add_argument("originalpdb", help="the original pdb uploaded to LigParGen")
  parser.add_argument("outgro", help="the .gro generated by LigParGen")
  parser.add_argument("outitp", help="the .itp generated by LigParGen")
  args = parser.parse_args()

  # map the indexes from one geometry to the others
  mapping, imapping = get_atom_correspondence(args.originalpdb, args.outgro)

  # based on the map return the reordered .gro
  reorder_gro(args.outgro, mapping)

  # now reorder the .itp
  reorder_itp(args.outitp, mapping, imapping)