LIN_FIX for fixing Linear molecule problems of LigParGen server

LIN_FIX.py

# THIS PROGRAM IS CREATED TO FIX PROBLEM WITH LINEAR MOLECULES
# WORKS ONLY FOR MOLECULES WITH ONE LINEAR BOND
import pandas as pd
import numpy as np
import os
import sys


def Refine_Zmat(zmat):
    for l in range(len(zmat)):
        if 'Geometry Variations follow' in zmat[l]:
            end_coos = l
        elif 'Additional Bonds follow' in zmat[l]:
            beg_addb = l
        elif 'Harmonic Constraints follow' in zmat[l]:
            end_addb = l
        elif 'Final blank line' in zmat[l]:
            nbd_part = l
    add_bonds = []
    for adb in zmat[beg_addb + 1:end_addb]: 
        add_bonds.append(adb.strip().split())
    addb_df = pd.DataFrame(add_bonds,dtype=np.int)
    addb_df.columns = ['I', 'J']
    zmat_data = np.loadtxt(zmat[1:end_coos], dtype=np.str, usecols=range(10))
    zmat_df = pd.DataFrame(zmat_data)
    zmat_df.columns = ['I', 'ATOM', 'EI', 'EJ','J', 'RIJ', 'K', 'TIJK', 'L', 'PIJKL']
    zmat_df[['I', 'J', 'K', 'L', 'EI', 'EJ']] = zmat_df[['I', 'J', 'K', 'L', 'EI', 'EJ']].astype(int)
    zmat_df[['RIJ', 'TIJK', 'PIJKL']] = zmat_df[['RIJ', 'TIJK', 'PIJKL']].astype(float)
    return(zmat_df, addb_df, zmat[nbd_part + 1:])


def WriteOUTZMAT(zname, zmat_df, add_bond, non_bond,xid):
    ofile = open(zname, 'w')
    tot_ats = zmat_df.I.max()
    ofile.write('LSD edited BOSS Z-matrix for LPG server \n')
    for i, Z in zmat_df.iterrows():
        ofile.write('%4d %-3s%5d%5d%5d%12.6f%4d%12.6f%4d%12.6f%4s%5d\n'
                    % (Z.I, Z.ATOM, Z.EI, Z.EJ, Z.J, Z.RIJ, Z.K, Z.TIJK, Z.L, Z.PIJKL, 'UNK', 1))
    ofile.write(
        '                    Geometry Variations follow    (2I4,F12.6)\n')
    ofile.write(
        '                    Variable Bonds follow         (I4)\n0004-%4d\n' % (tot_ats))
    ofile.write('                    Additional Bonds follow       (2I4)\n')
    for i, r in add_bond.iterrows():
        ofile.write('%4d%4d\n' % (r.I, r.J))
    ofile.write(
        '                    Harmonic Constraints follow   (2I4,4F10.4)\n')
    ofile.write(
        '                    Variable Bond Angles follow   (I4)\n0005-%4d\n%4d-%4d\n' % (xid-1,xid+1,tot_ats))
    ofile.write(
        '                    Additional Bond Angles follow (3I4)\nAUTO\n')
    ofile.write(
        '                    Variable Dihedrals follow     (3I4,F12.6)\n0006-%4d\n%4d-%4d\n' % (xid-1,xid+1,tot_ats))
    ofile.write(
        '                    Additional Dihedrals follow   (6I4)\nAUTO\n')
    ofile.write(
        '''                    Domain Definitions follow     (4I4)
                    Conformational Search (2I4,2F12.6)
                    Local Heating Residues follow (I4 or I4-I4)
                    Final blank line
    ''')
    for line in non_bond:
        ofile.write('%s' % line)
    ofile.close()
    return None


# In[6]:
def LinCheckZ(zmat_name):
    zmat = open(zmat_name, 'r').readlines()
    zmat_df, addb_df, nbd = Refine_Zmat(zmat)
    fin_zmat_df = zmat_df[~((zmat_df.ATOM =='X') & (zmat_df.I.astype(int)>2))]
    XID = list(zmat_df[(zmat_df.ATOM == 'X') & (zmat_df.I > 2)]['I'])
    if len(XID) > 0:
        for i, r in fin_zmat_df.iterrows():
            if r.I > XID[0]:
                fin_zmat_df.loc[i, 'I'] = r.I - 1
            if r.J > XID[0]:
                fin_zmat_df.loc[i, 'J'] = r.J - 1
            if r.K > XID[0]:
                fin_zmat_df.loc[i, 'K'] = r.K - 1
            if r.L > XID[0]:
                fin_zmat_df.loc[i, 'L'] = r.L - 1
        preds = zmat_df.set_index(['I'])['J'].to_dict()
        for i, r in addb_df.iterrows():
            if r.I > XID[0]:
                addb_df.loc[i, 'I'] = r.I - 1
            if r.J > XID[0]:
                addb_df.loc[i, 'J'] = r.J - 1
        for i, r in fin_zmat_df.iterrows():
            if r.K == XID[0]:
                fin_zmat_df.loc[i, 'K'] = preds[preds[XID[0]]]
                fin_zmat_df.loc[i, 'TIJK'] = 179.90
        for i, r in fin_zmat_df.iterrows():
            if (i > 1) and (r.K == r.L):
                fin_zmat_df.loc[i, 'L'] = preds[preds[r.K]]
        LINEAR = True
    else:
        LINEAR = False
    return(fin_zmat_df, addb_df, nbd, LINEAR,XID[0])
#zmat_name='LeBlanc.z'
zmat_name = sys.argv[1]
fin_zmat_df, addb_df, nbd, LINEAR,xid = LinCheckZ(zmat_name)
if LINEAR:
    WriteOUTZMAT(zname='LSD_TT.z', zmat_df=fin_zmat_df,
                 add_bond=addb_df, non_bond=nbd,xid = xid)
else: os.system('/bin/cp %s LSD_TT.z'%zmat_name)