train_lasso.py

# -*- coding: utf-8 -*-
"""
Created on Thu Aug 16 16:37:24 2018

@author: wangyf
"""


import lattice_functions as lf
import regression_tools as rtools
from set_ce_lattice import mother

from sklearn.preprocessing import StandardScaler 
from sklearn.linear_model import LassoCV, lasso_path, Lasso
from sklearn.model_selection import RepeatedKFold, cross_val_score, train_test_split
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import cross_validate

import numpy as np
import pickle 
import json
import os

model_name = 'lasso'
base_dir = os.getcwd()

#%% Import data
'''
fit_int_flag, explained: 
    
True - fit intercept automatically from LASSO regression
     - the intercept is obtained from LASSO.intercept_
False- fit intercept by adding additional first column of pi matrix 
     - the intercept is obtained from the first cluster interaction
     - the signficant cluster interactions need to be adjusted based on index
'''

fit_int_flag = False
iso_flag = True

selected_batches = [0,1,2, 3]

X_init, y, config = rtools.import_Xy(selected_batches  , fit_int_flag, iso_flag)
# the number of Pd atoms in each structure
NPd_list = lf.get_NPd_list(config)
n_sites = len(mother)

# Load clusters
with open('clusters.json') as f:
    Gcv = json.load(f)['Gcv']  


# Modify the model name with selected dataset
model_name = model_name + '_' + ''.join(str(i) for i in selected_batches) 
output_dir = os.path.join(base_dir, model_name)
if not os.path.exists(output_dir): os.makedirs(output_dir)    

#%% The date set is needed to be standardized for lasso, except for the intercept
X = X_init.copy()
if not fit_int_flag:
    scaler = StandardScaler().fit(X[:,1:])
    X[:,1:] = scaler.transform(X[:,1:])
else: 
    scaler = StandardScaler().fit(X)
    X= scaler.transform(X)
    
    
sv = scaler.scale_ # standard deviation for each x variable
mv = scaler.mean_ # mean for each x variable

n_unfilled_clusters = len(np.nonzero(mv)[0]) # the number of clusters that have not been filled

#%% Preparation before regression
# Train test split, save 10% of data point to the test set
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=0)

NPd_test = []
NPd_train = []
for i in y_test: NPd_test.append(NPd_list[np.where(y==i)[0][0]])
for i in y_train: NPd_train.append(NPd_list[np.where(y==i)[0][0]])
                               
NPd_test = np.array(NPd_test)
NPd_train = np.array(NPd_train)                              

# The alpha grid used for plotting path
alphas_grid = np.logspace(0, -3, 20)

# Cross-validation scheme                                  
rkf = RepeatedKFold(n_splits = 10, n_repeats = 3, random_state = 0)

# Explicitly take out the train/test set
X_cv_train, y_cv_train, X_cv_test, y_cv_test = [],[],[],[]
for train_index, test_index in rkf.split(X_train):
    X_cv_train.append(X_train[train_index])
    y_cv_train.append(y_train[train_index])
    X_cv_test.append(X_train[test_index])
    y_cv_test.append(y_train[test_index])

#%% LASSO regression
'''   
# LassoCV to obtain the best alpha, the proper training of Lasso
'''
lasso_cv  = LassoCV(cv = rkf,  max_iter = 1e7, tol = 0.001, fit_intercept=fit_int_flag, random_state=0)
lasso_cv.fit(X_train, y_train)

# the optimal alpha from lassocv
lasso_alpha = lasso_cv.alpha_
# Coefficients for each term
lasso_coefs = lasso_cv.coef_
# The original intercepts 
lasso_intercept = lasso_cv.intercept_

# Access the errors 
y_predict_test = lasso_cv.predict(X_test)
y_predict_train = lasso_cv.predict(X_train)

# error per cluster per site
lasso_RMSE_test_site = np.sqrt(mean_squared_error(y_test, y_predict_test))/n_sites
lasso_RMSE_train_site = np.sqrt(mean_squared_error(y_train, y_predict_train))/n_sites

# error per atom in the cluster
lasso_RMSE_test_atom = np.sqrt(mean_squared_error(y_test/NPd_test, y_predict_test/NPd_test))
lasso_RMSE_train_atom = np.sqrt(mean_squared_error(y_train/NPd_train, y_predict_train/NPd_train))



##Use alpha grid prepare for lassopath
lasso_RMSE_path, lasso_coef_path = rtools.cal_path(alphas_grid, Lasso, X_cv_train, y_cv_train, X_cv_test, y_cv_test, fit_int_flag)
##lasso_path to get alphas and coef_path, somehow individual CV does not work
#lasso_alphas, lasso_coef_path, _ = lasso_path(X_train, y_train, alphas = alphas_grid, fit_intercept=fit_int_flag)
rtools.plot_path(X, y, n_sites, NPd_list, lasso_alpha, alphas_grid, lasso_RMSE_path, lasso_coef_path, lasso_cv, model_name, output_dir)

'''
lasso coefficients needed to be tranformed back to the regular form
'''
lasso_coefs_regular = np.zeros(len(lasso_coefs))
lasso_coefs_regular[1:] = lasso_coefs[1:]/sv
lasso_coefs_regular[0] = lasso_coefs[0] - np.sum(mv/sv*lasso_coefs[1:])

#%% Select the significant cluster interactions 
'''
LASSO Post Processing
'''

# Set the tolerance for signficant interactions 
Tol = 1e-5
# The indices for non-zero coefficients/significant cluster interactions 
J_index = np.where(abs(lasso_coefs_regular)>Tol)[0]
# The number of non-zero coefficients/significant cluster interactions  
n_nonzero = len(J_index)
# The values of non-zero coefficients/significant cluster interactions  
J_nonzero = lasso_coefs_regular[J_index] 
pi_nonzero = X[:, J_index]

# Pick the significant clusters
Gcv_nonzero = []

# Adjust for the manual intercept fitting
if not fit_int_flag:
    
    intercept = J_nonzero[0]
    n_nonzero = n_nonzero - 1 
    J_nonzero = J_nonzero[1:]
    pi_nonzero = pi_nonzero[:,1:]

    
    for i in J_index[1:]:
        # take out the first one and adjust the indices by -1 
        Gcv_nonzero.append(Gcv[i-1]) 
else:       
    for i in J_index:
        Gcv_nonzero.append(Gcv[i]) 

'''
Save Gcv_nonzero and J_nonzero to pickle for further use
''' 
pickle.dump([Gcv_nonzero, J_nonzero, intercept, lasso_RMSE_test_atom, lasso_RMSE_test_site], 
            open(os.path.join(output_dir, model_name + '.p'),'wb'))
 			
			
#%% Energy test for lasso
from ase.io import read, write
def save_POV(atoms, index, output_dir):

    pov_args = {
    	'transparent': True, #Makes background transparent. I don't think I've had luck with this option though
        'canvas_width': 1800., #Size of the width. Height will automatically be calculated. This value greatly impacts POV-Ray processing times
        'display': False, #Whether you want to see the image rendering while POV-Ray is running. I've found it annoying
        'rotation': '0x, 0y,0z', #Position of camera. If you want different angles, the format is 'ax, by, cz' where a, b, and c are angles in degrees
        'celllinewidth': 0.02, #Thickness of cell lines
        'show_unit_cell': 0 #Whether to show unit cell. 1 and 2 enable it (don't quite remember the difference)
        #You can also color atoms by using the color argument. It should be specified by an list of length N_atoms of tuples of length 3 (for R, B, G)
        #e.g. To color H atoms white and O atoms red in H2O, it'll be:
        #colors: [(0, 0, 0), (0, 0, 0), (1, 0, 0)]
        }

    #Write to POV-Ray file
    filename = 'Pd'+'-' + str(index) + '.POV'
    if not os.path.exists(output_dir): os.makedirs(output_dir)    
    write(os.path.join(output_dir, filename), atoms, **pov_args)
    
    
       
#%% Seaborn lasso join plot
import matplotlib.pyplot as plt 
import seaborn as sns;   

y_predict_all = lasso_cv.predict(X)


lims = [
    np.min([y.min(), y_predict_all.min()]),  # min of both axes
    np.max([y.max(), y_predict_all.max()]),  # max of both axes
]
#    
fig = plt.figure(figsize=(10, 6))
sns.set(rc={'figure.figsize':(10,6)})
sns.set(font_scale=1.4)
sns.set_style("ticks", {"xtick.major.size": 8, "ytick.major.size": 8})    
g = sns.JointGrid(y, y_predict_all)
g = g.set_axis_labels("DFT Cluster Energy (eV)", "Predicted Cluster Energy (eV)")

g = g.plot_joint(plt.scatter, color="b", edgecolors = 'navy', s=60, alpha = 0.1) #, facecolors='white', edgecolors = 'b'
g.ax_joint.plot(lims, lims, '--k')

_ = g.ax_marg_x.scatter( y, (y_predict_all - y)/n_sites, s = 30, color ='b', edgecolors = 'navy', alpha = 0.1 ) #facecolors='white', edgecolors = 'b',
_ = g.ax_marg_y.hist(y_predict_all, color="b", alpha = 0.8, orientation="horizontal", bins=np.linspace(lims[0], lims[1], 12))
g.ax_marg_x.plot(lims, np.zeros(len(lims)), 'k--')
g.ax_marg_x.set_ylim([-0.02,0.02])

g.ax_marg_x.set_ylabel("Error/site (eV)")
g.ax_marg_y.set_xlabel("Frequency")
    
plt.setp(g.ax_marg_x.get_yticklabels(), visible=True)
plt.setp(g.ax_marg_y.get_xticklabels(), visible=True)    
plt.setp(g.ax_marg_x.yaxis.get_majorticklines(), visible=True)
plt.setp(g.ax_marg_x.yaxis.get_minorticklines(), visible=True)
plt.setp(g.ax_marg_y.xaxis.get_majorticklines(), visible=True)
plt.setp(g.ax_marg_y.xaxis.get_minorticklines(), visible=True)
plt.show()
#fig.savefig(os.path.join(output_dir, model_name + '_combo.png'))