1 """ plot wavefunction, together with atoms using VTK. 
   2 
   3 CO molecule is used as an example. 
   4 
   5 Usage: python -i plotwavefunction.py <bandno>
   6 
   7 """
   8 
   9 #Adjust the name of the nc output file here (if you have a different one)
  10 ncfilename = 'CO.nc'
  11 
  12 from ASE import Atom,ListOfAtoms
  13 from ASE.Visualization.VTK import VTKPlotWaveFunction
  14 from ASE.Visualization.VTK import VTKPlotAtoms
  15 from Dacapo import Dacapo
  16 import os.path,sys
  17 
  18 args = sys.argv[1:]
  19 if len(args) != 1:
  20    band = 0
  21 else:
  22    band = int(args[0])
  23 
  24 print 'plotting wavefunction for band number ',band
  25 
  26 atoms = Dacapo.ReadAtoms(filename=ncfilename)
  27 calc = atoms.GetCalculator()
  28 
  29 
  30 # make a combined plot of the wavefunction and the atoms
  31 atomplot = VTKPlotAtoms(atoms)
  32 wfplot = VTKPlotWaveFunction(atoms, band=band, parent=atomplot)
  33 atomplot.Update()
  34 
  35 # The appearence of the individual atomic elements can also be changed.
  36 # Finding aluminium atom avatar (found in the dictionary of species avatars)
  37 p1_O=atomplot.GetDictOfSpecies()['O']
  38 p1_C=atomplot.GetDictOfSpecies()['C']
  39 
  40 # Change the radius to 2.0 AA and set the color to blue
  41 p1_C.SetRadius(0.5)
  42 p1_O.SetRadius(0.7)
  43 
  44 #The color is given in an rgb (red,green,blue) scale
  45 p1_C.SetColor((1,0,0))
  46 p1_O.SetColor((0,1,0))
  47 atomplot.Render()
  48 
  49 # The unitcell can be removed
  50 unitcell=atomplot.unitcell
  51 atomplot.RemoveAvatar(unitcell)
  52 atomplot.Render()