cubediff.py

#!/usr/bin/python

###########################################################
#  
#  ./cubediff file1 file2
#  output=file2-file1
#  [file1] cube file (generated by regular Q-Chem plotter)
#  [file2] cube file (generated by libwfa)
#
#  !! Make sure that file1 is generated using N+1 points
#     (as compared to the libwfa plot) in the same box.
#     
###########################################################
import sys, re
import cubicle as cl
from string import Template

filename = sys.argv[1]
file_ref = sys.argv[2]
ext = filename[:-5]
ext_ref = file_ref[:-5]

cubelist = []

###### Parse cube headers ######
### (File 1) ENVIRONMENT
header = []
pointnums = []
head_pattern = '^\s*(\d+)\s+(-?\d+\.\d+\s*){3}\n'
natm = 1
with open(filename) as g:
    for j,line in enumerate(g):
        if j == 2:
            natm = int(line.split()[0])
        if j == natm+4+2:
            break
        header.append(line)
with open(filename) as f:
    cube = f.read()

head_str = ''.join(header) # string of total header (File 1)
for i,head in enumerate(re.finditer(head_pattern,head_str,flags=re.MULTILINE)):
     if i >= 1:
        pointnums.append(head.group(1))
pnums = list(map(int,pointnums))

### (File 2) TOTAL SYSTEM
header_ref = []
with open(file_ref) as ref:
    for j,line in enumerate(ref):
        if j == 2:
            natm = int(line.split()[0])
        if j == natm+4+2:
            break
        header_ref.append(line)
head_ref_str = ''.join(header_ref) # string of total header (File 2)

###### Parse cube data ######
pattern2 = '(-?\d+\.?\d+[Ee]\ *[-+]?\ *\d+\s*){' + re.escape(str(pnums[2])) + '}'
cubelist.append([])
xcnt = 0
ycnt = 0
for xy,m in enumerate(re.finditer(pattern2,cube)):
    if xy % pnums[1] == 0 and xy > 0:
        xcnt += 1
        ycnt = 0
        cubelist.append([]) # append empty list for x
    s = m.group().split()
    s = list(map(float,s))
    cubelist[xcnt].append(s)
    ycnt += 1

###### Manipulating cube file ######
del_last = True
# delete last
if del_last == True:
    for x in range(pnums[0]):
        if x == pnums[0]-1:
            del cubelist[x]
            continue
        for y in range(pnums[1]):
            if y == pnums[1]-1:
                del cubelist[x][y]
                continue
            for z in range(pnums[2]):
                if z == pnums[2]-1:
                    del cubelist[x][y][z]
                    continue
else:
#delete first
    for x in reversed(range(pnums[0])):
        if x == 0:
            del cubelist[x]
            continue
        for y in reversed(range(pnums[1])):
            if y == 0:
                del cubelist[x][y]
                continue
            for z in reversed(range(pnums[2])):
                if z == 0:
                    del cubelist[x][y][z]
                    continue
# new number of points
for n in range(3):
    pnums[n] = pnums[n]-1

###### Write new cube file ######
new_file = filename[:-5]+"_new.cube"
with open(new_file, "w") as out:
    out.write(head_ref_str)
    for x in range(pnums[0]):
        for y in range(pnums[1]):
            for z in range(pnums[2]):
                if (z+1) % 6 == 0 or z == pnums[2]-1:
                    out.write("  {0:.9E}\n".format(cubelist[x][y][z]))
                else:
                    out.write("  {0:.9E}".format(cubelist[x][y][z]))


###### Subtract from reference ######
calc = "1.0*"+file_ref+" -1.0*"+new_file
output = file_ref[:-5] + "_diff.cube"
with open(output, "w") as diff:
    diff.write(cl.compute_cube(calc))

###### Write Jmol input ######
jmol = Template("""load ${fn}
background white
isosurface ID diff cutoff 0.005 sign "${fn}"
color $$diff translucent 0.5
print "Pausing script to reorientate molecule. Resume with '&plot'."
throw context plot
write image 1024 786 png ${png}
""")
with open(output[:-5]+".spt", "w") as jm:
    jm.write(jmol.substitute(fn=output,png=output[:-5]+".png"))