#Saving the wave function as a .Cube file
#This does unfortunately only work for cubic unit cells. You should use VTK to plot wave functions in non cubic unit cells

import os
from Dacapo import Dacapo
from ASE import Atom, ListOfAtoms
from ASE.IO.Cube import WriteCube
atoms=Dacapo.ReadAtoms('filename')
calc=atoms.GetCalculator()
for band in range(6):  
 wavefunction_array=calc.GetWaveFunctionArray(band=band)
 #The wave function cube file is written as the density array, with the only difference being that we include the phase of the wave function as either + or - so that wave functions with opposite phases have opposite signs
 WriteCube(atoms,wavefunction_array,'wavefunctionarray_%d.cube' %band,real=True)