%matplotlib inline
import json
import numpy as np
import matplotlib.pyplot as plt

input_file = 'output.json'

# Parameters
input_file = "output.json"

with open(input_file, 'r') as f:
    output_dict = json.load(f)

plt.title('Energy Volume curve')
plt.plot(output_dict["volume_lst"], output_dict["energy_lst"], 'x-')
plt.axvline(output_dict["equilibrium_volume"], linestyle='--')
plt.axhline(output_dict["equilibrium_energy"], linestyle='--')
plt.xlabel('Volume $\AA ^ 3$')
plt.ylabel('Energy eV');

print({'equilibrium_b_prime': output_dict["equilibrium_b_prime"], 
       'equilibrium_bulk_modulus': output_dict["equilibrium_bulk_modulus"], 
       'equilibrium_energy': output_dict["equilibrium_energy"], 
       'equilibrium_volume': output_dict["equilibrium_volume"]})

for k in output_dict.keys(): 
    if 'struct' in k:
        plt.plot(output_dict[k][0], output_dict[k][1], label=str(np.round(output_dict["volume_lst"][int(k.split('_')[1])], 3)))
plt.legend()
plt.xlabel('Frequency [THz]')
plt.ylabel('DOS')
plt.title('Phonon DOS vs Energy')