velInit.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import ase, sys
import numpy as np
from ase.io import read, write
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation 

def init_vel(cml):
    '''
    '''
    arg = parse_cml_args(cml)

    # read in the initial structure
    init_pos = read(arg.poscar)
    # set the momenta corresponding to T
    MaxwellBoltzmannDistribution(init_pos, arg.temperature * ase.units.kB)
    Stationary(init_pos)
    ZeroRotation(init_pos)

    # scale the temperature to T 
    vel = init_pos.get_velocities()
    Tn = init_pos.get_temperature()
    vel *= np.sqrt(arg.temperature / Tn)
    init_pos.set_velocities(vel)

    # write the structure
    write(arg.out, init_pos, vasp5=True, direct=True)

    # units in VASP and ASE are different
    vel = init_pos.get_velocities() * ase.units.fs
    # np.savetxt('init_vel.dat', vel, fmt='%20.16f')
    # append the velocities to the POSCAR
    with open(arg.out, 'a+') as pos:
        pos.write('\n')
        pos.write(
                '\n'.join([
                    ''.join(["%20.16f" % x for x in row])
                    for row in vel
                   ])
                )

def parse_cml_args(cml):
    '''
    CML parser.
    '''
    arg = argparse.ArgumentParser(add_help=True)

    arg.add_argument('-i', dest='poscar', action='store', type=str,
                     default='POSCAR',
                     help='The slab POSCAR.')
    arg.add_argument('-o', dest='out', action='store', type=str,
                     default='out.vasp',
                     help='Default output filename.')
    arg.add_argument('-t', '--temperature', dest='temperature',
                     action='store', type=float,
                     default=300, 
                     help='The temperature.')

    return arg.parse_args(cml)

if __name__ == '__main__':
    init_vel(sys.argv[1:])