from QuantumATK import *

import pylab

# Read Y and Z coordinates
coords     = nlread('gnr_bandstructure.hdf5', BulkConfiguration)[0]
coords     = coords.cartesianCoordinates().inUnitsOf(Angstrom)[:,1:3]

# Read the electron density from the gnr_density-.hdf5 file
n_u = nlread('gnr_spin.hdf5', ElectronDensity)[-1]

# Extract the spin up and spin down component 
n_up = n_u.spinProjection(spin=Spin.Up)
n_dn = n_u.spinProjection(spin=Spin.Down)

# Calculate the spin density difference
n = (n_up - n_dn)

av = numpy.array(n[:,:,:].sum(axis=0))

# Set the horizontal/vertical axes by scaling the data indices with the unit cell
y = numpy.array(range(av.shape[0]))/float(av.shape[0])*n_u.unitCell()[1][1].inUnitsOf(Ang)
z = numpy.array(range(av.shape[1]))/float(av.shape[1])*n_u.unitCell()[2][2].inUnitsOf(Ang)

# A 'spin-dependent colorbar'
cdict = {
    'red':   ((0.0, 0.0, 0.0),
              (0.5, 0.0, 1.0),
              (1.0, 1.0, 1.0)),
    'green': ((0.0, 1.0, 1.0),
              (0.5, 0.0, 0.0),
              (1.0, 1.0, 1.0)),
    'blue':  ((0.0, 0.0, 0.0),
              (0.5, 1.0, 0.0),
              (1.0, 0.0, 0.0))
    }
blue_red2 = pylab.matplotlib.colors.LinearSegmentedColormap('BlueRed2', cdict)
levels = numpy.arange(-0.03, 0.03, 0.005)

# Build up the plot
pylab.figure(figsize=(5,y[-1]/z[-1]))
pylab.xlabel('z / Angstrom')
pylab.ylabel('y / Angstrom')

pylab.contourf(z,y,av,40,cmap=blue_red2)
pylab.contour(z,y,av,levels,colors='k')

pylab.plot(coords[:,1],coords[:,0],'ko',ms=15.0)
axis = pylab.axis('image')
# Limit the Z axis to the ribbon (default is the entire unit cell)
# The H-atom coordinates are z=10 and 19.3 Angstrom
v = [axis[0],axis[1],3.0,15.0]
pylab.axis(v)
pylab.savefig('av.png')
pylab.show()