/ase/io/vasp.py
Python | 422 lines | 338 code | 32 blank | 52 comment | 106 complexity | 606960165167cf1fff6f5db5059b535c MD5 | raw file
- """
- This module contains functionality for reading and writing an ASE
- Atoms object in VASP POSCAR format.
- """
- import os
- def get_atomtypes(fname):
- """Given a file name, get the atomic symbols.
- The function can get this information from OUTCAR and POTCAR
- format files. The files can also be compressed with gzip or
- bzip2.
- """
- atomtypes=[]
- if fname.find('.gz') != -1:
- import gzip
- f = gzip.open(fname)
- elif fname.find('.bz2') != -1:
- import bz2
- f = bz2.BZ2File(fname)
- else:
- f = open(fname)
- for line in f:
- if line.find('TITEL') != -1:
- atomtypes.append(line.split()[3].split('_')[0].split('.')[0])
- return atomtypes
- def atomtypes_outpot(posfname, numsyms):
- """Try to retreive chemical symbols from OUTCAR or POTCAR
-
- If getting atomtypes from the first line in POSCAR/CONTCAR fails, it might
- be possible to find the data in OUTCAR or POTCAR, if these files exist.
- posfname -- The filename of the POSCAR/CONTCAR file we're trying to read
-
- numsyms -- The number of symbols we must find
- """
- import os.path as op
- import glob
- # First check files with exactly same name except POTCAR/OUTCAR instead
- # of POSCAR/CONTCAR.
- fnames = [posfname.replace('POSCAR', 'POTCAR').replace('CONTCAR',
- 'POTCAR')]
- fnames.append(posfname.replace('POSCAR', 'OUTCAR').replace('CONTCAR',
- 'OUTCAR'))
- # Try the same but with compressed files
- fsc = []
- for fn in fnames:
- fsc.append(fn + '.gz')
- fsc.append(fn + '.bz2')
- for f in fsc:
- fnames.append(f)
- # Finally try anything with POTCAR or OUTCAR in the name
- vaspdir = op.dirname(posfname)
- fs = glob.glob(vaspdir + '*POTCAR*')
- for f in fs:
- fnames.append(f)
- fs = glob.glob(vaspdir + '*OUTCAR*')
- for f in fs:
- fnames.append(f)
- tried = []
- files_in_dir = os.listdir('.')
- for fn in fnames:
- if fn in files_in_dir:
- tried.append(fn)
- at = get_atomtypes(fn)
- if len(at) == numsyms:
- return at
- raise IOError('Could not determine chemical symbols. Tried files '
- + str(tried))
- def get_atomtypes_from_formula(formula):
- """Return atom types from chemical formula (optionally prepended
- with and underscore).
- """
- from ase.atoms import string2symbols
- symbols = string2symbols(formula.split('_')[0])
- atomtypes = [symbols[0]]
- for s in symbols[1:]:
- if s != atomtypes[-1]: atomtypes.append(s)
- return atomtypes
- def read_vasp(filename='CONTCAR'):
- """Import POSCAR/CONTCAR type file.
- Reads unitcell, atom positions and constraints from the POSCAR/CONTCAR
- file and tries to read atom types from POSCAR/CONTCAR header, if this fails
- the atom types are read from OUTCAR or POTCAR file.
- """
-
- from ase import Atoms, Atom
- from ase.constraints import FixAtoms, FixScaled
- from ase.data import chemical_symbols
- import numpy as np
- if isinstance(filename, str):
- f = open(filename)
- else: # Assume it's a file-like object
- f = filename
- # First line should contain the atom symbols , eg. "Ag Ge" in
- # the same order
- # as later in the file (and POTCAR for the full vasp run)
- atomtypes = f.readline().split()
- # Sometimes the first line in POSCAR/CONTCAR is of the form
- # "CoP3_In-3.pos". Check for this case and extract atom types
- if len(atomtypes) == 1 and '_' in atomtypes[0]:
- atomtypes = get_atomtypes_from_formula(atomtypes[0])
- lattice_constant = float(f.readline().split()[0])
- # Now the lattice vectors
- a = []
- for ii in range(3):
- s = f.readline().split()
- floatvect = float(s[0]), float(s[1]), float(s[2])
- a.append(floatvect)
- basis_vectors = np.array(a) * lattice_constant
- # Number of atoms. Again this must be in the same order as
- # in the first line
- # or in the POTCAR or OUTCAR file
- atom_symbols = []
- numofatoms = f.readline().split()
- #vasp5.1 has an additional line which gives the atom types
- #the following try statement skips this line
- try:
- int(numofatoms[0])
- except ValueError:
- numofatoms = f.readline().split()
- # check for comments in numofatoms line and get rid of them if necessary
- commentcheck = np.array(['!' in s for s in numofatoms])
- if commentcheck.any():
- # only keep the elements up to the first including a '!':
- numofatoms = numofatoms[:np.arange(len(numofatoms))[commentcheck][0]]
-
- numsyms = len(numofatoms)
- if len(atomtypes) < numsyms:
- # First line in POSCAR/CONTCAR didn't contain enough symbols.
- atomtypes = atomtypes_outpot(f.name, numsyms)
- else:
- try:
- for atype in atomtypes[:numsyms]:
- if not atype in chemical_symbols:
- raise KeyError
- except KeyError:
- atomtypes = atomtypes_outpot(f.name, numsyms)
- for i, num in enumerate(numofatoms):
- numofatoms[i] = int(num)
- [atom_symbols.append(atomtypes[i]) for na in xrange(numofatoms[i])]
- # Check if Selective dynamics is switched on
- sdyn = f.readline()
- selective_dynamics = sdyn[0].lower() == "s"
- # Check if atom coordinates are cartesian or direct
- if selective_dynamics:
- ac_type = f.readline()
- else:
- ac_type = sdyn
- cartesian = ac_type[0].lower() == "c" or ac_type[0].lower() == "k"
- tot_natoms = sum(numofatoms)
- atoms_pos = np.empty((tot_natoms, 3))
- if selective_dynamics:
- selective_flags = np.empty((tot_natoms, 3), dtype=bool)
- for atom in xrange(tot_natoms):
- ac = f.readline().split()
- atoms_pos[atom] = (float(ac[0]), float(ac[1]), float(ac[2]))
- if selective_dynamics:
- curflag = []
- for flag in ac[3:6]:
- curflag.append(flag == 'F')
- selective_flags[atom] = curflag
- # Done with all reading
- if type(filename) == str:
- f.close()
- if cartesian:
- atoms_pos *= lattice_constant
- atoms = Atoms(symbols = atom_symbols, cell = basis_vectors, pbc = True)
- if cartesian:
- atoms.set_positions(atoms_pos)
- else:
- atoms.set_scaled_positions(atoms_pos)
- if selective_dynamics:
- constraints = []
- indices = []
- for ind, sflags in enumerate(selective_flags):
- if sflags.any() and not sflags.all():
- constraints.append(FixScaled(atoms.get_cell(), ind, sflags))
- elif sflags.all():
- indices.append(ind)
- if indices:
- constraints.append(FixAtoms(indices))
- if constraints:
- atoms.set_constraint(constraints)
- return atoms
- def read_vasp_out(filename='OUTCAR',index = -1):
- """Import OUTCAR type file.
- Reads unitcell, atom positions, energies, and forces from the OUTCAR file
- and attempts to read constraints (if any) from CONTCAR/POSCAR, if present.
- """
- import os
- import numpy as np
- from ase.calculators.singlepoint import SinglePointCalculator
- from ase import Atoms, Atom
- try: # try to read constraints, first from CONTCAR, then from POSCAR
- constr = read_vasp('CONTCAR').constraints
- except:
- try:
- constr = read_vasp('POSCAR').constraints
- except:
- constr = None
- if isinstance(filename, str):
- f = open(filename)
- else: # Assume it's a file-like object
- f = filename
- data = f.readlines()
- natoms = 0
- images = []
- atoms = Atoms(pbc = True, constraint = constr)
- energy = 0
- species = []
- species_num = []
- symbols = []
- ecount = 0
- poscount = 0
- for n,line in enumerate(data):
- if 'POTCAR:' in line:
- temp = line.split()[2]
- for c in ['.','_','1']:
- if c in temp:
- temp = temp[0:temp.find(c)]
- species += [temp]
- if 'ions per type' in line:
- species = species[:len(species)/2]
- temp = line.split()
- for ispecies in range(len(species)):
- species_num += [int(temp[ispecies+4])]
- natoms += species_num[-1]
- for iatom in range(species_num[-1]): symbols += [species[ispecies]]
- if 'direct lattice vectors' in line:
- cell = []
- for i in range(3):
- temp = data[n+1+i].split()
- cell += [[float(temp[0]), float(temp[1]), float(temp[2])]]
- atoms.set_cell(cell)
- if 'FREE ENERGIE OF THE ION-ELECTRON SYSTEM' in line:
- energy = float(data[n+2].split()[4])
- if ecount < poscount:
- # reset energy for LAST set of atoms, not current one - VASP 5.11? and up
- images[-1].calc.energy = energy
- ecount += 1
- if 'POSITION ' in line:
- forces = []
- for iatom in range(natoms):
- temp = data[n+2+iatom].split()
- atoms += Atom(symbols[iatom],[float(temp[0]),float(temp[1]),float(temp[2])])
- forces += [[float(temp[3]),float(temp[4]),float(temp[5])]]
- atoms.set_calculator(SinglePointCalculator(energy,forces,None,None,atoms))
- images += [atoms]
- atoms = Atoms(pbc = True, constraint = constr)
- poscount += 1
- # return requested images, code borrowed from ase/io/trajectory.py
- if isinstance(index, int):
- return images[index]
- else:
- step = index.step or 1
- if step > 0:
- start = index.start or 0
- if start < 0:
- start += len(images)
- stop = index.stop or len(images)
- if stop < 0:
- stop += len(images)
- else:
- if index.start is None:
- start = len(images) - 1
- else:
- start = index.start
- if start < 0:
- start += len(images)
- if index.stop is None:
- stop = -1
- else:
- stop = index.stop
- if stop < 0:
- stop += len(images)
- return [images[i] for i in range(start, stop, step)]
- def write_vasp(filename, atoms, label='', direct=False, sort=None, symbol_count = None, long_format=True):
- """Method to write VASP position (POSCAR/CONTCAR) files.
- Writes label, scalefactor, unitcell, # of various kinds of atoms,
- positions in cartesian or scaled coordinates (Direct), and constraints
- to file. Cartesian coordiantes is default and default label is the
- atomic species, e.g. 'C N H Cu'.
- """
-
- import numpy as np
- from ase.constraints import FixAtoms, FixScaled
- if isinstance(filename, str):
- f = open(filename, 'w')
- else: # Assume it's a 'file-like object'
- f = filename
-
- if isinstance(atoms, (list, tuple)):
- if len(atoms) > 1:
- raise RuntimeError("Don't know how to save more than "+
- "one image to VASP input")
- else:
- atoms = atoms[0]
- # Write atom positions in scaled or cartesian coordinates
- if direct:
- coord = atoms.get_scaled_positions()
- else:
- coord = atoms.get_positions()
- if atoms.constraints:
- sflags = np.zeros((len(atoms), 3), dtype=bool)
- for constr in atoms.constraints:
- if isinstance(constr, FixScaled):
- sflags[constr.a] = constr.mask
- elif isinstance(constr, FixAtoms):
- sflags[constr.index] = [True, True, True]
- if sort:
- ind = np.argsort(atoms.get_chemical_symbols())
- symbols = np.array(atoms.get_chemical_symbols())[ind]
- coord = coord[ind]
- if atoms.constraints:
- sflags = sflags[ind]
- else:
- symbols = atoms.get_chemical_symbols()
- # Create a list sc of (symbol, count) pairs
- if symbol_count:
- sc = symbol_count
- else:
- sc = []
- psym = symbols[0]
- count = 0
- for sym in symbols:
- if sym != psym:
- sc.append((psym, count))
- psym = sym
- count = 1
- else:
- count += 1
- sc.append((psym, count))
- # Create the label
- if label == '':
- for sym, c in sc:
- label += '%2s ' % sym
- f.write(label + '\n')
- # Write unitcell in real coordinates and adapt to VASP convention
- # for unit cell
- # ase Atoms doesn't store the lattice constant separately, so always
- # write 1.0.
- f.write('%19.16f\n' % 1.0)
- if long_format:
- latt_form = ' %21.16f'
- else:
- latt_form = ' %11.6f'
- for vec in atoms.get_cell():
- f.write(' ')
- for el in vec:
- f.write(latt_form % el)
- f.write('\n')
- # Numbers of each atom
- for sym, count in sc:
- f.write(' %3i' % count)
- f.write('\n')
- if atoms.constraints:
- f.write('Selective dynamics\n')
- if direct:
- f.write('Direct\n')
- else:
- f.write('Cartesian\n')
- if long_format:
- cform = ' %19.16f'
- else:
- cform = ' %9.6f'
- for iatom, atom in enumerate(coord):
- for dcoord in atom:
- f.write(cform % dcoord)
- if atoms.constraints:
- for flag in sflags[iatom]:
- if flag:
- s = 'F'
- else:
- s = 'T'
- f.write('%4s' % s)
- f.write('\n')
- if type(filename) == str:
- f.close()