""" plot wavefunction, together with atoms using VTK. 

CO molecule is used as an example. 

"""
from ASE import Atom,ListOfAtoms
from ASE.Visualization.VTK import VTKPlotWaveFunction
from ASE.Visualization.VTK import VTKPlotAtoms
from Dacapo import Dacapo

band=0
myfile="al_bstr_gx.nc"

#Insert the name of the nc output file here (if you have a different one)
atoms = Dacapo.ReadAtoms(filename=myfile)
calc = atoms.GetCalculator()


# make a combined plot of the wavefunction and the atoms
atomplot = VTKPlotAtoms(atoms)
wfplot = VTKPlotWaveFunction(atoms, band=band, parent=atomplot)
atomplot.SetPeriods((3,3,3))
wfplot.SetPeriods((2,2,2))
atomplot.Update()

# The appearence of the individual atomic elements can also be changed.
# Finding aluminium atom avatar (found in the dictionary of species avatars)
p1_Al=atomplot.GetDictOfSpecies()['Al']

# Change the radius to 2.0 AA and set the color to blue
p1_Al.SetRadius(0.2)

#The color is given in an rgb (red,green,blue) scale
p1_Al.SetColor((0.7,0.7,0.7))
atomplot.Render()


# The unitcell can be removed
unitcell=atomplot.unitcell
atomplot.RemoveAvatar(unitcell)
atomplot.Render()