from ovito.data import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
def modify(frame, data):
grid = data.grids['binning[Potential Energy]']
array = grid['Potential Energy']
array_2d = np.reshape(array, (grid.shape[1], grid.shape[0]) )
plt.figure()
matplotlib.rcParams.update({'font.size': 18})
im = plt.imshow(array_2d, cmap='jet', aspect='equal', origin='lower')
plt.xlabel('coordinate 1')
plt.ylabel('coordinate 2')
plt.colorbar(shrink = 0.65)
plt.savefig('{}{}.pdf'.format(grid.title, frame), bbox='tight')