density_to_chgcar.py

#!/usr/bin/env python

from __future__ import print_function, division

import numpy as np
from gpaw import GPAW, restart
from ase.build import molecule
from ase import Atoms
from ase.io import write

try:
    myatoms, calc = restart('methane.gpw')
except:
    myatoms = molecule('CH4')
    myatoms.center(vacuum=2)
    calc = GPAW(xc='PBE',
                h=0.18,
                )
    myatoms.set_calculator(calc)

myatoms.get_potential_energy()
calc.write('methane.gpw')

nt = calc.get_pseudo_density()
n = calc.get_all_electron_density(gridrefinement=2)

write('methane.POSCAR', myatoms)

with open('methane.CHGCAR', 'w') as f:
    # Write name/label
    f.write('Methane\n')

    # Unit cell with scaling factor
    f.write('   1.00000000000000\n')
    for vec in myatoms.cell:
        f.write(' {:12.6f} {:12.6f} {:12.6f}\n'.format(*vec))

    # Order Atoms object so all atoms of the same element are contiguous
    by_symbol = {}
    for atom in myatoms:
        if atom.symbol not in by_symbol:
            by_symbol[atom.symbol] = Atoms()
        by_symbol[atom.symbol] += atom

    # List element names
    elements = list(by_symbol)
    elstring = ' ' + ' {:>3s}' * len(elements) + '\n'
    f.write(elstring.format(*elements))
    
    ordered_atoms = Atoms()
    for element in elements:
        ordered_atoms += by_symbol[element]
    
    # List how many of each element
    count = [len(by_symbol[element]) for element in elements]
    countstring = ' {:5d}' * len(elements) + '\n'
    f.write(countstring.format(*count))
    
    # Write scaled atomic positions
    f.write('Cartesian\n')
    for vec in ordered_atoms.get_scaled_positions():
        posstring = ' {:9.6f}' * 3 + '\n'
        f.write(posstring.format(*vec))

    f.write('\n')
    dimstring = ' {:4d}' * 3 + '\n'
    f.write(dimstring.format(*n.shape))
    nentries = np.prod(n.shape)
    nrows = nentries // 5

    nflat = n.ravel()
    
    rowstring = ' ' + '{:18.11E}' * 5 + '\n'
    for i in range(nrows):
        f.write(rowstring.format(*nflat[5*i:5*(i + 1)]))
    excess = nentries % 5
    if excess:
        rowstring = ' ' + '{:18.11E}' * excess + '\n'
        f.write(rowstring.format(*nflat[-excess:]))