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

# # Plotting Density of States with <br />  Plotly and Pymatgen

# ### About the Author:
# 
# This notebook was contributed by plotly user [Germain Salvato-Vallverdu](https://plotly.com/~gvallverdu).  You can follow him on [GitHub](https://github.com/gVallverdu?tab=activity) or visit [his webpage](http://gvallver.perso.univ-pau.fr/).

# ### Introduction

# This notebook will go over an example for plotting the density of states and the band diagram of Silicon using python with [pymatgen](http://pymatgen.org/) and [plotly](https://plotly.com/) packages.
# 
# * [pymatgen](http://pymatgen.org/) : (Python Materials Genomics) A robust, open-source Python library for materials analysis.
# * [plotly](https://plotly.com/) : A platform for publishing beautiful, interactive graphs from Python to the web.

# Firstly, we import numpy and pymatgen tools

# In[1]:


import numpy as np
from pymatgen.io.vaspio.vasp_output import Vasprun    # read vasprun.xml output file of VASP
from pymatgen.electronic_structure.core import Spin


# Next, we import plotly tools

# In[2]:


import plotly.plotly as pltly      # plotting functions
import plotly.tools as tls         # plotly tools
import plotly.graph_objs as go     # plot and configuration tools : Scatter, Line, Layout


# # Band Diagram and Denisty of States of Silicon

# ### 1) Plot the density of states

# First, read projected density of states (DOS) from a VASP calculation on ``"./DOS/vasprun.xml"`` file using pymatgen.

# In[3]:


dosrun = Vasprun("./DOS/vasprun.xml")
spd_dos = dosrun.complete_dos.get_spd_dos()


# Set up the scatter plots for the total DOS and the contribution of 3s and 3p atomic orbital to the total DOS.

# In[4]:


# total DOS
trace_tdos = go.Scatter(
    x=dosrun.tdos.densities[Spin.up],
    y=dosrun.tdos.energies - dosrun.efermi,
    mode="lines",
    name="total DOS",
    line=go.Line(color="#444444"),
    fill="tozeroy"
)
# 3s contribution to the total DOS
trace_3s = go.Scatter(
    x=spd_dos["S"].densities[Spin.up],
    y=dosrun.tdos.energies - dosrun.efermi,
    mode="lines",
    name="3s",
    line=go.Line(color="red")
)
# 3p contribution to the total DOS
trace_3p = go.Scatter(
    x=spd_dos["P"].densities[Spin.up],
    y=dosrun.tdos.energies - dosrun.efermi,
    mode="lines",
    name="3p",
    line=go.Line(color="green")
)
dosdata = go.Data([trace_tdos, trace_3s, trace_3p])


# Customize axes and general aspect of the plot.

# In[5]:


dosxaxis = go.XAxis(
    title="Density of states",
    showgrid=True,
    showline=True,
    range=[.01, 3],
    mirror="ticks",
    ticks="inside",
    linewidth=2,
    tickwidth=2
)
dosyaxis = go.YAxis(
    title="$E - E_f \quad / \quad \\text{eV}$",
    showgrid=True,
    showline=True,
    ticks="inside",
    mirror='ticks',
    linewidth=2,
    tickwidth=2,
    zerolinewidth=2
)
doslayout = go.Layout(
    title="Density of states of Silicon",
    xaxis=dosxaxis,
    yaxis=dosyaxis
)


# In[6]:


dosfig = go.Figure(data=dosdata, layout=doslayout)
plot_url = pltly.plot(dosfig, filename="DOS_Si", auto_open=False)
tls.embed(plot_url)


# ### 2) Plot the band diagram

# First, read bands from a VASP calculation on ``"./Bandes/vasprun.xml"`` file using pymatgen.

# In[7]:


run = Vasprun("./Bandes/vasprun.xml", parse_projected_eigen = True)
bands = run.get_band_structure("./Bandes/KPOINTS", line_mode=True, efermi=dosrun.efermi)


# Look for the boundaries of the band diagram in order to set up y axes range.

# In[8]:


emin = 1e100
emax = -1e100
for spin in bands.bands.keys():
    for band in range(bands.nb_bands):
        emin = min(emin, min(bands.bands[spin][band]))
        emax = max(emax, max(bands.bands[spin][band]))
emin = emin - bands.efermi - 1
emax = emax - bands.efermi + 1


# Each band is plotted using a scatter plot.

# In[20]:


kptslist = [k for k in range(len(bands.kpoints))]
bandTraces = list()
for band in range(bands.nb_bands):
    bandTraces.append(
        go.Scatter(
            x=kptslist,
            y=[e - bands.efermi for e in bands.bands[Spin.up][band]],
            mode="lines",
            line=go.Line(color="#666666"),
            showlegend=False
        )
    )


# Add vertical lines at each high symmetry k-points and a label at the bottom.

# In[21]:


labels = [r"$L$", r"$\Gamma$", r"$X$", r"$U,K$", r"$\Gamma$"]
step = len(bands.kpoints) / (len(labels) - 1)
# vertical lines
vlines = list()
for i, label in enumerate(labels):
    vlines.append(
        go.Scatter(
            x=[i * step, i * step],
            y=[emin, emax],
            mode="lines",
            line=go.Line(color="#111111", width=1),
            showlegend=False
        )
    )
# Labels of highsymetry k-points are added as Annotation object
annotations = list()
for i, label in enumerate(labels):
    annotations.append(
        go.Annotation(
            x=i * step, y=emin,
            xref="x1", yref="y1",
            text=label,
            xanchor="center", yanchor="top",
            showarrow=False
        )
    )


# Customize axes and general aspect of the plot.

# In[22]:


bandxaxis = go.XAxis(
    title="k-points",
    range=[0, len(bands.kpoints)],
    showgrid=True,
    showline=True,
    ticks="", 
    showticklabels=False,
    mirror=True,
    linewidth=2
)
bandyaxis = go.YAxis(
    title="$E - E_f \quad / \quad \\text{eV}$",
    range=[emin, emax],
    showgrid=True,
    showline=True,
    zeroline=True,
    mirror="ticks",
    ticks="inside",
    linewidth=2,
    tickwidth=2,
    zerolinewidth=2
)
bandlayout = go.Layout(
    title="Bands diagram of Silicon",
    xaxis=bandxaxis,
    yaxis=bandyaxis,
    annotations=go.Annotations(annotations)
)


# In[24]:


bandfig = go.Figure(data=bandTraces + vlines, layout=bandlayout)
plot_url = pltly.plot(bandfig, filename="Bands_Si", auto_open=False)
tls.embed(plot_url)


# ### 3) Use subplots to plot DOS and bands together

# We will now make a figure with both the band diagram and the density of states using the ``make_subplots`` facility. First, we set up a figure with two columns, one row. At the end, the two plots will share y axis.

# In[25]:


dosbandfig = tls.make_subplots(rows=1, cols=2, shared_yaxes=True)


# We use the previously defined traces for the band and the densities of state and add them to the figure object.
# 
# * The bands are plotted on the left subplot (1, 1)
# * The densities of states are plotted on the right subplot (1, 2)

# In[26]:


# add the bands
for btrace in bandTraces:
    dosbandfig.append_trace(btrace, 1, 1)
# add vlines for specific k-points
for vline in vlines:
    dosbandfig.append_trace(vline, 1, 1)
# add the densities
dosbandfig.append_trace(trace_tdos, 1, 2)
dosbandfig.append_trace(trace_3s, 1, 2)
dosbandfig.append_trace(trace_3p, 1, 2)    


# Customize axes and general aspect of the plot using previously defined axis and layout options. The ``"domain"`` of each subplot is also define.

# In[27]:


dosbandfig["layout"].update(
    go.Layout(
        title="Bands diagram and density of states of Silicon",
        xaxis1=bandxaxis,
        yaxis1=bandyaxis,
        xaxis2=dosxaxis,
        annotations=go.Annotations(annotations)
    )
)
# adjust size of subplots
dosbandfig["layout"]["xaxis1"]["domain"] = [0., 0.7]
dosbandfig["layout"]["xaxis2"]["domain"] = [0.702, 1.]
# add some specific options
dosbandfig["layout"]["yaxis1"]["mirror"] = "allticks"
dosbandfig["layout"]["xaxis2"]["mirror"] = "allticks"


# In[28]:


plot_url = pltly.plot(dosbandfig, filename="DOS_bands_Si", auto_open=False)
tls.embed(plot_url)


# ### 4) Atomic orbital contributions in bands using a color scale

# As previously done for the density of states, the contribution of 3s and 3p atomic orbital may be highlighted using a color scale. The main idea, from [this example](https://github.com/vossjo/ase-espresso/wiki/Band-structure-calculation-example), is to normalize atomic orbital contributions and build the RGB code of the color from these contributions.
# 
# Thus, we first compute atomic orbital normalized contributions from projected bands :

# In[29]:


# extract projected bands
name = "Si"
pbands = bands.get_projections_on_elts_and_orbitals({name: ["s", "p", "d"]})

# compute contributions
contrib = np.zeros((bands.nb_bands, len(bands.kpoints), 3))                        
for band in range(bands.nb_bands):
    for k in range(len(bands.kpoints)):
        sc = pbands[Spin.up][band][k][name]["s"]**2
        pc = pbands[Spin.up][band][k][name]["p"]**2
        dc = pbands[Spin.up][band][k][name]["d"]**2
        tot = sc + pc + dc
        if tot != 0.0:
            contrib[band, k, 0] = sc / tot
            contrib[band, k, 1] = pc / tot
            contrib[band, k, 2] = dc / tot


# Now for each band, for each couple of consecutive points, we plot a line plot with a specific color, defined from the atomic orbital contributions to the current band.

# In[30]:


colorBands = list() # will contain the list of all lines
nkpts = len(bands.kpoints)
for band in range(bands.nb_bands):
    eband = [e - bands.efermi for e in bands.bands[Spin.up][band]]
    for k in range(nkpts - 1):
        red, green, blue = [int(255 * (contrib[band, k, i] + contrib[band, k+1, i])/2) for i in range(3)]
        colorBands.append(
            go.Scatter(
                x=[k, k+1],
                y=[eband[k], eband[k+1]],
                mode="lines",
                line=go.Line(color="rgb({}, {}, {})".format(red, green, blue)),
                showlegend=False
            )
        )


# As previously, two subplots are used to plot the band diagram on the left and the density of states on the right.

# In[31]:


# set up a new figure with two subplots
colorbandfig = tls.make_subplots(rows=1, cols=2, shared_yaxes=True)
# add the bands in the first subplot
for btrace in colorBands:
    colorbandfig.append_trace(btrace, 1, 1)
# add vlines for specific k-points in the first subplot
for vline in vlines:
    colorbandfig.append_trace(vline, 1, 1)
# add the densities in the second subplot
colorbandfig.append_trace(trace_tdos, 1, 2)
colorbandfig.append_trace(trace_3s, 1, 2)
colorbandfig.append_trace(trace_3p, 1, 2)
# Layout configuration
colorbandfig["layout"].update(
    go.Layout(
        title="Bands diagram and density of states of Silicon",
        xaxis1=bandxaxis,
        yaxis1=bandyaxis,
        xaxis2=dosxaxis,
        annotations=go.Annotations(annotations)
    )
)
# adjust size of subplots
colorbandfig["layout"]["xaxis1"]["domain"] = [0., 0.7]
colorbandfig["layout"]["xaxis2"]["domain"] = [0.702, 1.]
# add some specific options
colorbandfig["layout"]["yaxis1"]["mirror"] = "allticks"
colorbandfig["layout"]["xaxis2"]["mirror"] = "allticks"
# add a custom legend
legend = go.Legend(
    x=.98, y=.98,
    xanchor="right", yanchor="top",
    bordercolor='#333', borderwidth=1
)
colorbandfig["layout"]["legend"] = legend
# do the plot
plot_url = pltly.plot(colorbandfig, filename="DOS_bands_Si_color", auto_open=False)
tls.embed(plot_url)


# In[2]:


from IPython.display import HTML, display

display(HTML('<link href="//fonts.googleapis.com/css?family=Open+Sans:600,400,300,200|Inconsolata|Ubuntu+Mono:400,700" rel="stylesheet" type="text/css" />'))
display(HTML('<link rel="stylesheet" type="text/css" href="https://help.plot.ly/documentation/all_static/css/ipython-notebook-custom.css">'))

import publisher
publisher.publish('density-of-states ', '/ipython-notebooks/density-of-states/', 
                  'plotting the density of states and the band diagram using pymatgen and plotly', 
                  'Plotting Density of States with Plotly and Pymatgen')


# In[ ]: