from ase.dft.stm import STM
from gpaw import GPAW

calc = GPAW('Al2x2x3.gpw')
str=calc.get_atoms()
stm=STM(str)
z=11.0
bias=1.0
c=stm.get_averaged_current(bias, z)
x,y,h=stm.scan(bias, c, repeat=(3, 3))

import matplotlib.pyplot as plt
import numpy as np
plt.gca(aspect='equal')
plt.contourf(x, y, h, 40)
plt.colorbar()
plt.savefig('2d.png')

plt.figure()
a = str.cell[0, 0]
x, y = stm.linescan(bias, c, [0, 0], [2 * a, 0])
plt.plot(x, y)
plt.savefig('line.png')