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

# # Matminer introduction - Predicting bulk modulus
# 
# This notebook demonstrates some of the basic features of matminer.
# 
# This notebook was last updated 06/07/21 for version 0.7.0 of matminer and verison 0.0.2 of figrecipes (which you can download by cloning the [figrecipes source repo](https://github.com/hackingmaterials/figrecipes)). 
# 
# 
# **Note that in order to get the in-line plotting to work, you might need to start Jupyter notebook with a higher data rate, e.g., ``jupyter notebook --NotebookApp.iopub_data_rate_limit=1.0e10``. We recommend you do this before starting.**
# 
# 
# ## Overview
# In this notebook, we will:
# 1. Load and examine a dataset in a pandas dataframe
# 2. Add descriptors to the dataframe using matminer
# 3. Train, compare, and visualize several machine learning methods with scikit-learn and matminer FigRecipes.
# 
# 
# 
# ## 1. Load and process data set
# 
# Matminer comes pre-loaded with several example data sets you can use. Below, we'll load a data set of computed elastic properties of materials which is sourced from the paper:  "Charting the complete elastic properties of inorganic crystalline compounds", M. de Jong *et al.*, Sci. Data. 2 (2015) 150009.

# In[1]:


from matminer.datasets.convenience_loaders import load_elastic_tensor
df = load_elastic_tensor()  # loads dataset in a pandas DataFrame object


# ## 1.1 A first look at the data set
# 
# The data set comes as a pandas DataFrame, which is a kind of "spreadsheet" object in Python. DataFrames have several useful methods you can use to explore and clean the data, some of which we'll explore below.

# In[2]:


df.head()


# ## 1.2 Removing unneeded columns from the data set
# 
# The data set above has many columns - we won't need all this data for our modeling. We'll mainly be trying to predict ``K_VRH`` and ``G_VRH`` (the Voight-Reuss-Hill average of the bulk and shear modulus, respectively) and the ``elastic_anisotropy``. We can drop most of the other output data. We also don't need various metadata such as the ``cif`` string of the structure since the crystal structure is already embedded in the ``structure`` column.

# In[3]:


df.columns


# In[4]:


unwanted_columns = ["volume", "nsites", "compliance_tensor", "elastic_tensor", 
                    "elastic_tensor_original", "K_Voigt", "G_Voigt", "K_Reuss", "G_Reuss"]
df = df.drop(unwanted_columns, axis=1)


# In[5]:


df.head()


# ## 1.3 Getting a feel for the data using descriptive statistics
# 
# A pandas DataFrame includes a function called ``describe()`` that helps determine statistics for the various numerical / categorical columns in the data.

# In[6]:


df.describe()


# Sometimes, the ``describe()`` function will reveal outliers that indicate mistakes in the data. For example, negative hence unphysical minimum bulk/shear moduli or maximum bulk/shear moduli that are too high. In this case, the data looks ok at first glance; meaing that there are no clear problems with the ranges of the various properties. Therefore, and we won't filter out any data.
# 
# Note that the ``describe()`` function only describes numerical columns by default.

# # 2. Add descriptors to the data ("featurization")
# 
# We are seeking to find relationships between the inputs (composition and crystal structure of a material) and outputs (elastic properties such as ``K_VRH``, ``G_VRH``, and ``elastic_anisotropy``). To find such relations, we need to "featurize" the input data such that they are numbers that meaningfully represent the underlying physical quantity. For example, one "feature" or "descriptor" of the composition of a material such as ``Nb4CoSi`` would be the standard deviation of the Pauling electronegativity of the elements in the compound (weighted by stoichiometry). Compositions with high values of this quantity would be more ionic and quantities with lower values would tend towards covalent or ionic. A descriptor of the crystal structure might be the average coordination number of sites; higher coordination numbers indicate more bonds and therefore might indicate stiffer materials. With matminer, we can start generating hundreds of possible descriptors using the descriptor library that is available. Data mining techniques can help narrow down which descriptors are the most relevant to the target problem using the available output data as a guide.

# ## 2.1 Add composition-based features
# 
# A major class of featurizers available in matminer uses the chemical composition to featurize the input data. Let's add some composition based features to our DataFrame.
# 
# The first step is to have a column representing the chemical composition as a pymatgen Composition object. One way to do this is to use the conversions Featurizers in matminer to turn a String composition (our ``formula`` column from before) into a pymatgen Composition.

# In[7]:


from matminer.featurizers.conversions import StrToComposition
df = StrToComposition().featurize_dataframe(df, "formula")
df.head()


# As you can see, we now have a new Composition column above. The visualization in the DataFrame is not so clear, but the column embeds both the elements *and* the amounts of each element in the composition (not shown).
# 
# Next, we'll use one of the featurizers in matminer to add a suite of descriptors to the DataFrame.

# In[8]:


from matminer.featurizers.composition import ElementProperty

ep_feat = ElementProperty.from_preset(preset_name="magpie")
df = ep_feat.featurize_dataframe(df, col_id="composition")  # input the "composition" column to the featurizer
df.head()


# As you can see, we now have many more columns in the DataFrame (a total of 141 columns - not all are shown). So we added *many* features to our data!
# 
# As an aside, note that each featurizer also has a ``citations()`` function that tells you where to find out more about the Featurizer.

# In[9]:


ep_feat.citations()


# ## 2.2 Add more composition-based features
# 
# There are many more Composition based featurizers apart from ``ElementProperty`` that are available in the ``matminer.featurizers.composition``. Let's try the ``ElectronegativityDiff`` featurizer which requires knowing the oxidation state of the various elements in the Composition. This information is not there currently, but one can use the ``conversions`` package to try to guess oxidation states and then apply the ``ElectronegativityDiff`` featurizer to this column.

# In[10]:


from matminer.featurizers.conversions import CompositionToOxidComposition
from matminer.featurizers.composition import OxidationStates

df = CompositionToOxidComposition().featurize_dataframe(df, "composition")

os_feat = OxidationStates()
df = os_feat.featurize_dataframe(df, "composition_oxid")
df.head()


# As you can see, the end of our data frame has now has some oxidation state based features!

# ## 2.3 Add some structure based features
# 
# Not all featurizers operate on compositions. Matminer can also analyze crystal structures and featurize those as well. Let's start by adding some simple density features.

# In[11]:


from matminer.featurizers.structure import DensityFeatures

df_feat = DensityFeatures()
df = df_feat.featurize_dataframe(df, "structure")  # input the structure column to the featurizer
df.head()


# In[12]:


df_feat.feature_labels()


# Again, we see more features in the last few columns: ``density``, ``vpa``, and ``packing fraction``.
# 
# There are many more structure based featurizers that are much more complex and detailed analysis of the crystal structure that are outside of the scope of this example. Let's move on to doing some machine learning predictions.

# # 3. Try some different machine learning models to relate input features to the bulk modulus
# 
# We now have enough data to do some machine learning! We'll need to first determine what columns we consider input (independent variables) and what column we consider output (dependent variable).
# 
# ## 3.1 Define input data and output data
# 
# For now, we'll use ``K_VRH`` (bulk modulus) as the output. You could retry this example with ``G_VRH``, ``elastic_anisotropy``, or ``poisson_ratio`` as the target output.
# 
# For the inputs, we'll use all the features we generated. That is, everything except the output data and non-numerical columns like ``composition`` and ``structure``.
# 
# 

# In[13]:


y = df['K_VRH'].values
excluded = ["G_VRH", "K_VRH", "elastic_anisotropy", "formula", "material_id", 
            "poisson_ratio", "structure", "composition", "composition_oxid"]
X = df.drop(excluded, axis=1)
print("There are {} possible descriptors:\n\n{}".format(X.shape[1], X.columns.values))


# ### 3.1 Try a linear regression model using scikit-learn
# 
# The scikit-learn library makes it easy to fit and cross-validate different types of regression models. Let's start with one of the simplest - a linear regression.

# In[14]:


from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import numpy as np

lr = LinearRegression()

lr.fit(X, y)

# get fit statistics
print('training R2 = ' + str(round(lr.score(X, y), 3)))
print('training RMSE = %.3f' % np.sqrt(mean_squared_error(y_true=y, y_pred=lr.predict(X))))


# This looks reasonable since linear regression is a simple (high bias) model. But to really validate that we are not over-fitting, we need to check the cross-validation score rather than the fitting score.

# In[16]:


from sklearn.model_selection import KFold, cross_val_score

# Use 10-fold cross validation (90% training, 10% test)
crossvalidation = KFold(n_splits=10, shuffle=True, random_state=1)
scores = cross_val_score(lr, X, y, scoring='neg_mean_squared_error', cv=crossvalidation, n_jobs=1)
rmse_scores = [np.sqrt(abs(s)) for s in scores]
r2_scores = cross_val_score(lr, X, y, scoring='r2', cv=crossvalidation, n_jobs=1)

print('Cross-validation results:')
print('Folds: %i, mean R2: %.3f' % (len(scores), np.mean(np.abs(r2_scores))))
print('Folds: %i, mean RMSE: %.3f' % (len(scores), np.mean(np.abs(rmse_scores))))


# A root-mean-squared error of ~22 GPa is not bad! Let's see what this looks like on a plot.
# 
# **Note that in order to get the code below to work, you might need to start Jupyter notebook with a higher data rate, e.g., ``jupyter notebook --NotebookApp.iopub_data_rate_limit=1.0e10``.**

# In[18]:


from figrecipes import PlotlyFig
from sklearn.model_selection import cross_val_predict

pf = PlotlyFig(x_title='DFT (MP) bulk modulus (GPa)',
               y_title='Predicted bulk modulus (GPa)',
               title='Linear regression',
               mode='notebook',
               filename="lr_regression.html")

pf.xy(xy_pairs=[(y, cross_val_predict(lr, X, y, cv=crossvalidation)), ([0, 400], [0, 400])], 
      labels=df['formula'], 
      modes=['markers', 'lines'],
      lines=[{}, {'color': 'black', 'dash': 'dash'}], 
      showlegends=False
     )


# Not too bad! However, there are definitely some outliers (you can hover over the points with your mouse to see what they are). We will see below (with a random forest model) how the accuracy changes when we try the model on a completely new test data.

# ## 3.2 Try a random forest model
# 
# Let's see if a more complex machine learning model does better. We can try a random forest model which is a good "starting" machine learning model, although one can usually do better. Let's repeat the steps for linear regression but for a random forest model.

# In[19]:


from sklearn.ensemble import RandomForestRegressor

rf = RandomForestRegressor(n_estimators=50, random_state=1)

rf.fit(X, y)
print('training R2 = ' + str(round(rf.score(X, y), 3)))
print('training RMSE = %.3f' % np.sqrt(mean_squared_error(y_true=y, y_pred=rf.predict(X))))


# At least on the training data, we have very low RMSE and very high R2 - this is good! But let's see if these numbers hold up on cross-validation.

# In[20]:


# compute cross validation scores for random forest model
r2_scores = cross_val_score(rf, X, y, scoring='r2', cv=crossvalidation, n_jobs=-1)
scores = cross_val_score(rf, X, y, scoring='neg_mean_squared_error', cv=crossvalidation, n_jobs=-1)
rmse_scores = [np.sqrt(abs(s)) for s in scores]

print('Cross-validation results:')
print('Folds: %i, mean R2: %.3f' % (len(scores), np.mean(np.abs(r2_scores))))
print('Folds: %i, mean RMSE: %.3f' % (len(scores), np.mean(np.abs(rmse_scores))))


# Looks like upon cross-validation, we do slightly better than a linear regression but not *too* much better. Let's plot this one as well.

# In[21]:


from figrecipes import PlotlyFig

pf_rf = PlotlyFig(x_title='DFT (MP) bulk modulus (GPa)',
                  y_title='Random forest bulk modulus (GPa)',
                  title='Random forest regression',
                  mode='notebook',
                  filename="rf_regression.html")

pf_rf.xy([(y, cross_val_predict(rf, X, y, cv=crossvalidation)), ([0, 400], [0, 400])], 
      labels=df['formula'], modes=['markers', 'lines'],
      lines=[{}, {'color': 'black', 'dash': 'dash'}], showlegends=False)


# Visually, this looks a little better but not *that* much better. The random forest did very well in training, but is only a little better than linear regression when we plot the cross-validation error as per above.
# 
# You could (optionally) visualize the *training* error by replacing the code ``cross_val_predict(rf, X, y, cv=crossvalidation)`` in the above cell with ``rf.predict(X)``. That would look a lot better, but would not be an accurate representation of your prediction error.

# In[22]:


from sklearn.model_selection import train_test_split
X['formula'] = df['formula']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=1)
train_formula = X_train['formula']
X_train = X_train.drop('formula', axis=1)
test_formula = X_test['formula']
X_test = X_test.drop('formula', axis=1)

rf_reg = RandomForestRegressor(n_estimators=50, random_state=1)
rf_reg.fit(X_train, y_train)

# get fit statistics
print('training R2 = ' + str(round(rf_reg.score(X_train, y_train), 3)))
print('training RMSE = %.3f' % np.sqrt(mean_squared_error(y_true=y_train, y_pred=rf_reg.predict(X_train))))
print('test R2 = ' + str(round(rf_reg.score(X_test, y_test), 3)))
print('test RMSE = %.3f' % np.sqrt(mean_squared_error(y_true=y_test, y_pred=rf_reg.predict(X_test))))


# In[24]:


from figrecipes import PlotlyFig
pf_rf = PlotlyFig(x_title='Bulk modulus prediction residual (GPa)',
                  y_title='Probability',
                  title='Random forest regression residuals',
                  mode="notebook",
                  filename="rf_regression_residuals.html")

hist_plot = pf_rf.histogram(data=[y_train-rf_reg.predict(X_train), 
                                  y_test-rf_reg.predict(X_test)],
                            histnorm='probability', colors=['blue', 'red'],
                            return_plot=True
                           )
hist_plot["data"][0]['name'] = 'train'
hist_plot["data"][1]['name'] = 'test'
pf_rf.create_plot(hist_plot)


# Finally, let's see what are the most important features used by the random forest model.

# In[25]:


importances = rf.feature_importances_
# included = np.asarray(included)
included = X.columns.values
indices = np.argsort(importances)[::-1]

pf = PlotlyFig(y_title='Importance (%)',
               title='Feature by importances',
               mode='notebook',
               fontsize=20,
               ticksize=15)

pf.bar(x=included[indices][0:10], y=importances[indices][0:10])


# Features relating to the melting point and to the volume per atom / density are the most important in the random forest model.
# 
# This concludes the tutorial! You are now familiar with some of the basic features of matminer.