import matplotlib.pyplot as plt
import matplotlib.cm as cm

################################################################
# Evolution of energies and volumes during 300 K equilibration
################################################################

# ------------------------------
# Initialize figure
# ------------------------------
fig = plt.figure(figsize=(8,8))
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)

# ------------------------------
# Read MD Trajectory
# ------------------------------
md_trajectory = nlread('Li0.4S_eq-300K.hdf5', MDTrajectory)[-1]

# ------------------------------
# Read total energies
# ------------------------------
energies = md_trajectory.kineticEnergies() + md_trajectory.potentialEnergies()
energies = energies.inUnitsOf(eV)
x = range(len(energies))

# ------------------------------
# Read cell volumes
# ------------------------------
volumes = []
for i in x:
	image = md_trajectory.image(i)
	lattice = image.bravaisLattice()
	v = lattice.unitCellVolume().inUnitsOf(Angstrom**3)
	volumes.append(v)

# ------------------------------
# Normalize energies and volumes and compute mean values
# ------------------------------
energies = numpy.asarray(energies)/len(image)
volumes  = numpy.asarray(volumes) /len(image)
mean1 = numpy.mean(energies[-100:])
mean2 = numpy.mean(volumes[-100:])

# ------------------------------
# Plot data
# ------------------------------
ax1.plot(x, energies, '-', color='g', label='x=0.40')
ax1.axhline(mean1, color='k', ls='--', lw=3, zorder=10)
ax2.plot(x, volumes, '-', color='g')
ax2.axhline(mean2, color='k', ls='--', lw=3, zorder=10)

# ------------------------------
# Finalize plot and save+show it
# ------------------------------
ax1.set_xlabel('MD step')
ax1.set_ylabel('Energy per atom (eV)')
ax2.set_xlabel('MD step')
ax2.set_ylabel('Volume per atom (Ang^3)')
ax1.legend(loc='upper right')
plt.tight_layout()
plt.savefig('evolution.png')
plt.show()