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

# # Gaussian input/output treatments

# ## using pymatgen library
# 
# Germain Salvato Vallverdu, March 23, 2023

# ### import used modules and objects

# Load some pymatgen specific classes and module

# In[1]:


import pymatgen.core


# In[2]:


pymatgen.core.__version__


# In[3]:


from pymatgen.core import Molecule, Site
from pymatgen.io.gaussian import GaussianInput, GaussianOutput


# Image class of IPython for rendering pictures.

# In[4]:


from IPython.display import Image


# ### Generate gaussian inputs

# Hereafter is the docstring of the ```GaussianInput``` class.

# In[3]:


help(GaussianInput)


# #### Examples

# In this example, we get a $PO_2N^{2-}$ molecule and look for the best place for two Li atoms in the plane of the molecule. This example shows how to set up a Gaussian input file for all possible combinations. Coordinates of this molecule are available in xyz format :

# In[5]:


mol = Molecule.from_file("data_nglview/PO2N.xyz")
print(mol.to(fmt="xyz"))
Image("data_nglview/PO2N.png")


# First we set up the three guessed positions for Li atoms in the plane of the molecule. Each position is along the bissector of two PO or PN bonds. We build a list of ```pymatgen.Site``` object which represent an atom with its symbol, coordinates ....

# In[6]:


posLi = [Site("Li", [-0.424, -0.001,  0.267]),
         Site("Li", [-3.789, -2.031, -0.046]),
         Site("Li", [-3.787,  2.032, -0.044])]
print(posLi)


# Now for each combinaison of two Li atoms among three possibilities we will write a gaussian input file. Look at the ```combinations``` method of ```itertools``` module.

# In[7]:


from itertools import combinations


# In[8]:


for sites in combinations(posLi, 2):
    print(sites, "\t indices : ", [posLi.index(site) for site in sites])


# Hereafter is the loop over combinations which will write all input files.

# In[9]:


# Gaussian keywords
route_parms = {"opt": "tight", "Freq": ""}
link0 = {"%Nproc": "5"}
DFT = "B3LYP"

# the loop
for sites in combinations(posLi, 2):
    # load the molecule
    mol = Molecule.from_file("data_nglview/PO2N.xyz")
    
    # set up a label according to the combination
    title = "combination_" + "".join([str(posLi.index(site)) for site in sites])
    print(title)
    
    # add the two Li atoms to the molecule
    for site in sites:
        mol.append(site.specie, site.coords)
        
    # set up the calculation
    gau = GaussianInput(mol, charge=0, spin_multiplicity=1, title=title, 
                        functional=DFT, route_parameters=route_parms, link0_parameters=link0)
    gau.write_file(title + ".com", cart_coords=True)


# We can also load an input file and get some information.

# In[10]:


gau = GaussianInput.from_file("combination_01.com")
print(gau.basis_set, gau.functional)


# In[11]:


print(gau.to_string(cart_coords=True))


# ### Post-treatments of Gaussian outputs

# #### Docstring of the ```GaussianOutput``` class

# In[11]:


help(GaussianOutput)


# #### Examples :

# Used the ```GaussianOutput``` class to read a Gaussian output file.

# In[12]:


logfile = GaussianOutput("data_nglview/config_0123.log")


# Is termination ```Normal``` ?

# In[13]:


logfile.properly_terminated


# Display final energy or mulliken charges :

# In[14]:


logfile.final_energy


# In[15]:


logfile.Mulliken_charges


# You can extract the coordinates of the final structure. All structures are ```pymatgen.Molecule``` objects, see [here](http://pymatgen.org/_static/Molecule.html) or [the reference guide](http://pymatgen.org/pymatgen.core.html#pymatgen.core.structure.Molecule). You can do lot of stuff with this object such as : neighbors list, compute distances or the distance matrix, symmetry operation, atomic substitution ...

# In[16]:


print(logfile.final_structure.to(fmt="xyz"))


# You would prefer a gaussian input format with a z-matrix of the 42th geometrical optimization step.

# In[17]:


print(logfile.structures[41].to(fmt="gjf"))


# You can plot the geometrical convergence (here using plotly).

# In[18]:


import plotly.graph_objs as go


# In[19]:


fig = go.Figure([
    go.Scatter(x=[i for i in range(len(logfile.energies))], y=logfile.energies)
],
    layout=go.Layout(
        xaxis=dict(title="step"), 
        yaxis=dict(title="Energy (ua)"),
        template="plotly_white",
        height=500,
    )
)
fig.show()


# In[20]:


import matplotlib.pyplot as plt


# In[21]:


plt.plot(logfile.energies, "o--")
plt.xlabel("step")
plt.ylabel("Energy (ua)")