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

# # PCE METAMODELING: TRUSS DATA SET
# 
# This example showcases how to perform PCE metamodeling using existing data sets.
# The data sets come from a finite element model of a truss structure and are retrieved from different MAT-files.
# The files consist of an experimental design of size $200$ and a validation set of size $10^4$.

# ## Package imports

# In[ ]:


from uqpylab import sessions, display_util
import numpy as np
import matplotlib.pyplot as plt
import scipy.io


# ## Initialize common plotting parameters

# In[ ]:


display_util.load_plt_defaults()
uq_colors = display_util.get_uq_color_order()


# ## Start a remote UQCloud session

# In[ ]:


# Start the session
mySession = sessions.cloud()
# (Optional) Get a convenient handle to the command line interface
uq = mySession.cli
# Reset the session
mySession.reset()


# ## Set random seed for reproducibility

# In[ ]:


uq.rng(0,'twister');


# ## Retrieve data sets

# The experimental design and the validation basis are stored in two separate files

# In[ ]:


# Read the binary files (stored in .mat format)
mat = scipy.io.loadmat('Truss_Experimental_Design.mat')
X = mat['X']
Y = mat['Y']

# Also read the validation basis data set file and store the contents in matrices:
mat = scipy.io.loadmat('Truss_Validation_Basis.mat')
X_val = mat['Xval']
Y_val = mat['Yval']


# ## Probabilistic input model
# 
# Because PCE requires a choice of polynomial basis, a probabilistic input model needs to be defined.
# 
# Specify the marginals of the probabilistic input model:
# 

# In[ ]:


InputOpts = {
    "Marginals": [
        # Young's modulus of the cross-sections
        {"Name" : "E1",
         "Type": "Lognormal",
         "Moments": [2.1e11, 2.1e10]
        },
        {"Name" : "E2",
         "Type": "Lognormal",
         "Moments": [2.1e11, 2.1e10]
        },
        # Cross-section of horizontal elements
        {"Name" : "A1",
         "Type": "Lognormal",
         "Moments": [2.0e-3, 2.0e-4]
        },
        # Cross-section of diagonal elements
        {"Name" : "A2",
         "Type": "Lognormal",
         "Moments": [1.0e-3, 1.0e-4]
        },
        # Loads
        {"Name" : "p1",
         "Type": "Gumbel",
         "Moments": [5.0e4, 7.5e3]
        },
        {"Name" : "p2",
         "Type": "Gumbel",
         "Moments": [5.0e4, 7.5e3]
        },
        {"Name" : "p3",
         "Type": "Gumbel",
         "Moments": [5.0e4, 7.5e3]
        },
        {"Name" : "p4",
         "Type": "Gumbel",
         "Moments": [5.0e4, 7.5e3]
        },
        {"Name" : "p5",
         "Type": "Gumbel",
         "Moments": [5.0e4, 7.5e3]
        },
        {"Name" : "p6",
         "Type": "Gumbel",
         "Moments": [5.0e4, 7.5e3]
        }
    ]
}


# Create an INPUT object based on the specified marginals:

# In[ ]:


myInput = uq.createInput(InputOpts)


# ## PCE metamodel

# Select PCE as the metamodeling tool:

# In[ ]:


MetaOpts = {
    'Type': 'Metamodel',
    'MetaType': 'PCE'
}


# Use experimental design loaded from the data files:

# In[ ]:


MetaOpts['ExpDesign'] = {
    'X': X.tolist(),
    'Y': Y.tolist()
}


# Set the maximum polynomial degree to 5:

# In[ ]:


MetaOpts["Degree"] = np.arange(1,5).tolist()


# Provide the validation data set to get the validation error:

# In[ ]:


MetaOpts['ValidationSet'] = {
    'X': X_val.tolist(),
    'Y': Y_val.tolist()
}


# Create the PCE metamodel:

# In[ ]:


myPCE = uq.createModel(MetaOpts)


# Print a summary of the resulting PCE metamodel:

# In[ ]:


uq.print(myPCE)


# ## Validation

# Evaluate the PCE metamodel at the validation set:

# In[ ]:


Y_PCE = uq.evalModel(myPCE, X_val)


# Plot histograms of the true output and the PCE prediction:

# In[ ]:


plt.hist(Y_val, 20, color=uq_colors[0], alpha = 0.8)
plt.hist(Y_PCE, 20, color=uq_colors[1], alpha = 0.8)
legend_text = ['True model response', 'PCE prediction']
plt.legend(legend_text, frameon=False, loc="best")
plt.xlabel('$\mathrm{Y}$')
plt.ylabel('Counts')
plt.show()


# Plot the true vs. predicted values:

# In[ ]:


fig = plt.figure(figsize=(6, 6))
plt.scatter(Y_val, Y_PCE, marker='o', s=3)
plt.plot([np.min(Y_val), np.max(Y_val)], [np.min(Y_val), np.max(Y_val)], 'k')
plt.axis([np.min(Y_val), np.max(Y_val), np.min(Y_val), np.max(Y_val)])
plt.axis('equal')
plt.grid(True)
plt.xlabel('$\mathrm{Y_{true}}$')
plt.ylabel('$\mathrm{Y_{PCE}}$')
plt.show()


# Print the validation and leave-one-out (LOO) cross-validation errors:

# In[ ]:


print('PCE metamodel validation error: {:5.4e}'.format(myPCE['Error']['Val']))
print('PCE metamodel LOO error:        {:5.4e}'.format(myPCE['Error']['LOO']))


# ## Terminate the remote UQCloud session

# In[ ]:


mySession.quit()

