#!/usr/bin/env python

from Dacapo import Dacapo
from ASE.Visualization.VTK import VTKPlotArray, VTKPlotAtoms

atoms = Dacapo.ReadAtoms("anti.nc")
calc = atoms.GetCalculator()
spin0 = calc.GetDensityArray(spin=0)
spin1 = calc.GetDensityArray(spin=1)
zeta = (spin0-spin1)/ (spin0+spin1)
plot = VTKPlotArray(zeta, atoms.GetUnitCell())

atomplot = VTKPlotAtoms(atoms, parent=plot)

p1_Fe=atomplot.GetDictOfSpecies()['Fe']
p1_Fe.SetRadius(0.5)
p1_Fe.SetColor((1,1,1))

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

plot.Render()