#!/usr/bin/env python
# coding: utf-8

# # Energy volume curve
# ## Theory 
# Fitting the energy volume curve allows to calculate the equilibrium energy $E_0$, the equilirbium volume $V_0$, the equilibrium bulk modulus $B_0$ and its derivative $B^{'}_0$. These quantities can then be used as part of the Einstein model to get an initial prediction for the thermodynamik properties, the heat capacity $C_v$ and the free energy $F$. 

# ## Initialisation

# We start by importing matplotlib, numpy and the pyiron project class. 

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
import numpy as np
from pyiron import Project


# In the next step we create a project, by specifying the name of the project.

# In[2]:


pr = Project(path='thermo')


# ## Atomistic structure

# To analyse the energy volume dependence a single super cell is sufficient, so we create an iron super cell as an example.

# In[3]:


basis = pr.create_structure(element='Fe', bravais_basis='bcc', lattice_constant=2.75)
basis.plot3d()


# ## Calculation
# Energy volume curves are commonly calculated with ab initio codes, so we use VASP in this example. But we focus on the generic commands so the same example works with any DFT code. We choose 'vasp' as job name prefix, select an energy cut off of $320 eV$ and assign the basis to the job. Afterwards we apply the corresponding strain.

# In[4]:


for strain in np.linspace(0.95, 1.05, 7):
    strain_str = str(strain).replace('.', '_')
    job_vasp_strain = pr.create_job(job_type=pr.job_type.Gpaw, job_name='gpaw_' + strain_str)
    job_vasp_strain.set_encut(320.0)
    job_vasp_strain.structure = basis.copy()
    job_vasp_strain.structure.set_cell(cell=basis.cell * strain ** (1/3), scale_atoms=True)
    job_vasp_strain.run()


# As these are simple calculation, there is no need to submit them to the queuing sytem. We can confirm the status of the calculation with the job_table. If the status of each job is marked as finished, then we can continue with the next step.

# In[5]:


pr.job_table()


# ## Analysis

# We aggregate the data for further processing in two separated lists, one for the volumes and one for the energies. To do so we iterate over the jobs within the project, filter the job names which contain the string 'vasp' and from those extract the final volume and the final energy. 

# In[6]:


volume_lst, energy_lst = zip(*[[job['output/generic/volume'][-1], job['output/generic/energy_pot'][-1]] 
                               for job in pr.iter_jobs(convert_to_object=False) if 'gpaw' in job.job_name])   


# We plot the aggregated data using matplotlib. 

# In[7]:


plt.plot(volume_lst, energy_lst, 'x-')
plt.xlabel('Volume ($\AA ^ 3$)')
plt.ylabel('Energy (eV)')


# ## Encut Dependence
# To extend the complexity of our simulation protocol we can not only iterate over different strains but also different energy cutoffs. For this we use multiple sub projects to structure the data. And we summarize the previous code in multiple functions to maintain a high level of readability. The first function calculates a specific strained configuration for an specifc energy cut off, while the second function analyses the different strained calculations for a specific energy cutoff and returns the list of energy volume pairs.

# ### Functions

# In[8]:


def vasp_calculation_for_strain(pr, basis, strain, encut):
    strain_str = str(strain).replace('.', '_')
    job_vasp_strain = pr.create_job(job_type=pr.job_type.Gpaw, job_name='gpaw_' + strain_str)
    job_vasp_strain.set_encut(encut)
    job_vasp_strain.structure = basis.copy()
    job_vasp_strain.structure.set_cell(cell=basis.cell * strain ** (1/3), scale_atoms=True)
    job_vasp_strain.run()


# In[9]:


def energy_volume_pairs(pr):
    volume_lst, energy_lst = zip(*[[job['output/generic/volume'][-1], job['output/generic/energy_pot'][-1]] 
                               for job in pr.iter_jobs(convert_to_object=False) if 'gpaw' in job.job_name])
    return volume_lst, energy_lst


# ### Calculation

# With these functions we can structure our code and implement the additional for loop to include multiple energy cutoffs.

# In[10]:


for encut in np.linspace(270, 320, 6):
    encut_str = 'encut_' + str(int(encut))
    pr_encut = pr.open(encut_str)
    for strain in np.linspace(0.95, 1.05, 7):
        vasp_calculation_for_strain(pr=pr_encut, 
                                    basis=basis, 
                                    strain=strain, 
                                    encut=encut)


# ### Analysis 

# The analysis is structured in a similar way. Here we use iter_groups() to iterate over the existing subprojects within our project and plot the individual energy volume curves using the functions defined above. 

# In[11]:


for pr_encut in pr.iter_groups():
    volume_lst, energy_lst = energy_volume_pairs(pr_encut)
    plt.plot(volume_lst, energy_lst, 'x-', label=pr_encut.base_name)
plt.xlabel('Volume ($\AA ^ 3$)')
plt.ylabel('Energy (eV)')
plt.legend()


# ## Fitting
# After we created multiple datasets we can now start to fit the converged results. While it is possible to fit the results using a simple polynomial fit we prefer to use the phyiscally motivated birch murnaghan equation or the vinet equation. For this we create the Murnaghan object and use it is fitting functionality:

# In[12]:


murn = pr.create_job(job_type=pr.job_type.Murnaghan, job_name='murn')


# ### Birch Marnaghan

# In[13]:


[e0, b0, bP, v0], [e0_error, b0_error, bP_error, v0_error] = murn._fit_leastsq(volume_lst=volume_lst,
                                                                               energy_lst=energy_lst,
                                                                               fittype='birchmurnaghan')
[e0, b0, bP, v0] 


# ### Vinet

# In[14]:


[e0, b0, bP, v0], [e0_error, b0_error, bP_error, v0_error] = murn._fit_leastsq(volume_lst=volume_lst,
                                                                               energy_lst=energy_lst,
                                                                               fittype='vinet')
[e0, b0, bP, v0] 


# We see that both equation of states give slightly different results, with overall good agreement. To validate the agreement we plot the with with the original data.

# In[15]:


vol_lst = np.linspace(np.min(volume_lst), np.max(volume_lst), 1000)
plt.plot(volume_lst, energy_lst, label='dft')
plt.plot(vol_lst, murn.fit_module.vinet_energy(vol_lst, e0, b0/ 160.21766208, bP, v0), label='vinet')
plt.xlabel('Volume ($\AA ^ 3$)')
plt.ylabel('Energy (eV)')
plt.legend()


# ## Murnaghan Module 
# Besides the fitting capabilities the Murnaghan module can also be used to run a set of calculations. For this we define a reference job, which can be either a Vasp calculation or any other pyiron job type and then specify the input parameters for the Murnaghan job.

# In[16]:


job_vasp_strain = pr.create_job(job_type=pr.job_type.Gpaw, job_name='gpaw')
job_vasp_strain.set_encut(320)
job_vasp_strain.structure = basis.copy()


# In[17]:


murn = pr.create_job(job_type=pr.job_type.Murnaghan, job_name='murn')
murn.ref_job = job_vasp_strain
murn.input


# We modify the input parameters to agree with the settings used in the examples above and execute the simulation by calling the run command on the murnaghan job object.

# In[18]:


murn.input['num_points'] = 7
murn.input['vol_range'] = 0.05


# In[19]:


type(murn.structure)


# In[20]:


pr.job_table()


# In[21]:


murn.run()


# Afterwards we can use the build in capabilites to plot the resulting energy volume curve and fit different equations of state to the calculated energy volume pairs.

# In[22]:


murn.output_to_pandas()


# In[23]:


murn.plot()


# In[24]:


murn.fit_vinet()


# ## Common mistakes 

# ### Not copying the basis 

# It is important to copy the basis before applying the strain, as the strain has to be applied on the initial structure, not the previous structure:

# In[25]:


volume_lst_with_copy = []
for strain in np.linspace(0.95, 1.05, 7):
    basis_copy = basis.copy()
    basis_copy.set_cell(cell=basis.cell * strain ** (1/3), scale_atoms=True)
    volume_lst_with_copy.append(basis_copy.get_volume())


# In[26]:


basis_copy = basis.copy()
volume_lst_without_copy = []
for strain in np.linspace(0.95, 1.05, 7):
    basis_copy.set_cell(cell=basis_copy.cell * strain ** (1/3), scale_atoms=True)
    volume_lst_without_copy.append(basis_copy.get_volume())


# In[27]:


volume_lst_with_copy, volume_lst_without_copy


# ### Rescaling the cell

# Another common issue is the rescaling of the supercell, there are multiple options to choose from. We used the option to scale the atoms with the supercell.

# In[28]:


basis_copy = basis.copy()
strain = 0.5
basis_copy.set_cell(cell=basis_copy.cell * strain ** (1/3), scale_atoms=True)
basis_copy.plot3d()


# A nother typical case is rescaling the cell to increase the distance between the atoms or add vacuum. But that is not what we want to fit an energy volume curve.

# In[29]:


basis_copy = basis.copy()
strain = 0.5
basis_copy.set_cell(cell=basis_copy.cell * strain ** (1/3), scale_atoms=False)
basis_copy.plot3d()


# The same can be achieved by setting the basis to relative coordinates.

# In[30]:


basis_copy = basis.copy()
strain = 0.5
basis_copy.set_relative()
basis_copy.cell *= strain ** (1/3)
basis_copy.plot3d()


# In[31]:


basis_copy = basis.copy()
strain = 0.5
basis_copy.cell *= strain ** (1/3)
basis_copy.plot3d()


# In[ ]: