#!/usr/bin/env python
from sys import exit
from optparse import OptionParser
 
def getIndexedDansity(idx, calc):
  if index_error(idx, calc):
    exit()
  (nb, nk, ns) = idx
  wf = calc.GetWaveFunctionArray(band = nb, kpt = nk, spin = ns)
  density = abs(wf)*abs(wf)
  return density
 
def index_error(idx, calc):
  (nband, nkpt, nspin) = idx
  NumberOfBands = calc.GetNumberOfBands()
  NumberOfKpoints = len(calc.GetIBZKPoints())
  isSpinPolarized = calc.GetSpinPolarized()
  error = False
  if not (nband < NumberOfBands):
    print "error: nband must be less than %d" % NumberOfBands
    error = True
  if not (nkpt < NumberOfKpoints):
    print "error: nkpt must be less than %d" % NumberOfKpoints 
    error = True
  if isSpinPolarized:
    if not ((nspin == 0) or (nspin == 1)):
      print "error: nspin must be 0 or 1"
      error = True
  else:
    if not (nspin == 0):
      print "error: nspin must be 0"
      error = True
  return error
 
 
def getNumberOfSpins(calc):
  if (calc.GetSpinPolarized()):
    return 2
  else:
    return 1
 
def getRangedDansity(ewin, calc):
  (min_energy, max_energy) = ewin
  if (min_energy > max_energy):
    (min_energy, max_energy) = (max_energy, min_energy)
  NumberOfBands = calc.GetNumberOfBands()
  NumberOfKpoints = len(calc.GetIBZKPoints())
  NumberOfSpins = getNumberOfSpins(calculator)
  density = calculator.GetDensityArray()*0.0
  out_of_range = True
  min_range = 0.0
  max_range = 0.0
  for ns in range(NumberOfSpins):
    for nk in range(NumberOfKpoints):
      energies = calculator.GetEigenvalues(kpt = nk, spin = ns)
      for nb in range(NumberOfBands):
        if (energies[nb] < min_range):
          min_range = energies[nb] 
        if (energies[nb] > max_range):
          max_range = energies[nb]
        if (min_energy <= energies[nb]) and (energies[nb] <= max_energy):
          out_of_range = False
          wf = calculator.GetWaveFunctionArray(band = nb, kpt = nk, spin = ns)
          density = density + abs(wf)*abs(wf)
  if out_of_range:
    print "error: energy range should be from %f to %f" % (min_range, max_range)
    exit()
  return density
 
 
from Dacapo import Dacapo
from ASE.IO.Cube import WriteCube
 
cmd = OptionParser(usage = '%prog [-i|-e|-a] input_nc_file output_cube_file')
 
cmd.add_option('-i', '--indecies', type = 'int', nargs = 3,
               help = 'density of N-th Band, K-th Kpoint, and S-th Spin state',
               metavar = 'N K S')
 
cmd.add_option('-e', '--energy', type = 'float', nargs = 2,
               help = 'density of states with energies from E1[eV] to E2[eV]',
               metavar = 'E1 E2')
 
(opt, argv) = cmd.parse_args()
 
if len(argv) != 2:
    cmd.print_help()
    raise SystemExit
 
if (opt.indecies is not None) and (opt.energy is not None):
  print 'error: do not specify both options -i and -e'
  exit()
 
ncfile = argv[0]
cubefile = argv[1]
 
model = Dacapo.ReadAtoms(ncfile)
calculator = model.GetCalculator()
 
if opt.indecies is not None:
  density = getIndexedDansity(opt.indecies, calculator)
elif opt.energy is not None:
  density = getRangedDansity(opt.energy, calculator)
else:
  density = calculator.GetDensityArray()
WriteCube(model, density, cubefile)

usage: getdensity [-i|-e|-a] input_nc_file output_cube_file
 
options:
  -h, --help            show this help message and exit
  -iN K S, --indecies=N K S
                        density of N-th Band, K-th Kpoint, and S-th Spin state
  -eE1 E2, --energy=E1 E2
                        density of states with energies from E1[eV] to E2[eV]