from QuantumATK import *

# Calculate the current density at the fermi level from right to left going states
cd = CurrentDensity(
    configuration,
    energies = [0.0]*eV,
    energy_weights = Right,
    kpoints=MonkhorstPackGrid(1,11),
    )

# Calculate the volume element
dX, dY, dZ = cd.volumeElement()
dx = dX.norm()
dy = dY.norm()
dz = dZ.norm()

# Calculate the current density along z integrated over x,y
shape = cd.shape()
# Take out spin up (index 0) z component (index 2) in every grid point.
cd_z = numpy.array( [ [ [ cd[i,j,k][0][2] for k in range(shape[2])] for j in range(shape[1])] for i in range(shape[0]) ] )

# Sum across y dimensions
cd_xz = numpy.array([ [cd_z[i,:,k].sum() * dy for k in range(shape[2])] for i in range(shape[0] )] )


# Make 2-d arrays for contour plot of the data
shape = cd.shape()
x = dx.inUnitsOf(Ang)*numpy.arange(shape[0])
z = dz.inUnitsOf(Ang)*numpy.arange(shape[2])
Z, X = numpy.meshgrid(z,x)
F = numpy.array(cd_xz).reshape(numpy.shape(Z))

#plot the current density in the (z,x) plane
import pylab
pylab.xlabel('z (Angstrom)',fontsize=12,family='sans-serif')
pylab.ylabel('x (Angstrom)',fontsize=12,family='sans-serif')
pylab.contourf(Z, X, F)
pylab.colorbar()
pylab.savefig('cd.png',dpi=100)
pylab.show()