import sys
sys.path.append('../..')

from SALib import ProblemSpec
from SALib.test_functions import Ishigami
from SALib.plotting.bar import plot as barplot

import matplotlib.pyplot as plt
import numpy as np


# By convention, we assign to "sp" (for "SALib Problem")
sp = ProblemSpec({
    'names': ['x1', 'x2', 'x3'],   # Name of each parameter
    'bounds': [[-np.pi, np.pi]]*3,  # bounds of each parameter
    'outputs': ['Y']               # name of outputs in expected order
})

(sp.sample_saltelli(512, calc_second_order=True)
    .evaluate(Ishigami.evaluate)
    .analyze_sobol())

# Display results in table format
print(sp)

# First-order indices expected with Saltelli sampling:
# x1: 0.3139
# x2: 0.4424
# x3: 0.0

# Basic plotting of results
sp.plot()

plt.title("Basic example plot")


# More advanced plotting

# Plot functions actually return matplotlib axes objects
# In the case of the Sobol' method if `calc_second_order=True`, there will be 
# 3 axes (one each for Total, First, and Second order indices)
axes = sp.plot()

# These can be modified as desired.
# Here, for example, we set the Y-axis to log scale
for ax in axes:
    ax.set_yscale('log')

axes[0].set_title("Example custom plot with log scale")

# Other custom layouts can be created in the usual matplotlib style
# with the basic bar plotter.

# Example: Direct control of plot elements
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(6, 16))

# Get result DataFrames
total, first, second = sp.to_df()

ax1 = barplot(total, ax=ax1)
ax2 = barplot(first, ax=ax2)
ax3 = barplot(second, ax=ax3)

ax1.set_yscale('log')
ax2.set_yscale('log')

ax1.set_title("Customized matplotlib plot")
plt.show()


# Plot sensitivity indices as a heatmap
# Note that plotting methods return a matplotlib axes object
ax = sp.heatmap()
ax.set_title("Basic heatmap")
plt.show()


# Another heatmap plot with more fine-grain control
# Displays Total and First-Order sensitivities in separate subplots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 6))
sp.heatmap('Y', 'ST', 'Total Order Sensitivity', ax1)
sp.heatmap('Y', 'S1', 'First Order Sensitivity', ax2)
plt.show()

# Yet another heatmap example
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10, 6), sharex=True, constrained_layout=True)
sp.heatmap('Y', 'ST', 'Total Order', ax=ax1)
sp.heatmap('Y', 'ST_conf', 'Total Order Conf.', ax=ax2)
sp.heatmap('Y', 'S1', 'First Order', ax=ax3)
sp.heatmap('Y', 'S1_conf', 'First Order Conf.', ax=ax4)
plt.show()

import sys
sys.path.append('../..')

from SALib import ProblemSpec
from SALib.test_functions import Ishigami
from SALib.plotting.bar import plot as barplot

import matplotlib.pyplot as plt
import numpy as np


# By convention, we assign to "sp" (for "SALib Problem")
sp = ProblemSpec({
    'names': ['x1', 'x2', 'x3'],   # Name of each parameter
    'bounds': [[-np.pi, np.pi]]*3,  # bounds of each parameter
    'outputs': ['Y']               # name of outputs in expected order
})

(sp.sample_saltelli(512, calc_second_order=True)
    .evaluate(Ishigami.evaluate)
    .analyze_sobol())

# Display results in table format
print(sp)

# First-order indices expected with Saltelli sampling:
# x1: 0.3139
# x2: 0.4424
# x3: 0.0

# Basic plotting of results
sp.plot()

plt.title("Basic example plot")


# More advanced plotting

# Plot functions actually return matplotlib axes objects
# In the case of the Sobol' method if `calc_second_order=True`, there will be 
# 3 axes (one each for Total, First, and Second order indices)
axes = sp.plot()

# These can be modified as desired.
# Here, for example, we set the Y-axis to log scale
for ax in axes:
    ax.set_yscale('log')

axes[0].set_title("Example custom plot with log scale")

# Other custom layouts can be created in the usual matplotlib style
# with the basic bar plotter.

# Example: Direct control of plot elements
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(6, 16))

# Get result DataFrames
total, first, second = sp.to_df()

ax1 = barplot(total, ax=ax1)
ax2 = barplot(first, ax=ax2)
ax3 = barplot(second, ax=ax3)

ax1.set_yscale('log')
ax2.set_yscale('log')

ax1.set_title("Customized matplotlib plot")
plt.show()


# Plot sensitivity indices as a heatmap
# Note that plotting methods return a matplotlib axes object
ax = sp.heatmap("Y")
ax.set_title("Basic heatmap")
plt.show()


# Another heatmap plot with more fine-grain control
# Displays Total and First-Order sensitivities in separate subplots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 6))
sp.heatmap('Y', 'ST', 'Total Order Sensitivity', ax1)
sp.heatmap('Y', 'S1', 'First Order Sensitivity', ax2)
plt.show()


# Yet another heatmap example
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10, 6), sharex=True, constrained_layout=True)
sp.heatmap('Y', 'ST', 'Total Order', ax=ax1)
sp.heatmap('Y', 'ST_conf', 'Total Order Conf.', ax=ax2)
sp.heatmap('Y', 'S1', 'First Order', ax=ax3)
sp.heatmap('Y', 'S1_conf', 'First Order Conf.', ax=ax4)
plt.show()

Make explicit which output is being analyzed