density2cube

#!/usr/bin/env python

import pickle
import optparse
import espresso.io
import sys
import os.path

import ase.io.cube
import ase.units

Bohr = ase.units.Bohr
Bohr = 0.5291772108

def density2cube(options):
    if 'trajfile' in options:
        traj = [ase.io.read(options['trajfile'])]
    elif 'logfile' in options:
        traj = espresso.io.read_log(options['logfile'])
    else:
        traj = [None]

    atoms = traj[-1]
    print(atoms)
    try:
        with open(options['pickle']) as infile:
            data = pickle.load(infile)
    except:
        pfilename = options['pickle']
        print "Error opening {pfilename}".format(**locals())
        return

    origin, cell, density = data
    density = density[:-1, :-1, :-1]

    cell /= Bohr

    print(density.shape)
    #density /= ase.units.Rydberg

    n_atoms = len(atoms)
    origin_s = ' '.join(map(lambda x: '   {: 10.6f}'.format(x), origin.tolist()))
    x, y, z = density.shape
    X = ' '.join(map(lambda x: '   {: 10.6f}'.format(x), cell[ 0]/float(x)))
    Y = ' '.join(map(lambda x: '   {: 10.6f}'.format(x), cell[ 1]/float(y)))
    Z = ' '.join(map(lambda x: '   {: 10.6f}'.format(x), cell[ 2]/float(z)))


    with open(options['outfile'], 'w') as outfile:
        if options['with_ase']:
            ase.io.cube.write_cube(outfile, atoms, density)
        else:
            outfile.write('# generated from {options[logfile]} and {options[pickle]}\n'.format(**locals()))
            outfile.write('OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n')
            outfile.write('{n_atoms: 5d} {origin_s}\n'.format(**locals()))
            outfile.write('{x: 5d} {X}\n'.format(**locals()))
            outfile.write('{y: 5d} {Y}\n'.format(**locals()))
            outfile.write('{z: 5d} {Z}\n'.format(**locals()))
            for i, atom in enumerate(atoms):
                atom_number = atoms[i].number
                position = ' '.join(map(lambda x: '   {: 10.6f}'.format(x), atom.position / Bohr))
                outfile.write('   {atom_number: 5d}   0.000000 {position}\n'.format(**locals()))
            for xi in range(int(x)):
                for yi in range(int(y)):
                    for zi in range(int(z)):
                        outfile.write('{: .5e} '.format(density[xi, yi, zi]))
                        if zi % 6 == 5 :
                            outfile.write('\n')
            outfile.write('\n')

if __name__ == '__main__':
    parser = optparse.OptionParser()

    parser.add_option('-l', '--logfile', dest='logfile', default=None)
    parser.add_option('-t', '--trajfile', dest='trajfile', default=None)
    parser.add_option('-p', '--pickle', dest='pickle', default=None)
    parser.add_option('-o', '--outfile', dest='outfile', default=None)
    parser.add_option('-a', '--with-ase', dest='with_ase', action='store_true', default=False)

    options, args = parser.parse_args()


    if options.logfile is None and options.trajfile is None:
        raise UserWarning('Please specify either a trajectory file with -t or a logfile with -l')

    if options.trajfile is None:
        print("Warning: if the geometry was optimized with an ASE optimizer, the */log file does not contain the update positions. Specify a traj file instead.")
    if options.pickle is None:
        raise UserWarning('Please specify pickle file with -p')
    if options.outfile is None:
        options.outfile = os.path.splitext(options.pickle)[0] + '.cub'
        print('Set output filename to {options.outfile}'.format(**locals()))

    options = dict(pickle=options.pickle,
             logfile=options.logfile,
             trajfile=options.trajfile,
             outfile=options.outfile,
             with_ase=options.with_ase)

    density2cube(options)