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

# # Helmholtz Free Energy
# The thermodynamic phase stability can be evaluated by comparing the Helmholtz free energy of different phases. So being able to calculate the Helmholtz free energy is an essential step towards calculating the full phase diagram. In the following there approximations are introduced to calculate the Helmholtz free energy, starting with the harmonic approximation, followed by the quasi-harmonic approximation which also includes the volume expansion and finally thermodynamic integration is used to quantify the anharmonic contributions. This addiabative approach to calculate the Helmholtz free energy is typically used in combination with ab-initio simulation methods like density functional theory, still here the approach is demonstrated with the [LAMMPS](https://lammps.org/) simulation code and the Morse potential.

# Starting with the import of the required python modules: 

# In[1]:


from ase.build import bulk
import numpy as np
from atomistics.workflows import optimize_positions_and_volume, LangevinWorkflow, PhonopyWorkflow, QuasiHarmonicWorkflow
from atomistics.calculators import evaluate_with_lammps_library, evaluate_with_lammps, get_potential_by_name, evaluate_with_hessian
from pylammpsmpi import LammpsASELibrary
from phonopy.units import VaspToTHz
from tqdm import tqdm
import matplotlib.pyplot as plt
import pandas
import scipy.constants 


# Followed by the selection of the reference element and the definition of the Morse interatomic potential:

# In[2]:


structure_bulk = bulk("Al", cubic=True)
df_pot_selected = pandas.DataFrame({
    "Config": [[
        "pair_style morse/smooth/linear 9.0",
        "pair_coeff * * 0.5 1.8 2.95"
    ]],
    "Filename": [[]],
    "Model": ["Morse"],
    "Name": ["Morse"],
    "Species": [["Al"]],
})


# As a prerequisite the volume and positions of the atomistic structure are optimized to calculate the Helmholtz of the equilibrium structure:

# In[3]:


task_dict = optimize_positions_and_volume(structure=structure_bulk)
result_dict = evaluate_with_lammps(
    task_dict=task_dict,
    potential_dataframe=df_pot_selected,
)
structure_opt = result_dict["structure_with_optimized_positions_and_volume"]


# Finally, the size of the atomistic structure is increased by repeating the structure in all three directions three times: 

# In[4]:


structure = structure_opt.repeat([3, 3, 3])


# The increased structure is required for all three approximations, for the harmonic and quasi-harmonic approximation it is required to evaluate the phonons up to an interaction range of 10 Angstrom and for the thermodynamic integration the larger supercell provides improved thermodynamic stability when calculating molecular dynamics trajectories at finite temperatures.

# ## Harmonic Approximation

# The harmonic approximation is calculated using the finite displacement method using the [phonopy](https://phonopy.github.io/phonopy/) package. In the `atomistics` package the [phonopy](https://phonopy.github.io/phonopy/) workflow is implemented in the `PhonopyWorkflow` object. Following the typical three step process of generating the structures `generate_structures()` evaluating them with the [LAMMPS](https://lammps.org/) simulation code using the `evaluate_with_lammps()` function and analysing the results using the `analyse_structures()` function the thermodynamic properties can be calculated using the `get_thermal_properties()` function. 

# In[5]:


workflow_fix = PhonopyWorkflow(
    structure=structure_opt,
    interaction_range=10,
    factor=VaspToTHz,
    displacement=0.01,
    dos_mesh=20,
    primitive_matrix=None,
    number_of_snapshots=None,
)
structure_dict = workflow_fix.generate_structures()
result_dict = evaluate_with_lammps(
    task_dict=structure_dict,
    potential_dataframe=df_pot_selected,
)
workflow_fix.analyse_structures(output_dict=result_dict)
term_base_dict = workflow_fix.get_thermal_properties(
    t_min=1,
    t_max=1500,
    t_step=250,
    temperatures=None,
    cutoff_frequency=None,
    pretend_real=False,
    band_indices=None,
    is_projection=False,
)
term_base_dict


# The dictionary returned by the `get_thermal_properties()` function contains the temperature dependence of the following thermodynamic properties: Helmholtz free energy `free_energy`, Entropy `entropy` and the heat capcity at constant volume `heat_capacity`. In addition, the temperature is also included in the result dictionary to simplify the plotting of the results. The results are compared to the other approximations in the following.

# ## Quasi-Harmonic Approximation

# The quasi-harmonic approximation extends the harmonix approximation by including the volume expansion. In the `atomistics` package this is implemented with the `QuasiHarmonicWorkflow` which internally is a combination of the `EnergyVolumeCurveWorkflow` and the `PhonopyWorkflow` introduced above. Consequently, the input parameters for the `QuasiHarmonicWorkflow` are a superset of the inputs of these workflows: 

# In[6]:


workflow_qh = QuasiHarmonicWorkflow(
    structure=structure_opt,
    num_points=11,
    vol_range=0.05,
    interaction_range=10,
    factor=VaspToTHz,
    displacement=0.01,
    dos_mesh=20,
    primitive_matrix=None,
    number_of_snapshots=None,
)
structure_dict = workflow_qh.generate_structures()
result_dict = evaluate_with_lammps(
    task_dict=structure_dict,
    potential_dataframe=df_pot_selected,
)
workflow_qh.analyse_structures(output_dict=result_dict)
term_qh_dict = workflow_qh.get_thermal_properties(
    t_min=1,
    t_max=1500,
    t_step=250,
    temperatures=None,
    cutoff_frequency=None,
    pretend_real=False,
    band_indices=None,
    is_projection=False,
    quantum_mechanical=True,
)
term_qh_dict


# In analogy to the `get_thermal_properties()` function of the `PhonopyWorkflow` the `get_thermal_properties()` function of the `QuasiHarmonicWorkflow` returns a dictionary with the temperature dependence of the thermodynamic properties. The volume at finite temperature is included as an additional output. 

# Furthermore, the `QuasiHarmonicWorkflow` adds the ability to calculate the volume expansion with the classical harmonic oscillator rather than the quantum mechanical osicllator which is used by default:

# In[7]:


term_qh_cl_dict = workflow_qh.get_thermal_properties(
    t_min=1,
    t_max=1500,
    t_step=250,
    temperatures=None,
    cutoff_frequency=None,
    pretend_real=False,
    band_indices=None,
    is_projection=False,
    quantum_mechanical=False,
)
term_qh_cl_dict


# In this case the dictionary of thermal properties calculated by the `get_thermal_properties()` function only contains the Helmholtz free energy `free_energy`, the volume at finite temperature `volumes` and the corresponding temperature `termpatures`.

# For the comparison of the harmonic approximation and the quasi-harmonic approximation the thermodynamic properties are visualized in the following. Given the coarse sampling of the temperature dependence these graphs look a bit rough. A better choice for the temperature step would be 50K rather than 250K. Here the larger temperature step is chosen to minimize the computational cost. 

# In[8]:


plt.title("Helmholtz Free Energy")
plt.plot(term_base_dict['temperatures'], term_base_dict['free_energy'], label="Harmonic")
plt.plot(term_qh_dict['temperatures'], term_qh_dict['free_energy'], label="Quasi-Harmonic (qm)")
plt.plot(term_qh_cl_dict['temperatures'], term_qh_cl_dict['free_energy'], label="Quasi-Harmonic (cl)")
plt.xlabel("Temperature (K)")
plt.ylabel("Helmholtz Free Energy (eV)")
plt.legend()


# In[9]:


plt.title("Entropy")
plt.plot(term_base_dict['temperatures'], term_base_dict['entropy'], label="harmonic")
plt.plot(term_qh_dict['temperatures'], term_qh_dict['entropy'], label="quasi-harmonic (qm)")
plt.xlabel("Temperature (K)")
plt.ylabel("Entropy (J/K/mol)")
plt.legend()


# In[10]:


plt.title("Heat Capacity")
plt.plot(term_base_dict['temperatures'], term_base_dict['heat_capacity'], label="harmonic")
plt.plot(term_qh_dict['temperatures'], term_qh_dict['heat_capacity'], label="quasi-harmonic")
plt.axhline(3 * scipy.constants.Boltzmann * scipy.constants.Avogadro * len(structure_opt), color="black", linestyle="--", label="$3Nk_B$")
plt.xlabel("Temperature (K)")
plt.ylabel("Heat Capacity (J/K/mol)")
plt.legend()


# ## Thermo Dynamic Integration

# To include the anharmonic contribution to the free energy thermo dynamic integration is used to integrate the internal energy between the quasi-harmonic reference and the full anharmonic molecular dynamics calculation. For the choice of computational efficiency molecular dynamics trajectories with 3000 steps are using and a set of five lambda values between 0.0 and 1.0. Again, increasing the number of lambda values from 5 to 11 can improve the approximation.

# In[11]:


steps = 3000
steps_lst = list(range(steps))
lambda_lst = np.linspace(0.0, 1.0, 5)


# From the finite temperature volume of the `QuasiHarmonicWorkflow` above the finite temperature lattice constant is calculated:

# In[12]:


lattice_constant_lst = np.array(term_qh_dict['volumes']) ** (1/3)
temperature_lst = term_qh_dict['temperatures']


# The thermodynamic integration workflow consists of two loops. The first loop iterates over the different temperatures and the corresponding finite temperature lattice constants while the second inner loop iterates over the different lambda parameters. In the outter loop the `PhonopyWorkflow` workflow is used to calculate the finite temperature force constants which can be accessed from the `PhonopyWorkflow` using the `get_hesse_matrix()` function. In addition the `evaluate_with_lammps()` function is used in the outter loop as well to calculate the equilibrium energy at finite volume. Finally, in the inner loop the `LangevinWorkflow` is used to calculate the molecular dynamics trajectory with both the forces and the energy being mixed in dependence of the lambda parameter from the Hessian calculation `evaluate_with_hessian()` and the [LAMMPS](https://lammps.org/) calculation `evaluate_with_lammps_library()`. Here the [LAMMPS](https://lammps.org/) library interface is used to evaluate multiple structures one after another. Finally the lambda dependence is fitted and the integral from 0 to 1 is calculated:

# In[13]:


free_energy_lst, eng_lambda_dependence_lst = [], []
for lattice_constant, temperature in zip(lattice_constant_lst, temperature_lst):
    structure = bulk("Al", a=lattice_constant, cubic=True).repeat([3, 3, 3])
    equilibrium_lammps = evaluate_with_lammps(task_dict={"calc_energy": structure}, potential_dataframe=df_pot_selected)['energy']
    workflow_fix = PhonopyWorkflow(
        structure=structure,
        interaction_range=10,
        factor=VaspToTHz,
        displacement=0.01,
        dos_mesh=20,
        primitive_matrix=None,
        number_of_snapshots=None,
    )
    structure_dict = workflow_fix.generate_structures()
    result_dict = evaluate_with_lammps(
        task_dict=structure_dict,
        potential_dataframe=df_pot_selected,
    )
    workflow_fix.analyse_structures(output_dict=result_dict)
    energy_pot_all_lst, energy_mean_lst, energy_kin_all_lst = [], [], []
    for lambda_parameter in lambda_lst: 
        thermo_eng_pot_lst, thermo_eng_kin_lst = [], []
        workflow_md_thermo = LangevinWorkflow(
            structure=structure,
            temperature=temperature,
            overheat_fraction=2.0,
            damping_timescale=100.0,
            time_step=1,
        )
        lmp = LammpsASELibrary(
            working_directory=None,
            cores=1,
            comm=None,
            logger=None,
            log_file=None,
            library=None,
            diable_log_file=True,
        )
        for i in tqdm(steps_lst):
            task_dict = workflow_md_thermo.generate_structures()
            hessian_dict = evaluate_with_hessian(
                task_dict=task_dict,
                structure_equilibrium=structure,
                force_constants=workflow_fix.get_hesse_matrix(),
                bulk_modulus=0,
                shear_modulus=0,
            )
            lammps_dict = evaluate_with_lammps_library(
                task_dict=task_dict,
                potential_dataframe=df_pot_selected,
                lmp=lmp,
            )
            result_dict = {
                "forces": {0: (1-lambda_parameter) * hessian_dict["forces"][0] + lambda_parameter * lammps_dict["forces"][0]},
                "energy": {0: (1-lambda_parameter) * hessian_dict["energy"][0] + lambda_parameter * (lammps_dict["energy"][0] - equilibrium_lammps)},
            }
            eng_pot, eng_kin = workflow_md_thermo.analyse_structures(output_dict=result_dict)
            thermo_eng_pot_lst.append(eng_pot)
            thermo_eng_kin_lst.append(eng_kin)
        lmp.close()
        thermo_energy = np.array(thermo_eng_pot_lst) + np.array(thermo_eng_kin_lst)
        energy_mean_lst.append(np.mean(thermo_energy[1000:]))
        energy_pot_all_lst.append(thermo_eng_pot_lst)
        energy_kin_all_lst.append(thermo_eng_kin_lst)
    eng_lambda_dependence_lst.append(np.array(energy_mean_lst) / len(structure) * 1000)
    fit = np.poly1d(np.polyfit(lambda_lst, np.array(energy_mean_lst) / len(structure) * 1000, 3))
    integral = np.polyint(fit)
    free_energy_lst.append(integral(1.0) - integral(0.0))


# ## Summary

# The Helmholtz free energy for all three approximations is compared in the following:

# In[14]:


plt.title("Helmholtz Free Energy")
plt.plot(term_base_dict['temperatures'], term_base_dict['free_energy'], label="Harmonic")
plt.plot(term_qh_dict['temperatures'], term_qh_dict['free_energy'], label="Quasi-Harmonic")
plt.plot(term_qh_dict['temperatures'], term_qh_dict['free_energy'] - np.array(free_energy_lst) / 1000, label="Thermodynamic Integration")
plt.xlabel("Temperature (K)")
plt.ylabel("Helmholtz Free Energy (eV)")
plt.legend()


# In[ ]: