import matplotlib.pyplot as plt
import numpy as np
from crystals import Crystal
from skued import pelectrostatic

extent = np.linspace(-5, 5, 256)
xx, yy = np.meshgrid(extent, extent)

potential = pelectrostatic(Crystal.from_database('C'), xx, yy)

fig = plt.figure()
ax = fig.add_subplot(111)
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
ax.set_title('Electrostatic potential of graphite')
im = ax.imshow(potential, cmap='inferno')
cbar = plt.colorbar(im)
cbar.set_label('Electrostatic potential ($V \cdot \AA$)')