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 from ASE import Atom,ListOfAtoms
   9 from ASE.Visualization.VTK import VTKPlotWaveFunction
  10 from ASE.Visualization.VTK import VTKPlotAtoms
  11 from Dacapo import Dacapo
  12 import os.path,sys
  13 
  14 args = sys.argv[1:]
  15 if len(args) != 1:
  16    band = 0
  17 else: 
  18    band = int(args[0])
  19 
  20 print 'plotting wavefunction for band number ',band
  21 
  22 # check if we have already calculated the nc file
  23 if os.path.isfile('CO_in_a_box.nc'): 
  24 
  25      # start calculator from file 
  26      atoms = Dacapo.ReadAtoms(filename='CO_in_a_box.nc')
  27      calc = atoms.GetCalculator()
  28 
  29 else: 
  30 
  31      # start from scratch
  32      atoms = ListOfAtoms([Atom('C', (2, 2, 2)),
  33                           Atom('O', (3.1, 2, 2))],
  34                           cell=(4, 4, 4), periodic=1)
  35 
  36      calc = Dacapo(planewavecutoff=300,    # in eV
  37                    nbands=6,               # 1 extra empty bands
  38                    out='CO_in_a_box.nc',txtout='CO_in_a_box.txt')
  39 
  40      atoms.SetCalculator(calc)
  41 
  42 # make a combined plot of the wavefunction and the atoms
  43 atomplot = VTKPlotAtoms(atoms)
  44 wfplot = VTKPlotWaveFunction(atoms,band=band,parent=atomplot)
  45 atomplot.Update()
  46 
  47 # The appearence of the individual atomic elements can also be changed.
  48 # Finding aluminium atom avatar (found in the dictionary of species avatars)
  49 p1_O=atomplot.GetDictOfSpecies()['O']
  50 p1_C=atomplot.GetDictOfSpecies()['C']
  51 
  52 # Change the radius to 2.0 AA and set the color to blue
  53 p1_C.SetRadius(0.5)
  54 p1_O.SetRadius(0.7)
  55 
  56 #The color is given in an rgb (red,green,blue) scale
  57 p1_C.SetColor((1,0,0))
  58 p1_O.SetColor((0,1,0))
  59 atomplot.Render()
  60 
  61 # The unitcell can be removed
  62 unitcell=atomplot.unitcell
  63 atomplot.RemoveAvatar(unitcell)
  64 atomplot.Render()