, :] = 0
B_Imag[:, 0, :, :] = 0
B_Imag[:, :, 0, :] = 0
B_Imag[:, :, :, 0] = 0

B_Complex = B_Real + 1j * B_Imag


#%% md
#
# Make sure that the input parameters are in order to prevent aliasing.

#%%

t_u = 2*np.pi/2/F

if dt>t_u:
    print('Error')

#%% md
#
# Defining the Power Spectral Density Function.

#%%

BSRM_object = BispectralRepresentation(n_sim, P, B_Complex, [dt, dt], [df, df], [nt, nt], [nf, nf])
samples = BSRM_object.samples

from pylab import *
t_list = [t for _ in range(dim)]
T = np.array(np.meshgrid(*t_list, indexing='ij'))

fig1 = plt.figure()
plt.title('2d random field with a prescribed Power Spectrum and Bispectrum')
pcm = pcolor(T[0], T[1], samples[0, 0], cmap='RdBu_r', vmin=-30, vmax=30)
plt.colorbar(pcm, extend='both', orientation='vertical')
plt.xlabel('$X_{1}$')
plt.ylabel('$X_{2}$')
plt.show()

print('The mean of the samples is ', np.mean(samples), 'whereas the expected mean is 0.000')
print('The variance of the samples is ', np.var(samples), 'whereas the expected variance is ', np.sum(P)*df*2)
print('The skewness of the samples is ', np.mean(skew(samples, axis=0)), 'whereas the expected skewness is ',
      np.sum(B_Real)*df**2*6/(np.sum(P)*df*2)**(3/2))

# docs/code/stochastic_processes/karhunen_loeve/plot_karhunen_loeve.py

"""

Karhunen Loeve Expansion
=================================================================

In this example, the KL Expansion is used to generate stochastic processes from a prescribed Autocorrelation Function.
This example illustrates how to use the :class:`.KarhunenLoeveExpansion` class for a one dimensional and compare the
statistics of the generated stochastic processes with the expected values.

"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.KarhunenLoeveExpansion` class from the :py:mod:`stochastic_processes` module of UQpy.

#%%

from UQpy.stochastic_process import KarhunenLoeveExpansion
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn')

#%% md
#
# The input parameters necessary for the generation of the stochastic processes are given below:

#%%

n_sim = 10000  # Num of samples

m = 400 + 1
T = 1000
dt = T / (m - 1)
t = np.linspace(0, T, m)

#%% md
#
# Defining the Autocorrelation Function.

#%%

# Target Covariance(ACF)
R = np.zeros([m, m])
for i in range(m):
    for j in range(m):
        R[i, j] = 2 * np.exp(-((t[j] - t[i]) / 281) ** 2)  # var = 2

KLE_Object = KarhunenLoeveExpansion(n_samples=n_sim, correlation_function=R, time_interval=dt)
samples = KLE_Object.samples

fig, ax = plt.subplots()
plt.title('Realisation of the Karhunen Loeve Expansion')
plt.plot(t, samples[0, 0])
ax.yaxis.grid(True)
ax.xaxis.grid(True)
plt.show()

print('The mean of the samples is ', np.mean(samples), 'whereas the expected mean is 0.000')
print('The variance of the samples is ', np.var(samples), 'whereas the expected variance is 2.000')

# docs/code/stochastic_processes/spectral/spectral_1d_1v.py

"""

One-dimensional & one variable
=================================================================

In this example, the Spectral Representation Method is used to generate stochastic processes from a prescribed Power
Spectrum. This example illustrates how to use the SRM class for a one dimensional and one variable case and compare the
statistics of the generated stochastic processes with the expected values.

"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the SRM class from the StochasticProcesses module of UQpy.

#%%

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn')

#%% md
#
# The input parameters necessary for the generation of the stochastic processes are given below:

#%%

n_sim = 10000  # Num of samples

n = 1 # Num of dimensions

m = 1 # Num of variables

T = 100  # Time(1 / T = dw)
nt = 256  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nw = 128  # Num of Discretized Freq.

# # Generation of Input Data(Stationary)
dt = T / nt
t = np.linspace(0, T - dt, nt)
dw = F / nw
w = np.linspace(0, F - dw, nw)

#%% md
#
# Make sure that the input parameters are in order to prevent aliasing.

#%%

t_u = 2*np.pi/2/F

if dt>t_u:
    print('Error')

#%% md
#
# Defining the Power Spectral Density Function.

#%%

S = 125 / 4 * w ** 2 * np.exp(-5 * w)

SRM_object = SpectralRepresentation(n_sim, S, dt, dw, nt, nw)
samples = SRM_object.samples

fig, ax = plt.subplots()
plt.title('Realisation of the Spectral Representation Method')
plt.plot(t, samples[0, 0])
ax.yaxis.grid(True)
ax.xaxis.grid(True)
plt.show()

print('The mean of the samples is ', np.mean(samples), 'whereas the expected mean is 0.000')
print('The variance of the samples is ', np.var(samples), 'whereas the expected variance is ', np.sum(S)*dw*2)

# docs/code/stochastic_processes/spectral/spectral_1d_mv.py

"""

One-dimensional & multiple variables
=================================================================

In this example, the Spectral Representation Method is used to generate stochastic processes from a prescribed Power
Spectrum and associated Cross Spectral Density. This example illustrates how to use the SRM class for a one dimensional
and 'm' variable case and compare the statistics of the generated stochastic processes with the expected values.

"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.SpectralRepresentation` class from the :class:`stochastic_processes` module of UQpy.

#%%

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn')

#%% md
#
# The input parameters necessary for the generation of the stochastic processes are given below:

#%%

n_sim = 10000  # Num of samples

n = 1 # Num of dimensions

m = 3 # Num of variables

T = 10  # Time(1 / T = dw)
nt = 256  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nf = 128  # Num of Discretized Freq.

# # Generation of Input Data(Stationary)
dt = T / nt
t = np.linspace(0, T - dt, nt)
df = F / nf
f = np.linspace(0, F - df, nf)

#%% md
#
# Make sure that the input parameters are in order to prevent aliasing

#%%

t_u = 2*np.pi/2/F

if dt>t_u:
    print('Error')

#%% md
#
# Defining the Power Spectral Density Function (S) and the Cross Spectral Density (g)

#%%

S_11 = 38.3 / (1 + 6.19 * f) ** (5 / 3)
S_22 = 43.4 / (1 + 6.98 * f) ** (5 / 3)
S_33 = 135 / (1 + 21.8 * f) ** (5 / 3)

g_12 = np.exp(-0.1757 * f)
g_13 = np.exp(-3.478 * f)
g_23 = np.exp(-3.392 * f)

S_list = np.array([S_11, S_22, S_33])
g_list = np.array([g_12, g_13, g_23])

# Assembly of S_jk
S_sqrt = np.sqrt(S_list)
S_jk = np.einsum('i...,j...->ij...', S_sqrt, S_sqrt)
# Assembly of g_jk
g_jk = np.zeros_like(S_jk)
counter = 0
for i in range(m):
    for j in range(i + 1, m):
        g_jk[i, j] = g_list[counter]
        counter = counter + 1
g_jk = np.einsum('ij...->ji...', g_jk) + g_jk

for i in range(m):
    g_jk[i, i] = np.ones_like(S_jk[0, 0])
S = S_jk * g_jk


SRM_object = SpectralRepresentation(n_sim, S, dt, df, nt, nf)
samples = SRM_object.samples

fig, ax = plt.subplots()
plt.title('Realisation of the Spectral Representation Method')
plt.plot(t, samples[0, 0], label='1st dimension')
plt.plot(t, samples[0, 1], label='2nd dimension')
plt.plot(t, samples[0, 2], label='3rd dimension')
ax.yaxis.grid(True)
ax.xaxis.grid(True)
plt.legend()
plt.show()

print('The mean of the samples is ', np.mean(samples), 'whereas the expected mean is 0.000')
print('The variance of the samples is ', np.var(samples), 'whereas the expected variance is ',
      np.sum(S_list)*np.prod(df)*2/m)

# docs/code/stochastic_processes/spectral/spectral_nd_1d.py

"""

N-dimensional & one-variable
=================================================================

In this example, the Spectral Representation Method is used to generate stochastic processes from a prescribed Power
Spectrum. This example illustrates how to use the :class:`.SpectralRepresentation` class for 'n' dimensional and one
variable case and compare the statistics of the generated stochastic processes with the expected values.

"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.SpectralRepresentation` class from the StochasticProcesses module of UQpy.

#%%

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np
import matplotlib.pyplot as plt
from pylab import *
plt.style.use('seaborn')

#%% md
#
# The input parameters necessary for the generating of the stochastic processes are given below:

#%%

n_sim = 1000 # Num of samples

n = 2 # Num of dimensions

m = 1 # Num of variables

T = 10
nt = 200
dt = T/nt
t = np.linspace(0, T-dt, nt)

# Frequency
W = np.array([1.0, 1.5])
nw = 100
dw = W / nw
x_list = [np.linspace(0, W[i] - dw[i], nw) for i in range(n)]
xy_list = np.array(np.meshgrid(*x_list, indexing='ij'))


#%% md
#
# Defining the Power Spectral Density(S)

#%%

S_nd_1v = 125 / 4 * np.linalg.norm(xy_list, axis=0) ** 2 * np.exp(-5 * np.linalg.norm(xy_list, axis=0))

#%% md
#
# Make sure that the input parameters are in order to prevent aliasing

#%%

t_u = 2*np.pi/2/W

if dt>t_u.all():
    print('Error')

SRM_object = SpectralRepresentation(n_sim, S_nd_1v, [dt, dt], dw, [nt, nt], [nw, nw], random_state=128)
samples_nd_1v = SRM_object.samples

t_list = [t for _ in range(n)]
tt_list = np.array(np.meshgrid(*t_list, indexing='ij'))

fig1 = plt.figure()
plt.title('2d random field with a prescribed Power Spectrum')
pcm = pcolor(tt_list[0], tt_list[1], samples_nd_1v[0, 0], cmap='RdBu_r', vmin=-6, vmax=6)
plt.colorbar(pcm, extend='both', orientation='vertical')
plt.xlabel('$X_{1}$')
plt.ylabel('$X_{2}$')
plt.show()

print('The mean of the samples is ', np.mean(samples_nd_1v), 'whereas the expected mean is 0.000')
print('The variance of the samples is ', np.var(samples_nd_1v), 'whereas the expected variance is ',
      np.sum(S_nd_1v)*np.prod(dw)*(2**n))

# docs/code/stochastic_processes/spectral/spectral_nd_mv.py

"""

N-dimensional & multiple-variables
=================================================================

In this example, the Spectral Representation Method is used to generate stochastic processes from a prescribed Power
Spectrum. This example illustrates how to use the SRM class for 'n' dimensional and one variable case and compare the
statistics of the generated stochastic processes with the expected values.

"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.SpectralRepresentation` class from the :py:mod:`stochastic_processes` module of UQpy.

#%%

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np
import matplotlib.pyplot as plt
from pylab import *
plt.style.use('seaborn')

#%% md
#
# The input parameters necessary for the generating of the stochastic processes are given below:

#%%

n_sim = 1000 # Num of samples

n = 2 # Num of dimensions

m = 1 # Num of variables

T = 10
nt = 200
dt = T/nt
t = np.linspace(0, T-dt, nt)

# Frequency
W = np.array([1.0, 1.5])
nw = 100
dw = W / nw
x_list = [np.linspace(0, W[i] - dw[i], nw) for i in range(n)]
xy_list = np.array(np.meshgrid(*x_list, indexing='ij'))

#%% md
#
# Defining the Power Spectral Density(S)

#%%

S_nd_1v = 125 / 4 * np.linalg.norm(xy_list, axis=0) ** 2 * np.exp(-5 * np.linalg.norm(xy_list, axis=0))

#%% md
#
# Make sure that the input parameters are in order to prevent aliasing

#%%

t_u = 2*np.pi/2/W

if dt>t_u.all():
    print('Error')

SRM_object = SpectralRepresentation(n_sim, S_nd_1v, [dt, dt], dw, [nt, nt], [nw, nw], random_state=128)
samples_nd_1v = SRM_object.samples

t_list = [t for _ in range(n)]
tt_list = np.array(np.meshgrid(*t_list, indexing='ij'))

fig1 = plt.figure()
plt.title('2d random field with a prescribed Power Spectrum')
pcm = pcolor(tt_list[0], tt_list[1], samples_nd_1v[0, 0], cmap='RdBu_r', vmin=-6, vmax=6)
plt.colorbar(pcm, extend='both', orientation='vertical')
plt.xlabel('$X_{1}$')
plt.ylabel('$X_{2}$')
plt.show()

print('The mean of the samples is ', np.mean(samples_nd_1v), 'whereas the expected mean is 0.000')
print('The variance of the samples is ', np.var(samples_nd_1v), 'whereas the expected variance is ',
      np.sum(S_nd_1v)*np.prod(dw)*(2**n))

# docs/code/stochastic_processes/translation/translation.py

"""

Translation
=================================================================

In this example, a Gaussian stochastic processes is translated into a stochastic processes of a number of distributions.
This example illustrates how to use the :class:`.Correlate` class to translate from Gaussian to other probability
distributions and compare how the statistics of the translated stochastic processes change along with distributions.

"""

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the class:`.Correlate` class from the :py:mod:`stochastic_processes` module of UQpy.

# %%

from UQpy.stochastic_process import Translation, SpectralRepresentation
import numpy as np
import matplotlib.pyplot as plt

plt.style.use('seaborn')

# %% md
#
# Firstly we generate Gaussian Stochastic Processes using the Spectral Representation Method.

# %%

n_sim = 10000  # Num of samples
T = 100  # Time(1 / T = dw)
nt = 256  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nw = 128  # Num of Discretized Freq.
dt = T / nt
t = np.linspace(0, T - dt, nt)
dw = F / nw
w = np.linspace(0, F - dw, nw)
S = 125 * w ** 2 * np.exp(-2 * w)

SRM_object = SpectralRepresentation(n_sim, S, dt, dw, nt, nw, random_state=128)
samples = SRM_object.samples


def S_to_R(S, w, t):
    dw = w[1] - w[0]
    fac = np.ones(len(w))
    fac[1: len(w) - 1: 2] = 4
    fac[2: len(w) - 2: 2] = 2
    fac = fac * dw / 3
    R = np.zeros(len(t))
    for i in range(len(t)):
        R[i] = 2 * np.dot(fac, S * np.cos(w * t[i]))
    return R


R_g = S_to_R(S, w, t)
r_g = R_g / R_g[0]

# %% md
#
# We translate the samples to be Uniform samples from :math:`1` to :math:`2`

# %%

from UQpy.distributions import Lognormal

distribution = Lognormal(0.5)
samples = samples.flatten()[:, np.newaxis]

Translate_object = Translation(distributions=distribution, time_interval=dt, frequency_interval=dw,
                               n_time_intervals=nt, n_frequency_intervals=nw,
                               correlation_function_gaussian=R_g, samples_gaussian=samples)
samples_ng = Translate_object.samples_non_gaussian

R_ng = Translate_object.scaled_correlation_function_non_gaussian
r_ng = Translate_object.correlation_function_non_gaussian

# %% md
#
# Plotting the actual and translated autocorrelation functions

# %%

fig1 = plt.figure()
plt.plot(r_g, label='Gaussian')
plt.plot(r_ng, label='non-Gaussian')
plt.title('Correlation Function (r)')
plt.legend()
plt.show()


# docs/code/surrogates/gpr/gpr_constraints.py

"""

Gaussian Process with Noise and Constraints
======================================================================

"""

# %% md
#
# This jupyter script shows the performance of GaussianProcessRegressor class in the UQpy. A training data is generated
# using a function (:math:`f(x)`, as defined below), which is used to train a surrogate model.

# %%

# %% md
#
# Import the necessary modules to run the example script. Notice that FminCobyla is used here, to solve the MLE
# optimization problem with constraints.

# %%

import numpy as np
import matplotlib.pyplot as plt
import warnings

from UQpy.surrogates.gaussian_process.regression_models.QuadraticRegression import QuadraticRegression

warnings.filterwarnings('ignore')
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from UQpy.utilities.FminCobyla import FminCobyla
from UQpy.surrogates import GaussianProcessRegression, NonNegative, RBF


# %% md
#
# Consider the following function :math:`f(x)`.
#
# .. math:: f(x) = \frac{1}{100} + \frac{5}{8}(2x-1)^4[(2x-1)^2 + 4\sin{(5 \pi x)^2}], \quad \quad x \in [0,1]

# %%

def funct(x):
    y = (1 / 100) + (5 / 8) * ((2 * x - 1) ** 4) * (((2 * x - 1) ** 2) + 4 * np.sin(5 * np.pi * x) ** 2)
    return y


# %% md
#
# Define the training data set. The following 13 points have been used to fit the GP.

# %%

X_train = np.array([0, 0.06, 0.08, 0.26, 0.27, 0.4, 0.52, 0.6, 0.68, 0.81, 0.9, 0.925, 1]).reshape(-1, 1)
y_train = funct(X_train)

# %% md
#
# Define the test data.

# %%

X_test = np.linspace(0, 1, 100).reshape(-1, 1)
y_test = funct(X_test)


# %% md
#
# Train GPR
# ~~~~~~~~~~~~~
# - Noise
# - Constraints
#
# Here, 30 equidistant point are selected over the domain of :math:`x`, lets call them constraint points. The idea is to
# train the surrogate model such that the probability of positive surrogates prediction is very high at these points.

# %%

X_c = np.linspace(0, 1, 31).reshape(-1,1)
y_c = funct(X_c)

# %% md
#
# In this approach, MLE problem is solved with the following constraints:
#
# .. math:: \hat{y}(x_c)-Z \sigma_{\hat{y}}(x_c) > 0  \quad \quad Z = 2
# .. math:: |\hat{y}(x_t) - y(x_t)| < \epsilon   \quad \quad \epsilon = 0.3
#
# where, :math:`x_c` and :math:`x_t` are the constraint and training sample points, respectively.
#
# Define kernel used to define the covariance matrix. Here, the application of Radial Basis Function (RBF) kernel is
# demonstrated.

# %%

kernel3 = RBF()

# %% md
#
# Define the optimizer used to identify the maximum likelihood estimate.

# %%

bounds_3 = [[10**(-6), 10**(-1)], [10**(-5), 10**(-1)], [10**(-13), 10**(-5)]]
optimizer3 = FminCobyla()

# %% md
#
# Define constraints for the Cobyla optimizer using UQpy's Nonnegatice class.

# %%

cons = NonNegative(constraint_points=X_c, observed_error=0.03, z_value=2)

# %% md
#
# Define the 'GaussianProcessRegressor' class object, the input attributes defined here are kernel, optimizer, initial
# estimates of hyperparameters and number of times MLE is identified using random starting point.

# %%

gpr3 = GaussianProcessRegression(kernel=kernel3, hyperparameters=[10**(-3), 10**(-2), 10**(-10)], optimizer=optimizer3,
                                 optimizations_number=10, optimize_constraints=cons, bounds=bounds_3, noise=True,
                                 regression_model=QuadraticRegression())

# %% md
#
# Call the 'fit' method to train the surrogate model (GPR).

# %%

gpr3.fit(X_train, y_train)

# %% md
#
# The maximum likelihood estimates of the hyperparameters are as follows:

# %%

print(gpr3.hyperparameters)

print('Length Scale: ', gpr3.hyperparameters[0])
print('Process Variance: ', gpr3.hyperparameters[1])
print('Noise Variance: ', gpr3.hyperparameters[2])


# %% md
#
# Use 'predict' method to compute surrogate prediction at the test samples. The attribute 'return_std' is a boolean
# indicator. If 'True', 'predict' method also returns the standard error at the test samples.

# %%

y_pred3, y_std3 = gpr3.predict(X_test, return_std=True)

# %% md
#
# The plot shows the test function in dashed red line and 13 training points are represented by blue dots. Also, blue
# curve shows the GPR prediction for $x \in (0, 1)$ and yellow shaded region represents 95% confidence interval.

# %%

fig, ax = plt.subplots(figsize=(8.5,7))
ax.plot(X_test,y_test,'r--',linewidth=2,label='Test Function')
ax.plot(X_train,y_train,'bo',markerfacecolor='b', markersize=10, label='Training Data')
ax.plot(X_test,y_pred3,'b-', lw=2, label='GP Prediction')
ax.plot(X_test, np.zeros((X_test.shape[0],1)))
ax.fill_between(X_test.flatten(), y_pred3-1.96*y_std3,
                y_pred3+1.96*y_std3,
                facecolor='yellow',label='95% Credibility Interval')
ax.tick_params(axis='both', which='major', labelsize=12)
ax.set_xlabel('x', fontsize=15)
ax.set_ylabel('f(x)', fontsize=15)
ax.set_ylim([-0.3,1.8])
ax.legend(loc="upper right",prop={'size': 12});
plt.grid()


# %% md
#
# Verify the constraints for the trained surrogate model. Notice that all values are positive, thus constraints are
# satisfied for the constraint points.

# %%

y_, ys_ = gpr3.predict(X_c, return_std=True)
y_ - 2*ys_

# %% md
#
# Notice that all values are negative, thus constraints are satisfied for the training points.

# %%

y_ = gpr3.predict(X_train, return_std=False)
np.abs(y_train[:, 0]-y_) - 0.03


# docs/code/surrogates/gpr/local_python_model_1Dfunction.py

def y_func(z):
    import numpy as np
    return np.sin(z)



# docs/code/surrogates/gpr/local_python_model_function.py

def y_func(z):
    return 1/(6.2727*(abs(0.3-z[:, 0]**2-z[:, 1]**2)+0.01))



# docs/code/surrogates/gpr/plot_gpr_custom2D.py

"""

Gaussian Process of a custom 2D function
======================================================================

"""

# %% md
#
# In this example, Gaussian Process Regression is used to generate a surrogate model for a given data. In this data,
# sample points are generated using TrueStratifiedSampling class and functional value at sample points are estimated
# using a model defined in python script ('python_model_function.py).

# %%

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the TrueStratifiedSampling, RunModel and GaussianProcessRegression class from UQpy.

# %%

import shutil

from UQpy import PythonModel
from UQpy.surrogates.gaussian_process.regression_models import ConstantRegression
from UQpy.sampling import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling
from UQpy.run_model.RunModel import RunModel
from UQpy.distributions import Uniform
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from UQpy.surrogates import GaussianProcessRegression, Matern

# %% md
#
# Create a distribution object.

# %%

marginals = [Uniform(loc=0., scale=1.), Uniform(loc=0., scale=1.)]

# %% md
#
# Create a strata object.

# %%

strata = RectangularStrata(strata_number=[10, 10])

# %% md
#
# Using UQpy TrueStratifiedSampling class to generate samples for two random variables,
# which are uniformly distributed between 0 and 1.

# %%

x = TrueStratifiedSampling(distributions=marginals, strata_object=strata,
                           nsamples_per_stratum=1, random_state=1)

# %% md
#
# RunModel is used to evaluate function values at sample points. Model is defined as a function in python file
# 'python_model_function.py'.

# %%

model = PythonModel(model_script='local_python_model_function.py', model_object_name="y_func")
rmodel = RunModel(model=model)

rmodel.run(samples=x.samples)

# %% md
#
# Using UQpy GaussianProcessRegression class to generate a surrogate for generated data. In this illustration, Quadratic regression model and
# Exponential correlation model are used.

# %%

regression_model = ConstantRegression()
kernel = Matern(nu=0.5)

from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

optimizer = MinimizeOptimizer(method="L-BFGS-B")
K = GaussianProcessRegression(regression_model=regression_model, optimizer=optimizer,
                              kernel=kernel,
                              optimizations_number=20,
                              hyperparameters=[1, 1, 0.1])
K.fit(samples=x.samples, values=rmodel.qoi_list)
print(K.hyperparameters)

# %% md
#
# This plot shows the actual model which is used to evaluate the samples to identify the function values.

# %%

num = 25
x1 = np.linspace(0, 1, num)
x2 = np.linspace(0, 1, num)

x1g, x2g = np.meshgrid(x1, x2)
x1gv, x2gv = x1g.reshape(x1g.size, 1), x2g.reshape(x2g.size, 1)

y2 = K.predict(np.concatenate([x1gv, x2gv], 1)).reshape(x1g.shape[0], x1g.shape[1])
model = PythonModel(model_script='local_python_model_function.py', model_object_name="y_func")
r2model = RunModel(model=model)
r2model.run(samples=np.concatenate([x1gv, x2gv], 1))
y_act = np.array(r2model.qoi_list).reshape(x1g.shape[0], x1g.shape[1])

fig1 = plt.figure()
ax = fig1.gca(projection='3d')
surf = ax.plot_surface(x1g, x2g, y_act, cmap=cm.coolwarm, linewidth=0, antialiased=False)
ax.set_zlim(-1, 15)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
# Add a color bar which maps values to colors.
fig1.colorbar(surf, shrink=0.5, aspect=5)
plt.show()

# %% md
#
# This plot shows the input data as red dot and green wireframe plot represent the kriging surrogate generated through
# Kriging class.

# %%

fig2 = plt.figure()
ax2 = plt.axes(projection='3d')
# Plot for estimated values
kr = ax2.plot_wireframe(x1g, x2g, y2, color='Green', label='Kriging interpolate')

# Plot for scattered data
ID = ax2.scatter3D(x.samples[:, 0], x.samples[:, 1], np.array(rmodel.qoi_list), color='Red', label='Input data')
plt.legend(handles=[kr, ID])
plt.show()

# docs/code/surrogates/gpr/plot_gpr_no_noise.py

"""

Gaussian Process without noise
======================================================================

"""

# %% md
#
# This jupyter script shows the performance of GaussianProcessRegressor class in the UQpy. A training data is generated
# using a function (:math:`f(x)`, as defined below), which is used to train a surrogate model.

# %%

# %% md
#
# Import the necessary modules to run the example script. Notice that FminCobyla is used here, to solve the MLE
# optimization problem with constraints.

# %%

import numpy as np
import matplotlib.pyplot as plt
import warnings

from UQpy.surrogates.gaussian_process.regression_models.LinearRegression import LinearRegression

warnings.filterwarnings('ignore')
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from UQpy.utilities.FminCobyla import FminCobyla
from UQpy.surrogates import GaussianProcessRegression, NonNegative, RBF


# %% md
#
# Consider the following function :math:`f(x)`.
#
# .. math:: f(x) = \frac{1}{100} + \frac{5}{8}(2x-1)^4[(2x-1)^2 + 4\sin{(5 \pi x)^2}], \quad \quad x \in [0,1]

# %%

def funct(x):
    y = (1 / 100) + (5 / 8) * ((2 * x - 1) ** 4) * (((2 * x - 1) ** 2) + 4 * np.sin(5 * np.pi * x) ** 2)
    return y


# %% md
#
# Define the training data set. The following 13 points have been used to fit the GP.

# %%

X_train = np.array([0, 0.06, 0.08, 0.26, 0.27, 0.4, 0.52, 0.6, 0.68, 0.81, 0.9, 0.925, 1]).reshape(-1, 1)
y_train = funct(X_train)

# %% md
#
# Define the test data.

# %%

X_test = np.linspace(0, 1, 100).reshape(-1, 1)
y_test = funct(X_test)

# %% md
#
# The plot shows the test function in dashed red line and 13 training points are represented by blue dots.

# %%

fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(X_test, y_test, 'r--', linewidth=2, label='Test Function')
ax.plot(X_train, y_train, 'bo', markerfacecolor='b', markersize=10, label='Training Data')
ax.plot(X_test, np.zeros((X_test.shape[0], 1)))
ax.tick_params(axis='both', which='major', labelsize=12)
ax.set_xlabel('x', fontsize=15)
ax.set_ylabel('f(x)', fontsize=15)
ax.set_ylim([-0.3, 1.8])
ax.legend(loc="upper right", prop={'size': 12});
plt.grid()

# %% md
#
# Train GPR
# ~~~~~~~~~~~~~
# - No Noise
# - No Constraints
#
# Define kernel used to define the covariance matrix. Here, the application of Radial Basis Function (RBF) kernel is
# demonstrated.

# %%

kernel1 = RBF()

# %% md
#
# Define the optimizer used to identify the maximum likelihood estimate.

# %%

bounds_1 = [[10 ** (-4), 10 ** 3], [10 ** (-3), 10 ** 2]]
optimizer1 = MinimizeOptimizer(method='L-BFGS-B', bounds=bounds_1)

# %% md
#
# Define the 'GaussianProcessRegressor' class object, the input attributes defined here are kernel, optimizer, initial
# estimates of hyperparameters and number of times MLE is identified using random starting point.

# %%

gpr1 = GaussianProcessRegression(kernel=kernel1, hyperparameters=[10 ** (-3), 10 ** (-2)], optimizer=optimizer1,
                                 optimizations_number=10, noise=False, regression_model=LinearRegression())

# %% md
#
# Call the 'fit' method to train the surrogate model (GPR).

# %%

gpr1.fit(X_train, y_train)


# %% md
#
# The maximum likelihood estimates of the hyperparameters are as follows:

# %%

gpr1.hyperparameters

print('Length Scale: ', gpr1.hyperparameters[0])
print('Process Variance: ', gpr1.hyperparameters[1])

# %% md
#
# Use 'predict' method to compute surrogate prediction at the test samples. The attribute 'return_std' is a boolean
# indicator. If 'True', 'predict' method also returns the standard error at the test samples.

# %%

y_pred1, y_std1 = gpr1.predict(X_test, return_std=True)


# %% md
#
# The plot shows the test function in dashed red line and 13 training points are represented by blue dots. Also, blue
# curve shows the GPR prediction for :math:`x \in (0, 1)` and yellow shaded region represents 95% confidence interval.

# %%

fig, ax = plt.subplots(figsize=(8.5,7))
ax.plot(X_test,y_test,'r--',linewidth=2,label='Test Function')
ax.plot(X_train,y_train,'bo',markerfacecolor='b', markersize=10, label='Training Data')
ax.plot(X_test,y_pred1,'b-', lw=2, label='GP Prediction')
ax.plot(X_test, np.zeros((X_test.shape[0],1)))
ax.fill_between(X_test.flatten(), y_pred1-1.96*y_std1,
                y_pred1+1.96*y_std1,
                facecolor='yellow',label='95% CI')
ax.tick_params(axis='both', which='major', labelsize=12)
ax.set_xlabel('x', fontsize=15)
ax.set_ylabel('f(x)', fontsize=15)
ax.set_ylim([-0.3,1.8])
plt.title('GP surrogate (No noise, No Constraints)')
ax.legend(loc="upper right",prop={'size': 12})
plt.grid()

# docs/code/surrogates/gpr/plot_gpr_noisy.py

"""

Gaussian Process with noisy output
======================================================================

"""

# %% md
#
# This jupyter script shows the performance of GaussianProcessRegressor class in the UQpy. A training data is generated
# using a function (:math:`f(x)`, as defined below), which is used to train a surrogate model.

# %%

# %% md
#
# Import the necessary modules to run the example script. Notice that FminCobyla is used here, to solve the MLE
# optimization problem with constraints.

# %%

import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from UQpy.surrogates.gaussian_process.regression_models.LinearRegression import LinearRegression
from UQpy.utilities.FminCobyla import FminCobyla
from UQpy.surrogates import GaussianProcessRegression, NonNegative, RBF


# %% md
#
# Consider the following function :math:`f(x)`.
#
# .. math:: f(x) = \frac{1}{100} + \frac{5}{8}(2x-1)^4[(2x-1)^2 + 4\sin{(5 \pi x)^2}], \quad \quad x \in [0,1]

# %%

def funct(x):
    y = (1 / 100) + (5 / 8) * ((2 * x - 1) ** 4) * (((2 * x - 1) ** 2) + 4 * np.sin(5 * np.pi * x) ** 2)
    return y


# %% md
#
# Define the training data set. The following 13 points have been used to fit the GP.

# %%

X_train = np.array([0, 0.06, 0.08, 0.26, 0.27, 0.4, 0.52, 0.6, 0.68, 0.81, 0.9, 0.925, 1]).reshape(-1, 1)
y_train = funct(X_train)

# %% md
#
# Define the test data.

# %%

X_test = np.linspace(0, 1, 100).reshape(-1, 1)
y_test = funct(X_test)


# %% md
#
# Train GPR
# ~~~~~~~~~~~~~
# - Noise
# - No Constraints
#
# Define kernel used to define the covariance matrix. Here, the application of Radial Basis Function (RBF) kernel is
# demonstrated.

# %%

kernel2 = RBF()


# %% md
#
# Define the optimizer used to identify the maximum likelihood estimate.

# %%

bounds_2 = [[10**(-3), 10**3], [10**(-3), 10**2], [10**(-3), 10**(2)]]
optimizer2 = MinimizeOptimizer(method='L-BFGS-B', bounds=bounds_2)


# %% md
#
# Define the 'GaussianProcessRegressor' class object, the input attributes defined here are kernel, optimizer, initial
# estimates of hyperparameters and number of times MLE is identified using random starting point.

# %%

gpr2 = GaussianProcessRegression(kernel=kernel2, hyperparameters=[1, 1, 0.1], optimizer=optimizer2,
                                 optimizations_number=10, noise=True, regression_model=LinearRegression())


# %% md
#
# Call the 'fit' method to train the surrogate model (GPR).

# %%

gpr2.fit(X_train, y_train)

# %% md
#
# The maximum likelihood estimates of the hyperparameters are as follows:

# %%

print(gpr2.hyperparameters)

print('Length Scale: ', gpr2.hyperparameters[0])
print('Process Variance: ', gpr2.hyperparameters[1])
print('Noise Variance: ', gpr2.hyperparameters[2])


# %% md
#
# Use 'predict' method to compute surrogate prediction at the test samples. The attribute 'return_std' is a boolean
# indicator. If 'True', 'predict' method also returns the standard error at the test samples.

# %%

y_pred2, y_std2 = gpr2.predict(X_test, return_std=True)


# %% md
#
# The plot shows the test function in dashed red line and 13 training points are represented by blue dots. Also, blue
# curve shows the GPR prediction for $x \in (0, 1)$ and yellow shaded region represents 95% confidence interval.

# %%

fig, ax = plt.subplots(figsize=(8.5,7))
ax.plot(X_test,y_test,'r--',linewidth=2,label='Test Function')
ax.plot(X_train,y_train,'bo',markerfacecolor='b', markersize=10, label='Training Data')
ax.plot(X_test,y_pred2,'b-', lw=2, label='GP Prediction')
ax.plot(X_test, np.zeros((X_test.shape[0],1)))
ax.fill_between(X_test.flatten(), y_pred2-1.96*y_std2,
                y_pred2+1.96*y_std2,
                facecolor='yellow',label='95% CI')
ax.tick_params(axis='both', which='major', labelsize=12)
ax.set_xlabel('x', fontsize=15)
ax.set_ylabel('f(x)', fontsize=15)
ax.set_ylim([-0.3,1.8])
plt.title('GP Surrogate (Noise, No Constraints)')
ax.legend(loc="upper right",prop={'size': 12});
plt.grid()


# docs/code/surrogates/gpr/plot_gpr_sine.py

"""

Gaussian Process of a sinusoidal function
======================================================================

"""

# %% md
#
# In this example, Gaussian Process Regression is used to generate a surrogate model for a given data. In this data,
# sample points are generated using TrueStratifiedSampling class and functional value at sample points are estimated
# using a model defined in python script ('python_model_function.py).

# %%

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the TrueStratifiedSampling, RunModel and GaussianProcessRegression class from UQpy.

# %%

import shutil

from UQpy import PythonModel
from UQpy.sampling.stratified_sampling.strata import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling
from UQpy.run_model.RunModel import RunModel
from UQpy.distributions import Gamma
import numpy as np
import matplotlib.pyplot as plt
from UQpy.surrogates import GaussianProcessRegression, RBF

# %% md
#
# Create a distribution object.

# %%

marginals = [Gamma(a=2., loc=1., scale=3.)]

# %% md
#
# Create a distribution object.

# %%

strata = RectangularStrata(strata_number=[20])

# %% md
#
# Run stratified sampling

# %%

x = TrueStratifiedSampling(distributions=marginals, strata_object=strata,
                           nsamples_per_stratum=1, random_state=2)


# %% md
#
# RunModel is used to evaluate function values at sample points. Model is defined as a function in python file
# 'python_model_function.py'.

# %%

model = PythonModel(model_script='local_python_model_1Dfunction.py', model_object_name='y_func',
                    delete_files=True)
rmodel = RunModel(model=model)
rmodel.run(samples=x.samples)

from UQpy.surrogates.gaussian_process.regression_models import LinearRegression
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

bounds = [[10**(-3), 10**3], [10**(-3), 10**2]]
optimizer = MinimizeOptimizer(method='L-BFGS-B', bounds=bounds)

K = GaussianProcessRegression(regression_model=LinearRegression(), kernel=RBF(),
                              optimizer=optimizer, optimizations_number=20, hyperparameters=[1, 0.1],
                              random_state=2)
K.fit(samples=x.samples, values=rmodel.qoi_list)
print(K.hyperparameters)

# %% md
#
# RunModel is used to evaluate function values at sample points. Model is defined as a function in python file
# 'python_model_function.py'.

# %%

num = 1000
x1 = np.linspace(min(x.samples), max(x.samples), num)

y, y_sd = K.predict(x1.reshape([num, 1]), return_std=True)


# %% md
#
# Actual model is evaluated at all points to compare it with kriging surrogate.

# %%

rmodel.run(samples=x1, append_samples=False)

# %% md
#
# This plot shows the input data as blue dot, blue curve is actual function and orange curve represents response curve.
# This plot also shows the gradient and 95% confidence interval of the kriging surrogate.

# %%

fig = plt.figure()
ax = plt.subplot(111)
plt.plot(np.squeeze(x1), np.squeeze(rmodel.qoi_list), label='Sine')
plt.plot(x1, y, label='Surrogate')
# plt.plot(x1, y_grad, label='Gradient')
plt.scatter(K.samples, K.values, label='Data')
plt.fill(np.concatenate([x1, x1[::-1]]), np.concatenate([y - 1.9600 * y_sd,
                                                         (y + 1.9600 * y_sd)[::-1]]),
         alpha=.5, fc='y', ec='None', label='95% CI')
box = ax.get_position()
ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])

# Put a legend to the right of the current axis
ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
plt.show()



# docs/code/surrogates/pce/pce_robot_arm.py

"""

Robot Arm function (8 random inputs, scalar output)
======================================================================

In this example, PCE is used to generate a surrogate model for a given set of 8D data.

**Dimensions:** 8

**Description:**  Models the position of a robot arm which has four segments.

**Input Domain:**  The input variables and their usual input ranges are: :math:`[0,1]` for the :math:`L_i` inputs
and :math:`[0,2\pi]` for the :math:`θ_i` inputs.

**Function:**

:math:`f(x) = \Big( \sum_{i=1}^{4}L_i cos\Big(\sum_{j=1}^{i} θ_j\Big) \Big)^{2} + \Big( \sum_{i=1}^{4}L_i sin\Big(\sum_{j=1}^{i} θ_j\Big) \Big)^{2}`

**Output:** The square distance from the end of the robot arm to the origin.

**Reference:**  An, J., & Owen, A. (2001). Quasi-regression. Journal of Complexity, 17(4), 588-607.

"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Uniform, JointIndependent
from UQpy.surrogates import *

# %% md
#
# Define the function.

# %%

def function(x):
    # without square root
    u1 = x[:, 4] * np.cos(x[:, 0])
    u2 = x[:, 4] * np.cos(x[:, 0]) + x[:, 5] * np.cos(np.sum(x[:, :2], axis=1))
    u3 = x[:, 4] * np.cos(x[:, 0]) + x[:, 5] * np.cos(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.cos(
        np.sum(x[:, :3], axis=1))
    u4 = x[:, 4] * np.cos(x[:, 0]) + x[:, 5] * np.cos(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.cos(
        np.sum(x[:, :3], axis=1)) + x[:, 7] * np.cos(np.sum(x[:, :4], axis=1))

    v1 = x[:, 4] * np.sin(x[:, 0])
    v2 = x[:, 4] * np.sin(x[:, 0]) + x[:, 5] * np.sin(np.sum(x[:, :2], axis=1))
    v3 = x[:, 4] * np.sin(x[:, 0]) + x[:, 5] * np.sin(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.sin(
        np.sum(x[:, :3], axis=1))
    v4 = x[:, 4] * np.sin(x[:, 0]) + x[:, 5] * np.sin(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.sin(
        np.sum(x[:, :3], axis=1)) + x[:, 7] * np.sin(np.sum(x[:, :4], axis=1))

    return (u1 + u2 + u3 + u4) ** 2 + (v1 + v2 + v3 + v4) ** 2

# %% md
#
# Create a distribution object, generate samples and evaluate the function at the samples.

# %%

np.random.seed(1)

dist_1 = Uniform(loc=0, scale=2*np.pi)
dist_2 = Uniform(loc=0, scale=1)

marg = [dist_1]*4
marg_1 = [dist_2]*4
marg.extend(marg_1)

joint = JointIndependent(marginals=marg)

n_samples = 9000
x = joint.rvs(n_samples)
y = function(x)

# %% md
#
# Create an object from the PCE class. Compute PCE coefficients using least squares regression.

# %%

max_degree = 6
polynomial_basis = TotalDegreeBasis(joint, max_degree)
least_squares = LeastSquareRegression()
pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

pce.fit(x,y)

# %% md
#
# Compute PCE coefficients using Lasso regression.

# %%

polynomial_basis = PolynomialBasis.create_total_degree_basis(joint, max_degree)
lasso = LassoRegression()
pce2 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=lasso)

pce2.fit(x,y)

# %% md
#
# Compute PCE coefficients using Ridge regression.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
ridge = RidgeRegression()
pce3 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=ridge)

pce3.fit(x,y)

# %% md
# Error Estimation
# -----------------
# Validation error.

# %%

n_samples = 500
x_val = joint.rvs(n_samples)
y_val = function(x_val)

y_pce = pce.predict(x_val).flatten()
y_pce2 = pce2.predict(x_val).flatten()
y_pce3 = pce3.predict(x_val).flatten()

error = np.sum(np.abs(y_pce - y_val)/np.abs(y_val))/n_samples
error2 = np.sum(np.abs(y_pce2 - y_val)/np.abs(y_val))/n_samples
error3 = np.sum(np.abs(y_pce3 - y_val)/np.abs(y_val))/n_samples

print('Validation error, LSTSQ-PCE:', error)
print('Validation error, LASSO-PCE:', error2)
print('Validation error, Ridge-PCE:', error3)

# %% md
# Moment Estimation
# -----------------
# Returns mean and variance of the PCE surrogate.

# %%

n_mc = 1000000
x_mc = joint.rvs(n_mc)
y_mc = function(x_mc)
mean_mc = np.mean(y_mc)
var_mc = np.var(y_mc)

print('Moments from least squares regression :', pce.get_moments())
print('Moments from LASSO regression :', pce2.get_moments())
print('Moments from Ridge regression :', pce3.get_moments())
print('Moments from MC integration: ', mean_mc, var_mc)

# docs/code/surrogates/pce/plot_pce_camel.py

"""

Camel function (3 random inputs, scalar output)
======================================================================

In this example, PCE is used to generate a surrogate model for a given set of 2D data.

Six-hump camel function

.. math:: f(x) = \Big(4-2.1x_1^2 + \frac{x_1^4}{3} \Big)x_1^2 + x_1x_2 + (-4 + 4x_2^2)x_2^2

**Description:**  Dimensions: 2

**Input Domain:**  This function is evaluated on the hypercube :math:`x_1 \in [-3, 3], x_2 \in [-2, 2]`.

**Global minimum:** :math:`f(x^*)=-1.0316,` at :math:`x^* = (0.0898, -0.7126)` and
:math:`(-0.0898, 0.7126)`.

**Reference:**  Molga, M., & Smutnicki, C. Test functions for optimization needs (2005). Retrieved June 2013, from http://www.zsd.ict.pwr.wroc.pl/files/docs/functions.pdf.
"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from UQpy.distributions import Uniform, JointIndependent
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from UQpy.surrogates import *


# %% md
#
# Define the function.

# %%

def function(x, y):
    return (4 - 2.1 * x ** 2 + x ** 4 / 3) * x ** 2 + x * y + (-4 + 4 * y ** 2) * y ** 2


# %% md
#
# Create a distribution object, generate samples and evaluate the function at the samples.

# %%

np.random.seed(1)

dist_1 = Uniform(loc=-2, scale=4)
dist_2 = Uniform(loc=-1, scale=2)

marg = [dist_1, dist_2]
joint = JointIndependent(marginals=marg)

n_samples = 250
x = joint.rvs(n_samples)
y = function(x[:, 0], x[:, 1])

# %% md
#
# Visualize the 2D function.

# %%

xmin, xmax = -2, 2
ymin, ymax = -1, 1
X1 = np.linspace(xmin, xmax, 50)
X2 = np.linspace(ymin, ymax, 50)
X1_, X2_ = np.meshgrid(X1, X2)  # grid of points
f = function(X1_, X2_)

fig = plt.figure(figsize=(10, 6))
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X1_, X2_, f, rstride=1, cstride=1, cmap='gnuplot2', linewidth=0, antialiased=False)
ax.set_title('True function')
ax.set_xlabel('$x_1$', fontsize=15)
ax.set_ylabel('$x_2$', fontsize=15)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
ax.view_init(20, 50)
fig.colorbar(surf, shrink=0.5, aspect=7)

plt.show()

# %% md
#
# Visualize training data.

# %%

fig = plt.figure(figsize=(10, 6))
ax = fig.gca(projection='3d')
ax.scatter(x[:, 0], x[:, 1], y, s=20, c='r')

ax.set_title('Training data')
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
ax.view_init(20, 50)
ax.set_xlabel('$x_1$', fontsize=15)
ax.set_ylabel('$x_2$', fontsize=15)
plt.show()

# %% md
#
# Visualize training data.

# %%

max_degree = 6

# %% md
#
# Compute a PCE using least squares regression.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
least_squares = LeastSquareRegression()
pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

pce.fit(x,y)


# %% md
#
# Compute PCE using LASSO regression.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
lasso = LassoRegression()
pce2 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=lasso)

pce2.fit(x,y)

# %% md
#
# Compute PCE using Ridge regression.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
ridge = RidgeRegression()
pce3 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=ridge)

pce3.fit(x,y)

# %% md
#
# PCE surrogate is used to predict the behavior of the function at new samples.

# %%

n_test_samples = 20000
x_test = joint.rvs(n_test_samples)
y_test = pce.predict(x_test)

# %% md
#
# Plot PCE prediction.

# %%

fig = plt.figure(figsize=(10,6))
ax = fig.gca(projection='3d')
ax.scatter(x_test[:,0], x_test[:,1], y_test, s=1)

ax.set_title('PCE predictor')
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
ax.view_init(20,50)
ax.set_xlim(-2,2)
ax.set_ylim(-1,1)
ax.set_xlabel('$x_1$', fontsize=15)
ax.set_ylabel('$x_2$', fontsize=15)
plt.show()

# %% md
#
# Error Estimation
# -----------------
# Validation error.

# %%

# validation sample
n_samples = 150
x_val = joint.rvs(n_samples)
y_val = function(x_val[:,0], x_val[:,1])

# PCE predictions
y_pce  = pce.predict(x_val).flatten()
y_pce2 = pce2.predict(x_val).flatten()
y_pce3 = pce3.predict(x_val).flatten()

# mean relative validation errors
error = np.sum(np.abs((y_val - y_pce)/y_val))/n_samples
error2 = np.sum(np.abs((y_val - y_pce2)/y_val))/n_samples
error3 = np.sum(np.abs((y_val - y_pce3)/y_val))/n_samples

print('Mean rel. error, LSTSQ:', error)
print('Mean rel. error, LASSO:', error2)
print('Mean rel. error, Ridge:', error3)

# %% md
#
# Moment Estimation
# -----------------
# Returns mean and variance of the PCE surrogate.

# %%

n_mc = 1000000
x_mc = joint.rvs(n_mc)
y_mc = function(x_mc[:,0], x_mc[:,1])
mean_mc = np.mean(y_mc)
var_mc = np.var(y_mc)

print('Moments from least squares regression :', pce.get_moments())
print('Moments from LASSO regression :', pce2.get_moments())
print('Moments from Ridge regression :', pce3.get_moments())
print('Moments from Monte Carlo integration: ', mean_mc, var_mc)

# docs/code/surrogates/pce/plot_pce_exponential.py

"""

Exponential function (3 random inputs, scalar output)
======================================================================

In this example, PCE is used to generate a surrogate model for a given set of 3D data.

Dette & Pepelyshev exponential function
----------------------------------------

.. math:: f(x) = 100(\exp{(-2/x_1^{1.75})} + \exp{(-2/x_2^{1.5})} + \exp{(-2/x_3^{1.25})})

**Description:**  Dimensions: 3

**Input Domain:**  This function is evaluated on the hypercube :math:`x_i \in [0,1]` for all :math:`i = 1,2,3`.

**Reference:**  Dette, H., & Pepelyshev, A. (2010). Generalized Latin hypercube design for computer experiments. Technometrics, 52(4).
"""

# %% md
#
# Import necessary libraries.

# %%

import math
import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Uniform, JointIndependent
from UQpy.surrogates import *

# %% md
#
# Define the function.

# %%

def function(x):
    return 100*(np.exp(-2/(x[:,0]**1.75)) + np.exp(-2/(x[:,1]**1.5)) + np.exp(-2/(x[:,2]**1.25)))

# %% md
#
# Define the input probability distributions.

# %%

# input distributions
dist = Uniform(loc=0, scale=1)
marg = [dist]*3
joint = JointIndependent(marginals=marg)

# %% md
#
# Compute reference mean and variance values using Monte Carlo sampling.

# %%

# reference moments via Monte Carlo Sampling
n_samples_mc = 1000000
xx = joint.rvs(n_samples_mc)
yy = function(xx)
mean_ref = yy.mean()
var_ref = yy.var()

# %% md
#
# Create validation data sets, to be used later to estimate the accuracy of the PCE.

# %%


# validation data sets
n_samples_val = 1000
xx_val = joint.rvs(n_samples_val)
yy_val = function(xx_val)

# %% md
#
# Assess the PCE in terms of approximation and moment estimation accuracy, for increasing maximum polynomial degree.

# %%

# construct PCE surrogate models
l2_err = []
mean_err = []
var_err = []
for max_degree in range(1, 10):
    print(' ')

    # PCE basis
    print('Total degree: ', max_degree)
    polynomial_basis = TotalDegreeBasis(joint, max_degree)
    print('Size of basis:', polynomial_basis.polynomials_number)

    # generate training data
    sampling_coeff = 5
    print('Sampling coefficient: ', sampling_coeff)
    np.random.seed(42)
    n_samples = math.ceil(sampling_coeff * polynomial_basis.polynomials_number)
    print('Training data: ', n_samples)
    xx_train = joint.rvs(n_samples)
    yy_train = function(xx_train)

    # fit model
    least_squares = LeastSquareRegression()
    pce_metamodel = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce_metamodel.fit(xx_train, yy_train)

    # validation errors
    yy_val_pce = pce_metamodel.predict(xx_val)
    errors = np.abs(yy_val - yy_val_pce.flatten())
    error_l2 = np.linalg.norm(errors)
    l2_err.append(error_l2)
    print('Validation error in the l2 norm:', error_l2)

    # moment errors
    pce_moments = pce_metamodel.get_moments()
    mean_pce = pce_moments[0]
    var_pce = pce_moments[1]
    mean_err.append(np.abs((mean_pce - mean_ref) / mean_ref))
    var_err.append(np.abs((var_pce - var_ref) / var_ref))
    print('Relative error in the mean:', mean_err[-1])
    print('Relative error in the variance:', var_err[-1])


# docs/code/surrogates/pce/plot_pce_friedman.py

"""

Friedman function (5 random inputs, scalar output)
======================================================================

In this example, PCE is used to generate a surrogate model for a given set of 5D data.

Friedman function
----------------------------------------

.. math:: f(x) = 10 sin(\pi x_1x_2) + 20(x_3 - 0.5)^2 + 10x_4 + 5x_5

**Description:**  Dimensions: 5

**Input Domain:**  This function is evaluated on the hypercube :math:`x_i \in [0,1]` for all i = 1, …, 5.

**Reference:**  Friedman, J. H. (1991). Multivariate adaptive regression splines. The annals of statistics, 19(1), 1-67.
"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Uniform, JointIndependent
from UQpy.surrogates import *

# %% md
#
# Define the function.

# %%

def function(x):
    return 10*np.sin(np.pi*x[:,0]*x[:,1]) + 20*(x[:,2]-0.5)**2 + 10*x[:,3] + 5*x[:,4]

# %% md
#
# Create a distribution object, generate samples and evaluate the function at the samples.

# %%

np.random.seed(1)

dist = Uniform(loc=0, scale=1)

marg = [dist]*5
joint = JointIndependent(marginals=marg)

n_samples = 200
x = joint.rvs(n_samples)
y = function(x)

# %% md
#
# Create an object from the PCE class. Then compute PCE coefficients using least squares regression.

# %%

max_degree = 3

polynomial_basis = TotalDegreeBasis(joint, max_degree)
least_squares = LeastSquareRegression()
pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
pce.fit(x,y)

# %% md
#
# Compute PCE coefficients using Lasso regression.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
lasso = LassoRegression()
pce2 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=lasso)
pce2.fit(x,y)

# %% md
#
# Compute PCE coefficients using Ridge regression.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
ridge = RidgeRegression()
pce3 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=ridge)
pce3.fit(x,y)

# %% md
# Error Estimation
# -----------------
# Validation error.

# %%

n_samples = 100
x_val = joint.rvs(n_samples)
y_val = function(x_val)

y_pce = pce.predict(x_val).flatten()
y_pce2 = pce2.predict(x_val).flatten()
y_pce3 = pce3.predict(x_val).flatten()

error = np.sum(np.abs(y_pce - y_val)/np.abs(y_val))/n_samples
error2 = np.sum(np.abs(y_pce2 - y_val)/np.abs(y_val))/n_samples
error3 = np.sum(np.abs(y_pce3 - y_val)/np.abs(y_val))/n_samples

print('Validation error, LSTSQ-PCE:', error)
print('Validation error, LASSO-PCE:', error2)
print('Validation error, Ridge-PCE:', error3)

# %% md
# Moment Estimation
# -----------------
# Returns mean and variance of the PCE surrogate.

# %%

n_mc = 1000000
x_mc = joint.rvs(n_mc)
y_mc = function(x_mc)
mean_mc = np.mean(y_mc)
var_mc = np.var(y_mc)

print('Moments from least squares regression :', pce.get_moments())
print('Moments from LASSO regression :', pce2.get_moments())
print('Moments from Ridge regression :', pce3.get_moments())
print('Moments from MC integration: ', mean_mc, var_mc)

# docs/code/surrogates/pce/plot_pce_helmholtz.py

"""

2D Helmholtz eigenvalues (2 random inputs, vector-valued output)
======================================================================

In this example, PCE is used to generate a surrogate model for a given set of 2D data for a numerical model with
multi-dimensional outputs.
"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import matplotlib.pyplot as plt
import math
from UQpy.distributions import Normal, JointIndependent
from UQpy.surrogates import *

# %% md
#
# The analytical function below describes the eigenvalues of the 2D Helmholtz equation on a square.

# %%

def analytical_eigenvalues_2d(Ne, lx, ly):
    """
    Computes the first Ne eigenvalues of a rectangular waveguide with
    dimensions lx, ly

    Parameters
    ----------
    Ne : integer
         number of eigenvalues.
    lx : float
         length in x direction.
    ly : float
         length in y direction.

    Returns
    -------
    ev : numpy 1d array
         the Ne eigenvalues
    """
    ev = [(m * np.pi / lx) ** 2 + (n * np.pi / ly) ** 2 for m in range(1, Ne + 1)
          for n in range(1, Ne + 1)]
    ev = np.array(ev)

    return ev[:Ne]


# %% md
#
# Create a distribution object.

# %%

pdf_lx = Normal(loc=2, scale=0.02)
pdf_ly = Normal(loc=1, scale=0.01)
margs = [pdf_lx, pdf_ly]
joint = JointIndependent(marginals=margs)

# %% md
#
# Define the number of input dimensions and choose the number of output dimensions (number of eigenvalues).

# %%

dim_in = 2
dim_out = 10

# %% md
#
# Construct PCE models by varying the maximum degree of polynomials (and therefore the number of polynomial basis) and
# compute the validation error for all resulting models.

# %%

errors = []
# construct PCE surrogate models
for max_degree in range(1, 6):
    print('Total degree: ', max_degree)
    polynomial_basis = TotalDegreeBasis(joint, max_degree)

    print('Size of basis:', polynomial_basis.polynomials_number)
    # training data
    sampling_coeff = 5
    print('Sampling coefficient: ', sampling_coeff)
    np.random.seed(42)
    n_samples = math.ceil(sampling_coeff * polynomial_basis.polynomials_number)
    print('Training data: ', n_samples)
    xx = joint.rvs(n_samples)
    yy = np.array([analytical_eigenvalues_2d(dim_out, x[0], x[1]) for x in xx])

    # fit model
    least_squares = LeastSquareRegression()
    pce_metamodel = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce_metamodel.fit(xx, yy)

    # coefficients
    # print('PCE coefficients: ', pce.C)

    # validation errors
    np.random.seed(999)
    n_samples = 1000
    x_val = joint.rvs(n_samples)
    y_val = np.array([analytical_eigenvalues_2d(dim_out, x[0], x[1]) for x in x_val])
    y_val_pce = pce_metamodel.predict(x_val)
    errors.append(np.linalg.norm((y_val - y_val_pce) / y_val, ord=1, axis=0))
    print('Relative absolute errors: ', errors[-1])
    print('')

# %% md
#
# Plot errors.

# %%

errors = np.array(errors)
plt.figure(1)
for i in range(np.shape(errors)[0]):
    plt.semilogy(np.linspace(1, dim_out, dim_out), errors[i], '--o', label='pol. degree: {}'.format(i+1))
plt.legend()
plt.show()

# %% md
#
# Moment estimation (directly estimated from the last PCE metamodel).

# %%

print('Mean PCE estimate:', pce_metamodel.get_moments()[0])
print('')
print('Variance PCE estimate:', pce_metamodel.get_moments()[1])

# docs/code/surrogates/pce/plot_pce_ishigami.py

"""

Ishigami function (3 random inputs, scalar output)
======================================================================

In this example, we approximate the well-known Ishigami function with a total-degree Polynomial Chaos Expansion.

"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import math
import numpy as np
from UQpy.distributions import Uniform, JointIndependent
from UQpy.surrogates import *

# %% md
#
# We then define the Ishigami function, which reads:
# :math:`f(x_1, x_2, x_3) = \sin(x_1) + a \sin^2(x_2) + b x_3^4 \sin(x_1)`

# %%

# function to be approximated
def ishigami(xx):
    """Ishigami function"""
    a = 7
    b = 0.1
    term1 = np.sin(xx[0])
    term2 = a * np.sin(xx[1])**2
    term3 = b * xx[2]**4 * np.sin(xx[0])
    return term1 + term2 + term3

# %% md
#
# The Ishigami function has three random inputs, which are uniformly distributed in :math:`[-\pi, \pi]`. Moreover, the
# input random variables are mutually independent, which simplifies the construction of the joint distribution. Let's
# define the corresponding distributions.

# %%

# input distributions
dist1 = Uniform(loc=-np.pi, scale=2*np.pi)
dist2 = Uniform(loc=-np.pi, scale=2*np.pi)
dist3 = Uniform(loc=-np.pi, scale=2*np.pi)
marg = [dist1, dist2, dist3]
joint = JointIndependent(marginals=marg)

# %% md
#
# We now define our PCE. Only thing we need is the joint distribution.
#
# We must now select a polynomial basis. Here we opt for a total-degree (TD) basis, such that the univariate
# polynomials have a maximum degree equal to :math:`P` and all multivariate polynomial have a total-degree
# (sum of degrees of corresponding univariate polynomials) at most equal to :math:`P`. The size of the basis is then
# given by :math:`\frac{(N+P)!}{N! P!}`
# where :math:`N` is the number of random inputs (here, :math:`N+3`).

# %%

# maximum polynomial degree
P = 6

# construct total-degree polynomial basis
polynomial_basis = TotalDegreeBasis(joint, P)

# check the size of the basis
print('Size of PCE basis:', polynomial_basis.polynomials_number)

# %% md
#
# We must now compute the PCE coefficients. For that we first need a training sample of input random variable
# realizations and the corresponding model outputs. These two data sets form what is also known as an
# ''experimental design''. It is generally advisable that the experimental design has :math:`2-10` times more data points
# than the number of PCE polynomials.

# %%

# create training data
sample_size = int(polynomial_basis.polynomials_number*5)
print('Size of experimental design:', sample_size)

# realizations of random inputs
xx_train = joint.rvs(sample_size)
# corresponding model outputs
yy_train = np.array([ishigami(x) for x in xx_train])


# %% md
#
# We now fit the PCE coefficients by solving a regression problem. There are multiple ways to do this, e.g. least
# squares regression, ridge regression, LASSO regression, etc. Here we opt for the _np.linalg.lstsq_ method, which
# is based on the _dgelsd_ solver of LAPACK.

# %%

# fit model
least_squares = LeastSquareRegression()
pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

pce.fit(xx_train, yy_train)

# %% md
#
# By simply post-processing the PCE's terms, we are able to get estimates regarding the mean and standard deviation
# of the model output.

# %%

mean_est = pce.get_moments()[0]
var_est = pce.get_moments()[1]
print('PCE mean estimate:', mean_est)
print('PCE variance estimate:', var_est)


# %% md
#
# Similarly to the mean and variance estimates, we can very simply estimate the Sobol sensitivity indices, which
# quantify the importance of the input random variables in terms of impact on the model output.

# %%

from UQpy.sensitivity import *
pce_sensitivity = PceSensitivity(pce)
pce_sensitivity.run()
sobol_first = pce_sensitivity.first_order_indices
sobol_total = pce_sensitivity.total_order_indices
print('First-order Sobol indices:')
print(sobol_first)
print('Total-order Sobol indices:')
print(sobol_total)

# %% md
#
# The PCE should become increasingly more accurate as the maximum polynomial degree :math:`P` increases. We will test
# that by computing the mean absolute error (MAE) between the PCE's predictions and the true model outputs, given a
# validation sample of :math:`10^5` data points.

# %%

# validation data sets
np.random.seed(999)  # fix random seed for reproducibility
n_samples_val = 100000
xx_val = joint.rvs(n_samples_val)
yy_val = np.array([ishigami(x) for x in xx_val])

mae = []  # to hold MAE for increasing polynomial degree
for degree in range(16):
    # define PCE
    polynomial_basis = TotalDegreeBasis(joint, degree)
    least_squares = LeastSquareRegression()
    pce_metamodel = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

    # create training data
    np.random.seed(1)  # fix random seed for reproducibility
    sample_size = int(pce_metamodel.polynomials_number * 5)
    xx_train = joint.rvs(sample_size)
    yy_train = np.array([ishigami(x) for x in xx_train])

    # fit PCE coefficients
    pce_metamodel.fit(xx_train, yy_train)

    # compute mean absolute validation error
    yy_val_pce = pce_metamodel.predict(xx_val).flatten()
    errors = np.abs(yy_val.flatten() - yy_val_pce)
    mae.append(np.linalg.norm(errors, 1) / n_samples_val)

    print('Polynomial degree:', degree)
    print('Mean absolute error:', mae[-1])
    print(' ')




# docs/code/surrogates/pce/plot_pce_oakley.py

"""

Oakley Function  (1 random input, scalar output)
======================================================================

In this example, PCE is used to generate a surrogate model of a sinusoidal function with a single random input and a
scalar output.

"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Normal
from UQpy.surrogates import *

# %% md
#
# Define the sinusoidal function to be approximated.
#
# Reference: Oakley, J. E., & O'Hagan, A. (2004). Probabilistic sensitivity analysis of complex models: a Bayesian
# approach. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 66(3), 751-769.

# %%

def oakley_function(x):
    return 5 + x + np.cos(x)

# %% md
#
# Create a distribution object, generate samples and evaluate the function at the samples.

# %%

np.random.seed(1)

dist = Normal(loc=0, scale=2)
n_samples = 100
x = dist.rvs(n_samples)
y = oakley_function(x)

# %% md
#
# Create an object from the PCE class, construct a total-degree polynomial basis given a maximum polynomial degree, and
# compute the PCE coefficients using least squares regression.

# %%

max_degree = 8
polynomial_basis = TotalDegreeBasis(dist, max_degree)
least_squares = LeastSquareRegression()
pce_lstsq = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

pce_lstsq.fit(x,y)

# %% md
#
# Create an object from the PCE class, construct a total-degree polynomial basis given a maximum polynomial degree, and
# compute the PCE coefficients using LASSO regression.

# %%

polynomial_basis = TotalDegreeBasis(dist, max_degree)
lasso = LassoRegression()
pce_lasso = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=lasso)

pce_lasso.fit(x,y)

# %% md
#
# Create an object from the PCE class, construct a total-degree polynomial basis given a maximum polynomial degree, and
# compute the PCE coefficients using ridge regression.

# %%

polynomial_basis = TotalDegreeBasis(dist, max_degree)
ridge = RidgeRegression()
pce_ridge = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=ridge)

pce_ridge.fit(x, y)

# %% md
#
# PCE surrogate is used to predict the behavior of the function at new samples.

# %%

x_test = dist.rvs(100)
x_test.sort(axis=0)

y_test_lstsq = pce_lstsq.predict(x_test)
y_test_lasso = pce_lasso.predict(x_test)
y_test_ridge = pce_ridge.predict(x_test)

# %% md
#
# Plot training data, true function and PCE surrogate

# %%

n_samples_ = 1000
x_ = np.linspace(min(x_test), max(x_test), n_samples_)
f = oakley_function(x_)

plt.figure()
plt.plot(x_test, y_test_lstsq, 'g', label='PCE predictor - LSTSQ')
plt.plot(x_test, y_test_lasso, 'r', label='PCE predictor - LASSO')
plt.plot(x_test, y_test_ridge, 'b', label='PCE predictor - Ridge')
plt.scatter(x, y, label='training data')
plt.plot(x_, f, 'm', label='function')
plt.title('PCE surrogate - prediction accuracy')
plt.legend(); plt.show()

# %% md
# Error Estimation
# -----------------
# Construct a validation dataset and get the validation error.

# %%

# validation sample
n_samples = 100000
x_val = dist.rvs(n_samples)
y_val = oakley_function(x_val).flatten()

# PCE predictions
y_pce_lstsq = pce_lstsq.predict(x_val).flatten()
y_pce_lasso = pce_lasso.predict(x_val).flatten()
y_pce_ridge = pce_ridge.predict(x_val).flatten()

# mean absolute errors
error_lstsq = np.sum(np.abs(y_val - y_pce_lstsq))/n_samples
error_lasso = np.sum(np.abs(y_val - y_pce_lasso))/n_samples
error_ridge = np.sum(np.abs(y_val - y_pce_ridge))/n_samples

print('Mean absolute error from least squares regression is: ', error_lstsq)
print('Mean absolute error from LASSO regression is: ', error_lasso)
print('Mean absolute error from ridge regression is: ', error_ridge)
print(' ')

# mean relative errors
error_lstsq = np.sum( np.abs((y_val - y_pce_lstsq)/y_val) )/n_samples
error_lasso = np.sum( np.abs((y_val - y_pce_lasso)/y_val) )/n_samples
error_ridge = np.sum( np.abs((y_val - y_pce_ridge)/y_val) )/n_samples

print('Mean relative error from least squares regression is: ', error_lstsq)
print('Mean relative error from LASSO regression is: ', error_lasso)
print('Mean relative error from ridge regression is: ', error_ridge)


# %% md
# Moment Estimation
# -----------------
# Returns mean and variance of the PCE surrogate.

# %%

n_mc = 1000000
x_mc = dist.rvs(n_mc)
y_mc = oakley_function(x_mc)
mean_mc = np.mean(y_mc)
var_mc = np.var(y_mc)

print('Moments from least squares regression :', pce_lstsq.get_moments())
print('Moments from LASSO regression :', pce_lasso.get_moments())
print('Moments from Ridge regression :', pce_ridge.get_moments())
print('Moments from Monte Carlo integration: ', mean_mc, var_mc)

# docs/code/surrogates/pce/plot_pce_sinusoidal.py

"""

Sinusoidal Function  (1 random input, scalar output)
======================================================================

In this example, PCE is used to generate a surrogate model of a sinusoidal function with a single random input and a
scalar output.

"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Uniform
from UQpy.surrogates import *


# %% md
#
# Define the sinusoidal function to be approximated.

# %%

def sinusoidal_function(x):
    return x * np.sin(x) / 10.0

# %% md
#
# Create a distribution object, generate samples and evaluate the function at the samples.

# %%

np.random.seed(1)

dist = Uniform(loc=0, scale=10)
n_samples = 200
x = dist.rvs(n_samples)
y = sinusoidal_function(x)


# %% md
#
# Create an object from the PCE class, construct a total-degree polynomial basis given a maximum polynomial degree, and
# compute the PCE coefficients using least squares regression.

# %%

max_degree = 15
polynomial_basis = TotalDegreeBasis(dist, max_degree)
least_squares = LeastSquareRegression()
pce_lstsq = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

pce_lstsq.fit(x,y)

# %% md
#
# Create an object from the PCE class, construct a total-degree polynomial basis given a maximum polynomial degree, and
# compute the PCE coefficients using LASSO regression.

# %%

polynomial_basis = TotalDegreeBasis(dist, max_degree)
lasso = LassoRegression()
pce_lasso = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=lasso)

pce_lasso.fit(x,y)

# %% md
#
# Create an object from the PCE class, construct a total-degree polynomial basis given a maximum polynomial degree, and
# compute the PCE coefficients using ridge regression.

# %%

polynomial_basis = TotalDegreeBasis(dist, max_degree)
ridge = RidgeRegression()
pce_ridge = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=ridge)

pce_ridge.fit(x,y)

# %% md
#
# PCE surrogate is used to predict the behavior of the function at new samples.

# %%

x_test = dist.rvs(100)
x_test.sort(axis=0)

y_test_lstsq = pce_lstsq.predict(x_test)
y_test_lasso = pce_lasso.predict(x_test)
y_test_ridge = pce_ridge.predict(x_test)


# %% md
#
# Plot training data, true function and PCE surrogate

# %%

n_samples_ = 1000
x_ = np.linspace(min(x_test), max(x_test), n_samples_)
f = sinusoidal_function(x_)

plt.figure()
plt.plot(x_test, y_test_lstsq, 'g', label='PCE predictor - LSTSQ')
plt.plot(x_test, y_test_lasso, 'r', label='PCE predictor - LASSO')
plt.plot(x_test, y_test_ridge, 'b', label='PCE predictor - Ridge')
plt.scatter(x, y, label='training data')
plt.plot(x_, f, 'm', label='function')
plt.title('PCE surrogate - prediction accuracy')
plt.legend(); plt.show()

# %% md
# Error Estimation
# -----------------
# Construct a validation dataset and get the validation error.

# %%

# validation sample
n_samples = 100000
x_val = dist.rvs(n_samples)
y_val = sinusoidal_function(x_val).flatten()

# PCE predictions
y_pce_lstsq = pce_lstsq.predict(x_val).flatten()
y_pce_lasso = pce_lasso.predict(x_val).flatten()
y_pce_ridge = pce_ridge.predict(x_val).flatten()

# mean absolute errors
error_lstsq = np.sum(np.abs(y_val - y_pce_lstsq))/n_samples
error_lasso = np.sum(np.abs(y_val - y_pce_lasso))/n_samples
error_ridge = np.sum(np.abs(y_val - y_pce_ridge))/n_samples

print('Mean absolute error from least squares regression is: ', error_lstsq)
print('Mean absolute error from LASSO regression is: ', error_lasso)
print('Mean absolute error from ridge regression is: ', error_ridge)
print(' ')

# mean relative errors
error_lstsq = np.sum( np.abs((y_val - y_pce_lstsq)/y_val) )/n_samples
error_lasso = np.sum( np.abs((y_val - y_pce_lasso)/y_val) )/n_samples
error_ridge = np.sum( np.abs((y_val - y_pce_ridge)/y_val) )/n_samples

print('Mean relative error from least squares regression is: ', error_lstsq)
print('Mean relative error from LASSO regression is: ', error_lasso)
print('Mean relative error from ridge regression is: ', error_ridge)

# %% md
# Moment Estimation
# -----------------
# Returns mean and variance of the PCE surrogate.

# %%

n_mc = 1000000
x_mc = dist.rvs(n_mc)
y_mc = sinusoidal_function(x_mc)
mean_mc = np.mean(y_mc)
var_mc = np.var(y_mc)

print('Moments from least squares regression :', pce_lstsq.get_moments())
print('Moments from LASSO regression :', pce_lasso.get_moments())
print('Moments from Ridge regression :', pce_ridge.get_moments())
print('Moments from Monte Carlo integration: ', mean_mc, var_mc)

# docs/code/surrogates/pce/plot_pce_sparsity_lars.py

"""

Hyperbolic truncation and Least Angle Regression
======================================================================

In this example, we approximate the well-known Ishigami function with a total-degree Polynomial Chaos Expansion further
reduced by hyperbolic truncation. In order to reduce the number of basis functions, we use the best-model selection
algorithm based on Least Angle Regression.
"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import math
import numpy as np
from UQpy.distributions import Uniform, JointIndependent
from UQpy.surrogates import *


# %% md
#
# We then define the Ishigami function, which reads:
# :math:` f(x_1, x_2, x_3) = \sin(x_1) + a \sin^2(x_2) + b x_3^4 \sin(x_1)`

# %%


def ishigami(xx):
    """Ishigami function"""
    a = 7
    b = 0.1
    term1 = np.sin(xx[0])
    term2 = a * np.sin(xx[1]) ** 2
    term3 = b * xx[2] ** 4 * np.sin(xx[0])
    return term1 + term2 + term3


# %% md
#
# The Ishigami function has three indepdent random inputs, which are uniformly distributed in
# interval :math:`[-\pi, \pi]`.

# %%

# input distributions
dist1 = Uniform(loc=-np.pi, scale=2 * np.pi)
dist2 = Uniform(loc=-np.pi, scale=2 * np.pi)
dist3 = Uniform(loc=-np.pi, scale=2 * np.pi)
marg = [dist1, dist2, dist3]
joint = JointIndependent(marginals=marg)

# %% md
#
# We now define our complete PCE, which will be further used for the best model selection algorithm.

# %%

# %% md
#
# We must now select a polynomial basis. Here we opt for a total-degree (TD) basis, such that the univariate
# polynomials have a maximum degree equal to :math:`P` and all multivariate polynomial have a total-degree
# (sum of degrees of corresponding univariate polynomials) at most equal to :math:`P`. The size of the basis
# is then given by
#
# .. math:: \frac{(N+P)!}{N! P!}
#
# where :math:`N` is the number of random inputs (here, :math:`N=3`).
#
# Note that the size of the basis is highly dependent both on :math:`N` and :math:`P:math:`. It is generally advisable
# that the experimental design has :math:`2-10` times more data points than the number of PCE polynomials. This might
# lead to curse of dimensionality and thus we will utilize the best model selection algorithm based on
# Least Angle Regression.

# %%

# maximum polynomial degree
P = 15
# construct total-degree polynomial basis
polynomial_basis = TotalDegreeBasis(joint, P)

# %% md
#
# We must now compute the PCE coefficients. For that we first need a training sample of input random variable
# realizations and the corresponding model outputs. These two data sets form what is also known as an
# ''experimental design''. In case of adaptive construction of PCE by the best model selection algorithm, size of
# ED is given apriori and the most suitable basis functions are adaptively selected.

# %%

# create training data
sample_size = 500
print('Size of experimental design:', sample_size)

# realizations of random inputs
xx_train = joint.rvs(sample_size)
# corresponding model outputs
yy_train = np.array([ishigami(x) for x in xx_train])

# %% md
#
# We now fit the PCE coefficients by solving a regression problem. Here we opt for the _np.linalg.lstsq_ method,
# which is based on the _dgelsd_ solver of LAPACK. This original PCE class will be used for further selection of
# the best basis functions.

# %%

# fit model
least_squares = LeastSquareRegression()
pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
pce.fit(xx_train, yy_train)

# %% md
#
# Once we have created the PCE containing all basis functions generated by TD algorithm, it is possible to reduce
# the number of basis functions by LAR algorithm. The best model selection algorithm in UQPy is based on results
# of LAR adding basis functions to active set one-by-one until the target accuracy is obtained. Approximation
# error is measured by leave-one-out error :math:`Q^2` on given ED in :math:`[0,1]`. Target error represents the target
# accuracy measured by :math:`Q^2`.
#
# Note that if the target error is too high (close to 1), there is a risk of over-fitting. Therefore, we must check
# the over-fitting by empirical rule: if the three steps of LAR in row lead to decreasing accuracy - stop the
# algorithm. It is recommended to always check the over-fitting.

# %%

# check the size of the basis
print('Size of the full set of PCE basis:', polynomial_basis.polynomials_number)

target_error = 1
CheckOverfitting = True
pceLAR = polynomial_chaos.regressions.LeastAngleRegression.model_selection(pce, target_error, CheckOverfitting)

print('Size of the LAR PCE basis:', pceLAR.polynomial_basis.polynomials_number)

# %% md
#
# By simply post-processing the PCE's terms, we are able to get estimates regarding the mean and standard deviation of
# the model output.

# %%

mean_est, var_est = pceLAR.get_moments(higher=False)
print('PCE mean estimate:', mean_est)
print('PCE variance estimate:', var_est)

# %% md
#
# It is possible to obtain skewness and kurtosis (3rd and 4th moments), though it might be computationally demanding
# for high :math:`N` and :math:`P`.

# %%

mean_est, var_est, skew_est, kurt_est = pceLAR.get_moments(True)
print('PCE mean estimate:', mean_est)
print('PCE variance estimate:', var_est)
print('PCE skewness estimate:', skew_est)
print('PCE kurtosis estimate:', kurt_est)

# %% md
#
# Similarly to the statistical moments, we can very simply estimate the Sobol sensitivity indices, which quantify the
# importance of the input random variables in terms of impact on the model output.

# %%

from UQpy.sensitivity import *

pce_sensitivity = PceSensitivity(pceLAR)
pce_sensitivity.run()
sobol_first = pce_sensitivity.first_order_indices
sobol_total = pce_sensitivity.total_order_indices
print('First-order Sobol indices:')
print(sobol_first)
print('Total-order Sobol indices:')
print(sobol_total)

# %% md
#
# The accuracy of PCE is typically measured by leave-one-out error :math:`Q^2` on given ED. Moreover, we will test that
# also by computing the mean absolute error (MAE) between the PCE's predictions and the true model outputs, given a
# validation sample of :math:`10^5` data points.

# %%

n_samples_val = 100000
xx_val = joint.rvs(n_samples_val)
yy_val = np.array([ishigami(x) for x in xx_val])

yy_val_pce = pceLAR.predict(xx_val).flatten()
errors = np.abs(yy_val.flatten() - yy_val_pce)
MAE = (np.linalg.norm(errors, 1) / n_samples_val)

print('Mean absolute error:', MAE)
print('Leave-one-out cross validation on ED:', pceLAR.leaveoneout_error())

# %% md
#
# For the comparison, we can check results of PCE solved by OLS without the model selection algorithm. Note that, it is
# necessary to use 2−10 times more data points than the number of PCE polynomials.

# %%

# validation data sets
np.random.seed(999)  # fix random seed for reproducibility
n_samples_val = 100000
xx_val = joint.rvs(n_samples_val)
yy_val = np.array([ishigami(x) for x in xx_val])

mae = []  # to hold MAE for increasing polynomial degree
for degree in range(16):
    # define PCE
    polynomial_basis = TotalDegreeBasis(joint, degree)
    least_squares = LeastSquareRegression()
    pce_metamodel = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

    # create training data
    np.random.seed(1)  # fix random seed for reproducibility
    sample_size = int(pce_metamodel.polynomials_number * 5)
    xx_train = joint.rvs(sample_size)
    yy_train = np.array([ishigami(x) for x in xx_train])

    # fit PCE coefficients
    pce_metamodel.fit(xx_train, yy_train)

    # compute mean absolute validation error
    yy_val_pce = pce_metamodel.predict(xx_val).flatten()
    errors = np.abs(yy_val.flatten() - yy_val_pce)
    mae.append(np.linalg.norm(errors, 1) / n_samples_val)
    print('Size of ED:', sample_size)
    print('Polynomial degree:', degree)
    print('Mean absolute error:', mae[-1])
    print(' ')

# %% md
#
# In case of high-dimensional input and/or high :math:P` it is also beneficial to reduce the TD basis set by hyperbolic
# trunction. The hyperbolic truncation reduces higher-order interaction terms in dependence to parameter :math:`q` in
# interval :math:`(0,1)`. The set of multi indices :math:`\alpha` is reduced as follows:
#
# :math:`\alpha\in \mathbb{N}^{N}: || \boldsymbol{\alpha}||_q \equiv \Big( \sum_{i=1}^{N} \alpha_i^q \Big)^{1/q} \leq P`
#
# Note that :math:`q=1` leads to full TD set.

# %%

print('Size of the full set of PCE basis:', TotalDegreeBasis(joint, P).polynomials_number)
q = 0.8
polynomial_basis_hyperbolic = TotalDegreeBasis(joint, P, q)
# check the size of the basis
print('Size of the hyperbolic full set of PCE basis:', polynomial_basis_hyperbolic.polynomials_number)

# %% md
#
# The reduction of the full set size significantly reduces the necessary number of data points in ED for non-adaptive
# PCE. However, it is suitable only for mathematical models without significant higher-order interaction terms.

# %%

pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis_hyperbolic, regression_method=least_squares)
pce.fit(xx_train, yy_train)
yy_val_pce = pce.predict(xx_val).flatten()
errors = np.abs(yy_val.flatten() - yy_val_pce)
MAE = (np.linalg.norm(errors, 1) / n_samples_val)

print('Mean absolute error:', MAE)
print('Leave-one-out cross validation on ED:', pce.leaveoneout_error())


# docs/code/surrogates/pce/plot_pce_sphere.py

"""

Sinusoidal Function  Sphere function (2 random inputs, scalar output)
======================================================================

In this example, PCE is used to generate a surrogate model for a given set of 2D data.

.. math:: f(x) = x_1^2 + x_2^2

**Description:**  Dimensions: 2

**Input Domain:**  This function is evaluated on the hypercube :math:`x_i \in [-5.12, 5.12]` for all :math:`i = 1,2`.

**Global minimum:** :math:`f(x^*)=0,` at :math:`x^* = (0,0)`.

**Reference:**  Dixon, L. C. W., & Szego, G. P. (1978). The global optimization problem: an introduction. Towards global optimization, 2, 1-15.
"""

# %% md
#
# Import necessary libraries.

# %%

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from UQpy.surrogates import *
from UQpy.distributions import Uniform, JointIndependent

# %% md
#
# Define the function.

# %%

def function(x,y):
    return x**2 + y**2

# %% md
#
# Create a distribution object, generate samples and evaluate the function at the samples.

# %%

np.random.seed(1)

dist_1 = Uniform(loc=-5.12, scale=10.24)
dist_2 = Uniform(loc=-5.12, scale=10.24)

marg = [dist_1, dist_2]
joint = JointIndependent(marginals=marg)

n_samples = 100
x = joint.rvs(n_samples)
y = function(x[:,0], x[:,1])

# %% md
#
# Visualize the 2D function.

# %%

xmin, xmax = -6,6
ymin, ymax = -6,6
X1 = np.linspace(xmin, xmax, 50)
X2 = np.linspace(ymin, ymax, 50)
X1_, X2_ = np.meshgrid(X1, X2) # grid of points
f = function(X1_, X2_)

fig = plt.figure(figsize=(10,6))
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X1_, X2_, f, rstride=1, cstride=1, cmap='gnuplot2', linewidth=0, antialiased=False)
ax.set_title('True function')
ax.set_xlabel('$x_1$', fontsize=15)
ax.set_ylabel('$x_2$', fontsize=15)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
ax.view_init(20, 140)
fig.colorbar(surf, shrink=0.5, aspect=7)

plt.show()

# %% md
#
# Visualize training data.

# %%

fig = plt.figure(figsize=(10,6))
ax = fig.gca(projection='3d')
ax.scatter(x[:,0], x[:,1], y, s=20, c='r')

ax.set_title('Training data')
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
ax.view_init(20,140)
ax.set_xlabel('$x_1$', fontsize=15)
ax.set_ylabel('$x_2$', fontsize=15)
plt.show()

# %% md
#
# Create an object from the PCE class. Compute PCE coefficients using least squares regression.

# %%

max_degree = 3
polynomial_basis = TotalDegreeBasis(joint, max_degree)
least_squares = LeastSquareRegression()
pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)

pce.fit(x,y)

# %% md
#
# Compute PCE coefficients using LASSO.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
lasso = LassoRegression()
pce2 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=lasso)

pce2.fit(x,y)

# %% md
#
# Compute PCE coefficients with Ridge regression.

# %%

polynomial_basis = TotalDegreeBasis(joint, max_degree)
ridge = RidgeRegression()
pce3 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=ridge)

pce3.fit(x,y)

# %% md
#
# PCE surrogate is used to predict the behavior of the function at new samples.

# %%

n_test_samples = 10000
x_test = joint.rvs(n_test_samples)
y_test = pce.predict(x_test)


# %% md
#
# Plot PCE prediction.

# %%

fig = plt.figure(figsize=(10,6))
ax = fig.gca(projection='3d')
ax.scatter(x_test[:,0], x_test[:,1], y_test, s=1)

ax.set_title('PCE predictor')
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
ax.view_init(20,140)
ax.set_xlim(-6,6)
ax.set_ylim(-6,6)
ax.set_xlabel('$x_1$', fontsize=15)
ax.set_ylabel('$x_2$', fontsize=15)
plt.show()

# %% md
# Error Estimation
# -----------------
# Construct a validation dataset and get the validation error.

# %%

# validation sample
n_samples = 150
x_val = joint.rvs(n_samples)
y_val = function(x_val[:,0], x_val[:,1])

# PCE predictions
y_pce  = pce.predict(x_val).flatten()
y_pce2 = pce2.predict(x_val).flatten()
y_pce3 = pce3.predict(x_val).flatten()

# mean relative validation errors
error = np.sum(np.abs((y_val - y_pce)/y_val))/n_samples
error2 = np.sum(np.abs((y_val - y_pce2)/y_val))/n_samples
error3 = np.sum(np.abs((y_val - y_pce3)/y_val))/n_samples

print('Mean rel. error, LSTSQ:', error)
print('Mean rel. error, LASSO:', error2)
print('Mean rel. error, Ridge:', error3)

# %% md
# Moment Estimation
# -----------------
# Returns mean and variance of the PCE surrogate.

# %%

n_mc = 1000000
x_mc = joint.rvs(n_mc)
y_mc = function(x_mc[:,0], x_mc[:,1])
mean_mc = np.mean(y_mc)
var_mc = np.var(y_mc)

print('Moments from least squares regression :', pce.get_moments())
print('Moments from LASSO regression :', pce2.get_moments())
print('Moments from Ridge regression :', pce3.get_moments())
print('Moments from Monte Carlo integration: ', mean_mc, var_mc)

# docs/code/surrogates/srom/local_eigenvalue_model.py

"""

Auxiliary File
======================================================================

"""
import numpy as np


class RunPythonModel:

    def __init__(self, samples=None, dimension=None):

        self.samples = samples
        self.dimension = dimension
        self.qoi = np.zeros_like(self.samples)
        for i in range(self.samples.shape[0]):
            p = np.array([[self.samples[i, 0]+self.samples[i, 1], -self.samples[i, 1], 0], 
                          [-self.samples[i, 1], self.samples[i, 1]+self.samples[i, 2], -self.samples[i, 2]], 
                          [0, -self.samples[i, 2], self.samples[i, 2]]])
            w, v = np.linalg.eig(p)
            self.qoi[i, :] = w



# docs/code/surrogates/srom/plot_srom_eigenvalues.py

"""

SROM on the eigenvalues of a system
======================================================================

In this example, Uncertainty in eigenvalues of a system is studied using :class:`.SROM` and it is compared with the
Monte Carlo Simulation results. Stiffness of each element (i.e. k1, k2 and k3) are treated as random variables which
follows gamma distribution. :class:`.SROM` is created for all three random variables and distribution of eigenvalues are
identified using :class:`.SROM`.
"""

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.MonteCarloSampling`, :class:`.TrueStratifiedSampling` and :class:`.SROM` class from UQpy.

# %%
import shutil

from UQpy import PythonModel
from UQpy.sampling import MonteCarloSampling, TrueStratifiedSampling
from UQpy.sampling import RectangularStrata
from UQpy.distributions import Gamma
from UQpy.surrogates import SROM
from UQpy.run_model.RunModel import RunModel
from scipy.stats import gamma
import numpy as np
import matplotlib.pyplot as plt

# %% md
#
# Create a distribution object for :class:`.Gamma` distribution with shape, shift and scale parameters as :math:`2`,
# :math:`1` and :math:`3`.

# %%

marginals = [Gamma(a=2., loc=1., scale=3.), Gamma(a=2., loc=1., scale=3.), Gamma(a=2., loc=1., scale=3.)]

# %% md
#
# Create a strata object.

# %%

strata = RectangularStrata(strata_number=[3, 3, 3])

# %% md
#
# Using UQpy :class:`.TrueStratifiedSampling` class to generate samples for two random variables having :class:`.Gamma`
# distribution.

# %%

x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1)


# %% md
#
# Run :class:`.SROM` to minimize the error in distribution, first order and second order moment about origin.

# %%

y = SROM(samples=x.samples, target_distributions=marginals, moments=[[6, 6, 6], [54, 54, 54]],
         weights_errors=[1, 1, 0], properties=[True, True, True, False])

# %% md
#
# Plot the sample sets and weights from :class:`.SROM` class. Also, compared with the CDF of gamma distribution of k1.

# %%

# Arrange samples in increasing order and sort samples accordingly
com = np.append(y.samples, np.transpose(np.matrix(y.sample_weights)), 1)
srt = com[np.argsort(com[:, 0].flatten())]
s = np.array(srt[0, :, 0])
a = srt[0, :, 3]
a0 = np.array(np.cumsum(a))
# Plot the SROM approximation and compare with actual gamma distribution
l = 3
fig = plt.figure()
plt.rcParams.update({'font.size': 12})
plt.plot(s[0], a0[0], linewidth=l)
plt.plot(np.arange(3, 12, 0.05), gamma.cdf(np.arange(3, 12, 0.05), a=2, loc=3, scale=1), linewidth=l)
plt.legend(['SROM Approximation', 'Gamma CDF'], loc=5, prop={'size': 12}, bbox_to_anchor=(1, 0.75))
plt.show()

# %% md
#
# Run the model 'eigenvalue_model.py' for each sample generated through :class:`.TrueStratifiedSampling` class. This
# model defines the stiffness matrix corresponding to each sample and estimate the eigenvalues of the matrix.

# %%

m = PythonModel(model_script='local_eigenvalue_model.py',model_object_name="RunPythonModel" )
model = RunModel(model=m)
# model = RunModel(model_script='local_eigenvalue_model.py')
model.run(samples=y.samples)
r_srom = model.qoi_list

# %% md
#
# :class:`MonteCarloSampling` class is used to generate 1000 samples.

# %%


x_mcs = MonteCarloSampling(distributions=marginals, nsamples=1000)

# %% md
#
# Run the model 'eigenvalue_model.py' for each sample generated through :class:`.MonteCarloSampling` class.

# %%

model.run(samples=x_mcs.samples, append_samples=False)
r_mcs = model.qoi_list

# %% md
#
# Plot the distribution of each eigenvalue, estimated using :class:`.SROM` and :class:`.MonteCarloSampling` weights.

# %%

# Plot SROM and MCS approximation for first eigenvalue
r = np.array([x[:,0] for x in r_srom])
r_mcs = np.squeeze(np.array(r_mcs))
com = np.append(np.atleast_2d(r), np.transpose(np.matrix(y.sample_weights)), 1)
srt = com[np.argsort(com[:, 0].flatten())]
s = np.array(srt[0, :, 0])
a = srt[0, :, 1]
a0 = np.array(np.cumsum(a))
fig1 = plt.figure()
plt.plot(s[0], a0[0], linewidth=l)
r_mcs0 = r_mcs[np.argsort(r_mcs[:, 0].flatten())]
plt.plot(r_mcs0[:, 0], np.cumsum(0.001*np.ones([1, 1000])), linewidth=l)
plt.title('Eigenvalue, $\lambda_1$')
plt.legend(['SROM', 'MCS'], loc=1, prop={'size': 12}, bbox_to_anchor=(1, 0.92))
plt.show()

# Plot SROM and MCS approximation for second eigenvalue
r = np.array([x[:,1] for x in r_srom])
com = np.append(np.atleast_2d(r), np.transpose(np.matrix(y.sample_weights)), 1)
srt = com[np.argsort(com[:, 0].flatten())]
s = np.array(srt[0, :, 0])
a = srt[0, :, 1]
a0 = np.array(np.cumsum(a))
fig2 = plt.figure()
plt.plot(s[0], a0[0], linewidth=l)
r_mcs0 = r_mcs[np.argsort(r_mcs[:, 1].flatten())]
plt.plot(r_mcs0[:, 1], np.cumsum(0.001*np.ones([1, 1000])), linewidth=l)
plt.title('Eigenvalue, $\lambda_2$')
plt.legend(['SROM', 'MCS'], loc=1, prop={'size': 12}, bbox_to_anchor=(1, 0.92))
plt.show()

# Plot SROM and MCS approximation for third eigenvalue
r = np.array([x[:,2] for x in r_srom])
com = np.append(np.atleast_2d(r), np.transpose(np.matrix(y.sample_weights)), 1)
srt = com[np.argsort(com[:, 0].flatten())]
s = np.array(srt[0, :, 0])
a = srt[0, :, 1]
a0 = np.array(np.cumsum(a))
fig3 = plt.figure()
plt.plot(s[0], a0[0], linewidth=l)
r_mcs0 = r_mcs[np.argsort(r_mcs[:, 2].flatten())]
plt.plot(r_mcs0[:, 2], np.cumsum(0.001*np.ones([1, 1000])), linewidth=l)
plt.title('Eigenvalue, $\lambda_3$')
plt.legend(['SROM', 'MCS'], loc=1, prop={'size': 12}, bbox_to_anchor=(1, 0.92))
plt.show()

# %% md
#
# Note: Monte Carlo Simulation used 1000 samples, whereas SROM used 27 samples.


# docs/code/surrogates/srom/plot_srom_gamma.py

"""

SROM on a Gamma distribution
======================================================================

In this example, Stratified sampling is used to generate samples from Gamma distribution and weights are defined using
Stochastic Reduce Order Model (SROM).
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.TrueStratifiedSampling` and :class:`.SROM` class from UQpy.

#%%

from UQpy.surrogates import SROM
from UQpy.sampling import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling
from UQpy.distributions import Gamma
import scipy.stats as stats
import matplotlib.pyplot as plt
import numpy as np

#%% md
#
# Create a distribution object for :class:`.Gamma` distribution with shape, shift and scale parameters as :math:`2`,
# :math:`1` and :math:`3`.

#%%

marginals = [Gamma(a=2., loc=1., scale=3.), Gamma(a=2., loc=1., scale=3.)]

#%% md
#
# Create a strata object.

#%%

strata = RectangularStrata(strata_number=[4, 4])

#%% md
#
# Using UQpy :class:`.TrueStratifiedSampling` class to generate samples for two random variables having :class:`.Gamma`
# distribution.

#%%

x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1)

#%% md
#
# Run :class:`.SROM` using the defined :class:`.Gamma` distribution. Here we use the following parameters.
#
# - :class:`.Gamma` distribution with shape, shift and scale parameters as :math:`2`, :math:`1` and :math:`3`.
# - First and second order moments about origin are :math:`6` and :math:`54`.
# - Notice that :code:`pdf_target` references the :class:`.Gamma` function directly and does not designate it as a string.
# - Samples are uncorrelated, i.e. also default value of correlation.

#%%

y = SROM(samples=x.samples,
         target_distributions=marginals,
         moments=np.array([[6., 6.], [54., 54.]]))
y.run(properties=[True, True, True, True])

#%% md
#
# Plot the samples and weights from :class:`.SROM` class. Also, compared with the CDF of gamma distribution.

#%%

c = np.concatenate((y.samples, y.sample_weights.reshape(y.sample_weights.shape[0], 1)), axis=1)
d = c[c[:, 0].argsort()]
plt.plot(d[:, 0], np.cumsum(d[:, 2], axis=0), 'o')
plt.plot(np.arange(1, 15, 0.1), stats.gamma.cdf(np.arange(1, 15, 0.1), 2, loc=1, scale=3))
plt.legend(['RV1_SROM', 'CDF'])
plt.show()
e = c[c[:, 1].argsort()]
plt.plot(e[:, 1], np.cumsum(e[:, 2], axis=0), 'o')
plt.plot(np.arange(1, 15, 0.1), stats.gamma.cdf(np.arange(1, 15, 0.1), 2, loc=1, scale=3))
plt.legend(['RV2_SROM', 'CDF'])
plt.show()


#%% md
#
# A note on the weights corresponding to error in distribution, moments and correlation of random variables:
#
# - For this illustration, error_weights are not defined and default value is :math:`[1, 0.2, 0]`. These weights can be changed to obtain desired accuracy in certain properties.

#%%

print(y.sample_weights)


# docs/code/surrogates/srom/plot_srom_gamma_2.py

"""

SROM on a Gamma distribution 2
======================================================================

In this example, Stratified sampling is used to generate samples from Gamma distribution and weights are defined using
Stochastic Reduce Order Model (SROM). This example illustrate how to define same weights for each sample of a
random variable.
"""

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.TrueStratifiedSampling` and :class:`.SROM` class from UQpy.

# %%

from UQpy.surrogates import SROM
from UQpy.sampling import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling
from UQpy.distributions import Gamma
import scipy.stats as stats
import matplotlib.pyplot as plt
import numpy as np

# %% md
#
# Create a distribution object for :class:`.Gamma` distribution with shape, shift and scale parameters as :math:`2`,
# :math:`1` and :math:`3`.

# %%

marginals = [Gamma(a=2., loc=1., scale=3.), Gamma(a=2., loc=1., scale=3.)]

# %% md
#
# Create a strata object.

# %%

strata = RectangularStrata(strata_number=[4, 4])

# %% md
#
# Using UQpy :class:`.TrueStratifiedSampling` class to generate samples for two random variables having :class:`.Gamma`
# distribution.

# %%

x = TrueStratifiedSampling(distributions=marginals,
                           strata_object=strata,
                           nsamples_per_stratum=1)

# %% md
#
# Run :class:`.SROM` using the defined Gamma distribution. Here we use the following parameters.
#
# - :class:`.Gamma` distribution with shape, shift and scale parameters as :math:`2`, :math:`1` and :math:`3`.
# - First and second order moments about origin are :math:`6` and :math:`54`.
# - Notice that :code:`pdf_target` references the :class:`.Gamma` function directly and does not designate it as a string.
# - Samples are uncorrelated, i.e. also default value of correlation.

# %%

y1 = SROM(samples=x.samples, target_distributions=marginals, moments=[[6., 6.], [54., 54.]],
          properties=[True, True, True, False])

# %% md
#
# In this case, sample_weights are generated using default values of :code:`weights_distribution`,
# :code:`weights_moments` and :code:`weights_correlation`. Default values are:

# %%

print('weights_distribution', '\n', y1.weights_distribution, '\n', 'weights_moments', '\n', y1.weights_moments, '\n',
      'weights_correlation', '\n', y1.weights_correlation)

y2 = SROM(samples=x.samples, target_distributions=marginals, moments=[[6., 6.], [54., 54.]],
          properties=[True, True, True, False],
          weights_distribution=[[0.4, 0.5]], weights_moments=[[0.2, 0.7]],
          weights_correlation=[[0.3, 0.4], [0.4, 0.6]])

# %% md
#
# In second case, :code:`weights_distribution` is modified by :class:`SROM` class. First, it defines an array of size
# :math:`2×16` with all elements equal to :math:`1` and then multiply first column by :math:`0.4` and second column by
# :math:`0.5` . Similarly, :code:`weights_moments` and :code:`weights_correlation` are modified.

# %%

print('weights_distribution', '\n', y2.weights_distribution, '\n', 'weights_moments', '\n', y2.weights_moments, '\n',
      'weights_correlation', '\n', y2.weights_correlation)

# %% md
#
# Plot below shows the comparison of samples weights generated using two different weights with the actual CDF of
# gamma distribution.

# %%

c1 = np.concatenate((y1.samples, y1.sample_weights.reshape(y1.sample_weights.shape[0], 1)), axis=1)
d1 = c1[c1[:, 0].argsort()]
c2 = np.concatenate((y2.samples, y2.sample_weights.reshape(y2.sample_weights.shape[0], 1)), axis=1)
d2 = c2[c2[:, 0].argsort()]
plt.plot(d1[:, 0], np.cumsum(d1[:, 2], axis=0), 'x')
plt.plot(d2[:, 0], np.cumsum(d2[:, 2], axis=0), 'o')
plt.plot(np.arange(1,15,0.1), stats.gamma.cdf(np.arange(1,15,0.1), 2, loc=1, scale=3))
plt.legend(['Case 1','Case 2','CDF'])
plt.title('1st random variable')
plt.show()

# %% md
#
# A note on the weights corresponding to distribution, moments and correlation of random variables:
#
# - For this illustration, default weights_moments are square of reciprocal of moments. Thus, moments should be of :any:'list[float]' type.


# docs/code/transformations/nataf/plot_nataf.py

"""

Nataf
==============================================

"""

# %% md
#
# We'll be using UQpy's Nataf transformation functionalities. We also use Matplotlib to display results graphically.
#
# Additionally, this demonstration opts to use Numpy's random state management to ensure that results are reproducible
# between notebook runs.

# %%

from UQpy.transformations import Nataf
import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Normal, Gamma, Lognormal
from UQpy.transformations import Decorrelate, Correlate


# %% md
#
# We will start by constructing two non-Gaussian random variables. The first one, follows a :code:`Gamma` distribution
# :code:`dist1`,  while the second one follows a :code:`Lognormal` distribution  :code:`dist2`.
# The two distributions are correlated according to the correlation matrix :code:`Rx`.

# %%

dist1 = Gamma(4.0, loc=0.0, scale=1.0)
dist2 = Lognormal(s=2., loc=0., scale=np.exp(1))
Rx = np.array([[1.0, 0.9], [0.9, 1.0]])

# %% md
#
# Next, we'll construct a :code:`Nataf` object :code:`nataf_obj`. Here, we provide the distribution of the random
# variables :code:`distributions` and the correlation matrix :code:`corr_x` in the parameter space.

# %%

nataf_obj = Nataf(distributions=[dist1, dist2], corr_x=Rx)

# %% md
#
# We can use the :code:`rvs` method of the :code:`nataf_obj` object to draw random samples from the two distributions.

# %%

samples_x = nataf_obj.rvs(1000)

# %% md
#
# We can visualize the samples by plotting them on axes of each distribution's range.

# %%


plt.figure()
plt.title('non-Gaussian random variables')
plt.scatter(samples_x[:, 0], samples_x[:, 1])
plt.grid(True)
plt.xlabel('$X_1$')
plt.ylabel('$X_2$')
plt.show()

# %% md
#
# We can invoke the :code:`run` method of the :code:`nataf_obj` object to transform the non-Gaussian random variables
# from space :code:`X`to the correlated standard normal space :code:`Z`. Here, we provide as :code:`boolean (True)` an
# optional attribute :code:`jacobian` to calculate the Jacobian of the transformation. The distorted correlation matrix
# can be accessed in the attribute :code:`corr_z`

# %%

nataf_obj.run(samples_x=samples_x, jacobian=True)
print(nataf_obj.corr_z)


# %% md
#
# We can visualize the correlated (standard normal) samples by plotting them on axes of each distribution's range.

# %%


plt.figure()
plt.title('Correlated standard normal samples')
plt.scatter(nataf_obj.samples_z[:, 0], nataf_obj.samples_z[:, 1])
plt.grid(True)
plt.xlabel('$Z_1$')
plt.ylabel('$Z_2$')
plt.show()

# %% md
#
# We can use the :code:`Decorrelate` class to transform the correlated standard normal samples to the uncorrelated
# standard normal space :code:`U`.

# %%

samples_u = Decorrelate(nataf_obj.samples_z, nataf_obj.corr_z).samples_u

# %% md
#
# We can visualize the uncorrelated (standard normal) samples by plotting them on axes of each distribution's range.

# %%

plt.figure()
plt.title('Uncorrelated standard normal samples')
plt.scatter(samples_u[:, 0], samples_u[:, 1])
plt.grid(True)
plt.xlabel('$U_1$')
plt.ylabel('$U_2$')
plt.show()

# %% md
#
# We can use the :code:`Correlate` class to transform the uncorrelated standard normal samples back to the correlated
# standard normal space :code:`Z`.

# %%

samples_z = Correlate(samples_u, nataf_obj.corr_z).samples_z

# %% md
#
# We can visualize the correlated (standard normal) samples by plotting them on axes of each distribution's range.

# %%

plt.figure()
plt.title('Correlated standard normal samples')
plt.scatter(samples_z[:, 0], samples_z[:, 1])
plt.grid(True)
plt.xlabel('$U_1$')
plt.ylabel('$U_2$')
plt.show()


# %% md
#
# In the second example, we will calculate the distortion of the correlation coefficient (in the standard normal space)
# as a function of the correlation coefficient (in the parameter space).  To this end, we consider N=20 values for the
# latter, ranging from -0.999 to 0.999. We do not consider the values -1 and 1 since they result in numerical
# instabilities.

# %%

N = 20
w3 = np.zeros(N)
rho = np.linspace(-0.999, 0.999, N)
for i in range(N):
    Rho1 = np.array([[1.0, rho[i]], [rho[i], 1.0]])
    ww = Nataf([dist1, dist2], corr_x=Rho1).corr_z
    w3[i] = ww[0, 1]

plt.plot(rho, w3)
plt.xlabel('rho_X')
plt.ylabel('rho_Z')
plt.show()


# %% md
#
# In the third example, we will calculate the distortion of the correlation coefficient (in the parameter space) as a
# function of the correlation coefficient (in the standard normal space).  To this end, we consider N=20 values for the
# latter, ranging from -0.999 to 0.999. We do not consider the values -1 and 1 since they result in numerical
# instabilities.

# %%

w4 = np.zeros(N)
rho = np.linspace(-0.999, 0.999, N)
for i in range(N):
    Rho1 = np.array([[1.0, rho[i]], [rho[i], 1.0]])
    ww = Nataf(distributions=[dist1, dist2], corr_z=Rho1).corr_x
    w4[i] = ww[0, 1]

plt.plot(rho, w4)
plt.plot(rho, rho)
plt.show()


# src/UQpy/dimension_reduction/diffusion_maps/DiffusionMaps.py

import itertools

import scipy.sparse as sps
import scipy as sp
from scipy.sparse.linalg import eigsh, eigs
import scipy.spatial.distance as sd
import matplotlib.pyplot as plt
import scipy
import numpy as np
from itertools import combinations

from UQpy.utilities import GaussianKernel, GrassmannPoint
from UQpy.utilities.Utilities import *
from UQpy.utilities.Utilities import _nn_coord
from beartype import beartype
from typing import Annotated, Union
from beartype.vale import Is
from UQpy.utilities.ValidationTypes import Numpy2DFloatArray, NumpyFloatArray
from UQpy.utilities.kernels.baseclass import Kernel


class DiffusionMaps:
    AlphaType = Annotated[Union[float, int], Is[lambda number: 0 <= number <= 1]]
    IntegerLargerThanUnityType = Annotated[int, Is[lambda number: number >= 1]]

    @beartype
    def __init__(
            self,
            kernel_matrix: Numpy2DFloatArray = None,
            data: Union[Numpy2DFloatArray, list[GrassmannPoint]] = None,
            kernel: Kernel = None,
            alpha: AlphaType = 0.5,
            n_eigenvectors: IntegerLargerThanUnityType = 2,
            is_sparse: bool = False,
            n_neighbors: IntegerLargerThanUnityType = 1,
            random_state: RandomStateType = None,
            t: int = 1
    ):
        """

        :param kernel_matrix: Kernel matrix defining the similarity between the data points. Either `kernel_matrix` or
            both `data` and `kernel` parameters must be provided. In the former case, `kernel_matrix` is precomputed
            using a :class:`Kernel` class. In the second case the `kernel_matrix` is internally and used for the
            evaluation of the :class:`.DiffusionMaps`. In case all three of the aforementioned parameters are provided,
            then :class:`.DiffusionMaps` will be fitted only using the `kernel_matrix`
        :param data: Set of data points. Either `kernel_matrix` or both `data` and `kernel` parameters must be
            provided.
        :param kernel: Kernel object used to compute the kernel matrix defining similarity between the data points.
            Either `kernel_matrix` or both `data` and `kernel` parameters must be provided.
        :param alpha: A scalar that corresponds to different diffusion operators. `alpha` should be between zero and
            one.
        :param n_eigenvectors: Number of eigenvectors to retain.
        :param is_sparse:  Work with sparse matrices to improve computational performance.
        :param n_neighbors: If :code:`is_sparse is True`, defines the number of nearest neighbors to use when making
            matrices sparse.
        :param random_state: Random seed used to initialize the pseudo-random number generator. If an :any:`int` is
         provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise, the
         object itself can be passed directly.
        :param t: Time exponent.
        """
        self.parsimonious_residuals = None
        """Residuals calculated from the Parsimonious Representation. This attribute will only be populated if the 
        :py:meth:`parsimonious`  method is invoked."""
        self.parsimonious_indices = None
        """Indices of the most important eigenvectors. This attribute will only be populated if the 
        :py:meth:`parsimonious`  method is invoked."""

        if kernel_matrix is not None:
            self.kernel_matrix = kernel_matrix
        elif data is not None and kernel is not None:
            kernel.calculate_kernel_matrix(points=data)
            self.kernel_matrix = kernel.kernel_matrix
        else:
            raise ValueError("Either `kernel_matrix` or both `data` and `kernel` must be provided")

        self.alpha = alpha
        self.eigenvectors_number = n_eigenvectors
        self.is_sparse = is_sparse
        self.neighbors_number = n_neighbors

        self.random_state = random_state,
        self.t = t

        self.transition_matrix: np.ndarray = None
        '''
        Markov Transition Probability Matrix.
        '''
        self.diffusion_coordinates: np.ndarray = None
        '''
        Coordinates of the data in the diffusion space.
        '''
        self.eigenvectors: np.ndarray = None
        '''
        Eigenvectors of the transition probability matrix.
        '''
        self.eigenvalues: np.ndarray = None
        '''
        Eigenvalues of the transition probability matrix.
        '''
        self.cut_off = None

        self._fit()

    def _fit(self):
        """
        Perform diffusion map embedding.
        """

        if self.is_sparse:
            self.kernel_matrix = self.__sparse_kernel(self.kernel_matrix, self.neighbors_number)

        alpha = self.alpha
        # Compute the diagonal matrix D(i,i) = sum(Kernel(i,j)^alpha,j) and its inverse.
        # d, d_inv = self._d_matrix(self.kernel_matrix, self.alpha)
        d = np.array(self.kernel_matrix.sum(axis=1)).flatten()
        d_inv = np.power(d, -alpha)

        # Compute L^alpha = D^(-alpha)*L*D^(-alpha).
        m = d_inv.shape[0]
        d_alpha = sps.spdiags(d_inv, 0, m, m) if self.is_sparse else np.diag(d_inv)
        l_star = d_alpha.dot(self.kernel_matrix.dot(d_alpha))

        # d_star, d_star_inv = self._d_matrix(l_star, 1.0)
        d_star = np.array(l_star.sum(axis=1)).flatten()
        d_star_inv = np.power(d_star, -1)

        if self.is_sparse:
            d_star_inv_d = sps.spdiags(d_star_inv, 0, d_star_inv.shape[0], d_star_inv.shape[0])
        else:
            d_star_inv_d = np.diag(d_star_inv)

        # Compute the transition matrix.
        transition_matrix = d_star_inv_d.dot(l_star)

        if self.is_sparse:
            is_symmetric = sp.sparse.linalg.norm(transition_matrix - transition_matrix.T, sp.inf) < 1e-08
        else:
            is_symmetric = np.allclose(transition_matrix, transition_matrix.T, rtol=1e-5, atol=1e-08)

        # Find the eigenvalues and eigenvectors of Ps.
        eigenvalues, eigenvectors = DiffusionMaps.eig_solver(transition_matrix, is_symmetric,
                                                             (self.eigenvectors_number + 1))

        ix = np.argsort(np.abs(eigenvalues))
        ix = ix[::-1]
        s = np.real(eigenvalues[ix])
        u = np.real(eigenvectors[:, ix])

        eigenvalues = s[:self.eigenvectors_number]
        eigenvectors = u[:, :self.eigenvectors_number]

        # Compute the diffusion coordinates
        eig_values_time = np.power(eigenvalues, self.t)
        diffusion_coordinates = eigenvectors @ np.diag(eig_values_time)

        self.transition_matrix = transition_matrix

        self.diffusion_coordinates = diffusion_coordinates
        self.eigenvalues = eigenvalues
        self.eigenvectors = eigenvectors

    # Private method
    @staticmethod
    def __sparse_kernel(kernel_matrix, neighbors_number):
        rows = np.shape(kernel_matrix)[0]
        for i in range(rows):
            row_data = kernel_matrix[i, :]
            index = _nn_coord(row_data, neighbors_number)
            kernel_matrix[i, index] = 0
            if sum(kernel_matrix[i, :]) <= 0:
                raise ValueError("UQpy: Consider increasing `n_neighbors` to have a connected graph.")

        return sps.csc_matrix(kernel_matrix)

    @staticmethod
    def diffusion_distance(diffusion_coordinates: Numpy2DFloatArray) -> Numpy2DFloatArray:
        distance_matrix = np.zeros((diffusion_coordinates.shape[0], diffusion_coordinates.shape[0]))
        pairs = list(itertools.combinations(diffusion_coordinates.shape[0], 2))
        for id_pair in range(np.shape(pairs)[0]):
            i = pairs[id_pair][0]
            j = pairs[id_pair][1]

            xi = diffusion_coordinates[i, :]
            xj = diffusion_coordinates[j, :]

            distance_matrix[i, j] = np.linalg.norm(xi - xj)
            distance_matrix[j, i] = distance_matrix[i, j]

        return distance_matrix

    def __normalize_kernel_matrix(self, kernel_matrix, inverse_diagonal_matrix):

        m = inverse_diagonal_matrix.shape[0]
        if self.is_sparse:
            d_alpha = sps.spdiags(inverse_diagonal_matrix, 0, m, m)
        else:
            d_alpha = np.diag(inverse_diagonal_matrix)

        normalized_kernel = d_alpha.dot(kernel_matrix.dot(d_alpha))

        return normalized_kernel

    def parsimonious(self, dim: int):
        """
        Selection of independent vectors for parsimonious data manifold embedding, based on
        local regression.  The eigenvectors with the largest residuals are considered for the
        embedding. The scale of the kernel used for the local linear regression is:

        .. code::

            scale = median(distances) / 3

        :param dim: Number of eigenvectors to select with largest residuals.
        """

        residuals = np.zeros(self.eigenvectors.shape[1])
        residuals[0] = np.nan
        # residual 1 for the first eigenvector.
        residuals[1] = 1.0

        # Get the residuals of each eigenvector.
        for i in range(2, self.eigenvectors.shape[1]):
            residuals[i] = DiffusionMaps.__get_residual(f_mat=self.eigenvectors[:, 1:i], f=self.eigenvectors[:, i])

        # Get the index of the eigenvalues associated with each residual.
        indices = np.argsort(residuals)[::-1][1:dim + 1]
        self.parsimonious_indices = indices
        self.parsimonious_residuals = residuals

    @staticmethod
    def __get_residual(f_mat, f):
        n_samples = np.shape(f_mat)[0]
        distance_matrix = sd.squareform(sd.pdist(f_mat))
        m = 3
        epsilon = (np.median(np.square(distance_matrix.flatten())) / m)
        kernel_matrix = np.exp(-1 * np.square(distance_matrix) / epsilon)
        coefficients = np.zeros((n_samples, n_samples))

        vec_1 = np.ones((n_samples, 1))

        for i in range(n_samples):
            # Weighted least squares:
            mat_x = np.hstack([vec_1, f_mat - f_mat[i, :]])
            mat_x_k = mat_x.T * kernel_matrix[i, :]
            u, _, _, _ = np.linalg.lstsq((mat_x_k @ mat_x), mat_x_k, rcond=1e-6)
            coefficients[i, :] = u[0, :]

        estimated_f = coefficients @ f

        # normalized leave-one-out cross-validation error.
        residual = np.sqrt(np.sum(np.square((f - estimated_f))) / np.sum(np.square(f)))
        return residual

    @staticmethod
    def eig_solver(kernel_matrix: Numpy2DFloatArray, is_symmetric: bool, n_eigenvectors: int) -> tuple[
        NumpyFloatArray, Numpy2DFloatArray]:

        n_samples, n_features = kernel_matrix.shape

        if n_eigenvectors == n_features:
            solver = sp.linalg.eigh if is_symmetric else sp.linalg.eig
            solver_kwargs = {"check_finite": False}

        else:
            solver = eigsh if is_symmetric else eigs
            solver_kwargs = {
                "sigma": None,
                "k": n_eigenvectors,
                "which": "LM",
                "v0": np.ones(n_samples),
                "tol": 1e-14,
            }

        eigenvalues, eigenvectors = solver(kernel_matrix, **solver_kwargs)

        eigenvectors /= np.linalg.norm(eigenvectors, axis=0)[np.newaxis, :]

        return eigenvalues, eigenvectors

    @staticmethod
    def _plot_eigen_pairs(eigenvectors: Numpy2DFloatArray,
                          trivial: bool = False, pair_indices: list = None, **kwargs):  # pragma: no cover
        """
        Plot scatter plot of n-th eigenvector on x-axis and remaining eigenvectors on
        y-axis.

        :param eigenvectors: Eigenvectors of the kernel matrix of shape `(n_samples, n_eigenvectors)`.
        :param  trivial: When trivial constant eigenvectors are ignored is set to `False`.
        :param pair_indices: Indices of the pair of eigenvectors to plot.
        :param kwargs:
            color: visualize the points.
            figure_size: Size of the figure to be passed as keyword argument to `matplotlib.pyplot.figure()`.
            font_size: Size of the font to be passed as keyword argument to `matplotlib.pyplot.figure()`.
        """
        figure_size = kwargs.get('figure_size', None)
        font_size = kwargs.get('font_size', None)
        color = kwargs.get('color', None)

        plt.figure()

        if figure_size is None and font_size is None:
            plt.rcParams["figure.figsize"] = (10, 10)
            plt.rcParams.update({'font.size': 18})
        elif figure_size is not None and font_size is None:
            plt.rcParams["figure.figsize"] = kwargs['figure_size']
            plt.rcParams.update({'font.size': 18})
        elif figure_size is None:
            plt.rcParams["figure.figsize"] = (10, 10)
            plt.rcParams.update({'font.size': kwargs['font_size']})
        else:
            plt.rcParams["figure.figsize"] = kwargs['figure_size']
            plt.rcParams.update({'font.size': kwargs['font_size']})

        if color is None:
            color = 'b'

        n_eigenvectors = eigenvectors.shape[1]

        start = 1 if not trivial else 0
        num_pairs = sum(1 for _ in combinations(range(start, n_eigenvectors), 2))
        if num_pairs % 2 == 0:
            n_rows = int(np.ceil(num_pairs / 2))
        else:
            n_rows = int(np.ceil(num_pairs / 2)) + 1

        if pair_indices is None:
            _, _ = plt.subplots(
                nrows=n_rows, ncols=2, sharex=True, sharey=True,
            )

            for count, arg in enumerate(combinations(range(start, n_eigenvectors), 2), start=1):
                i = arg[0]
                j = arg[1]
                plt.subplot(n_rows, 2, count)
                plt.scatter(eigenvectors[:, i], eigenvectors[:, j], c=color, cmap=plt.cm.Spectral)
                plt.title(
                    r"$\Psi_{{{}}}$ vs. $\Psi_{{{}}}$".format(i, j))

        else:
            _, _ = plt.subplots(
                nrows=1, ncols=1, sharex=True, sharey=True)
            plt.scatter(eigenvectors[:, pair_indices[0]], eigenvectors[:, pair_indices[1]], c=color,
                        cmap=plt.cm.Spectral)
            plt.title(
                r"$\Psi_{{{}}}$ vs. $\Psi_{{{}}}$".format(pair_indices[0], pair_indices[1]))


# src/UQpy/dimension_reduction/diffusion_maps/__init__.py

from UQpy.dimension_reduction.diffusion_maps.DiffusionMaps import DiffusionMaps


# src/UQpy/dimension_reduction/grassmann_manifold/GrassmannInterpolation.py

import copy
from typing import Union

import numpy as np
from scipy.interpolate import LinearNDInterpolator

from UQpy.dimension_reduction.grassmann_manifold import GrassmannOperations
from UQpy.surrogates.baseclass import Surrogate
from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.utilities.ValidationTypes import NumpyFloatArray
from UQpy.utilities.distances import GrassmannianDistance


class GrassmannInterpolation:

    def __init__(self, interpolation_method: Union[Surrogate, callable, None],
                 manifold_data: list[GrassmannPoint],
                 coordinates: Union[np.ndarray, list[NumpyFloatArray]],
                 distance: GrassmannianDistance,
                 optimization_method: str = "GradientDescent"):
        """
        A class to perform interpolation of points on the Grassmann manifold.

        :param interpolation_method: Type of interpolation to perform. This may be specified as a :class:`Surrogate`
            object or a callable function. If :any:`None`, then multi-linear interpolation is performed.
        :param manifold_data: Data points on the Grassmann manifold.
        :param optimization_method: Optimization method for calculating the Karcher mean. See
            :py:meth:`.GrassmannOperations.karcher_mean`.
        :param coordinates: Nodes of the interpolant.
        :param distance:  Distance measure.
        """
        self.interpolation_method = interpolation_method

        self.mean = GrassmannOperations.karcher_mean(grassmann_points=manifold_data,
                                                     optimization_method=optimization_method,
                                                     distance=distance)

        self.tangent_points = GrassmannOperations.log_map(grassmann_points=manifold_data,
                                                          reference_point=self.mean)

        if self.interpolation_method is Surrogate:
            self.surrogates = [[]]
            nargs = len(self.tangent_points)
            n_rows = self.tangent_points[0].shape[0]
            n_cols = self.tangent_points[0].shape[1]
            for j in range(n_rows):
                for k in range(n_cols):
                    val_data = [[self.tangent_points[i][j, k]] for i in range(nargs)]
                    # if all the elements of val_data are the same.
                    val_data = np.array(val_data)
                    self.surrogates[j][k] = copy.copy(self.interpolation_method)
                    self.surrogates[j][k].fit(coordinates, val_data)

        else:
            self.coordinates = coordinates

    def interpolate_manifold(self, point: np.ndarray):
        """
        :param point:  Point at which to interpolate.
        :return: Interpolated point on the Grassmann manifold.
        """

        shape_ref = np.shape(self.tangent_points[0])
        interp_point = np.zeros(shape_ref)
        n_rows = self.tangent_points[0].shape[0]
        n_cols = self.tangent_points[0].shape[1]

        nargs = len(self.tangent_points)
        shape_ref = np.shape(self.tangent_points[0])
        for i in range(1, nargs):
            if np.shape(self.tangent_points[i]) != shape_ref:
                raise TypeError("UQpy: Input matrices have different shape.")

        if self.interpolation_method is Surrogate:
            for j in range(n_rows):
                for k in range(n_cols):
                    y = self.surrogates[j][k].predict(point, return_std=False)
                    interp_point[j, k] = y
        elif self.interpolation_method is callable:
            interp_point = self.interpolation_method(self.coordinates, self.tangent_points, point)
        else:
            interp = LinearNDInterpolator(self.coordinates, self.tangent_points)
            interp_point = interp(point)

        return GrassmannOperations.exp_map(tangent_points=[interp_point.squeeze()], reference_point=self.mean)[0]


# src/UQpy/dimension_reduction/grassmann_manifold/GrassmannOperations.py

import copy
from typing import Union

import numpy as np
from beartype import beartype

from UQpy.utilities.distances import GeodesicDistance
from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.utilities.ValidationTypes import Numpy2DFloatArray, Numpy2DFloatArrayOrthonormal
from UQpy.utilities.distances.baseclass import GrassmannianDistance
from UQpy.utilities.kernels import GrassmannianKernel, ProjectionKernel


class GrassmannOperations:
    @beartype
    def __init__(self, grassmann_points: Union[list[Numpy2DFloatArrayOrthonormal], list[GrassmannPoint]],
                 kernel: GrassmannianKernel = ProjectionKernel(),
                 p: Union[int, str] = "max", optimization_method: str = "GradientDescent",
                 distance: GrassmannianDistance = GeodesicDistance()):
        """
        The :class:`.GrassmannOperations` class can used in two ways. In the first case, the user can invoke the
        initializer by providing all the required input data. The class will then automatically calculate the
        kernel matrix, distance matrix, karcher mean, and frechet variance of the input Grassmann points. Alternatively,
        the user may invoke each of the static methods individually and calculate only the required quantities, without
        having to instantiate a new object.

        :param grassmann_points: Data points projected on the Grassmann manifold
        :param kernel: Kernel to be used to evaluate the kernel matrix of the given Grassmann points
        :param p: Rank of the Grassmann projected points.
        :param optimization_method: String that defines the optimization method for computation of the Karcher mean.

            Options: `"GradientDescent"`, `"StochasticGradientDescent"`
        :param distance: Distance measure to be used for the optimization in computation of the Karcher mean and Frechet
            variance.
        """
        self.grassmann_points = grassmann_points
        self.p = p
        self.kernel_matrix = kernel.calculate_kernel_matrix(self.grassmann_points)
        self.distance_matrix = distance.calculate_distance_matrix(self.grassmann_points)
        self.karcher_mean = GrassmannOperations.karcher_mean(self.grassmann_points, optimization_method, distance)
        self.frechet_variance = GrassmannOperations.frechet_variance(self.grassmann_points, self.karcher_mean, distance)

    @staticmethod
    def calculate_kernel_matrix(grassmann_points: Union[list[Numpy2DFloatArrayOrthonormal], list[GrassmannPoint]],
                                kernel: GrassmannianKernel = ProjectionKernel()):
        return kernel.calculate_kernel_matrix(grassmann_points)

    @staticmethod
    @beartype
    def log_map(grassmann_points: Union[list[Numpy2DFloatArrayOrthonormal], list[GrassmannPoint]],
                reference_point: Union[Numpy2DFloatArrayOrthonormal, GrassmannPoint]) -> list[Numpy2DFloatArray]:
        """
        :param grassmann_points: Point(s) on the Grassmann manifold.
        :param reference_point: Origin of the tangent space.
        :return: Point(s) on the tangent space.
        """
        reference_point = reference_point.data if isinstance(reference_point, GrassmannPoint) else reference_point
        from UQpy.utilities.distances.baseclass.GrassmannianDistance import GrassmannianDistance
        number_of_points = len(grassmann_points)
        for i in range(number_of_points):
            GrassmannianDistance.check_rows(reference_point, grassmann_points[i])
            if reference_point.data.shape[1] != grassmann_points[i].data.shape[1]:
                raise ValueError("UQpy: Point {0} is on G({1},{2}) - Reference is on"
                                 " G({1},{2})".format(i, grassmann_points[i].data.shape[1],
                                                      grassmann_points[i].data.shape[0]))

        # Multiply ref by its transpose.
        reference_point_transpose = reference_point.T
        m_ = np.dot(reference_point.data, reference_point_transpose)

        tangent_points = []
        for i in range(number_of_points):
            u_trunc = grassmann_points[i]
            # compute: M = ((I - psi0*psi0')*psi1)*inv(psi0'*psi1)
            m_inv = np.linalg.inv(np.dot(reference_point_transpose, u_trunc.data))
            m = np.dot(u_trunc.data - np.dot(m_, u_trunc.data), m_inv)
            ui, si, vi = np.linalg.svd(m, full_matrices=False)
            tangent_points.append(np.dot(np.dot(ui, np.diag(np.arctan(si))), vi))

        return tangent_points

    @staticmethod
    @beartype
    def exp_map(tangent_points: list[Numpy2DFloatArray],
                reference_point: Union[np.ndarray, GrassmannPoint]) -> list[GrassmannPoint]:
        """
        :param tangent_points: Tangent vector(s).
        :param reference_point: Origin of the tangent space.
        :return: Point(s) on the Grassmann manifold.
        """
        number_of_points = len(tangent_points)
        for i in range(number_of_points):
            if reference_point.data.shape[1] != tangent_points[i].shape[1]:
                raise ValueError("UQpy: Point {0} is on G({1},{2}) - Reference is on"
                                 " G({1},{2})".format(i, tangent_points[i].shape[1], tangent_points[i].shape[0]))

        # Map the each point back to the manifold.
        manifold_points = []
        for i in range(number_of_points):
            u_trunc = tangent_points[i]
            ui, si, vi = np.linalg.svd(u_trunc, full_matrices=False)

            x0 = np.dot(
                np.dot(np.dot(reference_point.data, vi.T), np.diag(np.cos(si)))
                + np.dot(ui, np.diag(np.sin(si))),
                vi,
            )

            if not np.allclose(x0.T @ x0, np.eye(u_trunc.shape[1])):
                x0, _ = np.linalg.qr(x0)

            manifold_points.append(GrassmannPoint(x0))

        return manifold_points

    @staticmethod
    @beartype
    def frechet_variance(grassmann_points: Union[list[Numpy2DFloatArrayOrthonormal], list[GrassmannPoint]],
                         reference_point: Union[Numpy2DFloatArrayOrthonormal, GrassmannPoint],
                         distance: GrassmannianDistance) -> float:
        """
        :param grassmann_points: Point(s) on the Grassmann manifold
        :param reference_point: Reference point for the Frechet variance (:math:`\mu`). Typically assigned as the
            Karcher mean.
        :param distance: Distance measure to be used in the variance calculation.
        """
        p_dim = [min(np.shape(grassmann_points[i].data)) for i in range(len(grassmann_points))]

        points_number = len(grassmann_points)

        variance_nominator = 0
        for i in range(points_number):
            distance.calculate_distance_matrix([reference_point, grassmann_points[i]], p_dim)
            distances = distance.distance_matrix
            variance_nominator += distances[0] ** 2

        return variance_nominator / points_number

    @staticmethod
    @beartype
    def karcher_mean(grassmann_points: Union[list[Numpy2DFloatArrayOrthonormal], list[GrassmannPoint]],
                     optimization_method: str, distance: GrassmannianDistance,
                     acceleration: bool = False, tolerance: float = 1e-3,
                     maximum_iterations: int = 1000) -> GrassmannPoint:
        """
        :param maximum_iterations: Maximum number of iterations performed by the optimization algorithm.
        :param tolerance: Tolerance used as the convergence criterion of the optimization.
        :param acceleration: Boolean flag used in combination with :code:`GradientDescent` optimization method which
         activates the Nesterov acceleration scheme
        :param grassmann_points: Point(s) on the Grassmann manifold.
        :param optimization_method: String that defines the optimization method.

            Options: `"GradientDescent"`, `"StochasticGradientDescent"`
        :param distance: Distance measure to be used for the optimization.
        """
        # Compute and test the number of input matrices necessary to compute the Karcher mean.
        nargs = len(grassmann_points)
        if nargs < 2:
            raise ValueError("UQpy: At least two matrices must be provided.")

        if optimization_method == "GradientDescent":
            return GrassmannOperations._gradient_descent(grassmann_points, distance, acceleration, tolerance, maximum_iterations)
        else:
            return GrassmannOperations._stochastic_gradient_descent(grassmann_points, distance, tolerance, maximum_iterations)

    @staticmethod
    def _gradient_descent(data_points, distance_fun, acceleration, tolerance, maximum_iterations):
        # acc is a boolean variable to activate the Nesterov acceleration scheme.
        acc = acceleration
        # Error tolerance
        tol = tolerance
        # Maximum number of iterations.
        maxiter = maximum_iterations
        # Number of points.
        n_mat = len(data_points)

        # =========================================
        alpha = 0.5
        rnk = [min(np.shape(data_points[i].data)) for i in range(n_mat)]
        max_rank = max(rnk)
        fmean = [GrassmannOperations.frechet_variance(data_points, data_points[i], distance_fun) for i in range(n_mat)]

        index_0 = fmean.index(min(fmean))
        mean_element = data_points[index_0].data.tolist()

        avg_gamma = np.zeros([np.shape(data_points[0].data)[0], np.shape(data_points[0].data)[1]])

        itera = 0

        l = 0
        avg = []
        _gamma = []
        if acc:
            _gamma = GrassmannOperations.log_map(grassmann_points=data_points,
                                                 reference_point=np.asarray(mean_element))

            avg_gamma.fill(0)
            for i in range(n_mat):
                avg_gamma += _gamma[i] / n_mat
            avg.append(avg_gamma)

        # Main loop
        while itera <= maxiter:
            _gamma = GrassmannOperations.log_map(grassmann_points=data_points,
                                                 reference_point=np.asarray(mean_element))
            avg_gamma.fill(0)

            for i in range(n_mat):
                avg_gamma += _gamma[i] / n_mat

            test_0 = np.linalg.norm(avg_gamma, 'fro')
            if test_0 < tol and itera == 0:
                break

            # Nesterov: Accelerated Gradient Descent
            if acc:
                avg.append(avg_gamma)
                l0 = l
                l1 = 0.5 * (1 + np.sqrt(1 + 4 * l * l))
                ls = (1 - l0) / l1
                step = (1 - ls) * avg[itera + 1] + ls * avg[itera]
                l = copy.copy(l1)
            else:
                step = alpha * avg_gamma

            x = GrassmannOperations.exp_map(tangent_points=[step],
                                            reference_point=np.asarray(mean_element))

            test_1 = np.linalg.norm(x[0].data - mean_element, 'fro')

            if test_1 < tol:
                break

            mean_element = []
            mean_element = x[0].data.tolist()

            itera += 1

        # return the Karcher mean.
        return GrassmannPoint(np.asarray(mean_element))

    @staticmethod
    def _stochastic_gradient_descent(data_points, distance_fun, tolerance, maximum_iterations):

        tol = tolerance
        maxiter = maximum_iterations
        n_mat = len(data_points)

        rnk = [min(np.shape(data_points[i].data)) for i in range(n_mat)]
        max_rank = max(rnk)

        fmean = [GrassmannOperations.frechet_variance(data_points, data_points[i], distance_fun) for i in range(n_mat)]

        index_0 = fmean.index(min(fmean))

        mean_element = data_points[index_0].data.tolist()
        itera = 0
        _gamma = []
        k = 1
        while itera < maxiter:

            indices = np.arange(n_mat)
            np.random.shuffle(indices)

            melem = mean_element
            for i in range(len(indices)):
                alpha = 0.5 / k
                idx = indices[i]
                _gamma = GrassmannOperations.log_map(grassmann_points=[data_points[idx]],
                                                     reference_point=np.asarray(mean_element))

                step = 2 * alpha * _gamma[0]

                X = GrassmannOperations.exp_map(tangent_points=[step],
                                                reference_point=np.asarray(mean_element))

                _gamma = []
                mean_element = X[0].data

                k += 1

            test_1 = np.linalg.norm(mean_element - melem, 'fro')
            if test_1 < tol:
                break

            itera += 1

        return GrassmannPoint(np.asarray(mean_element))


# src/UQpy/dimension_reduction/grassmann_manifold/__init__.py

from UQpy.dimension_reduction.grassmann_manifold.GrassmannOperations import GrassmannOperations
from UQpy.dimension_reduction.grassmann_manifold.GrassmannInterpolation import GrassmannInterpolation
from UQpy.dimension_reduction.grassmann_manifold.projections import *


# src/UQpy/dimension_reduction/hosvd/HigherOrderSVD.py

import logging

from beartype import beartype

from UQpy.utilities.ValidationTypes import *
import numpy as np


class HigherOrderSVD:
    @beartype
    def __init__(
            self,
            solution_snapshots: Union[np.ndarray, list],
            modes: PositiveInteger = 10 ** 10,
            reconstruction_percentage: Union[PositiveFloat, PositiveInteger] = 10 ** 10,
    ):
        """

        :param solution_snapshots: Second order tensor or list containing the solution snapshots. Third dimension or
         length of list corresponds to the number of snapshots.
        :param modes: Number of modes to keep for dataset reconstruction.
        :param reconstruction_percentage: Dataset reconstruction percentage.
        """
        self.reconstruction_error = None
        self.s3hat: Numpy2DFloatArray = None
        """Normalized core tensor produced by the HOSVD decomposition."""
        self.u3hat: Numpy2DFloatArray = None
        """Normalized unitary array produced by the HOSVD decomposition"""
        self.reduced_solutions = None
        self.s3 = None
        self.v3 = None
        self.sig_3 = None
        self.u3 = None
        self.v2 = None
        self.sig_2 = None
        self.u2: Numpy2DFloatArray = None
        """Unitary array of the SVD of the second unfolded matrix."""
        self.v1 = None
        self.sig_1 = None
        self.u1: Numpy2DFloatArray = None
        """Unitary array of the SVD of the first unfolded matrix"""
        self.solution_snapshots = solution_snapshots
        self.logger = logging.getLogger(__name__)
        self.modes = modes
        self.reconstruction_percentage = reconstruction_percentage

        if self.modes != 10 ** 10 and self.reconstruction_percentage != 10 ** 10:
            raise ValueError("Either a number of modes or a reconstruction percentage must be chosen, not both.")

        if self.solution_snapshots is not None:
            self.factorize(get_error=True)

    def factorize(self, get_error: bool = False):
        """
        Executes the HOSVD method.

        :param get_error: A boolean declaring whether to return the reconstruction error.
        """
        rows = self.solution_snapshots[0].shape[0]

        a1, a2, a3 = HigherOrderSVD.unfold3d(self.solution_snapshots)

        self.u1, self.sig_1, self.v1 = np.linalg.svd(a1, full_matrices=True)
        self.u2, self.sig_2, self.v2 = np.linalg.svd(a2, full_matrices=True)
        self.u3, self.sig_3, self.v3 = np.linalg.svd(a3, full_matrices=True)

        sig_3_ = np.diag(self.sig_3)
        hold = np.dot(np.linalg.inv(self.u3), a3)
        kronecker_product = np.kron(self.u1, self.u2)
        self.s3 = np.array(np.dot(hold, np.linalg.inv(kronecker_product.T)))

        if self.modes == 10 ** 10:
            error_ = []
            for i in range(rows):
                error_.append(np.sqrt(((self.sig_3[i + 1:]) ** 2).sum()) / np.sqrt((self.sig_3 ** 2).sum()))
                if i == rows:
                    error_.append(0)
            error = [i * 100 for i in error_]
            error.reverse()
            perc = error.copy()
            percentage = min(perc, key=lambda x: abs(x - self.reconstruction_percentage))
            self.modes = perc.index(percentage) + 1

        elif self.modes > rows:
            self.logger.warning(
                "A number of modes greater than the number of temporal dimensions was given."
                "Number of temporal dimensions is {}.".format(rows))

        self.reduced_solutions = np.dot(self.u3, sig_3_)
        self.u3hat = np.dot(self.u3[:, : self.modes], sig_3_[: self.modes, : self.modes])
        self.s3hat = np.dot(np.linalg.inv(sig_3_[: self.modes, : self.modes]), self.s3[: self.modes, :])

        if self.modes == 10 ** 10:
            self.logger.info("Dataset reconstruction: {0:.3%}".format(self.reconstruction_percentage / 100))

        elif get_error:
            self.reconstruction_error = np.sqrt(((self.s3hat[self.modes:]) ** 2).sum()) / np.sqrt(
                (self.sig_3 ** 2).sum())
            self.logger.warning("Reduced-order reconstruction error: {0:.3%}".format(self.reconstruction_error))

    @staticmethod
    def unfold3d(second_order_tensor: np.ndarray):
        """

        :param second_order_tensor: A three-dimensional :class:`ndarray` which contains the data to be unfolded.
        :returns: Three unfolded matrices in the form of two-dimensional :class:`ndarray`.
        """
        if type(second_order_tensor) == list:
            rows = second_order_tensor[0].shape[0]
            columns = second_order_tensor[0].shape[1]
            number_of_data = len(second_order_tensor)
            tensor_of_list = np.zeros((rows, columns, number_of_data))
            for i in range(number_of_data):
                tensor_of_list[:, :, i] = second_order_tensor[i]
            del second_order_tensor
            second_order_tensor = np.copy(tensor_of_list)

        permutation1 = [0, 2, 1]
        permutation2 = [1, 2, 0]
        permutation3 = [2, 0, 1]

        permuted_tensor1 = np.transpose(second_order_tensor, permutation1)
        permuted_tensor2 = np.transpose(second_order_tensor, permutation2)
        permuted_tensor3 = np.transpose(second_order_tensor, permutation3)

        matrix1 = permuted_tensor1.reshape(second_order_tensor.shape[0],
                                           second_order_tensor.shape[2] * second_order_tensor.shape[1], )
        matrix2 = permuted_tensor2.reshape(second_order_tensor.shape[1],
                                           second_order_tensor.shape[2] * second_order_tensor.shape[0], )
        matrix3 = permuted_tensor3.reshape(second_order_tensor.shape[2],
                                           second_order_tensor.shape[0] * second_order_tensor.shape[1], )

        return matrix1, matrix2, matrix3

    @staticmethod
    def reconstruct(u1: Numpy2DFloatArray,
                    u2: Numpy2DFloatArray,
                    u3hat: Numpy2DFloatArray,
                    s3hat: Numpy2DFloatArray):
        """
        Reconstructs the approximated solution.

        :param u1: Unitary array of the SVD of the first unfolded matrix.
        :param u2: Unitary array of the SVD of the second unfolded matrix.
        :param u3hat: Normalized unitary array produced by the HOSVD decomposition.
        :param s3hat: Normalized core tensor produced by the HOSVD decomposition
        :return: An :class:`ndarray` containing the reconstruction of the approximated solution.
        """
        b = np.kron(u1, u2)
        c = np.dot(s3hat, b.T)
        d = np.dot(u3hat[:, :], c)

        rows = u1.shape[0]
        columns = u2.shape[0]
        snapshot_number = u3hat.shape[0]

        reconstructed_solutions = np.zeros((rows, columns, snapshot_number))
        for i in range(snapshot_number):
            reconstructed_solutions[0:rows, 0:columns, i] = d[i, :].reshape((rows, columns))

        return reconstructed_solutions


# src/UQpy/dimension_reduction/hosvd/__init__.py

from UQpy.dimension_reduction.hosvd.HigherOrderSVD import HigherOrderSVD

# src/UQpy/dimension_reduction/pod/DirectPOD.py

import numpy as np
from UQpy.dimension_reduction.pod.baseclass.POD import POD


class DirectPOD(POD):

    def run(self, solution_snapshots):
        """
        Executes proper orthogonal decomposition using the :class:`.DirectPOD` algorithm.
        """
        return super().run(solution_snapshots)

    def _calculate_c_and_iterations(self, u, snapshot_number, rows, columns):
        if snapshot_number < rows * columns and rows * columns > 1000:
            self.logger.warning("Snapshot POD is recommended.")

        c = np.dot(u.T, u) / (snapshot_number - 1)
        n_iterations = rows * columns
        return c, n_iterations

    def _calculate_reduced_and_reconstructed_solutions(self, u, phi, rows, columns, snapshot_number):
        a = np.dot(u, phi)
        reconstructed_solutions_ = np.dot(a[:, : self.modes], phi[:, : self.modes].T)
        reduced_solutions = np.dot(u, phi[:, : self.modes])
        reconstructed_solutions = np.zeros((rows, columns, snapshot_number))
        for i in range(snapshot_number):
            reconstructed_solutions[0:rows, 0:columns, i] = reconstructed_solutions_[i, :].reshape((rows, columns))
        return reconstructed_solutions, reduced_solutions


# src/UQpy/dimension_reduction/pod/SnapshotPOD.py

import numpy as np
from UQpy.dimension_reduction.pod.baseclass.POD import POD


class SnapshotPOD(POD):

    def run(self, solution_snapshots):
        """
        Executes proper orthogonal decomposition using the :class:`.SnapshotPOD` algorithm.
        """
        return super().run(solution_snapshots)

    def _calculate_reduced_and_reconstructed_solutions(self, u, phi, rows, columns, snapshot_number):
        phi_s = np.dot(u.T, phi)
        reconstructed_solutions_ = np.dot(phi[:, : self.modes], phi_s[:, : self.modes].T)
        reduced_solutions_ = (np.dot(u.T, phi[:, : self.modes])).T

        reconstructed_solutions = np.zeros((rows, columns, snapshot_number))
        reduced_solutions = np.zeros((rows, columns, self.modes))

        for i in range(snapshot_number):
            reconstructed_solutions[0:rows, 0:columns, i] = reconstructed_solutions_[i, :].reshape((rows, columns))

        for i in range(self.modes):
            reduced_solutions[0:rows, 0:columns, i] = reduced_solutions_[i, :].reshape((rows, columns))

        return reconstructed_solutions, reduced_solutions

    def _calculate_c_and_iterations(self, u, snapshot_number, rows, columns):
        c = np.dot(u, u.T) / (snapshot_number - 1)
        n_iterations = snapshot_number
        return c, n_iterations


# src/UQpy/dimension_reduction/pod/__init__.py

from UQpy.dimension_reduction.pod.DirectPOD import DirectPOD
from UQpy.dimension_reduction.pod.SnapshotPOD import SnapshotPOD

from UQpy.dimension_reduction.pod.baseclass import *

# src/UQpy/distributions/baseclass/Copula.py

from UQpy.distributions.baseclass.DistributionContinuous1D import (
    DistributionContinuous1D,
)
from abc import ABC


class Copula(ABC):

    def __init__(self, ordered_parameters: dict = None, **kwargs: dict):
        """
        Define a copula for a multivariate distribution whose dependence structure is defined with a copula.
        This class is used in support of the :class:`.JointCopula` class.

        :param ordered_parameters: List of parameter names
        :param kwargs: Parameters of the copula.
        """
        self.parameters: dict = kwargs
        """Parameters of the copula."""
        self.ordered_parameters: dict = ordered_parameters
        """List of parameter names"""
        if self.ordered_parameters is None:
            self.ordered_parameters = tuple(kwargs.keys())
        if len(self.ordered_parameters) != len(self.parameters):
            raise ValueError(
                "Inconsistent dimensions between ordered_parameters tuple and parameters dictionary."
            )

    def get_parameters(self) -> dict:
        """
        :return: A dictionary containing the parameter names.
        """
        return self.parameters

    def update_parameters(self, **kwargs: dict):
        """
        Given a dictionary with keys the names and values the new parameter values,
        the method updates the current values.

        :param kwargs: Dictionary containing the updated parameter values.
        """
        for key in kwargs.keys():
            if key not in self.parameters.keys():
                raise ValueError("Wrong parameter name.")
            self.parameters[key] = kwargs[key]

    @staticmethod
    def check_marginals(marginals: list[DistributionContinuous1D]):
        """
        Perform some checks on the marginals, raise errors if necessary.

        As an example, Archimedian copulas are only defined for bi-variate continuous distributions, thus this method
        checks that marginals is of length 2 and continuous, and raise an error if that is not the case.

        :param marginals: List of 1D continuous distributions.
        """
        if len(marginals) != 2:
            raise ValueError("Maximum dimension for the Copula is 2.")
        if not all(isinstance(m, DistributionContinuous1D) for m in marginals):
            raise ValueError("Marginals should be 1d continuous distributions.")


# src/UQpy/distributions/baseclass/Distribution.py

from abc import ABC


class Distribution(ABC):

    def __init__(self, ordered_parameters: list = None, **kwargs: dict):
        """
        A parent class to all :class:`.Distribution` classes.

        All distributions possess a number of methods to perform basic probabilistic operations. For most of the
        predefined distributions in :py:mod:`UQpy` these methods are inherited from the :py:mod:`scipy.stats` package.
        These include standard operations such as computing probability density/mass functions, cumulative distribution
        functions and their inverse, drawing random samples, computing moments and parameter fitting. However, for
        user-defined distributions, any desired method can be constructed into the child class structure.

        For bookkeeping purposes, all :class:`.Distribution` objects possesses :meth:`get_parameters` and
        :meth:`update_parameters` methods. These are described in more detail below.

        Any :class:`.Distribution` further inherits from one of the following classes:

        - :class:`DistributionContinuous1D`: Parent class to 1-dimensional continuous probability distributions.
        - :class:`DistributionDiscrete1D`: Parent class to 1-dimensional discrete probability distributions.
        - :class:`DistributionND`: Parent class to multivariate probability distributions.

        :param ordered_parameters: Ordered list of parameter names, useful when parameter values are stored in vectors
         and must be passed to the :meth:`update_parameters` method.
        :param kwargs: Parameters of the distribution. Note: this attribute is not defined for certain
         :class:`.Distribution` objects such as those of type :class:`.JointIndependent` or :class:`.JointCopula`. The
         user is advised to use the :meth:`get_parameters` method to access the parameters.
        """
        self.parameters: dict = kwargs
        """ Ordered list of parameter names, useful when parameter values are stored in vectors and must be passed to 
        the :meth:`update_params` method."""
        self.ordered_parameters: list = ordered_parameters
        """Parameters of the distribution. 
        
        This attribute is not defined for certain :class:`.Distribution` 
        objects such as those of type :class:`.JointIndependent` or :class:`.JointCopula`. The user is advised to use 
        the :meth:`get_parameters` method to access the parameters."""
        if self.ordered_parameters is None:
            self.ordered_parameters = tuple(kwargs.keys())
        if len(self.ordered_parameters) != len(self.parameters):
            raise ValueError(
                "Inconsistent dimensions between order_params tuple and params dictionary."
            )

    def update_parameters(self, **kwargs: dict):
        """
        Update the parameters of a :class:`.Distribution` object.

        To update the parameters of a :class:`.JointIndependent` or a :class:`.JointCopula` distribution, each parameter
        is assigned a unique string identifier as :code:`key_index` - where :code:`key` is the parameter name and
        :code:`index` the index of the marginal (e.g., location parameter of the 2nd marginal is identified as
        :code:`loc_1`).

        :param kwargs: Parameters to be updated, designated by their respective keywords.
        """
        for key in kwargs.keys():
            if key not in self.get_parameters().keys():
                raise ValueError("Wrong parameter name.")
            self.parameters[key] = kwargs[key]

    def get_parameters(self):
        """
        Return the parameters of a :class:`.Distribution` object.

        To update the parameters of a :class:`.JointIndependent` or a :class:`.JointCopula` distribution, each parameter
        is assigned a unique string identifier as :code:`key_index` - where :code:`key` is the parameter name and
        :code:`index` the index of the marginal (e.g., location parameter of the 2nd marginal is identified as
        :code:`loc_1`).

        :return: Parameters of the distribution.
        """
        return self.parameters


# src/UQpy/distributions/baseclass/Distribution1D.py

import numpy as np
import scipy.stats as stats
from UQpy.distributions.baseclass.Distribution import Distribution
from abc import ABC


class Distribution1D(Distribution, ABC):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @staticmethod
    def check_x_dimension(x):
        """
        Check the dimension of input x - must be an ndarray of shape (npoints,) or (npoints, 1)
        """
        array = np.atleast_1d(x)
        if len(array.shape) > 2 or (len(array.shape) == 2 and array.shape[1] != 1):
            raise ValueError("Wrong dimension in x.")
        return array.reshape((-1,))

    def _retrieve_1d_data_from_scipy(self, scipy_name=stats.rv_continuous, is_continuous=True):
        if is_continuous:
            self.cdf = lambda x: scipy_name.cdf(x=self.check_x_dimension(x), **self.parameters)
        else:
            self.cdf = lambda x: scipy_name.cdf(k=self.check_x_dimension(x), **self.parameters)
        self.icdf = lambda x: scipy_name.ppf(q=self.check_x_dimension(x), **self.parameters)
        self.moments = lambda moments2return="mvsk": scipy_name.stats(moments=moments2return, **self.parameters)
        self.rvs = lambda nsamples=1, random_state=None: scipy_name.rvs(
            size=nsamples, random_state=random_state, **self.parameters).reshape((nsamples, 1))


# src/UQpy/distributions/baseclass/DistributionContinuous1D.py

import scipy.stats as stats
from UQpy.distributions.baseclass.Distribution1D import Distribution1D
from abc import ABC


class DistributionContinuous1D(Distribution1D, ABC):
    """
    Parent class for univariate continuous probability distributions.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def _construct_from_scipy(self, scipy_name=stats.rv_continuous):
        self.pdf = lambda x: scipy_name.pdf(x=self.check_x_dimension(x), **self.parameters)
        self.log_pdf = lambda x: scipy_name.logpdf(x=self.check_x_dimension(x), **self.parameters)
        self._retrieve_1d_data_from_scipy(scipy_name)

        def tmp_fit(dist, data):
            data = self.check_x_dimension(data)
            fixed_params = {}
            for key, value in dist.parameters.items():
                if value is not None:
                    fixed_params["f" + key] = value
            params_fitted = scipy_name.fit(data=data, **fixed_params)
            return dict(zip(dist.ordered_parameters, params_fitted))

        self.fit = lambda data: tmp_fit(self, data)


# src/UQpy/distributions/baseclass/DistributionDiscrete1D.py

import scipy.stats as stats
from UQpy.distributions.baseclass.Distribution1D import Distribution1D
from abc import ABC


class DistributionDiscrete1D(Distribution1D, ABC):
    def __init__(self, **kwargs):
        """
        Parent class for univariate discrete distributions.
        """
        super().__init__(**kwargs)

    def _construct_from_scipy(self, scipy_name=stats.rv_discrete):
        self.pmf = lambda x: scipy_name.pmf(k=self.check_x_dimension(x), **self.parameters)
        self.log_pmf = lambda x: scipy_name.logpmf(k=self.check_x_dimension(x), **self.parameters)
        self._retrieve_1d_data_from_scipy(scipy_name, is_continuous=False)


# src/UQpy/distributions/baseclass/DistributionND.py

import numpy as np
from .Distribution import Distribution
from abc import ABC


class DistributionND(Distribution, ABC):
    """
    Parent class for multivariate probability distributions.

    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @staticmethod
    def check_x_dimension(x: np.ndarray, d: int = None):
        """
        Check the dimension of input x - must be an ndarray of shape (npoints, d)
        """
        x = np.array(x)
        if len(x.shape) != 2:
            raise ValueError("Wrong dimension in x.")
        if (d is not None) and (x.shape[1] != d):
            raise ValueError("Wrong dimension in x.")
        return x


# src/UQpy/distributions/baseclass/__init__.py

"""Collection of baseclass files."""
from UQpy.distributions.baseclass.Copula import Copula
from UQpy.distributions.baseclass.Distribution import Distribution
from UQpy.distributions.baseclass.Distribution1D import Distribution1D
from UQpy.distributions.baseclass.DistributionContinuous1D import (
    DistributionContinuous1D,
)
from UQpy.distributions.baseclass.DistributionDiscrete1D import DistributionDiscrete1D
from UQpy.distributions.baseclass.DistributionND import DistributionND


# src/UQpy/distributions/collection/Beta.py

from typing import Union

import scipy.stats as stats
from UQpy.distributions.baseclass import DistributionContinuous1D
from beartype import beartype


class Beta(DistributionContinuous1D):

    @beartype
    def __init__(
        self,
        a: Union[None, float, int],
        b: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param a: first shape parameter
        :param b: second shape parameter
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            a=a,
            b=b,
            loc=loc,
            scale=scale,
            ordered_parameters=("a", "b", "loc", "scale"),
        )
        self._construct_from_scipy(scipy_name=stats.beta)


# src/UQpy/distributions/collection/Binomial.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionDiscrete1D


class Binomial(DistributionDiscrete1D):
    @beartype
    def __init__(
        self,
        n: Union[None, int],
        p: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
    ):
        """

        :param n: number of trials, integer :math:`\ge 0`
        :param p: success probability for each trial, real number in :math:`[0, 1]`.
        :param loc: location parameter
        """
        super().__init__(n=n, p=p, loc=loc, ordered_parameters=("n", "p", "loc"))
        self._construct_from_scipy(scipy_name=stats.binom)


# src/UQpy/distributions/collection/Cauchy.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Cauchy(DistributionContinuous1D):

    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.cauchy)


# src/UQpy/distributions/collection/ChiSquare.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class ChiSquare(DistributionContinuous1D):
    @beartype
    def __init__(
        self,
        df: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param df: shape parameter (degrees of freedom) (given by k in the equation)
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            df=df, loc=loc, scale=scale, ordered_parameters=("df", "loc", "scale")
        )
        self._construct_from_scipy(scipy_name=stats.chi2)


# src/UQpy/distributions/collection/Exponential.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Exponential(DistributionContinuous1D):

    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.expon)


# src/UQpy/distributions/collection/Gamma.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Gamma(DistributionContinuous1D):

    @beartype
    def __init__(
        self,
        a: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param a: shape parameter
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            a=a, loc=loc, scale=scale, ordered_parameters=("a", "loc", "scale")
        )
        self._construct_from_scipy(scipy_name=stats.gamma)


# src/UQpy/distributions/collection/GeneralizedExtreme.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class GeneralizedExtreme(DistributionContinuous1D):

    @beartype
    def __init__(
        self,
        c: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param c: shape parameter
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            c=c, loc=loc, scale=scale, ordered_parameters=("c", "loc", "scale")
        )
        self._construct_from_scipy(scipy_name=stats.genextreme)


# src/UQpy/distributions/collection/InverseGaussian.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class InverseGauss(DistributionContinuous1D):

    @beartype
    def __init__(
        self,
        mu: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param mu: shape parameter :math:`\mu`
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            mu=mu, loc=loc, scale=scale, ordered_parameters=("mu", "loc", "scale")
        )
        self._construct_from_scipy(scipy_name=stats.invgauss)


# src/UQpy/distributions/collection/JointCopula.py

from types import MethodType
from typing import Union

import numpy as np
from beartype import beartype

from UQpy.distributions.baseclass import Copula
from UQpy.distributions.baseclass import (
    DistributionContinuous1D,
    DistributionND,
    DistributionDiscrete1D,
)


class JointCopula(DistributionND):
    @beartype
    def __init__(
        self,
        marginals: Union[list[DistributionContinuous1D], list[DistributionDiscrete1D]],
        copula: Copula,
    ):
        """

        :param marginals: list of distribution objects that define the marginals
        :param copula: copula object
        """
        super().__init__()
        self.ordered_parameters = []
        for i, m in enumerate(marginals):
            self.ordered_parameters.extend(
                [key + "_" + str(i) for key in m.ordered_parameters]
            )
        self.ordered_parameters.extend(
            [key + "_c" for key in copula.ordered_parameters]
        )

        # Check and save the marginals
        self.marginals = marginals
        if not (
            isinstance(self.marginals, list)
            and all(
                isinstance(d, (DistributionContinuous1D, DistributionDiscrete1D))
                for d in self.marginals
            )
        ):
            raise ValueError(
                "Input marginals must be a list of 1d continuous Distribution objects."
            )

        # Check the copula. Also, all the marginals should have a cdf method
        self.copula = copula
        if not isinstance(self.copula, Copula):
            raise ValueError("The input copula should be a Copula object.")
        if not all(hasattr(m, "cdf") for m in self.marginals):
            raise ValueError(
                "All the marginals should have a cdf method in order to define a joint with copula."
            )
        Copula.check_marginals(marginals=self.marginals)

        # Check if methods should exist, if yes define them bound them to the object
        if hasattr(self.copula, "evaluate_cdf"):

            def joint_cdf(dist, x):
                x = dist.check_x_dimension(x)
                # Compute cdf of independent marginals
                unif = np.array(
                    [marg.cdf(x[:, ind_m]) for ind_m, marg in enumerate(dist.marginals)]
                ).T
                # Compute copula
                cdf_val = dist.copula.evaluate_cdf(unit_uniform_samples=unif)
                return cdf_val

            self.cdf = MethodType(joint_cdf, self)

        if all(hasattr(m, "pdf") for m in self.marginals) and hasattr(
            self.copula, "evaluate_pdf"
        ):

            def joint_pdf(dist, x):
                x = dist.check_x_dimension(x)
                # Compute pdf of independent marginals
                pdf_val = np.prod(
                    np.array(
                        [
                            marg.pdf(x[:, ind_m])
                            for ind_m, marg in enumerate(dist.marginals)
                        ]
                    ),
                    axis=0,
                )
                # Add copula term
                unif = np.array(
                    [marg.cdf(x[:, ind_m]) for ind_m, marg in enumerate(dist.marginals)]
                ).T
                c_ = dist.copula.evaluate_pdf(unit_uniform_samples=unif)
                return c_ * pdf_val

            self.pdf = MethodType(joint_pdf, self)

        if all(hasattr(m, "log_pdf") for m in self.marginals) and hasattr(
            self.copula, "evaluate_pdf"
        ):

            def joint_log_pdf(dist, x):
                x = dist.check_x_dimension(x)
                # Compute pdf of independent marginals
                logpdf_val = np.sum(
                    np.array(
                        [
                            marg.log_pdf(x[:, ind_m])
                            for ind_m, marg in enumerate(dist.marginals)
                        ]
                    ),
                    axis=0,
                )
                # Add copula term
                unif = np.array(
                    [marg.cdf(x[:, ind_m]) for ind_m, marg in enumerate(dist.marginals)]
                ).T
                c_ = dist.copula.evaluate_pdf(unit_uniform_samples=unif)
                return np.log(c_) + logpdf_val

            self.log_pdf = MethodType(joint_log_pdf, self)

    def get_parameters(self) -> dict:
        """
        Return the parameters of a :class:`.Distributions` object.

        To update the parameters of a :class:`.JointIndependent` or a :class:`.JointCopula` distribution, each parameter
        is assigned a unique string identifier as :code:`key_index` - where :code:`key` is the parameter name and
        :code:`index` the index of the marginal (e.g., location parameter of the 2nd marginal is identified as
        :code:`loc_1`).

        :return: Parameters of the distribution.
        """
        params = {}
        for i, m in enumerate(self.marginals):
            for key, value in m.get_parameters().items():
                params[key + "_" + str(i)] = value
        for key, value in self.copula.get_parameters().items():
            params[key + "_c"] = value
        return params

    def update_parameters(self, **kwargs: dict):
        """
        Update the parameters of a :class:`.Distributions` object.

        To update the parameters of a :class:`.JointIndependent` or a :class:`.JointCopula` distribution, each
        parameter is assigned a unique string identifier as :code:`key_index` - where :code:`key` is the parameter name
        and :code:`index` the index of the marginal (e.g., location parameter of the 2nd marginal is identified as
        :code:`loc_1`).

        :param kwargs: Parameters to be updated
        :raises ValueError: if *kwargs* contains key that does not already exist.
        """
        # check arguments
        all_keys = self.get_parameters().keys()
        # update the marginal parameters
        for key_indexed, value in kwargs.items():
            if key_indexed not in all_keys:
                raise ValueError("Unrecognized keyword argument " + key_indexed)
            key_split = key_indexed.split("_")
            key, index = "_".join(key_split[:-1]), key_split[-1]
            if index == "c":
                self.copula.parameters[key] = value
            else:
                self.marginals[int(index)].parameters[key] = value


# src/UQpy/distributions/collection/JointIndependent.py

from types import MethodType
from typing import Union

import numpy as np
from beartype import beartype

from UQpy.distributions.baseclass import (
    DistributionContinuous1D,
    DistributionND,
    DistributionDiscrete1D,
)


class JointIndependent(DistributionND):
    @beartype
    def __init__(
            self,
            marginals: Union[list[DistributionContinuous1D], list[DistributionDiscrete1D]],
    ):
        """
        :param marginals: list of distribution objects that define the marginals.
        """
        super().__init__()
        self.ordered_parameters = []
        for i, m in enumerate(marginals):
            self.ordered_parameters.extend(
                [key + "_" + str(i) for key in m.ordered_parameters])

        # Check and save the marginals
        if not (isinstance(marginals, list)
                and all(isinstance(d, (DistributionContinuous1D, DistributionDiscrete1D)) for d in marginals)):
            raise ValueError("Input marginals must be a list of Distribution1d objects.")
        self.marginals = marginals

        # If all marginals have a method, the joint has it to
        if all(hasattr(m, "pdf") or hasattr(m, "pmf") for m in self.marginals):

            def joint_pdf(dist, x):
                x = dist.check_x_dimension(x)
                # Compute pdf of independent marginals
                pdf_val = np.ones((x.shape[0],))
                for ind_m in range(len(self.marginals)):
                    if hasattr(self.marginals[ind_m], "pdf"):
                        pdf_val *= marginals[ind_m].pdf(x[:, ind_m])
                    else:
                        pdf_val *= marginals[ind_m].pmf(x[:, ind_m])
                return pdf_val

            if any(hasattr(m, "pdf") for m in self.marginals):
                self.pdf = MethodType(joint_pdf, self)
            else:
                self.pmf = MethodType(joint_pdf, self)

        if all(hasattr(m, "log_pdf") or hasattr(m, "log_pmf") for m in self.marginals):

            def joint_log_pdf(dist, x):
                x = dist.check_x_dimension(x)
                # Compute pdf of independent marginals
                pdf_val = np.zeros((x.shape[0],))
                for ind_m in range(len(self.marginals)):
                    if hasattr(self.marginals[ind_m], "log_pdf"):
                        pdf_val += marginals[ind_m].log_pdf(x[:, ind_m])
                    else:
                        pdf_val += marginals[ind_m].log_pmf(x[:, ind_m])
                return pdf_val

            if any(hasattr(m, "log_pdf") for m in self.marginals):
                self.log_pdf = MethodType(joint_log_pdf, self)
            else:
                self.log_pmf = MethodType(joint_log_pdf, self)

        if all(hasattr(m, "cdf") for m in self.marginals):
            def joint_cdf(dist, x):
                x = dist.check_x_dimension(x)
                # Compute cdf of independent marginals
                cdf_val = np.prod(
                    np.array(
                        [
                            marg.cdf(x[:, ind_m])
                            for ind_m, marg in enumerate(dist.marginals)
                        ]
                    ),
                    axis=0,
                )
                return cdf_val

            self.cdf = MethodType(joint_cdf, self)

        if all(hasattr(m, "rvs") for m in self.marginals):

            def joint_rvs(dist, nsamples=1, random_state=None):
                # Go through all marginals
                rv_s = np.zeros((nsamples, len(dist.marginals)))
                for ind_m, marg in enumerate(dist.marginals):
                    rv_s[:, ind_m] = marg.rvs(
                        nsamples=nsamples, random_state=random_state
                    ).reshape((-1,))
                return rv_s

            self.rvs = MethodType(joint_rvs, self)

        if all(hasattr(m, "fit") for m in self.marginals):

            def joint_fit(dist, data):
                data = dist.check_x_dimension(data)
                # Compute ml estimates of independent marginal parameters
                mle_all = {}
                for ind_m, marg in enumerate(dist.marginals):
                    if any(
                            param_value is None
                            for param_value in marg.get_parameters().values()
                    ):
                        mle_i = marg.fit(data[:, ind_m])
                    else:
                        mle_i = marg.get_parameters().copy()
                    mle_all.update(
                        {key + "_" + str(ind_m): val for key, val in mle_i.items()}
                    )
                return mle_all

            self.fit = MethodType(joint_fit, self)

        if all(hasattr(m, "moments") for m in self.marginals):

            def joint_moments(dist, moments2return="mvsk"):
                # Go through all marginals
                if len(moments2return) == 1:
                    return np.array([marg.moments(moments2return=moments2return) for marg in dist.marginals])
                moments_ = [np.empty((len(dist.marginals),)) for _ in range(len(moments2return))]
                for ind_m, marg in enumerate(dist.marginals):
                    moments_i = marg.moments(moments2return=moments2return)
                    for j in range(len(moments2return)):
                        moments_[j][ind_m] = moments_i[j]
                return tuple(moments_)

            self.moments = MethodType(joint_moments, self)

    def get_parameters(self) -> dict:
        """
        Return the parameters of a :class:`.Distributions` object.

        To update the parameters of a :class:`.JointIndependent` or a :class:`.JointCopula` distribution, each parameter
        is assigned a unique string identifier as :code:`key_index` - where :code:`key` is the parameter name and
        :code:`index` the index of the marginal (e.g., location parameter of the 2nd marginal is identified as
        :code:`loc_1`).

        :return: Parameters of the distribution
        """
        params = {}
        for i, m in enumerate(self.marginals):
            params_m = m.get_parameters()
            for key, value in params_m.items():
                params[key + "_" + str(i)] = value
        return params

    def update_parameters(self, **kwargs: dict):
        """
        Update the parameters of a :class:`.Distributions` object.

        To update the parameters of a :class:`.JointIndependent` or a :class:`.JointCopula` distribution, each parameter
        is assigned a unique string identifier as :code:`key_index` - where :code:`key` is the parameter name and
        :code:`index` the index of the marginal (e.g., location parameter of the 2nd marginal is identified as
        :code:`loc_1`).

        :param kwargs: Parameters to be updated
        :raises ValueError: if *kwargs* contains key that does not already exist.
        """
        # check arguments
        all_keys = self.get_parameters().keys()
        # update the marginal parameters
        for key_indexed, value in kwargs.items():
            if key_indexed not in all_keys:
                raise ValueError("Unrecognized keyword argument " + key_indexed)
            key_split = key_indexed.split("_")
            key, index = "_".join(key_split[:-1]), int(key_split[-1])
            self.marginals[index].parameters[key] = value


# src/UQpy/distributions/collection/Laplace.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Laplace(DistributionContinuous1D):

    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.laplace)


# src/UQpy/distributions/collection/Levy.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Levy(DistributionContinuous1D):

    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.levy)


# src/UQpy/distributions/collection/Logistic.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Logistic(DistributionContinuous1D):

    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.logistic)


# src/UQpy/distributions/collection/Lognormal.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Lognormal(DistributionContinuous1D):

    @beartype
    def __init__(
        self,
        s: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param s: shape parameter
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            s=s, loc=loc, scale=scale, ordered_parameters=("s", "loc", "scale")
        )
        self._construct_from_scipy(scipy_name=stats.lognorm)


# src/UQpy/distributions/collection/Maxwell.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Maxwell(DistributionContinuous1D):

    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.maxwell)


# src/UQpy/distributions/collection/Multinomial.py

from typing import Union
import numpy as np
import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionND


class Multinomial(DistributionND):

    @beartype
    def __init__(self, n: Union[None, int], p: Union[list[float], np.ndarray]):
        """

        :param n: number of trials
        :param p: probability of a trial falling into each category; should sum to 1
        """
        super().__init__(n=n, p=p)

    def pmf(self, x):
        pdf_val = stats.multinomial.pmf(x=x, **self.parameters)
        return np.atleast_1d(pdf_val)

    def log_pmf(self, x):
        logpdf_val = stats.multinomial.logpmf(x=x, **self.parameters)
        return np.atleast_1d(logpdf_val)

    def rvs(self, nsamples=1, random_state=None):
        if not (isinstance(nsamples, int) and nsamples >= 1):
            raise ValueError("Input nsamples must be an integer > 0.")
        return stats.multinomial.rvs(
            size=nsamples, random_state=random_state, **self.parameters
        ).reshape((nsamples, -1))

    def moments(self, moments2return="mv"):
        n = self.get_parameters()["n"]
        p = np.array(self.get_parameters()["p"])
        d = len(p)
        if moments2return == "m":
            mean = n * np.array(p)
            return mean
        elif moments2return == "v":
            cov = (
                -n
                * np.tile(p[np.newaxis, :], [d, 1])
                * np.tile(p[:, np.newaxis], [1, d])
            )
            np.fill_diagonal(cov, n * p * (1.0 - p))
            return cov
        elif moments2return == "mv":
            cov = (
                -n
                * np.tile(p[np.newaxis, :], [d, 1])
                * np.tile(p[:, np.newaxis], [1, d])
            )
            np.fill_diagonal(cov, n * p * (1.0 - p))
            mean = n * np.array(p)
            return mean, cov
        else:
            raise ValueError('UQpy: moments2return must be "m", "v" or "mv".')


# src/UQpy/distributions/collection/MultivariateNormal.py

from typing import Union

import numpy as np
import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionND


class MultivariateNormal(DistributionND):

    @beartype
    def __init__(
        self,
        mean: Union[None, np.ndarray, list],
        cov: Union[None, int, float, np.ndarray, list[list]] = 1.0,
    ):
        """

        :param mean: mean vector, :class:`numpy.ndarray` of shape :code:`(dimension, )`
        :param cov: covariance, :any:`float` or :class:`numpy.ndarray` of shape :code:`(dimension, )` or
         :code:`(dimension, dimension)`. Default is 1.
        """
        if mean is not None and cov is not None:
            if isinstance(cov, (int, float)):
                pass
            else:
                if not (len(np.array(cov).shape) in [1, 2] and all(sh == len(mean) for sh in np.array(cov).shape)):
                    raise ValueError("Input covariance must be a float or ndarray of appropriate dimensions.")
        super().__init__(mean=mean, cov=cov, ordered_parameters=["mean", "cov"])

    def cdf(self, x):
        cdf_val = stats.multivariate_normal.cdf(x=x, **self.parameters)
        return np.atleast_1d(cdf_val)

    def pdf(self, x):
        pdf_val = stats.multivariate_normal.pdf(x=x, **self.parameters)
        return np.atleast_1d(pdf_val)

    def log_pdf(self, x):
        logpdf_val = stats.multivariate_normal.logpdf(x=x, **self.parameters)
        return np.atleast_1d(logpdf_val)

    def rvs(self, nsamples=1, random_state=None):
        if not (isinstance(nsamples, int) and nsamples >= 1):
            raise ValueError("Input nsamples must be an integer > 0.")
        return stats.multivariate_normal.rvs(size=nsamples, random_state=random_state, **self.parameters
                                             ).reshape((nsamples, -1))

    def fit(self, data):
        data = self.check_x_dimension(data)
        mle_mu, mle_cov = self.parameters["mean"], self.parameters["cov"]
        if mle_mu is None:
            mle_mu = np.mean(data, axis=0)
        if mle_cov is None:
            mle_cov = np.cov(data, rowvar=False)
        return {"mean": mle_mu, "cov": mle_cov}

    def moments(self, moments2return="mv"):
        if moments2return == "m":
            return self.get_parameters()["mean"]
        elif moments2return == "v":
            return self.get_parameters()["cov"]
        elif moments2return == "mv":
            return self.get_parameters()["mean"], self.get_parameters()["cov"]
        else:
            raise ValueError('UQpy: moments2return must be "m", "v" or "mv".')


# src/UQpy/distributions/collection/Normal.py

from typing import Union
import scipy.stats as stats
from beartype import beartype
from UQpy.distributions.baseclass import DistributionContinuous1D


class Normal(DistributionContinuous1D):

    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.norm)


# src/UQpy/distributions/collection/Pareto.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Pareto(DistributionContinuous1D):
    @beartype
    def __init__(
        self,
        b: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param b: shape parameter
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            b=b, loc=loc, scale=scale, ordered_parameters=("b", "loc", "scale")
        )
        self._construct_from_scipy(scipy_name=stats.pareto)


# src/UQpy/distributions/collection/Poisson.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionDiscrete1D


class Poisson(DistributionDiscrete1D):

    @beartype
    def __init__(self, mu: Union[None, float, int], loc: Union[None, float, int] = 0.0):
        """

        :param mu: shape parameter
        :param loc: location parameter
        """
        super().__init__(mu=mu, loc=loc, ordered_parameters=("mu", "loc"))
        self._construct_from_scipy(scipy_name=stats.poisson)


# src/UQpy/distributions/collection/Rayleigh.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Rayleigh(DistributionContinuous1D):
    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.rayleigh)


# src/UQpy/distributions/collection/TruncatedNormal.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class TruncatedNormal(DistributionContinuous1D):

    @beartype
    def __init__(
        self,
        a: Union[None, float, int],
        b: Union[None, float, int],
        loc: Union[None, float, int] = 0.0,
        scale: Union[None, float, int] = 1.0,
    ):
        """

        :param a: shape parameter
        :param b: shape parameter
        :param loc: location parameter
        :param scale: scale parameter
        """
        super().__init__(
            a=a,
            b=b,
            loc=loc,
            scale=scale,
            ordered_parameters=("a", "b", "loc", "scale"),
        )
        self._construct_from_scipy(scipy_name=stats.truncnorm)


# src/UQpy/distributions/collection/Uniform.py

from typing import Union

import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D


class Uniform(DistributionContinuous1D):
    @beartype
    def __init__(
        self, loc: Union[None, float, int] = 0.0, scale: Union[None, float, int] = 1.0
    ):
        """

        :param loc: lower bound
        :param scale: range
        """
        super().__init__(loc=loc, scale=scale, ordered_parameters=("loc", "scale"))
        self._construct_from_scipy(scipy_name=stats.uniform)


# src/UQpy/distributions/collection/__init__.py

"""distributions module."""
from UQpy.distributions.collection.Beta import Beta
from UQpy.distributions.collection.Binomial import Binomial
from UQpy.distributions.collection.Cauchy import Cauchy
from UQpy.distributions.collection.ChiSquare import ChiSquare
from UQpy.distributions.collection.Exponential import Exponential
from UQpy.distributions.collection.Gamma import Gamma
from UQpy.distributions.collection.GeneralizedExtreme import GeneralizedExtreme
from UQpy.distributions.collection.InverseGaussian import InverseGauss
from UQpy.distributions.collection.Laplace import Laplace
from UQpy.distributions.collection.Levy import Levy
from UQpy.distributions.collection.Logistic import Logistic
from UQpy.distributions.collection.Lognormal import Lognormal
from UQpy.distributions.collection.Maxwell import Maxwell
from UQpy.distributions.collection.Multinomial import Multinomial
from UQpy.distributions.collection.MultivariateNormal import MultivariateNormal
from UQpy.distributions.collection.Normal import Normal
from UQpy.distributions.collection.Pareto import Pareto
from UQpy.distributions.collection.Poisson import Poisson
from UQpy.distributions.collection.Rayleigh import Rayleigh
from UQpy.distributions.collection.TruncatedNormal import TruncatedNormal
from UQpy.distributions.collection.Uniform import Uniform
from UQpy.distributions.collection.JointIndependent import JointIndependent
from UQpy.distributions.collection.JointCopula import JointCopula


# src/UQpy/distributions/copulas/Clayton.py

import numpy
from beartype import beartype

from UQpy.utilities.ValidationTypes import Numpy2DFloatArray
from UQpy.distributions.baseclass import Copula
import numpy as np


class Clayton(Copula):

    @beartype
    def __init__(self, theta: float):
        """

        :param theta: Parameter of the Clayton copula, real number in :math:`[1, +\infty)`
        """
        super().__init__(theta=theta)

    def evaluate_cdf(self, unit_uniform_samples: Numpy2DFloatArray) -> numpy.ndarray:
        """
        Compute the copula cdf :math:`C(u_1, u_2, ..., u_d)` for a `d`-variate uniform distribution.

        For a generic multivariate distribution with marginal cdfs :math:`F_1, ..., F_d` the joint cdf is computed as:

        :math:`F(x_1, ..., x_d) = C(u_1, u_2, ..., u_d)`

        where :math:`u_i = F_i(x_i)` is uniformly distributed. This computation is performed in the
        :meth:`.JointCopula.cdf` method.

        :param unit_uniform_samples: Points (uniformly distributed) at which to evaluate the copula cdf, must be of
         shape :code:`(npoints, dimension)`.

        :return: Values of the cdf.
        """
        theta, u, v = self.extract_data(unit_uniform_samples)
        cdf_val = (np.maximum(u ** (-theta) + v ** (-theta) - 1.0, 0.0)) ** (-1.0 / theta)
        return cdf_val

    def extract_data(self, unit_uniform_samples: Numpy2DFloatArray):
        u = unit_uniform_samples[:, 0]
        v = unit_uniform_samples[:, 1]
        theta = self.parameters["theta"]
        return theta, u, v


# src/UQpy/distributions/copulas/Frank.py

import numpy
import numpy as np
from beartype import beartype

from UQpy.utilities.ValidationTypes import Numpy2DFloatArray
from UQpy.distributions.baseclass import Copula


class Frank(Copula):
    @beartype
    def __init__(self, theta: float):
        """

        :param theta: Parameter of the copula, real number in :math:`\mathbb{R}`
        """
        super().__init__(theta=theta)

    def evaluate_cdf(self, unit_uniform_samples: Numpy2DFloatArray) -> numpy.ndarray:
        """
        Compute the copula cdf :math:`C(u_1, u_2, ..., u_d)` for a `d`-variate uniform distribution.

        For a generic multivariate distribution with marginal cdfs :math:`F_1, ..., F_d` the joint cdf is computed as:

        :math:`F(x_1, ..., x_d) = C(u_1, u_2, ..., u_d)`

        where :math:`u_i = F_i(x_i)` is uniformly distributed. This computation is performed in the
        :meth:`.JointCopula.cdf` method.

        :param unit_uniform_samples: Points (uniformly distributed) at which to evaluate the copula cdf, must be of
         shape :code:`(npoints, dimension)`.

        :return: Values of the cdf.
        """
        theta, u, v = self.extract_data(unit_uniform_samples)
        tmp_ratio = ((np.exp(-theta * u) - 1.0) * (np.exp(-theta * v) - 1.0) / (np.exp(-theta) - 1.0))
        cdf_val = -1.0 / theta * np.log(1.0 + tmp_ratio)
        return cdf_val

    def extract_data(self, unit_uniform_samples: Numpy2DFloatArray):
        u = unit_uniform_samples[:, 0]
        v = unit_uniform_samples[:, 1]
        theta = self.parameters["theta"]
        return theta, u, v


# src/UQpy/distributions/copulas/Gumbel.py

from typing import Union

import numpy
import numpy as np
from beartype import beartype

from UQpy.utilities.ValidationTypes import Numpy2DFloatArray
from UQpy.distributions.baseclass import Copula
from numpy import log, exp


class Gumbel(Copula):
    @beartype
    def __init__(self, theta: Union[None, float]):
        """

        :param theta: Parameter of the Gumbel copula, real number in :math:`[1, +\infty)`.
        """
        super().__init__(theta=theta)

    def evaluate_cdf(self, unit_uniform_samples: Numpy2DFloatArray) -> np.ndarray:
        """
        Compute the copula cdf :math:`C(u_1, u_2, ..., u_d)` for a `d`-variate uniform distribution.

        For a generic multivariate distribution with marginal cdfs :math:`F_1, ..., F_d` the joint cdf is computed as:

        :math:`F(x_1, ..., x_d) = C(u_1, u_2, ..., u_d)`

        where :math:`u_i = F_i(x_i)` is uniformly distributed. This computation is performed in the
        :meth:`.JointCopula.cdf` method.

        :param unit_uniform_samples: Points (uniformly distributed) at which to evaluate the copula cdf, must be of
         shape :code:`(npoints, dimension)`.
        :return: Values of the cdf.
        """
        if unit_uniform_samples.shape[1] > 2:
            raise ValueError("Maximum dimension for the Gumbel Copula is 2")

        theta, u, v = self.extract_data(unit_uniform_samples)

        cdf_val = exp(-(((-log(u)) ** theta + (-log(v)) ** theta) ** (1 / theta)))

        return cdf_val

    def evaluate_pdf(self, unit_uniform_samples: Numpy2DFloatArray) -> numpy.ndarray:
        """
        Compute the copula pdf :math:`c(u_1, u_2, ..., u_d)` for a `d`-variate uniform distribution.

        For a generic multivariate distribution with marginals pdfs :math:`f_1, ..., f_d` and marginals cdfs
        :math:`F_1, ..., F_d`, the joint pdf is computed as:

        :math:`f(x_1, ..., x_d) = c(u_1, u_2, ..., u_d) f_1(x_1) ... f_d(x_d)`

        where :math:`u_i = F_i(x_i)` is uniformly distributed. This computation is performed in the
        :meth:`.JointCopula.pdf` method.

        :param unit_uniform_samples: Points (uniformly distributed) at which to evaluate the copula pdf, must be of
         shape :code:`(npoints, dimension)`.

        :return: Values of the copula pdf term.
        """
        theta, u, v = self.extract_data(unit_uniform_samples)
        c = exp(-(((-log(u)) ** theta + (-log(v)) ** theta) ** (1 / theta)))

        pdf_val = (c * 1 / u * 1 / v
                   * ((-log(u)) ** theta + (-log(v)) ** theta) ** (-2 + 2 / theta)
                   * (log(u) * log(v)) ** (theta - 1)
                   * (1 + (theta - 1) * ((-log(u)) ** theta + (-log(v)) ** theta) ** (-1 / theta)))
        return pdf_val

    def extract_data(self, unit_uniform_samples):
        u = unit_uniform_samples[:, 0]
        v = unit_uniform_samples[:, 1]
        theta = self.parameters["theta"]
        return theta, u, v


# src/UQpy/distributions/copulas/__init__.py

from UQpy.distributions.copulas.Gumbel import Gumbel
from UQpy.distributions.copulas.Clayton import Clayton
from UQpy.distributions.copulas.Frank import Frank


# src/UQpy/inference/evidence_methods/HarmonicMean.py

from UQpy.inference.evidence_methods.baseclass.EvidenceMethod import EvidenceMethod
import numpy as np


class HarmonicMean(EvidenceMethod):
    """
    Class used for the computation of model evidence using the harmonic mean method.
    """
    def estimate_evidence(self, inference_model, posterior_samples, log_posterior_values):
        log_likelihood_values = (log_posterior_values - inference_model.prior.log_pdf(x=posterior_samples))
        temp = np.mean(1.0 / np.exp(log_likelihood_values))
        return 1.0 / temp


# src/UQpy/inference/evidence_methods/__init__.py

from UQpy.inference.evidence_methods.baseclass import *
from UQpy.inference.evidence_methods.HarmonicMean import HarmonicMean


# src/UQpy/inference/inference_models/ComputationalModel.py

from typing import Callable

from beartype import beartype

from UQpy.run_model.RunModel import RunModel
from UQpy.distributions.baseclass import Distribution
from UQpy.inference.inference_models.baseclass.InferenceModel import *
from UQpy.distributions.collection.MultivariateNormal import MultivariateNormal
from UQpy.utilities.ValidationTypes import *
from UQpy.distributions.collection import Normal

import warnings

warnings.filterwarnings('ignore')


class ComputationalModel(InferenceModel):
    @beartype
    def __init__(self, n_parameters: PositiveInteger, runmodel_object: RunModel,
                 error_covariance: Union[np.ndarray, float] = 1.0, name: str = "", prior: Distribution = None,
                 log_likelihood: Callable = None):
        """
        Define a (non-)Gaussian error model for inference.

        :param n_parameters: Number of parameters to be estimated.
        :param runmodel_object: :class:`.RunModel` class object that defines the forward model. This input is required
         for **cases 1a and 1b**.
        :param error_covariance: Covariance for Gaussian error model **(case 1a)**. It can be a scalar (in which case
         the covariance matrix is the identity times that value), a 1D :class:`numpy.ndarray` in which case the
         covariance is assumed to be diagonal, or a full covariance matrix (2D :class:`numpy.ndarray`).
         Default value is :math:`1`.
        :param name: Name of model - optional but useful in a model selection setting.
        :param prior: Prior distribution, must have a :py:meth:`log_pdf` or :py:meth:`pdf` method.
        :param log_likelihood: Function that defines the log-likelihood model, possibly in conjunction with the
         `runmodel_object` **(cases 1b and 2)**. Default is :any:`None`, and a Gaussian-error model is considered
         **(case 1a)**.
        """
        super().__init__(n_parameters, name)
        self.runmodel_object = runmodel_object
        self.error_covariance = error_covariance
        self.log_likelihood = log_likelihood

        self.prior = prior
        if self.prior is not None:
            if not isinstance(self.prior, Distribution):
                raise TypeError("UQpy: Input prior should be an object of class Distribution.")
            if not hasattr(self.prior, "log_pdf"):
                if not hasattr(self.prior, "pdf"):
                    raise AttributeError("UQpy: Input prior should have a log_pdf or pdf method.")
                self.prior.log_pdf = lambda x: np.log(self.prior.pdf(x))

    def evaluate_log_likelihood(self, parameters: NumpyFloatArray, data: NumpyFloatArray):
        self.runmodel_object.run(samples=parameters, append_samples=False)
        model_outputs = self.runmodel_object.qoi_list

        # Case 1.a: Gaussian error model
        if self.log_likelihood is None:
            if isinstance(self.error_covariance, (float, int)):
                norm = Normal(loc=0.0, scale=np.sqrt(self.error_covariance))
                log_like_values = np.array(
                    [np.sum([norm.log_pdf(data_i - outpt_i) for data_i, outpt_i in zip(data, output)])
                     for output in model_outputs])
            else:
                multivariate_normal = MultivariateNormal(data, cov=self.error_covariance)
                log_like_values = np.array(
                    [multivariate_normal.log_pdf(x=np.array(output).reshape((-1,))) for output in model_outputs])

        # Case 1.b: likelihood is user-defined
        else:
            log_like_values = self.log_likelihood(data=data, model_outputs=model_outputs, params=parameters)
            if not isinstance(log_like_values, np.ndarray):
                log_like_values = np.array(log_like_values)
            if log_like_values.shape != (parameters.shape[0],):
                raise ValueError(
                    "UQpy: Likelihood function should output a (nsamples, ) ndarray of likelihood "
                    "values.")
        return log_like_values


# src/UQpy/inference/inference_models/DistributionModel.py

from typing import Union

import numpy as np
from beartype import beartype

from UQpy.distributions.baseclass import Distribution
from UQpy.inference.inference_models.baseclass.InferenceModel import *


class DistributionModel(InferenceModel):
    @beartype
    def __init__(self, distributions: Union[Distribution, list[Distribution]],
                 n_parameters: PositiveInteger, name: str = "", prior: Distribution = None):
        """
        Define a probability distribution model for inference.

        :param distributions: Distribution :math:`\pi` for which to learn parameters from iid data **(case 3)**.
         When creating this :class:`.Distribution` object, the parameters to be learned should be set to :any:`None`.
         Any parameter that is assigned a value will not be learned.
        :param n_parameters: Number of parameters to be estimated.
        :param name: Name of model - optional but useful in a model selection setting.
        :param prior: Prior distribution, must have a :py:meth:`log_pdf` or :py:meth:`pdf` method.
        """
        super().__init__(n_parameters, name)
        self.distributions = distributions

        if self.distributions is not None:
            if not isinstance(self.distributions, Distribution):
                raise TypeError("UQpy: Input dist_object should be an object of class Distribution.")
            if not hasattr(self.distributions, "log_pdf"):
                if not hasattr(self.distributions, "pdf"):
                    raise AttributeError("UQpy: dist_object should have a log_pdf or pdf method.")
                self.distributions.log_pdf = lambda x: np.log(self.distributions.pdf(x))
            init_params = self.distributions.get_parameters()
            self.list_params = [
                key for key in self.distributions.ordered_parameters if init_params[key] is None]
            if len(self.list_params) != self.n_parameters:
                raise TypeError("UQpy: Incorrect dimensions between nparams and number of inputs set to None.")

        self.prior = prior
        if self.prior is not None:
            if not isinstance(self.prior, Distribution):
                raise TypeError("UQpy: Input prior should be an object of class Distribution.")
            if not hasattr(self.prior, "log_pdf"):
                if not hasattr(self.prior, "pdf"):
                    raise AttributeError("UQpy: Input prior should have a log_pdf or pdf method.")
                self.prior.log_pdf = lambda x: np.log(self.prior.pdf(x))

    def evaluate_log_likelihood(self, parameters, data):
        log_like_values = []
        for params_ in parameters:
            self.distributions.update_parameters(**dict(zip(self.list_params, params_)))
            log_like_values.append(np.sum(self.distributions.log_pdf(x=data)))
        log_like_values = np.array(log_like_values)
        return log_like_values


# src/UQpy/inference/inference_models/LogLikelihoodModel.py

from typing import Callable

import numpy as np
from beartype import beartype
import warnings

warnings.filterwarnings('ignore')

from UQpy.inference.inference_models.baseclass.InferenceModel import *


class LogLikelihoodModel(InferenceModel):
    @beartype
    def __init__(self, n_parameters: PositiveInteger, log_likelihood: Callable, name: str = ""):
        """
        Define a log-likelihood model for inference.

        :param n_parameters: Number of parameters to be estimated.
        :param log_likelihood: Function that defines the log-likelihood model.
        :param name: Name of model - optional but useful in a model selection setting.
        """
        super().__init__(n_parameters, name)
        self.name = name
        self.log_likelihood = log_likelihood
        self.n_parameters = n_parameters

    def evaluate_log_likelihood(self, parameters: np.ndarray, data: np.ndarray):
        log_like_values = self.log_likelihood(data=data, params=parameters)
        if not isinstance(log_like_values, np.ndarray):
            log_like_values = np.array(log_like_values)
        if log_like_values.shape != (parameters.shape[0],):
            raise ValueError("UQpy: Likelihood function should output a (nsamples, ) ndarray of likelihood values.")
        return log_like_values


# src/UQpy/inference/inference_models/__init__.py

from UQpy.inference.inference_models.baseclass import *

from UQpy.inference.inference_models.ComputationalModel import ComputationalModel
from UQpy.inference.inference_models.DistributionModel import DistributionModel
from UQpy.inference.inference_models.LogLikelihoodModel import LogLikelihoodModel



# src/UQpy/inference/information_criteria/AIC.py

from UQpy.inference import MLE
from UQpy.utilities.ValidationTypes import NumpyFloatArray
from UQpy.inference.information_criteria.baseclass.InformationCriterion import InformationCriterion


class AIC(InformationCriterion):
    def minimize_criterion(self,
                           data: NumpyFloatArray,
                           parameter_estimator: MLE,
                           return_penalty: bool = False):
        inference_model = parameter_estimator.inference_model
        max_log_like = parameter_estimator.max_log_like
        n_parameters = inference_model.n_parameters

        penalty_term = self._calculate_penalty_term(n_parameters)
        if return_penalty:
            return -2 * max_log_like + penalty_term, penalty_term
        return -2 * max_log_like + penalty_term

    def _calculate_penalty_term(self, n_parameters):
        return 2 * n_parameters


# src/UQpy/inference/information_criteria/AICc.py

from UQpy.inference.information_criteria.baseclass.InformationCriterion import InformationCriterion
from UQpy.inference import MLE
from UQpy.utilities.ValidationTypes import NumpyFloatArray


class AICc(InformationCriterion):

    def minimize_criterion(self,
                           data: NumpyFloatArray,
                           parameter_estimator: MLE,
                           return_penalty: bool = False):
        inference_model = parameter_estimator.inference_model
        max_log_like = parameter_estimator.max_log_like
        n_parameters = inference_model.n_parameters
        n_data = len(data)

        penalty_term = self._calculate_penalty_term(n_data, n_parameters)
        if return_penalty:
            return -2 * max_log_like + penalty_term, penalty_term
        return -2 * max_log_like + penalty_term

    def _calculate_penalty_term(self, n_data, n_parameters):
        return 2 * n_parameters + (2 * n_parameters ** 2 + 2 * n_parameters) / (n_data - n_parameters - 1)



# src/UQpy/inference/information_criteria/BIC.py

from UQpy.inference.information_criteria.baseclass.InformationCriterion import InformationCriterion
import numpy as np
from UQpy.inference import MLE
from UQpy.utilities.ValidationTypes import NumpyFloatArray


class BIC(InformationCriterion):

    def minimize_criterion(self,
                           data: NumpyFloatArray,
                           parameter_estimator: MLE,
                           return_penalty: bool = False):
        inference_model = parameter_estimator.inference_model
        max_log_like = parameter_estimator.max_log_like
        n_parameters = inference_model.n_parameters
        n_data = len(data)

        penalty_term = self._calculate_penalty_term(n_data, n_parameters)
        if return_penalty:
            return -2 * max_log_like + penalty_term, penalty_term
        return -2 * max_log_like + penalty_term

    def _calculate_penalty_term(self, n_data, n_parameters):
        return np.log(n_data) * n_parameters


# src/UQpy/inference/information_criteria/__init__.py

from UQpy.inference.information_criteria.baseclass import *

from UQpy.inference.information_criteria.AIC import AIC
from UQpy.inference.information_criteria.BIC import BIC
from UQpy.inference.information_criteria.AICc import AICc

# src/UQpy/reliability/taylor_series/FORM.py

import logging
from typing import Union
import numpy as np
import scipy.stats as stats
from beartype import beartype

from UQpy.run_model.RunModel import RunModel
from UQpy.transformations import *
from UQpy.distributions import *
from UQpy.reliability.taylor_series.baseclass.TaylorSeries import TaylorSeries
from UQpy.utilities.ValidationTypes import PositiveInteger
from UQpy.transformations import Decorrelate
import warnings

warnings.filterwarnings('ignore')


class FORM(TaylorSeries):

    @beartype
    def __init__(
        self,
        distributions: Union[None, Distribution, list[Distribution]],
        runmodel_object: RunModel,
        seed_x: Union[list, np.ndarray] = None,
        seed_u: Union[list, np.ndarray] = None,
        df_step: Union[int, float] = 0.01,
        corr_x: Union[list, np.ndarray] = None,
        corr_z: Union[list, np.ndarray] = None,
        n_iterations: PositiveInteger = 100,
        tol1: Union[float, int] = None,
        tol2: Union[float, int] = None,
        tol3: Union[float, int] = None,
    ):
        """
        A class perform the First Order reliability Method. The :meth:`run` method of the :class:`.FORM` class can be invoked many
        times and each time the results are appended to the existing ones.
        This is a child class of the :class:`.TaylorSeries` class.

        :param distributions: Marginal probability distributions of each random variable. Must be an object of
         type :class:`.DistributionContinuous1D` or :class:`.JointIndependent`.
        :param runmodel_object: The computational model. It should be of type :class:`RunModel`.
        :param seed_u: The initial starting point in the uncorrelated standard normal space **U** for the `Hasofer-Lind`
         algorithm.
         Either `seed_u` or `seed_x` must be provided.
         Default: :code:`seed_u = (0, 0, ..., 0)`
         If either `seed_u` or `seed_x` is provided, then the :py:meth:`run` method will be executed automatically.
         Otherwise, the the :py:meth:`run` method must be executed by the user.
        :param seed_x: The initial starting point in the parameter space **X** for the `Hasofer-Lind` algorithm.
         Either `seed_u` or `seed_x` must be provided.
         If either `seed_u` or `seed_x` is provided, then the :py:meth:`run` method will be executed automatically.
         Otherwise, the the :py:meth:`run` method must be executed by the user.
        :param df_step: Finite difference step in standard normal space. Default: :math:`0.01`
        :param corr_z: Covariance matrix
         If `corr_z` is provided, it is the correlation matrix (:math:`\mathbf{C_Z}`) of the standard normal random
         vector **Z** .
         If `corr_x` is provided, it is the correlation matrix (:math:`\mathbf{C_X}`) of the random vector **X** .
        :param corr_z: Covariance matrix
         If `corr_x` is provided, it is the correlation matrix (:math:`\mathbf{C_X}`) of the random vector **X** .
         If `corr_z` is provided, it is the correlation matrix (:math:`\mathbf{C_Z}`) of the standard normal random
         vector **Z** .
         Default: `corr_z` is specified as the identity matrix.
        :param n_iterations: Maximum number of iterations for the `HLRF` algorithm. Default: :math:`100`
        :param tol1: Convergence threshold for criterion `e1` of the `HLRF` algorithm. Default: :math:`1.0e-3`
        :param tol2: Convergence threshold for criterion `e2` of the `HLRF` algorithm. Default: :math:`1.0e-3`
        :param tol3: Convergence threshold for criterion `e3` of the  `HLRF` algorithm. Default: :math:`1.0e-3`

        Any number of tolerances can be provided. Only the provided tolerances will be considered for the convergence
        of the algorithm. In case none of the tolerances is provided then they are considered equal to :math:`1e-3` and
        all are checked for the convergence.
        """
        if isinstance(distributions, list):
            self.dimension = len(distributions)
            self.dimension = 0
            for i in range(len(distributions)):
                if isinstance(distributions[i], DistributionContinuous1D):
                    self.dimension += 1
                elif isinstance(distributions[i], JointIndependent):
                    self.dimension += len(distributions[i].marginals)
                else:
                    raise TypeError(
                        "UQpy: A  ``DistributionContinuous1D`` or ``JointIndependent`` object must be "
                        "provided.")
        elif isinstance(distributions, DistributionContinuous1D):
            self.dimension = 1
        elif isinstance(distributions, JointIndependent):
            self.dimension = len(distributions.marginals)
        else:
            raise TypeError("UQpy: A  ``DistributionContinuous1D``  or ``JointIndependent`` object must be provided.")

        self.nataf_object = Nataf(distributions=distributions, corr_z=corr_z, corr_x=corr_x)

        self.corr_x = corr_x
        self.corr_z = corr_z
        self.dist_object = distributions
        self.n_iterations = n_iterations
        self.runmodel_object = runmodel_object
        self.tol1 = tol1
        self.tol2 = tol2
        self.tol3 = tol3
        self.seed_u = seed_u
        self.seed_x = seed_x

        self.logger = logging.getLogger(__name__)

        # Initialize output
        self.beta: float = None
        """Hasofer-Lind reliability index."""
        self.DesignPoint_U: list = None
        """Design point in the uncorrelated standard normal space U."""
        self.DesignPoint_X: list = None
        """Design point in the parameter space X."""
        self.alpha: float = None
        """Direction cosine."""
        self.failure_probability = None
        self.x = None
        self.g0 = None
        self.iterations: int = None
        """Number of model evaluations."""
        self.df_step = df_step
        self.error_record: list = None
        """Record of the error defined by criteria `e1, e2, e3`."""

        self.tol1 = tol1
        self.tol2 = tol2
        self.tol3 = tol3

        self.u_record: list = None
        """Record of all iteration points in the standard normal space **U**."""
        self.x_record: list = None
        """Record of all iteration points in the parameter space **X**."""
        self.g_record: list = None
        """Record of the performance function."""
        self.dg_u_record: list = None
        """Record of the model’s gradient in the standard normal space."""
        self.alpha_record: list = None
        """Record of the alpha (directional cosine)."""
        self.beta_record: list = None
        """Record of all Hasofer-Lind reliability index values."""
        self.jacobian_zx = None
        """Jacobian of the transformation from correlated standard normal space to the parameter space."""
        self.call = None

        if self.seed_u is not None:
            self.run(seed_u=self.seed_u)
        elif self.seed_x is not None:
            self.run(seed_x=self.seed_x)

    def run(self, seed_x: Union[list, np.ndarray] = None,
            seed_u: Union[list, np.ndarray] = None):
        """
        Runs FORM.

        :param seed_u: Either `seed_u` or `seed_x` must be provided.
         If `seed_u` is provided, it should be a point in the uncorrelated standard normal space of **U**.
         If `seed_x` is provided, it should be a point in the parameter space of **X**.
        :param seed_x: The initial starting point for the `Hasofer-Lind` algorithm.
         Either `seed_u` or `seed_x` must be provided.
         If `seed_u` is provided, it should be a point in the uncorrelated standard normal space of **U**.
         If `seed_x` is provided, it should be a point in the parameter space of **X**.
        """
        self.logger.info("UQpy: Running FORM...")
        if seed_u is None and seed_x is None:
            seed = np.zeros(self.dimension)
        elif seed_u is None and seed_x is not None:
            self.nataf_object.run(samples_x=seed_x.reshape(1, -1), jacobian=False)
            seed_z = self.nataf_object.samples_z
            seed = Decorrelate(seed_z, self.nataf_object.corr_z)
        elif seed_u is not None and seed_x is None:
            seed = np.squeeze(seed_u)
        else:
            raise ValueError("UQpy: Only one seed (seed_x or seed_u) must be provided")

        u_record = list()
        x_record = list()
        g_record = list()
        alpha_record = list()
        error_record = list()

        converged = False
        k = 0
        beta = np.zeros(shape=(self.n_iterations + 1,))
        u = np.zeros([self.n_iterations + 1, self.dimension])
        u[0, :] = seed
        g_record.append(0.0)
        dg_u_record = np.zeros([self.n_iterations + 1, self.dimension])

        while not converged:
            self.logger.info("Number of iteration: %i", k)
            # FORM always starts from the standard normal space
            if k == 0:
                if seed_x is not None:
                    x = seed_x
                else:
                    seed_z = Correlate(samples_u=seed.reshape(1, -1), corr_z=self.nataf_object.corr_z).samples_z
                    self.nataf_object.run(samples_z=seed_z.reshape(1, -1), jacobian=True)
                    x = self.nataf_object.samples_x
                    self.jacobian_zx = self.nataf_object.jxz
            else:
                z = Correlate(u[k, :].reshape(1, -1), self.nataf_object.corr_z).samples_z
                self.nataf_object.run(samples_z=z, jacobian=True)
                x = self.nataf_object.samples_x
                self.jacobian_zx = self.nataf_object.jxz

            self.x = x
            u_record.append(u)
            x_record.append(x)
            self.logger.info(
                "Design point Y: {0}\n".format(u[k, :])
                + "Design point X: {0}\n".format(self.x)
                + "Jacobian Jzx: {0}\n".format(self.jacobian_zx))

            # 2. evaluate Limit State Function and the gradient at point u_k and direction cosines
            dg_u, qoi, _ = self._derivatives(
                point_u=u[k, :],
                point_x=self.x,
                runmodel_object=self.runmodel_object,
                nataf_object=self.nataf_object,
                df_step=self.df_step,
                order="first")
            g_record.append(qoi)

            dg_u_record[k + 1, :] = dg_u
            norm_grad = np.linalg.norm(dg_u_record[k + 1, :])
            alpha = dg_u / norm_grad
            self.logger.info(
                "Directional cosines (alpha): {0}\n".format(alpha)
                + "Gradient (dg_y): {0}\n".format(dg_u_record[k + 1, :])
                + "norm dg_y:",
                norm_grad,)

            self.alpha = alpha.squeeze()
            alpha_record.append(self.alpha)
            beta[k] = -np.inner(u[k, :].T, self.alpha)
            beta[k + 1] = beta[k] + qoi / norm_grad
            self.logger.info(
                "Beta: {0}\n".format(beta[k])
                + "Pf: {0}".format(stats.norm.cdf(-beta[k])))

            u[k + 1, :] = -beta[k + 1] * self.alpha

            if (self.tol1 is not None) and (self.tol2 is not None) and (self.tol3 is not None):
                error1 = np.linalg.norm(u[k + 1, :] - u[k, :])
                error2 = np.linalg.norm(beta[k + 1] - beta[k])
                error3 = np.linalg.norm(dg_u_record[k + 1, :] - dg_u_record[k, :])
                error_record.append([error1, error2, error3])
                if error1 <= self.tol1 and error2 <= self.tol2 and error3 < self.tol3:
                    converged = True
                else:
                    k = k + 1

            if (self.tol1 is None) and (self.tol2 is None) and (self.tol3 is None):
                error1 = np.linalg.norm(u[k + 1, :] - u[k, :])
                error2 = np.linalg.norm(beta[k + 1] - beta[k])
                error3 = np.linalg.norm(dg_u_record[k + 1, :] - dg_u_record[k, :])
                error_record.append([error1, error2, error3])
                if error1 <= 1e-3 or error2 <= 1e-3 or error3 < 1e-3:
                    converged = True
                else:
                    k = k + 1

            elif (self.tol1 is not None) and (self.tol2 is None) and (self.tol3 is None):
                error1 = np.linalg.norm(u[k + 1, :] - u[k, :])
                error_record.append(error1)
                if error1 <= self.tol1:
                    converged = True
                else:
                    k = k + 1

            elif (
                (self.tol1 is None) and (self.tol2 is not None) and (self.tol3 is None)
            ):
                error2 = np.linalg.norm(beta[k + 1] - beta[k])
                error_record.append(error2)
                if error2 <= self.tol2:
                    converged = True
                else:
                    k = k + 1

            elif (self.tol1 is None) and (self.tol2 is None) and (self.tol3 is not None):
                error3 = np.linalg.norm(dg_u_record[k + 1, :] - dg_u_record[k, :])
                error_record.append(error3)
                if error3 < self.tol3:
                    converged = True
                else:
                    k = k + 1

            elif (self.tol1 is not None) and (self.tol2 is not None) and (self.tol3 is None):
                error1 = np.linalg.norm(u[k + 1, :] - u[k, :])
                error2 = np.linalg.norm(beta[k + 1] - beta[k])
                error_record.append([error1, error2])
                if error1 <= self.tol1 and error2 <= self.tol1:
                    converged = True
                else:
                    k = k + 1

            elif (self.tol1 is not None) and (self.tol2 is None) and (self.tol3 is not None):
                error1 = np.linalg.norm(u[k + 1, :] - u[k, :])
                error3 = np.linalg.norm(dg_u_record[k + 1, :] - dg_u_record[k, :])
                error_record.append([error1, error3])
                if error1 <= self.tol1 and error3 < self.tol3:
                    converged = True
                else:
                    k = k + 1

            elif (self.tol1 is None) and (self.tol2 is not None) and (self.tol3 is not None):
                error2 = np.linalg.norm(beta[k + 1] - beta[k])
                error3 = np.linalg.norm(dg_u_record[k + 1, :] - dg_u_record[k, :])
                error_record.append([error2, error3])
                if error2 <= self.tol2 and error3 < self.tol3:
                    converged = True
                else:
                    k = k + 1

            self.logger.error("Error: %s", error_record[-1])

            if converged is True or k > self.n_iterations:
                break

        if k > self.n_iterations:
            self.logger.info("UQpy: Maximum number of iterations {0} was reached before convergence."
                             .format(self.n_iterations))
            self.error_record = error_record
            self.u_record = [u_record]
            self.x_record = [x_record]
            self.g_record = [g_record]
            self.dg_u_record = [dg_u_record[:k]]
            self.alpha_record = [alpha_record]
        else:
            if self.call is None:
                self.beta_record = [beta[:k]]
                self.error_record = error_record
                self.beta = [beta[k]]
                self.DesignPoint_U = [u[k, :]]
                self.DesignPoint_X = [np.squeeze(self.x)]
                self.failure_probability = [stats.norm.cdf(-self.beta[-1])]
                self.iterations = [k]
                self.u_record = [u_record[:k]]
                self.x_record = [x_record[:k]]
                self.g_record = [g_record]
                self.dg_u_record = [dg_u_record[:k]]
                self.alpha_record = [alpha_record]
            else:
                self.beta_record = self.beta_record + [beta[:k]]
                self.beta = self.beta + [beta[k]]
                self.error_record = self.error_record + error_record
                self.DesignPoint_U = self.DesignPoint_U + [u[k, :]]
                self.DesignPoint_X = self.DesignPoint_X + [np.squeeze(self.x)]
                self.failure_probability = self.failure_probability + [stats.norm.cdf(-beta[k])]
                self.iterations = self.iterations + [k]
                self.u_record = self.u_record + [u_record[:k]]
                self.x_record = self.x_record + [x_record[:k]]
                self.g_record = self.g_record + [g_record]
                self.dg_u_record = self.dg_u_record + [dg_u_record[:k]]
                self.alpha_record = self.alpha_record + [alpha_record]
            self.call = True


# src/UQpy/reliability/taylor_series/SORM.py

import logging
from typing import Union

import numpy as np
import scipy.stats as stats
from beartype import beartype

from UQpy.run_model.RunModel import RunModel
from UQpy.distributions.baseclass import Distribution
from UQpy.utilities.ValidationTypes import PositiveInteger
from UQpy.reliability.taylor_series.FORM import FORM
from UQpy.reliability.taylor_series.baseclass.TaylorSeries import TaylorSeries


class SORM(TaylorSeries):

    @beartype
    def __init__(
            self,
            form_object: FORM,
            df_step: Union[float, int] = 0.01,
    ):
        """
        :class:`.SORM` is a child class of the :class:`.TaylorSeries` class. Input: The :class:`.SORM` class requires an
        object of type :class:`.FORM` as input.

        :param form_object: Object of type :class:`.FORM`
        :param df_step: Finite difference step in standard normal space. Default: :math:`0.01`
        """

        self.logger = logging.getLogger(__name__)
        """Design point in the parameter space **X**."""
        self.df_step = df_step
        """Record of the error defined by criteria `e1, e2, e3`."""

        self.failure_probability: float = None
        self.beta_second_order = None
        """Hasofer-Lind reliability index using the SORM method."""
        self.call = None

        self.form_object = form_object
        self._run()

    @classmethod
    @beartype
    def build_from_first_order(
            cls,
            distributions: Union[None, Distribution, list[Distribution]],
            runmodel_object: RunModel,
            seed_x: Union[list, np.ndarray] = None,
            seed_u: Union[list, np.ndarray] = None,
            df_step: Union[int, float] = 0.01,
            corr_x: Union[list, np.ndarray] = None,
            corr_z: Union[list, np.ndarray] = None,
            n_iterations: PositiveInteger = 100,
            tol1: Union[float, int] = None,
            tol2: Union[float, int] = None,
            tol3: Union[float, int] = None,
    ):
        """
        :param distributions: Marginal probability distributions of each random variable. Must be an object of
         type :class:`.DistributionContinuous1D` or :class:`.JointIndependent`.
        :param runmodel_object: The computational model. It should be of type :class:`RunModel`.
        :param seed_u: The initial starting point in the uncorrelated standard normal space **U** for the `Hasofer-Lind`
         algorithm.
         Either `seed_u` or `seed_x` must be provided.
         Default: :code:`seed_u = (0, 0, ..., 0)`
         If either `seed_u` or `seed_x` is provided, then the :py:meth:`run` method will be executed automatically.
         Otherwise, the the :py:meth:`run` method must be executed by the user.
        :param seed_x: The initial starting point in the parameter space **X** for the `Hasofer-Lind` algorithm.
         Either `seed_u` or `seed_x` must be provided.
         If either `seed_u` or `seed_x` is provided, then the :py:meth:`run` method will be executed automatically.
         Otherwise, the the :py:meth:`run` method must be executed by the user.
        :param df_step: Finite difference step in standard normal space. Default: :math:`0.01`
        :param corr_z: Covariance matrix
         If `corr_z` is provided, it is the correlation matrix (:math:`\mathbf{C_Z}`) of the standard normal random
         vector **Z** .
         If `corr_x` is provided, it is the correlation matrix (:math:`\mathbf{C_X}`) of the random vector **X** .
        :param corr_z: Covariance matrix
         If `corr_x` is provided, it is the correlation matrix (:math:`\mathbf{C_X}`) of the random vector **X** .
         If `corr_z` is provided, it is the correlation matrix (:math:`\mathbf{C_Z}`) of the standard normal random
         vector **Z** .
         Default: `corr_z` is specified as the identity matrix.
        :param n_iterations: Maximum number of iterations for the `HLRF` algorithm. Default: :math:`100`
        :param tol1: Convergence threshold for criterion `e1` of the `HLRF` algorithm. Default: :math:`1.0e-3`
        :param tol2: Convergence threshold for criterion `e2` of the `HLRF` algorithm. Default: :math:`1.0e-3`
        :param tol3: Convergence threshold for criterion `e3` of the  `HLRF` algorithm. Default: :math:`1.0e-3`

        Any number of tolerances can be provided. Only the provided tolerances will be considered for the convergence
        of the algorithm. In case none of the tolerances is provided then they are considered equal to :math:`1e-3` and
        all are checked for the convergence.
        """
        f = FORM(distributions, runmodel_object, seed_x, seed_u, df_step,
                 corr_x, corr_z, n_iterations, tol1, tol2, tol3)

        return cls(f, df_step)

    def _run(self):
        """
        Runs SORM
        """
        self.logger.info("UQpy: Calculating SORM correction...")
        beta = self.form_object.beta[-1]

        self.dimension = self.form_object.dimension
        model = self.form_object.runmodel_object
        dg_u_record = self.form_object.dg_u_record

        matrix_a = np.fliplr(np.eye(self.dimension))
        matrix_a[:, 0] = self.form_object.alpha

        def normalize(v):
            return v / np.sqrt(v.dot(v))

        q = np.zeros(shape=(self.dimension, self.dimension))
        q[:, 0] = normalize(matrix_a[:, 0])

        for i in range(1, self.dimension):
            ai = matrix_a[:, i]
            for j in range(0, i):
                aj = matrix_a[:, j]
                t = ai.dot(aj)
                ai = ai - t * aj
            q[:, i] = normalize(ai)

        r1 = np.fliplr(q).T
        self.logger.info("UQpy: Calculating the hessian for SORM..")

        hessian_g = self._derivatives(point_u=self.form_object.DesignPoint_U[-1],
                                      point_x=self.form_object.DesignPoint_X[-1],
                                      runmodel_object=model,
                                      nataf_object=self.form_object.nataf_object,
                                      order="second",
                                      df_step=self.df_step,
                                      point_qoi=self.form_object.g_record[-1][-1])

        matrix_b = np.dot(np.dot(r1, hessian_g), r1.T) / np.linalg.norm(dg_u_record[-1])
        kappa = np.linalg.eig(matrix_b[: self.dimension - 1, : self.dimension - 1])
        if self.call is None:
            self.failure_probability = [stats.norm.cdf(-1 * beta) * np.prod(1 / (1 + beta * kappa[0]) ** 0.5)]
            self.beta_second_order = [-stats.norm.ppf(self.failure_probability)]
        else:
            self.failure_probability += [stats.norm.cdf(-1 * beta) * np.prod(1 / (1 + beta * kappa[0]) ** 0.5)]
            self.beta_second_order += [-stats.norm.ppf(self.failure_probability)]

        self.call = True


# src/UQpy/reliability/taylor_series/__init__.py

from UQpy.reliability.taylor_series.FORM import FORM
from UQpy.reliability.taylor_series.SORM import SORM
from UQpy.reliability.taylor_series.baseclass.TaylorSeries import TaylorSeries


# src/UQpy/run_model/model_execution/ParallelExecution.py

# pragma: no cover
from __future__ import print_function

import math
import sys

import numpy as np
from mpi4py import MPI
import os
import pickle

try:
    comm = MPI.COMM_WORLD

    model = None
    samples = None
    samples_per_process = 0
    samples_shape = None
    samples_list = None
    ranges_list = None
    local_ranges = None
    local_samples = None
    if comm.rank == 0:
        n_existing_simulations = int(sys.argv[1])
        n_new_simulations = int(sys.argv[2])

        with open('model.pkl', 'rb') as filehandle:
            model = pickle.load(filehandle)

        with open('samples.pkl', 'rb') as filehandle:
            samples = pickle.load(filehandle)

        print(len(samples))
        print(samples[0].shape)

        samples_shape = list(samples.shape)
        n_samples = len(samples)
        samples_per_process = math.floor(n_samples / comm.size if n_samples / comm.size == 0 else n_samples / comm.size + 1)

        samples_list = []
        ranges_list = []
        for i in range(comm.size):
            start_index = samples_per_process * i
            if start_index > n_samples:
                samples_shape[0] = 0
                samples_list.append(np.empty(samples_shape))
                ranges_list.append(range(start_index, start_index))
                continue
            end_index = min(n_samples, samples_per_process * (i + 1))
            samples_list.append(samples.take(indices=range(start_index, end_index), axis=0))
            ranges_list.append(range(start_index, end_index))

    local_ranges = comm.scatter(ranges_list, root=0)
    print(list(local_ranges))
    local_samples = comm.scatter(samples_list, root=0)

    # broadcast model among processes
    model = comm.bcast(model, root=0)

    results = []

    if len(local_ranges) != 0:
        print(f"I am process {comm.rank} out of {comm.size} on node {MPI.Get_processor_name()} and my range is {list(local_ranges)}")
        index_start = local_ranges[0]
        print(index_start)
        for i in local_ranges:
            index = i - local_ranges[0]
            sample = model.preprocess_single_sample(i, sample=local_samples[index])

            execution_output = model.execute_single_sample(i, sample)

            results.append(model.postprocess_single_file(i, execution_output))

    qoi = comm.gather(results, root=0)

    if comm.rank == 0:
        result = []
        [result.extend(el) for el in qoi]
        with open('qoi.pkl', 'wb') as filehandle:
            pickle.dump(result, filehandle)

    comm.Barrier()  # wait for everybody to synchronize _here_

except Exception as e:
    print(e)



# src/UQpy/run_model/model_execution/PythonModel.py

import logging
import pathlib
import platform

import numpy as np
from beartype import beartype

from UQpy.utilities.ValidationTypes import Numpy2DFloatArray


class PythonModel:
    @beartype
    def __init__(self, model_script: str, model_object_name: str, var_names: list[str] = None,
                 delete_files: bool = False, **model_object_name_kwargs):
        """

        :param model_script: The filename (with .py extension) of the Python script which contains commands to
         execute the model.

         The named file must be present in the current working directory from which :class:`.RunModel` is called.
        :param model_object_name: In the Python workflow, `model_object_name` specifies the name of the function or
         class within `model_script' that executes the model. If there is only one function or class in the
         `model_script`, then it is not necessary to specify the model_object_name. If there are multiple objects within
         the `model_script`, then `model_object_name` must be specified.

         `model_object_name` is not used in the third-party software model workflow.
        :param var_names: A list containing the names of the variables present in `input_template`.

         If `input template` is provided and  `var_names` is not passed, i.e. if ``var_names=None``, then the default
         variable names `x0`, `x1`, `x2`,..., `xn` are created and used by :class:`.RunModel`, where `n` is the number of
         variables (`n_vars`).

         The number of variables is equal to the second dimension of `samples` (i.e. ``n_vars=len(samples[0])``).

         `var_names` is not used in the Python model workflow.
        :param delete_files: Specifies whether or not to delete individual run output files after model execution
         and output processing.

         If `delete_files = True`, :class:`.RunModel` will remove all `run_i...` directories in the `model_dir`.
        :param model_object_name_kwargs: Additional inputs to the Python object specified by `model_object_name` in the
         Python model workflow.
        """
        if var_names is None:
            var_names = []
        self.var_names = var_names
        self._model_output = None
        self.logger = logging.getLogger(__name__)

        if platform.system() in ["Windows"]:
            self.python_command = "python"
        else:
            self.python_command = "python3"

        self.model_object_name = model_object_name
        self.model_object_name_kwargs = model_object_name_kwargs

        self.delete_files = delete_files

        # Check if the model script is a python script
        model_extension = pathlib.Path(model_script).suffix
        if model_extension == ".py":
            self.model_script = model_script
        else:
            raise ValueError("\nUQpy: The model script must be the name of a python script, with extension '.py'.")

        # Import the python module
        python_model = __import__(self.model_script[:-3])
        self.model_object = getattr(python_model, self.model_object_name)
        # Run function which checks if the python model has the model object
        self._check_python_model(python_model)
        self.logger.info('\nUQpy: Performing serial execution of a Python model.\n')

    def initialize(self, samples):
        pass

    def finalize(self):
        pass

    def preprocess_single_sample(self, index, sample) -> Numpy2DFloatArray:
        return np.atleast_2d(sample)

    def execute_single_sample(self, index, sample_to_send):
        if len(self.model_object_name_kwargs) == 0:
            return self.model_object(sample_to_send)
        else:
            return self.model_object(sample_to_send, **self.model_object_name_kwargs)

    def postprocess_single_file(self, index, model_output):
        return model_output.qoi if self.model_is_class else model_output

    def _check_python_model(self, python_model):
        """
        Check if python model name is valid

        This function gets the name of the classes and functions in the imported python module whose names is passed in
        as the python model to RunModel. There should be at least one class or function in the module - if not there,
        then the function exits raising a ValueError. If there is at least one class or function in the module,
        if the model object name is not given as input and there is only one class or function, that class name or
        function name is used to run the model. If there is a model_object_name given, check if it is a valid name.
        Else, a ValueError is raised.
        """
        # Get the names of the classes and functions in the imported module
        import inspect

        class_list = []
        function_list = []
        for name, obj in inspect.getmembers(python_model):
            if inspect.isclass(obj):
                class_list.append(name)
            elif inspect.isfunction(obj):
                function_list.append(name)

        # There should be at least one class or function in the module - if not there, exit with error.
        if len(class_list) == 0 and len(function_list) == 0:
            raise ValueError("\nUQpy: A python model should be defined as a function or class in the script.\n")

        else:  # If there is at least one class or function in the module
            # If the model object name is not given as input and there is only one class or function,
            # take that class name or function name to run the model.
            if self.model_object_name is None and len(class_list) + len(function_list) == 1:
                if len(class_list) == 1:
                    self.model_object_name = class_list[0]
                elif len(function_list) == 1:
                    self.model_object_name = function_list[0]

            # If there is a model_object_name given, check if it is in the list.
            if self.model_object_name in class_list:
                self.logger.info("\nUQpy: The model class that will be run: " + self.model_object_name)
                self.model_is_class = True
            elif self.model_object_name in function_list:
                self.logger.info("\nUQpy: The model function that will be run: " + self.model_object_name)
                self.model_is_class = False
            else:
                if self.model_object_name is None:
                    raise ValueError("\nUQpy: There are more than one objects in the module. Specify the name of the "
                                     "function or class which has to be executed.\n")
                else:
                    print("\nUQpy: You specified the model_object_name as: " + str(self.model_object_name))
                    raise ValueError("\nUQpy: The file does not contain an object which was specified as the model.\n")


# src/UQpy/run_model/model_execution/SerialExecution.py

import logging

import numpy as np


class SerialExecution:
    def __init__(self):
        self.logger = logging.getLogger(__name__)

    def run(self, model, n_existing_simulations, n_new_simulations, samples):
        results = []
        for i in range(n_existing_simulations, n_existing_simulations + n_new_simulations):
            sample = model.preprocess_single_sample(i, samples)

            execution_output = model.execute_single_sample(i, sample)

            results.append(model.postprocess_single_file(i, execution_output))

        self.logger.info("\nUQpy: Serial execution of the python model complete.\n")
        return results


# src/UQpy/run_model/model_execution/ThirdPartyModel.py

import collections
import datetime
import logging
import os
import pathlib
import platform
import re
import shutil
import subprocess

import numpy as np


class ThirdPartyModel:

    def __init__(self, var_names: list[str], input_template: str, model_script: str, model_object_name: str,
                 output_script: str, output_object_name: str, fmt: str = None, separator: str = ', ',
                 delete_files: bool = False, model_dir: str = "Model_Runs"):
        """

        :param var_names: A list containing the names of the variables present in `input_template`.

         If `input template` is provided and  `var_names` is not passed, i.e. if ``var_names=None``, then the default
         variable names `x0`, `x1`, `x2`,..., `xn` are created and used by :class:`.RunModel`, where `n` is the number of
         variables (`n_vars`).

         The number of variables is equal to the second dimension of `samples` (i.e. ``n_vars=len(samples[0])``).

         `var_names` is not used in the Python model workflow.
        :param input_template: The name of the template input file that will be used to generate input files for
         each run of the model. When operating :class:`.RunModel` with a third-party software model, ``input_template`` must
         be specified.

         The named file must be present in the current working directory from which :class:`.RunModel` is called.
        :param model_script: The filename (with .py extension) of the Python script which contains commands to
         execute the model.

         The named file must be present in the current working directory from which :class:`.RunModel` is called.
        :param model_object_name: In the Python workflow, `model_object_name` specifies the name of the function or
         class within `model_script' that executes the model. If there is only one function or class in the
         `model_script`, then it is not necessary to specify the model_object_name. If there are multiple objects within
         the `model_script`, then `model_object_name` must be specified.
        :param output_script: The filename of the Python script that contains the commands to process the output from
         third-party software model evaluation. `output_script` is used to extract quantities of interest from model
         output files and return the quantities of interest to :class:`.RunModel` for subsequent :py:mod:`UQpy` processing (e.g. for
         adaptive methods that utilize the results of previous simulations to initialize new simulations).

         If, in the third-party software model workflow, ``output_script = None`` (the default), then the
         :py:attr:`qoi_list` attribute is empty and postprocessing must be handled outside of :py:mod:`UQpy`.

         If used, the named file must be present in the current working directory from which :class:`.RunModel` is called.

         `output_script` is not used in the Python model workflow. In the Python model workflow, all model postprocessing
         is handled directly within `model_script`.
        :param output_object_name: The name of the function or class within `output_script` that is used to collect
         and process the output values from third-party software model output files. If the object is a class, the
         output must be saved as an attribute called :py:attr:`qoi`. If it is a function, it should return the output
         quantity of interest.

         If there is only one function or only one class in `output_script`, then it is not necessary to specify
         `output_object_name`. If there are multiple objects in `output_script`, then output_object_name must be
         specified.
        :param fmt: If the `template_input` requires variables to be written in specific format, this format can be
         specified here.

         Format specification follows standard Python conventions for the str.format() command described at:
         https://docs.python.org/3/library/stdtypes.html#str.format. For additional details, see the Format String Syntax
         description at: https://docs.python.org/3/library/string.html#formatstrings.

         For example, ls-dyna .k files require each card is to be exactly 10 characters. The following format string
         syntax can be used, "{:>10.4f}".

        :param separator: A string used to delimit values when printing arrays to the `template_input`.
        :param delete_files: Specifies whether or not to delete individual run output files after model execution
         and output processing.

         If `delete_files = True`, :class:`.RunModel` will remove all `run_i...` directories in the `model_dir`.
        :param model_dir: Specifies the name of the sub-directory from which the model will be executed and to which
         output files will be saved.  A new directory is created by :class:`.RunModel` within the current directory whose name
         is `model_dir` appended with a timestamp.
        """
        self.template_text = None
        self.logger = logging.getLogger(__name__)

        if platform.system() in ["Windows"]:
            self.python_command = "python"
        else:
            self.python_command = "python3"

        self.separator = separator
        self.fmt = fmt
        self.check_formatting(fmt)
        self.delete_files = delete_files

        self.input_template = input_template
        self.var_names = var_names
        self.n_variables: int = 0

        if self.var_names is not None and not ThirdPartyModel._is_list_of_strings(self.var_names):
            raise ValueError("\nUQpy: Variable names should be passed as a list of strings.\n")

        # Establish parent directory for simulations
        self.parent_dir = os.getcwd()

        # Create a list of all of the files and directories in the working directory. Do not include any other
        # directories containing the same name as model_dir
        model_files = self.create_model_files_list(model_dir)

        # Create a new directory where the model will be executed
        self.create_model_execution_directory(model_dir, model_files)

        # Check if the model script is a python script
        model_extension = pathlib.Path(model_script).suffix
        if model_extension == '.py':
            self.model_script = model_script
        else:
            raise ValueError("\nUQpy: The model script must be the name of a python script, with extension '.py'.")

        self.model_object_name = model_object_name
        self.output_script = output_script
        self.output_object_name = output_object_name

        self._check_python_model()

        # Import the output script
        if self.output_script is not None:
            # self.output_module = __import__(self.output_script[:-3])
            # Run function which checks if the output module has the output object
            self._check_output_module()

    def create_model_execution_directory(self, model_dir, model_files):
        ts = datetime.datetime.now().strftime("%Y_%m_%d_%I_%M_%f")
        self.model_dir = os.path.join(self.parent_dir, model_dir + "_" + ts)
        os.makedirs(self.model_dir)

        os.chdir(self.model_dir)

        # self.logger.info("\nUQpy: The following directory has been created for model evaluations: \n" + self.model_dir)
        # Copy files from the model list to model run directory
        for file_name in model_files:
            full_file_name = os.path.join(self.parent_dir, file_name)
            if not os.path.isdir(full_file_name):
                shutil.copy(full_file_name, self.model_dir)
            else:
                new_dir_name = os.path.join(self.model_dir, os.path.basename(full_file_name))
                shutil.copytree(full_file_name, new_dir_name)
        # self.logger.info("\nUQpy: The model files have been copied to the following directory for evaluation: \n"
        #                  + self.model_dir)

    def create_model_files_list(self, model_dir):
        model_files = []
        for f_name in os.listdir(self.parent_dir):
            path = os.path.join(self.parent_dir, f_name)
            if model_dir not in path:
                model_files.append(path)
        self.model_files = model_files
        return model_files

    def check_formatting(self, fmt):
        if self.fmt is None:
            pass
        elif isinstance(self.fmt, str):
            if (self.fmt[0] != "{") or (self.fmt[-1] != "}") or (":" not in self.fmt):
                raise ValueError("\nUQpy: fmt should be a string in brackets indicating a standard Python format.\n")
        else:
            raise TypeError("\nUQpy: fmt should be a str.\n")

    @staticmethod
    def _is_list_of_strings(list_of_strings):
        """
        Check if input list contains only strings

        ** Input: **

        :param list_of_strings: A list whose entries should be checked to see if they are strings
        :type list_of_strings: list
        """
        return (bool(list_of_strings) and isinstance(list_of_strings, list)
                and all(isinstance(element, str) for element in list_of_strings))

    def initialize(self, samples):
        os.chdir(self.model_dir)
        self.logger.info("\nUQpy: All model evaluations will be executed from the following directory: \n"
                         + self.model_dir)

        self.n_variables = len(samples[0])

        if self.input_template is not None:
            if self.var_names is None:
                # If var_names is not passed and there is an input template, create default variable names
                self.var_names = []
                for i in range(self.n_variables):
                    self.var_names.append('x%d' % i)

            elif len(self.var_names) != self.n_variables:
                raise ValueError("\nUQpy: var_names must have the same length as the number of variables (i.e. "
                                 "len(var_names) = len(samples[0]).\n")
        assert os.path.isfile(self.input_template) and os.access(self.input_template, os.R_OK), \
            "\nUQpy: File {} doesn't exist or isn't readable".format(self.input_template)
        # Read in the text from the template files
        with open(self.input_template, "r") as f:
            self.template_text = str(f.read())

    def finalize(self):
        parent_dir = os.path.dirname(self.model_dir)
        os.chdir(parent_dir)

    def preprocess_single_sample(self, i, sample):
        work_dir = os.path.join(self.model_dir, "run_" + str(i))
        self._copy_files(work_dir=work_dir)

        # Change current working directory to model run directory
        os.chdir(work_dir)
        self.logger.info("\nUQpy: Running model number " + str(i) + " in the following directory: \n" + work_dir)

        # Call the input function
        self._input_serial(i, sample)

    def execute_single_sample(self, index, sample_to_send):
        # os.system(f"{self.python_command} {self.model_script} {index}")
        python_model = __import__(self.model_script[:-3])
        model_object = getattr(python_model, self.model_object_name)
        model_object(index)

    def postprocess_single_file(self, index, model_output):
        if self.output_script is not None:
            output = self._output_serial(index)

        work_dir = os.path.join(self.model_dir, "run_" + str(index))
        # Remove the copied files and folders
        self._remove_copied_files(work_dir)

        # Return to the model directory
        os.chdir(self.model_dir)
        self.logger.info("\nUQpy: Model evaluation " + str(index) + " complete.\n")
        self.logger.info("\nUQpy: Returning to the model directory:\n" + self.model_dir)
        return output

    def _input_serial(self, index, sample):
        """
        Create one input file using the template and attach the index to the filename

        ** Input: **

        :param index: The simulation number
        :type index: int
        """
        self.new_text = self._find_and_replace_var_names_with_values(sample=sample)
        # Write the new text to the input file
        self._create_input_files(file_name=self.input_template, num=index, text=self.new_text,
                                 new_folder="InputFiles", )

    def _create_input_files(self, file_name, num, text, new_folder="InputFiles"):
        """
        Create input files using filename, index, text

        ** Input: **

        :param file_name: Name of input file
        :type file_name: str

        :param num: The simulation number
        :type num: int

        :param text: Contents of the input file
        :type text: str

        :param new_folder: Name of directory where the created input files are placed

                           Default: 'InputFiles'
        :type new_folder: str
        """
        if not os.path.exists(new_folder):
            os.makedirs(new_folder)
        base_name = os.path.splitext(os.path.basename(file_name))
        new_name = os.path.join(
            new_folder, base_name[0] + "_" + str(num) + base_name[1]
        )
        with open(new_name, "w") as f:
            f.write(text)
        return

    def _find_and_replace_var_names_with_values(self, sample):
        """
        Replace placeholders containing variable names in template input text with sample values.

        ** Input: **

        :param index: The sample number
        :type index: int
        """

        template_text = self.template_text
        var_names = self.var_names

        new_text = template_text
        for j in range(self.n_variables):
            string_regex = re.compile(r"<" + var_names[j] + r".*?>")
            count = 0
            for string in string_regex.findall(template_text):
                temp_check = string[1:-1].split("[")[0]
                pattern_check = re.compile(var_names[j])
                if pattern_check.fullmatch(temp_check):
                    temp = string[1:-1].replace(var_names[j], "sample[" + str(j) + "]")
                    try:
                        temp = eval(temp)
                    except IndexError as err:
                        print("\nUQpy: Index Error: {0}\n".format(err))
                        raise IndexError("{0}".format(err))

                    if isinstance(temp, collections.Iterable):
                        # If it is iterable, flatten and write as text file with designated separator
                        temp = np.array(temp).flatten()
                        to_add = ""
                        for i in range(len(temp) - 1):
                            if self.fmt is None:
                                to_add += str(temp[i]) + self.separator
                            else:
                                to_add += self.fmt.format(temp[i]) + self.separator
                        if self.fmt is None:
                            to_add += str(temp[-1])
                        else:
                            to_add += self.fmt.format(temp[-1])
                    else:
                        if self.fmt is None:
                            to_add = str(temp)
                        else:
                            to_add = self.fmt.format(temp)
                    new_text = (new_text[0: new_text.index(string)] + to_add
                                + new_text[(new_text.index(string) + len(string)):])
                    count += 1
        return new_text

    def _output_serial(self, index):
        """
        Execute the output script, obtain the output qoi and save it in qoi_list

        ** Input: **

        :param index: The simulation number
        :type index: int
        """
        # Run output module
        output_module = __import__(self.output_script[:-3])
        output_object = getattr(output_module, self.output_object_name)
        model_output = output_object(index)
        return model_output.qoi if self.model_is_class else model_output

    def _copy_files(self, work_dir):
        os.makedirs(work_dir)

        # Copy files from the model list to model run directory
        for file_name in self.model_files:
            full_file_name = os.path.join(self.model_dir, file_name)
            if not os.path.isdir(full_file_name):
                shutil.copy(full_file_name, work_dir)
            else:
                new_dir_name = os.path.join(work_dir, os.path.basename(full_file_name))
                shutil.copytree(full_file_name, new_dir_name)

    def _remove_copied_files(self, work_dir):
        """
        Remove the copied files from each run directory to avoid having many redundant files.

        ** Input: **

        :param work_dir: The working directory of the current run.
        :type work_dir: str
        """

        for file_name in self.model_files:
            full_file_name = os.path.join(work_dir, os.path.basename(file_name))
            if not os.path.isdir(full_file_name):
                os.remove(full_file_name)
            else:
                shutil.rmtree(full_file_name)

    def _check_output_module(self):
        """
        Check if output script name is valid

        This function get the names of the classes and functions in the imported module. There should be at least one
        class or function in the module - if not there, exit with ValueError. If there is at least one class or
        function in the module, if the output object name is not given as input and there is only one class or function,
        take that class name or function name to extract output. If there is a output_object_name given, check if it is
        a valid name. Else, a ValueError is raised.
        """
        # Get the names of the classes and functions in the imported module
        import inspect
        output_module = __import__(self.output_script[:-3])

        class_list = []
        function_list = []
        for name, obj in inspect.getmembers(output_module):
            if inspect.isclass(obj):
                class_list.append(name)
            elif inspect.isfunction(obj):
                function_list.append(name)

        # There should be at least one class or function in the module - if not there, exit with error.
        if len(class_list) == 0 and len(function_list) == 0:
            raise ValueError("\nUQpy: The output object should be defined as a function or class in the script.\n")

        else:  # If there is at least one class or function in the module
            # If the model object name is not given as input and there is only one class or function,
            # take that class name or function name to run the model.
            if self.output_object_name is None and len(class_list) + len(function_list) == 1:
                if len(class_list) == 1:
                    self.output_object_name = class_list[0]
                elif len(function_list) == 1:
                    self.output_object_name = function_list[0]

            # If there is a model_object_name given, check if it is in the list.
            if self.output_object_name in class_list:
                # self.logger.info("\nUQpy: The output class that will be run: " + self.output_object_name)
                self.output_is_class = True
            elif self.output_object_name in function_list:
                # self.logger.info("\nUQpy: The output function that will be run: " + self.output_object_name)
                self.output_is_class = False
            else:
                if self.output_object_name is None:
                    raise ValueError("\nUQpy: There are more than one objects in the module. Specify the name of the "
                                     "function or class which has to be executed.\n")
                else:
                    print("\nUQpy: You specified the output object name as: " + str(self.output_object_name))
                    raise ValueError("\nUQpy: The file does not contain an object which was specified as the output "
                                     "processor.\n")

    def _check_python_model(self):
        """
        Check if python model name is valid

        This function gets the name of the classes and functions in the imported python module whose names is passed in
        as the python model to RunModel. There should be at least one class or function in the module - if not there,
        then the function exits raising a ValueError. If there is at least one class or function in the module,
        if the model object name is not given as input and there is only one class or function, that class name or
        function name is used to run the model. If there is a model_object_name given, check if it is a valid name.
        Else, a ValueError is raised.
        """
        # Get the names of the classes and functions in the imported module
        import inspect

        class_list = []
        function_list = []
        python_model = __import__(self.model_script[:-3])
        for name, obj in inspect.getmembers(python_model):
            if inspect.isclass(obj):
                class_list.append(name)
            elif inspect.isfunction(obj):
                function_list.append(name)

        # There should be at least one class or function in the module - if not there, exit with error.
        if len(class_list) == 0 and len(function_list) == 0:
            raise ValueError("\nUQpy: A python model should be defined as a function or class in the script.\n")

        else:  # If there is at least one class or function in the module
            # If the model object name is not given as input and there is only one class or function,
            # take that class name or function name to run the model.
            if self.model_object_name is None and len(class_list) + len(function_list) == 1:
                if len(class_list) == 1:
                    self.model_object_name = class_list[0]
                elif len(function_list) == 1:
                    self.model_object_name = function_list[0]

            # If there is a model_object_name given, check if it is in the list.
            if self.model_object_name in class_list:
                # self.logger.info("\nUQpy: The model class that will be run: " + self.model_object_name)
                self.model_is_class = True
            elif self.model_object_name in function_list:
                # self.logger.info("\nUQpy: The model function that will be run: " + self.model_object_name)
                self.model_is_class = False
            else:
                if self.model_object_name is None:
                    raise ValueError("\nUQpy: There are more than one objects in the module. Specify the name of the "
                                     "function or class which has to be executed.\n")
                else:
                    print("\nUQpy: You specified the model_object_name as: " + str(self.model_object_name))
                    raise ValueError("\nUQpy: The file does not contain an object which was specified as the model.\n")


# src/UQpy/run_model/model_execution/__init__.py

# from UQpy.utilities.model_execution.ParallelExecution import *
from UQpy.run_model.model_execution.SerialExecution import *
from UQpy.run_model.model_execution.PythonModel import *
from UQpy.run_model.model_execution.ThirdPartyModel import *


# src/UQpy/sampling/adaptive_kriging_functions/ExpectedFeasibility.py

from typing import Union

from beartype import beartype

from UQpy.sampling.adaptive_kriging_functions.baseclass.LearningFunction import (
    LearningFunction,
)
import scipy.stats as stats


class ExpectedFeasibility(LearningFunction):

    @beartype
    def __init__(
        self,
        eff_a: Union[float, int] = 0,
        eff_epsilon: Union[float, int] = 2,
        eff_stop: Union[float, int] = 0.001,
    ):
        """
        Expected Feasibility Function (EFF) for reliability analysis, see :cite:`AKMCS4` for a detailed explanation.

        :param eff_a: Reliability threshold.
        :param eff_epsilon: EGRA method epsilon
        :param eff_stop: Stopping threshold
        """
        self.eff_a = eff_a
        self.eff_epsilon = eff_epsilon
        self.eff_stop = eff_stop

    def evaluate_function(self, distributions, n_add, surrogate, population, qoi=None, samples=None):

        g, sig = surrogate.predict(population, True)

        # Remove the inconsistency in the shape of 'g' and 'sig' array
        g = g.reshape([population.shape[0], 1])
        sig = sig.reshape([population.shape[0], 1])
        # reliability threshold: a_ = 0
        # EGRA method: epsilon = 2*sigma(x)
        a_, ep = self.eff_a, self.eff_epsilon * sig
        t1 = (a_ - g) / sig
        t2 = (a_ - ep - g) / sig
        t3 = (a_ + ep - g) / sig
        eff = (g - a_) * (
            2 * stats.norm.cdf(t1) - stats.norm.cdf(t2) - stats.norm.cdf(t3)
        )
        eff += -sig * (2 * stats.norm.pdf(t1) - stats.norm.pdf(t2) - stats.norm.pdf(t3))
        eff += ep * (stats.norm.cdf(t3) - stats.norm.cdf(t2))
        rows = eff[:, 0].argsort()[-n_add:]

        stopping_criteria_indicator = max(eff[:, 0]) <= self.eff_stop
        new_samples = population[rows, :]
        learning_function_values = eff[rows, :]
        return new_samples, learning_function_values, stopping_criteria_indicator


# src/UQpy/sampling/adaptive_kriging_functions/ExpectedImprovement.py

from typing import Union
from beartype import beartype
from UQpy.sampling.adaptive_kriging_functions.baseclass.LearningFunction import (
    LearningFunction,
)
import scipy.stats as stats
import numpy as np


class ExpectedImprovement(LearningFunction):

    @beartype
    def __init__(self, eif_stop: Union[float, int] = 0.01):
        """
        Expected Improvement Function (EIF) for Efficient Global Optimization (EFO). See :cite:`AKMCS2` for a detailed
        explanation.

        :param eif_stop: Stopping threshold
        """
        self.eif_stop = eif_stop

    def evaluate_function(self, distributions, n_add, surrogate, population, qoi=None, samples=None):
        g, sig = surrogate.predict(population, True)

        # Remove the inconsistency in the shape of 'g' and 'sig' array
        g = g.reshape([population.shape[0], 1])
        sig = sig.reshape([population.shape[0], 1])

        fm = min(qoi)
        eif = (fm - g) * stats.norm.cdf((fm - g) / sig) + sig * stats.norm.pdf((fm - g) / sig)
        rows = eif[:, 0].argsort()[(np.size(g) - n_add) :]

        stopping_criteria_indicator = max(eif[:, 0]) / abs(fm) <= self.eif_stop
        new_samples = population[rows, :]
        learning_function_values = eif[rows, :]
        return new_samples, learning_function_values, stopping_criteria_indicator


# src/UQpy/sampling/adaptive_kriging_functions/ExpectedImprovementGlobalFit.py

from UQpy.sampling.adaptive_kriging_functions.baseclass.LearningFunction import (
    LearningFunction,
)
import numpy as np
from sklearn.neighbors import NearestNeighbors


class ExpectedImprovementGlobalFit(LearningFunction):
    """
    Expected Improvement for Global Fit (EIGF) learning function. See :cite:`AKMCS5` for a detailed explanation.
    """

    def evaluate_function(
        self, distributions, n_add, surrogate, population, qoi=None, samples=None
    ):
        g, sig = surrogate.predict(population, True)

        # Remove the inconsistency in the shape of 'g' and 'sig' array
        g = g.reshape([population.shape[0], 1])
        sig = sig.reshape([population.shape[0], 1])

        # Evaluation of the learning function
        # First, find the nearest neighbor in the training set for each point in the population.

        knn = NearestNeighbors(n_neighbors=1)
        knn.fit(np.atleast_2d(samples))
        neighbors = knn.kneighbors(np.atleast_2d(population), return_distance=False)

        # noinspection PyTypeChecker
        qoi_array = np.array([qoi[x] for x in np.squeeze(neighbors)])

        # Compute the learning function at every point in the population.
        u = np.square(g - qoi_array) + np.square(sig)
        rows = u[:, 0].argsort()[(np.size(g) - n_add) :]

        stopping_criteria_indicator = False
        new_samples = population[rows, :]
        learning_function_evaluations = u[rows, :]

        return new_samples, learning_function_evaluations, stopping_criteria_indicator


# src/UQpy/sampling/adaptive_kriging_functions/UFunction.py

from beartype import beartype

from UQpy.sampling.adaptive_kriging_functions.baseclass.LearningFunction import (
    LearningFunction,
)


class UFunction(LearningFunction):

    @beartype
    def __init__(self, u_stop: int = 2):
        """
        U-function for reliability analysis. See :cite:`AKMCS1` for a detailed explanation.

        :param u_stop: U-Function stopping parameter
        """
        self.u_stop = u_stop

    def evaluate_function(self, distributions, n_add, surrogate, population, qoi=None, samples=None):

        g, sig = surrogate.predict(population, True)

        # Remove the inconsistency in the shape of 'g' and 'sig' array
        g = g.reshape([population.shape[0], 1])
        sig = sig.reshape([population.shape[0], 1])

        u = abs(g) / sig
        rows = u[:, 0].argsort()[:n_add]

        stopping_criteria_indicator = min(u[:, 0]) >= self.u_stop
        new_samples = population[rows, :]
        learning_function_values = u[rows, 0]
        return new_samples, learning_function_values, stopping_criteria_indicator


# src/UQpy/sampling/adaptive_kriging_functions/WeightedUFunction.py

from beartype import beartype

from UQpy.sampling.adaptive_kriging_functions.baseclass.LearningFunction import (
    LearningFunction,
)
import numpy as np


class WeightedUFunction(LearningFunction):

    @beartype
    def __init__(self, weighted_u_stop: int):
        """
        Probability Weighted U-function for reliability analysis. See :cite:`AKMCS3` for a detailed explanation.

        :param weighted_u_stop: Stopping parameter required for the WeightedU learning function
        """
        self.weighted_u_stop = weighted_u_stop

    def evaluate_function(self, distributions, n_add, surrogate, population, qoi=None, samples=None):
        g, sig = surrogate.predict(population, True)

        # Remove the inconsistency in the shape of 'g' and 'sig' array
        g = g.reshape([population.shape[0], 1])
        sig = sig.reshape([population.shape[0], 1])

        u = abs(g) / sig
        p1 = np.ones([population.shape[0], population.shape[1]])
        p2 = np.ones([samples.shape[0], population.shape[1]])

        for j in range(samples.shape[1]):
            p1[:, j] = distributions[j].pdf(np.atleast_2d(population[:, j]).T)
            p2[:, j] = distributions[j].pdf(np.atleast_2d(samples[:, j]).T)

        p1 = p1.prod(1).reshape(u.size, 1)
        max_p = max(p2.prod(1))
        u_ = u * ((max_p - p1) / max_p)
        rows = u_[:, 0].argsort()[:n_add]

        stopping_criteria_indicator = min(u[:, 0]) >= self.weighted_u_stop
        new_samples = population[rows, :]
        learning_function_values = u_[rows, :]
        return new_samples, learning_function_values, stopping_criteria_indicator


# src/UQpy/sampling/adaptive_kriging_functions/__init__.py

from UQpy.sampling.adaptive_kriging_functions.ExpectedFeasibility import *
from UQpy.sampling.adaptive_kriging_functions.ExpectedImprovement import *
from UQpy.sampling.adaptive_kriging_functions.ExpectedImprovementGlobalFit import *
from UQpy.sampling.adaptive_kriging_functions.baseclass import *
from UQpy.sampling.adaptive_kriging_functions.UFunction import *
from UQpy.sampling.adaptive_kriging_functions.WeightedUFunction import *


# src/UQpy/sampling/mcmc/DRAM.py

import logging
from typing import Callable
import warnings

import numpy as np

warnings.filterwarnings('ignore')

from beartype import beartype
from UQpy.sampling.mcmc.baseclass.MCMC import MCMC
from UQpy.distributions import *
from UQpy.utilities.ValidationTypes import *


class DRAM(MCMC):

    @beartype
    def __init__(
            self,
            pdf_target: Union[Callable, list[Callable]] = None,
            log_pdf_target: Union[Callable, list[Callable]] = None,
            args_target: tuple = None,
            burn_length: Annotated[int, Is[lambda x: x >= 0]] = 0,
            jump: int = 1,
            dimension: int = None,
            seed: list = None,
            save_log_pdf: bool = False,
            concatenate_chains: bool = True,
            initial_covariance: float = None,
            covariance_update_rate: float = 100,
            scale_parameter: float = None,
            delayed_rejection_scale: float = 1 / 5,
            save_covariance: bool = False,
            random_state: RandomStateType = None,
            n_chains: int = None,
            nsamples: int = None,
            nsamples_per_chain: int = None,
    ):
        """
        Delayed Rejection Adaptive Metropolis algorithm :cite:`Dram1` :cite:`MCMC2`

        In this algorithm, the proposal density is Gaussian and its covariance C is being updated from samples as
        :code:`C = scale_parameter * C_sample` where :code:`C_sample` is the sample covariance. Also, the delayed
        rejection scheme is applied, i.e, if a candidate is not accepted another one is generated from the proposal with
        covariance :code:`delayed_rejection_scale ** 2 * C`.

        :param pdf_target: Target density function from which to draw random samples. Either `pdf_target` or
         `log_pdf_target` must be provided (the latter should be preferred for better numerical stability).

         If `pdf_target` is a callable, it refers to the joint pdf to sample from, it must take at least one input
         **x**, which are the point(s) at which to evaluate the pdf. Within :class:`.MCMC` the `pdf_target` is evaluated
         as:
         :code:`p(x) = pdf_target(x, \*args_target)`

         where **x** is a :class:`numpy.ndarray of shape :code:`(nsamples, dimension)` and `args_target` are additional
         positional arguments that are provided to :class:`.MCMC` via its `args_target` input.

         If `pdf_target` is a list of callables, it refers to independent marginals to sample from. The marginal in
         dimension :code:`j` is evaluated as: :code:`p_j(xj) = pdf_target[j](xj, \*args_target[j])` where **x** is a
         :class:`numpy.ndarray` of shape :code:`(nsamples, dimension)`
        :param log_pdf_target: Logarithm of the target density function from which to draw random samples.
         Either pdf_target or log_pdf_target must be provided (the latter should be preferred for better numerical
         stability).

         Same comments as for input `pdf_target`.
        :param args_target: Positional arguments of the pdf / log-pdf target function. See `pdf_target`
        :param burn_length: Length of burn-in - i.e., number of samples at the beginning of the chain to discard (note:
         no thinning during burn-in). Default is :math:`0`, no burn-in.
        :param jump: Thinning parameter, used to reduce correlation between samples. Setting :code:`jump=n` corresponds
         to skipping n-1 states between accepted states of the chain. Default is :math:`1` (no thinning).
        :param dimension: A scalar value defining the dimension of target density function. Either `dimension` and
         `n_chains` or `seed` must be provided.
        :param seed: Seed of the Markov chain(s), shape :code:`(n_chains, dimension)`.
         Default: :code:`zeros(n_chains x dimension)`.

         If `seed` is not provided, both `n_chains` and `dimension` must be provided.
        :param save_log_pdf: Boolean that indicates whether to save log-pdf values along with the samples.
         Default: :any:`False`
        :param concatenate_chains: Boolean that indicates whether to concatenate the chains after a run, i.e., samples
         are stored as a :class:`numpy.ndarray` of shape :code:`(nsamples * n_chains, dimension)` if :any:`True`,
         :code:`(nsamples, n_chains, dimension)` if :any:`False`.
         Default: :any:`True`
        :param n_chains: The number of Markov chains to generate. Either `dimension` and `n_chains` or `seed` must be
         provided.
        :param initial_covariance: Initial covariance for the gaussian proposal distribution. Default: I(dim)
        :param covariance_update_rate: Rate at which covariance is being updated, i.e., every k0 iterations.
         Default: :math:`100`
        :param scale_parameter: Scale parameter for covariance updating. Default: :math:`2.38^2/dim`
        :param delayed_rejection_scale: Scale parameter for delayed rejection. Default: :math:`1/5`
        :param save_covariance: If :any:`True`, updated covariance is saved in attribute :py:attr:`adaptive_covariance`.
         Default: :any:`False`
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is
         :any:`None`.
        :param nsamples: Number of samples to generate.
        :param nsamples_per_chain: Number of samples to generate per chain.
        """
        self.nsamples = nsamples
        self.nsamples_per_chain = nsamples_per_chain
        super().__init__(
            pdf_target=pdf_target,
            log_pdf_target=log_pdf_target,
            args_target=args_target,
            dimension=dimension,
            seed=seed,
            burn_length=burn_length,
            jump=jump,
            save_log_pdf=save_log_pdf,
            concatenate_chains=concatenate_chains,
            random_state=random_state,
            n_chains=n_chains,
        )

        self.logger = logging.getLogger(__name__)
        # Check the initial covariance
        self.initial_covariance = initial_covariance
        if self.initial_covariance is None:
            self.initial_covariance = np.eye(self.dimension)
        elif not (isinstance(self.initial_covariance, np.ndarray)
                  and self.initial_covariance == (self.dimension, self.dimension)):
            raise TypeError(
                "UQpy: Input initial_covariance should be a 2D ndarray of shape (dimension, dimension)")

        self.covariance_update_rate = covariance_update_rate
        self.scale_parameter = scale_parameter
        if self.scale_parameter is None:
            self.scale_parameter = 2.38 ** 2 / self.dimension
        self.delayed_rejection_scale = delayed_rejection_scale
        self.save_covariance = save_covariance
        for key, typ in zip(
                [
                    "covariance_update_rate",
                    "scale_parameter",
                    "delayed_rejection_scale",
                    "save_covariance",
                ],
                [int, float, float, bool],
        ):
            if not isinstance(getattr(self, key), typ):
                raise TypeError("Input " + key + " must be of type " + typ.__name__)

        # initialize the sample mean and sample covariance that you need
        self.current_covariance = np.tile(
            self.initial_covariance[np.newaxis, ...], (self.n_chains, 1, 1))
        self.sample_mean = np.zeros((self.n_chains, self.dimension,))
        self.sample_covariance = np.zeros((self.n_chains, self.dimension, self.dimension))
        if self.save_covariance:
            self.adaptive_covariance = [self.current_covariance.copy(), ]

        self.logger.info("\nUQpy: Initialization of " + self.__class__.__name__ + " algorithm complete.")

        if (nsamples is not None) or (nsamples_per_chain is not None):
            self.run(nsamples=nsamples, nsamples_per_chain=nsamples_per_chain)

    def run_one_iteration(self, current_state: np.ndarray, current_log_pdf: np.ndarray):
        """
        Run one iteration of the mcmc chain for DRAM algorithm, starting at current state -
        see :class:`MCMC` class.
        """
        from UQpy.distributions import MultivariateNormal

        multivariate_normal = MultivariateNormal(mean=np.zeros(self.dimension, ), cov=1.0)

        # Sample candidate
        candidate = np.zeros_like(current_state)
        for nc, current_cov in enumerate(self.current_covariance):
            multivariate_normal.update_parameters(cov=current_cov)
            candidate[nc, :] = current_state[nc, :] + \
                               multivariate_normal.rvs(nsamples=1, random_state=self.random_state) \
                                   .reshape((self.dimension,))

        # Compute log_pdf_target of candidate sample
        log_p_candidate = self.evaluate_log_target(candidate)

        # Compare candidate with current sample and decide or not to keep the candidate (loop over nc chains)
        accept_vec = np.zeros((self.n_chains,))
        delayed_chains_indices = ([])  # indices of chains that will undergo delayed rejection
        unif_rvs = (Uniform().rvs(nsamples=self.n_chains, random_state=self.random_state)
                    .reshape((-1,)))
        for nc, (cand, log_p_cand, log_p_curr) in enumerate(
                zip(candidate, log_p_candidate, current_log_pdf)):
            accept = np.log(unif_rvs[nc]) < log_p_cand - log_p_curr
            if accept:
                current_state[nc, :] = cand
                current_log_pdf[nc] = log_p_cand
                accept_vec[nc] += 1.0
            else:  # enter delayed rejection
                delayed_chains_indices.append(nc)  # these indices will enter the delayed rejection part

        # Delayed rejection
        if delayed_chains_indices:  # performed delayed rejection for some chains
            current_states_delayed = np.zeros(
                (len(delayed_chains_indices), self.dimension))
            candidates_delayed = np.zeros((len(delayed_chains_indices), self.dimension))
            candidate2 = np.zeros((len(delayed_chains_indices), self.dimension))
            # Sample other candidates closer to the current one
            for i, nc in enumerate(delayed_chains_indices):
                current_states_delayed[i, :] = current_state[nc, :]
                candidates_delayed[i, :] = candidate[nc, :]
                multivariate_normal.update_parameters(
                    cov=self.delayed_rejection_scale ** 2 * self.current_covariance[nc])
                candidate2[i, :] = current_states_delayed[i, :] + \
                                   multivariate_normal.rvs(nsamples=1, random_state=self.random_state) \
                                       .reshape((self.dimension,))
            # Evaluate their log_target
            log_p_candidate2 = self.evaluate_log_target(candidate2)
            log_prop_cand_cand2 = multivariate_normal.log_pdf(candidates_delayed - candidate2)
            log_prop_cand_curr = multivariate_normal.log_pdf(candidates_delayed - current_states_delayed)
            # Accept or reject
            unif_rvs = (Uniform().rvs(nsamples=len(delayed_chains_indices),
                                      random_state=self.random_state).reshape((-1,)))
            for (nc, cand2, log_p_cand2, j1, j2, u_rv) in zip(
                    delayed_chains_indices,
                    candidate2,
                    log_p_candidate2,
                    log_prop_cand_cand2,
                    log_prop_cand_curr,
                    unif_rvs,
            ):
                alpha_cand_cand2 = min(1.0, np.exp(log_p_candidate[nc] - log_p_cand2))
                alpha_cand_curr = min(1.0, np.exp(log_p_candidate[nc] - current_log_pdf[nc]))
                log_alpha2 = (log_p_cand2 - current_log_pdf[nc] + j1 - j2
                              + np.log(max(1.0 - alpha_cand_cand2, 10 ** (-320)))
                              - np.log(max(1.0 - alpha_cand_curr, 10 ** (-320))))
                accept = np.log(u_rv) < min(0.0, log_alpha2)
                if accept:
                    current_state[nc, :] = cand2
                    current_log_pdf[nc] = log_p_cand2
                    accept_vec[nc] += 1.0

        # Adaptive part: update the covariance
        for nc in range(self.n_chains):
            # update covariance
            self.sample_mean[nc], self.sample_covariance[nc], = self._recursive_update_mean_covariance(
                nsamples=self.iterations_number,
                new_sample=current_state[nc, :],
                previous_mean=self.sample_mean[nc],
                previous_covariance=self.sample_covariance[nc], )
            if (self.iterations_number > 1) and (self.iterations_number % self.covariance_update_rate == 0):
                self.current_covariance[nc] = self.scale_parameter * self.sample_covariance[nc] + \
                                              1e-6 * np.eye(self.dimension)
        if self.save_covariance and \
                ((self.iterations_number > 1) and (self.iterations_number % self.covariance_update_rate == 0)):
            self.adaptive_covariance.append(self.current_covariance.copy())

        # Update the acceptance rate
        self._update_acceptance_rate(accept_vec)
        return current_state, current_log_pdf

    @staticmethod
    def _recursive_update_mean_covariance(
            nsamples, new_sample, previous_mean, previous_covariance=None
    ):
        """
        Iterative formula to compute a new sample mean and covariance based on previous ones and new sample.

        New covariance is computed only of previous_covariance is provided.

        **Inputs:**

        * n (int): Number of samples used to compute the new mean
        * new_sample (ndarray (dim, )): new sample
        * previous_mean (ndarray (dim, )): Previous sample mean, to be updated with new sample value
        * previous_covariance (ndarray (dim, dim)): Previous sample covariance, to be updated with new sample value

        **Output/Returns:**

        * new_mean (ndarray (dim, )): Updated sample mean
        * new_covariance (ndarray (dim, dim)): Updated sample covariance

        """
        new_mean = (nsamples - 1) / nsamples * previous_mean + 1 / nsamples * new_sample
        if previous_covariance is None:
            return new_mean
        dimensions = new_sample.size
        if nsamples == 1:
            new_covariance = np.zeros((dimensions, dimensions))
        else:
            delta_n = (new_sample - previous_mean).reshape((dimensions, 1))
            new_covariance = (nsamples - 2) / (nsamples - 1) \
                             * previous_covariance + 1 / nsamples * np.matmul(delta_n, delta_n.T)
        return new_mean, new_covariance


# src/UQpy/sampling/mcmc/DREAM.py

import logging
from typing import Callable
import warnings

import numpy as np

warnings.filterwarnings('ignore')

from beartype import beartype
from UQpy.sampling.mcmc.baseclass.MCMC import MCMC
from UQpy.distributions import *
from UQpy.utilities.ValidationTypes import *


class DREAM(MCMC):

    @beartype
    def __init__(
            self,
            pdf_target: Union[Callable, list[Callable]] = None,
            log_pdf_target: Union[Callable, list[Callable]] = None,
            args_target: tuple = None,
            burn_length: Annotated[int, Is[lambda x: x >= 0]] = 0,
            jump: PositiveInteger = 1,
            dimension: int = None,
            seed: list = None,
            save_log_pdf: bool = False,
            concatenate_chains: bool = True,
            jump_rate: int = 3,
            c: float = 0.1,
            c_star: float = 1e-6,
            crossover_probabilities_number: int = 3,
            gamma_probability: float = 0.2,
            crossover_adaptation: tuple = (-1, 1),
            check_chains: tuple = (-1, 1),
            random_state: RandomStateType = None,
            n_chains: int = None,
            nsamples: int = None,
            nsamples_per_chain: int = None,
    ):
        """
        DiffeRential Evolution Adaptive Metropolis algorithm :cite:`Dream1` :cite:`Dream2`

        :param pdf_target: Target density function from which to draw random samples. Either `pdf_target` or
         `log_pdf_target` must be provided (the latter should be preferred for better numerical stability).

         If `pdf_target` is a callable, it refers to the joint pdf to sample from, it must take at least one input
         **x**, which are the point(s) at which to evaluate the pdf. Within :class:`.MCMC` the `pdf_target` is evaluated
         as:
         :code:`p(x) = pdf_target(x, \*args_target)`

         where **x** is a :class:`numpy.ndarray  of shape :code:`(nsamples, dimension)` and `args_target` are additional
         positional arguments that are provided to :class:`.MCMC` via its `args_target` input.

         If `pdf_target` is a list of callables, it refers to independent marginals to sample from. The marginal in
         dimension :code:`j` is evaluated as: :code:`p_j(xj) = pdf_target[j](xj, \*args_target[j])` where **x** is a
         :class:`numpy.ndarray` of shape :code:`(nsamples, dimension)`
        :param log_pdf_target: Logarithm of the target density function from which to draw random samples.
         Either `pdf_target` or `log_pdf_target` must be provided (the latter should be preferred for better numerical
         stability).

         Same comments as for input `pdf_target`.
        :param args_target: Positional arguments of the pdf / log-pdf target function. See `pdf_target`
        :param burn_length: Length of burn-in - i.e., number of samples at the beginning of the chain to discard (note:
         no thinning during burn-in). Default is :math:`0`, no burn-in.
        :param jump: Thinning parameter, used to reduce correlation between samples. Setting :code:`jump=n` corresponds
         to skipping :code:`n-1` states between accepted states of the chain. Default is :math:`1` (no thinning).
        :param dimension: A scalar value defining the dimension of target density function. Either `dimension` and
         `n_chains` or `seed` must be provided.
        :param seed: Seed of the Markov chain(s), shape :code:`(n_chains, dimension)`.
         Default: :code:`zeros(n_chains x dimension)`.

         If `seed` is not provided, both `n_chains` and `dimension` must be provided.
        :param save_log_pdf: Boolean that indicates whether to save log-pdf values along with the samples.
         Default: :any:`False`
        :param concatenate_chains: Boolean that indicates whether to concatenate the chains after a run, i.e., samples
         are stored as a :class:`numpy.ndarray` of shape :code:`(nsamples * n_chains, dimension)` if :any:`True`,
         :code:`(nsamples, n_chains, dimension)` if :any:`False`.
         Default: :any:`True`
        :param n_chains: The number of Markov chains to generate. Either `dimension` and `n_chains` or `seed` must be
         provided.
        :param jump_rate: Jump rate. Default: :math:`3`
        :param c: Differential evolution parameter. Default: :math:`0.1`
        :param c_star: Differential evolution parameter, should be small compared to width of target.
         Default: :math:`1e-6`
        :param crossover_probabilities_number: Number of crossover probabilities. Default: :math:`3`
        :param gamma_probability: :code:`Prob(gamma=1)`. Default: :math:`0.2`
        :param crossover_adaptation: :code:`(iter_max, rate)` governs adaptation of crossover probabilities (adapts
         every rate iterations if :code:`iter<iter_max`). Default: :code:`(-1, 1)`, i.e., no adaptation
        :param check_chains: :code:`(iter_max, rate)` governs discarding of outlier chains (discard every rate
         iterations if :code:`iter<iter_max`). Default: :code:`(-1, 1)`, i.e., no check on outlier chains
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is
         :any:`None`.

        :param nsamples: Number of samples to generate.
        :param nsamples_per_chain: Number of samples to generate per chain.
        """
        self.nsamples = nsamples
        self.nsamples_per_chain = nsamples_per_chain
        super().__init__(
            pdf_target=pdf_target,
            log_pdf_target=log_pdf_target,
            args_target=args_target,
            dimension=dimension,
            seed=seed,
            burn_length=burn_length,
            jump=jump,
            save_log_pdf=save_log_pdf,
            concatenate_chains=concatenate_chains,
            random_state=random_state,
            n_chains=n_chains,
        )

        self.logger = logging.getLogger(__name__)
        # Check nb of chains
        if self.n_chains < 2:
            raise ValueError("UQpy: For the DREAM algorithm, a seed must be provided with at least two samples.")

        # Check user-specific algorithms
        self.jump_rate = jump_rate
        self.c = c
        self.c_star = c_star
        self.crossover_probabilities_number = crossover_probabilities_number
        self.gamma_probability = gamma_probability
        self.crossover_adaptation = crossover_adaptation
        self.check_chains = check_chains

        for key, typ in zip(
                [
                    "jump_rate",
                    "c",
                    "c_star",
                    "crossover_probabilities_number",
                    "gamma_probability",
                ],
                [int, float, float, int, float],
        ):
            if not isinstance(getattr(self, key), typ):
                raise TypeError("Input " + key + " must be of type " + typ.__name__)
        if (self.dimension is not None and self.crossover_probabilities_number > self.dimension):
            self.crossover_probabilities_number = self.dimension
        for key in ["crossover_adaptation", "check_chains"]:
            p = getattr(self, key)
            if not (isinstance(p, tuple) and len(p) == 2 and all(isinstance(i, (int, float)) for i in p)):
                raise TypeError("Inputs " + key + " must be a tuple of 2 integers.")
        if (not self.save_log_pdf) and (self.check_chains[0] > 0):
            raise ValueError("UQpy: Input save_log_pdf must be True in order to check outlier chains")

        # Initialize a few other variables
        self.j_ind = np.zeros((self.crossover_probabilities_number,))
        self.n_id = np.zeros((self.crossover_probabilities_number,))
        self.cross_prob = (
                np.ones((self.crossover_probabilities_number,))
                / self.crossover_probabilities_number)

        self.logger.info("UQpy: Initialization of " + self.__class__.__name__ + " algorithm complete.\n")

        if (nsamples is not None) or (nsamples_per_chain is not None):
            self.run(nsamples=nsamples, nsamples_per_chain=nsamples_per_chain, )

    def run_one_iteration(self, current_state: np.ndarray, current_log_pdf: np.ndarray):
        """
        Run one iteration of the mcmc chain for DREAM algorithm, starting at current state -
        see :class:`MCMC` class.
        """
        r_diff = np.array([np.setdiff1d(np.arange(self.n_chains), j) for j in range(self.n_chains)])
        cross = (np.arange(1, self.crossover_probabilities_number + 1) / self.crossover_probabilities_number)

        # Dynamic part: evolution of chains
        unif_rvs = (Uniform().rvs(nsamples=self.n_chains * (self.n_chains - 1),
                                  random_state=self.random_state, )
                    .reshape((self.n_chains - 1, self.n_chains)))
        draw = np.argsort(unif_rvs, axis=0)
        dx = np.zeros_like(current_state)
        lmda = (Uniform(scale=2 * self.c).rvs(nsamples=self.n_chains, random_state=self.random_state)
                .reshape((-1,)))
        std_x_tmp = np.std(current_state, axis=0)

        multi_rvs = Multinomial(n=1, p=[1.0 / self.jump_rate, ] * self.jump_rate).rvs(
            nsamples=self.n_chains, random_state=self.random_state)
        d_ind = np.nonzero(multi_rvs)[1]
        as_ = [r_diff[j, draw[slice(d_ind[j]), j]] for j in range(self.n_chains)]
        bs_ = [r_diff[j, draw[slice(d_ind[j], 2 * d_ind[j], 1), j]] for j in range(self.n_chains)]
        multi_rvs = Multinomial(n=1, p=self.cross_prob).rvs(
            nsamples=self.n_chains, random_state=self.random_state)
        id_ = np.nonzero(multi_rvs)[1]
        # id = np.random.choice(self.n_CR, size=(self.nchains, ), replace=True, trial_probability=self.pCR)
        z = (Uniform().rvs(nsamples=self.n_chains * self.dimension,
                           random_state=self.random_state, )
             .reshape((self.n_chains, self.dimension)))
        subset_a = [np.where(z_j < cross[id_j])[0] for (z_j, id_j) in zip(z, id_)]  # subset A of selected dimensions
        d_star = np.array([len(a_j) for a_j in subset_a])
        for j in range(self.n_chains):
            if d_star[j] == 0:
                subset_a[j] = np.array([np.argmin(z[j])])
                d_star[j] = 1
        gamma_d = 2.38 / np.sqrt(2 * (d_ind + 1) * d_star)
        g = (Binomial(n=1, p=self.gamma_probability).rvs(nsamples=self.n_chains, random_state=self.random_state)
             .reshape((-1,)))
        g[g == 0] = gamma_d[g == 0]
        norm_vars = (Normal(loc=0.0, scale=1.0).rvs(nsamples=self.n_chains ** 2, random_state=self.random_state)
                     .reshape((self.n_chains, self.n_chains)))
        for j in range(self.n_chains):
            for i in subset_a[j]:
                dx[j, i] = self.c_star * norm_vars[j, i] + (1 + lmda[j]) * g[j] * np.sum(current_state[as_[j], i] -
                                                                                         current_state[bs_[j], i])
        candidates = current_state + dx

        # Evaluate log likelihood of candidates
        logp_candidates = self.evaluate_log_target(candidates)

        # Accept or reject
        accept_vec = np.zeros((self.n_chains,))
        unif_rvs = (Uniform().rvs(nsamples=self.n_chains, random_state=self.random_state)
                    .reshape((-1,)))
        for nc, (lpc, candidate, log_p_curr) in enumerate(
                zip(logp_candidates, candidates, current_log_pdf)):
            accept = np.log(unif_rvs[nc]) < lpc - log_p_curr
            if accept:
                current_state[nc, :] = candidate
                current_log_pdf[nc] = lpc
                accept_vec[nc] = 1.0
            else:
                dx[nc, :] = 0
            self.j_ind[id_[nc]] = self.j_ind[id_[nc]] + np.sum((dx[nc, :] / std_x_tmp) ** 2)
            self.n_id[id_[nc]] += 1

        # Save the acceptance rate
        self._update_acceptance_rate(accept_vec)

        # update selection cross prob
        if (self.iterations_number < self.crossover_adaptation[0]
                and self.iterations_number % self.crossover_adaptation[1] == 0):
            self.cross_prob = self.j_ind / self.n_id
            self.cross_prob /= sum(self.cross_prob)
        # check outlier chains (only if you have saved at least 100 values already)
        if ((self.samples_counter >= 100)
                and (self.iterations_number < self.check_chains[0])
                and (self.iterations_number % self.check_chains[1] == 0)):
            self.check_outlier_chains(replace_with_best=True)

        return current_state, current_log_pdf

    def check_outlier_chains(self, replace_with_best: bool = False):
        if not self.save_log_pdf:
            raise ValueError("UQpy: Input save_log_pdf must be True in order to check outlier chains")
        start_ = self.nsamples_per_chain // 2
        avgs_logpdf = np.mean(self.log_pdf_values[start_: self.nsamples_per_chain], axis=0)
        best_ = np.argmax(avgs_logpdf)
        avg_sorted = np.sort(avgs_logpdf)
        ind1, ind3 = (1 + round(0.25 * self.n_chains), 1 + round(0.75 * self.n_chains),)
        q1, q3 = avg_sorted[ind1], avg_sorted[ind3]
        qr = q3 - q1

        outlier_num = 0
        for j in range(self.n_chains):
            if avgs_logpdf[j] < q1 - 2.0 * qr:
                outlier_num += 1
                if replace_with_best:
                    self.samples[start_:, j, :] = self.samples[start_:, best_, :].copy()
                    self.log_pdf_values[start_:, j] = self.log_pdf_values[
                                                      start_:, best_
                                                      ].copy()
                else:
                    self.logger.info("UQpy: Chain {} is an outlier chain".format(j))
        if outlier_num > 0:
            self.logger.info("UQpy: Detected {} outlier chains".format(outlier_num))


# src/UQpy/sampling/mcmc/MetropolisHastings.py

import logging
from typing import Callable
from beartype import beartype
from UQpy.sampling.mcmc.baseclass.MCMC import MCMC
from UQpy.distributions import *
from UQpy.utilities.ValidationTypes import *
import warnings

warnings.filterwarnings('ignore')


class MetropolisHastings(MCMC):

    @beartype
    def __init__(
        self,
        pdf_target: Union[Callable, list[Callable]] = None,
        log_pdf_target: Union[Callable, list[Callable]] = None,
        args_target: tuple = None,
        burn_length: Annotated[int, Is[lambda x: x >= 0]] = 0,
        jump: int = 1,
        dimension: int = None,
        seed: list = None,
        save_log_pdf: bool = False,
        concatenate_chains: bool = True,
        n_chains: int = None,
        proposal: Distribution = None,
        proposal_is_symmetric: bool = False,
        random_state: RandomStateType = None,
        nsamples: PositiveInteger = None,
        nsamples_per_chain: PositiveInteger = None,
    ):
        """
        Metropolis-Hastings algorithm :cite:`MCMC1` :cite:`MCMC2`

        :param pdf_target: Target density function from which to draw random samples. Either `pdf_target` or
         `log_pdf_target` must be provided (the latter should be preferred for better numerical stability).

         If `pdf_target` is a callable, it refers to the joint pdf to sample from, it must take at least one input
         **x**, which are the point(s) at which to evaluate the pdf. Within :class:`.MCMC` the pdf_target is evaluated
         as:
         :code:`p(x) = pdf_target(x, \*args_target)`

         where **x** is a :class:`numpy.ndarray  of shape :code:`(nsamples, dimension)` and `args_target` are additional
         positional arguments that are provided to :class:`.MCMC` via its `args_target` input.

         If `pdf_target` is a list of callables, it refers to independent marginals to sample from. The marginal in
         dimension :code:`j` is evaluated as:
         :code:`p_j(xj) = pdf_target[j](xj, \*args_target[j])` where **x** is a :class:`numpy.ndarray` of shape
         :code:`(nsamples, dimension)`
        :param log_pdf_target: Logarithm of the target density function from which to draw random samples.
         Either `pdf_target` or `log_pdf_target` must be provided (the latter should be preferred for better numerical
         stability).

         Same comments as for input `pdf_target`.
        :param args_target: Positional arguments of the pdf / log-pdf target function. See `pdf_target`
        :param burn_length: Length of burn-in - i.e., number of samples at the beginning of the chain to discard (note:
         no thinning during burn-in). Default is :math:`0`, no burn-in.
        :param jump: Thinning parameter, used to reduce correlation between samples. Setting :code:`jump=n` corresponds
         to skipping :code:`n-1` states between accepted states of the chain. Default is :math:`1` (no thinning).
        :param dimension: A scalar value defining the dimension of target density function. Either `dimension` and
         `n_chains` or `seed` must be provided.
        :param seed: Seed of the Markov chain(s), shape :code:`(n_chains, dimension)`.
         Default: :code:`zeros(n_chains x dimension)`.

         If seed is not provided, both n_chains and dimension must be provided.
        :param save_log_pdf: Boolean that indicates whether to save log-pdf values along with the samples.
         Default: :any:`False`
        :param concatenate_chains: Boolean that indicates whether to concatenate the chains after a run, i.e., samples
         are stored as an :class:`numpy.ndarray` of shape :code:`(nsamples * n_chains, dimension)` if :any:`True`,
         :code:`(nsamples, n_chains, dimension)` if :any:`False`.
         Default: :any:`True`
        :param n_chains: The number of Markov chains to generate. Either dimension and `n_chains` or `seed` must be
         provided.
        :param proposal: Proposal distribution, must have a log_pdf/pdf and rvs method. Default: standard
         multivariate normal
        :param proposal_is_symmetric: Indicates whether the proposal distribution is symmetric, affects computation of
         acceptance probability alpha Default: :any:`False`, set to :any:`True` if default proposal is used
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is
         :any:`None`.

        :param nsamples: Number of samples to generate.
        :param nsamples_per_chain: Number of samples to generate per chain.
        """
        self.nsamples = nsamples
        self.nsamples_per_chain = nsamples_per_chain
        super().__init__(
            pdf_target=pdf_target,
            log_pdf_target=log_pdf_target,
            args_target=args_target,
            dimension=dimension,
            seed=seed,
            burn_length=burn_length,
            jump=jump,
            save_log_pdf=save_log_pdf,
            concatenate_chains=concatenate_chains,
            random_state=random_state,
            n_chains=n_chains,
        )

        self.logger = logging.getLogger(__name__)
        # Initialize algorithm specific inputs
        self.proposal = proposal
        self.proposal_is_symmetric = proposal_is_symmetric
        if self.proposal is None:
            if self.dimension is None:
                raise ValueError("UQpy: Either input proposal or dimension must be provided.")
            from UQpy.distributions import JointIndependent, Normal

            self.proposal = JointIndependent([Normal()] * self.dimension)
            self.proposal_is_symmetric = True
        else:
            self._check_methods_proposal(self.proposal)

        self.logger.info("\nUQpy: Initialization of " + self.__class__.__name__ + " algorithm complete.")

        if (nsamples is not None) or (nsamples_per_chain is not None):
            self.run(nsamples=nsamples, nsamples_per_chain=nsamples_per_chain,)

    def run_one_iteration(self, current_state: np.ndarray, current_log_pdf: np.ndarray):
        """
        Run one iteration of the mcmc chain for MH algorithm, starting at current state -
        see :class:`MCMC` class.
        """
        # Sample candidate
        candidate = current_state + self.proposal.rvs(
            nsamples=self.n_chains, random_state=self.random_state)

        # Compute log_pdf_target of candidate sample
        log_p_candidate = self.evaluate_log_target(candidate)

        # Compute acceptance ratio
        if self.proposal_is_symmetric:  # proposal is symmetric
            log_ratios = log_p_candidate - current_log_pdf
        else:  # If the proposal is non-symmetric, one needs to account for it in computing acceptance ratio
            log_proposal_ratio = self.proposal.log_pdf(
                candidate - current_state
            ) - self.proposal.log_pdf(current_state - candidate)
            log_ratios = log_p_candidate - current_log_pdf - log_proposal_ratio

        # Compare candidate with current sample and decide or not to keep the candidate (loop over nc chains)
        accept_vec = np.zeros(
            (self.n_chains,)
        )  # this vector will be used to compute accept_ratio of each chain
        unif_rvs = (
            Uniform()
            .rvs(nsamples=self.n_chains, random_state=self.random_state)
            .reshape((-1,))
        )
        for nc, (cand, log_p_cand, r_) in enumerate(
            zip(candidate, log_p_candidate, log_ratios)
        ):
            accept = np.log(unif_rvs[nc]) < r_
            if accept:
                current_state[nc, :] = cand
                current_log_pdf[nc] = log_p_cand
                accept_vec[nc] = 1.0
        # Update the acceptance rate
        self._update_acceptance_rate(accept_vec)

        return current_state, current_log_pdf


# src/UQpy/sampling/mcmc/ModifiedMetropolisHastings.py

import logging
from typing import Callable
import warnings

warnings.filterwarnings('ignore')

import numpy as np
from beartype import beartype
from UQpy.sampling.mcmc.baseclass.MCMC import MCMC
from UQpy.distributions import *
from UQpy.utilities.ValidationTypes import *


class ModifiedMetropolisHastings(MCMC):
    @beartype
    def __init__(
            self,
            pdf_target: Union[Callable, list[Callable]] = None,
            log_pdf_target: Union[Callable, list[Callable]] = None,
            args_target: tuple = None,
            burn_length: Annotated[int, Is[lambda x: x >= 0]] = 0,
            jump: PositiveInteger = 1,
            dimension: int = None,
            seed: list = None,
            save_log_pdf: bool = False,
            concatenate_chains: bool = True,
            proposal: Union[Distribution, list[Distribution]] = None,
            proposal_is_symmetric: Union[bool, list[bool]] = False,
            random_state: RandomStateType = None,
            n_chains: int = None,
            nsamples: PositiveInteger = None,
            nsamples_per_chain: PositiveInteger = None,
    ):
        """
        Component-wise Modified Metropolis-Hastings algorithm. :cite:`SubsetSimulation`

        In this algorithm, candidate samples are drawn separately in each dimension, thus the proposal consists of a
        list of 1D distributions. The target pdf can be given as a joint pdf or a list of marginal pdfs in all
        dimensions. This will trigger two different algorithms. If a list of marginals is provided, the acceptance ratio
        is computed for every dimension independently using the marginal densities. If a joint pdf is provided, the
        acceptance ratio for each component is computed in a loop using this joint pdf.

        :param pdf_target: Target density function from which to draw random samples. Either `pdf_target` or
         `log_pdf_target` must be provided (the latter should be preferred for better numerical stability).

         If `pdf_target` is a callable, it refers to the joint pdf to sample from, it must take at least one input
         **x**, which are the point(s) at which to evaluate the pdf. Within :class:`.MCMC` the `pdf_target` is evaluated
         as:
         :code:`p(x) = pdf_target(x, \*args_target)`

         where **x** is a :class:`numpy.ndarray` of shape :code:`(nsamples, dimension)` and `args_target` are additional
         positional arguments that are provided to :class:`.MCMC` via its args_target input.

         If `pdf_target` is a list of callables, it refers to independent marginals to sample from. The marginal in
         dimension :code:`j` is evaluated as: :code:`p_j(xj) = pdf_target[j](xj, \*args_target[j])` where **x** is a
         :class:`numpy.ndarray` of shape :code:`(nsamples, dimension)`
        :param log_pdf_target: Logarithm of the target density function from which to draw random samples.
         Either `pdf_target` or `log_pdf_target` must be provided (the latter should be preferred for better numerical
         stability).

         Same comments as for input `pdf_target`.
        :param args_target: Positional arguments of the pdf / log-pdf target function. See `pdf_target`
        :param burn_length: Length of burn-in - i.e., number of samples at the beginning of the chain to discard (note:
         no thinning during burn-in). Default is :math:`0`, no burn-in.
        :param jump: Thinning parameter, used to reduce correlation between samples. Setting :code:`jump=n` corresponds
         to skipping :code:`n-1` states between accepted states of the chain. Default is :math:`1` (no thinning).
        :param dimension: A scalar value defining the dimension of target density function. Either `dimension` and
         `n_chains` or `seed` must be provided.
        :param seed: Seed of the Markov chain(s), shape :code:`(n_chains, dimension)`.
         Default: :code:`zeros(n_chains x dimension)`.

         If `seed` is not provided, both `n_chains` and `dimension` must be provided.
        :param save_log_pdf: Boolean that indicates whether to save log-pdf values along with the samples.
         Default: :any:`False`
        :param concatenate_chains: Boolean that indicates whether to concatenate the chains after a run, i.e., samples
         are stored as an :class:`numpy.ndarray` of shape :code:`(nsamples * n_chains, dimension)` if :any:`True`,
         :code:`(nsamples, n_chains, dimension)` if :any:`False`.
         Default: True
        :param n_chains: The number of Markov chains to generate. Either `dimension` and `n_chains` or `seed` must be
         provided.
        :param proposal: Proposal distribution, must have a log_pdf/pdf and rvs method. Default: standard
         multivariate normal
        :param proposal_is_symmetric: Indicates whether the proposal distribution is symmetric, affects computation of
         acceptance probability alpha Default: :any:`False`, set to :any:`True` if default proposal is used
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is
         :any:`None`.
        :param nsamples: Number of samples to generate.
        :param nsamples_per_chain: Number of samples to generate per chain.
        """
        self.nsamples = nsamples
        self.nsamples_per_chain = nsamples_per_chain
        super().__init__(
            pdf_target=pdf_target,
            log_pdf_target=log_pdf_target,
            args_target=args_target,
            dimension=dimension,
            seed=seed,
            burn_length=burn_length,
            jump=jump,
            save_log_pdf=save_log_pdf,
            concatenate_chains=concatenate_chains,
            random_state=random_state,
            n_chains=n_chains,
        )

        self.target_type = None
        self.logger = logging.getLogger(__name__)
        # If proposal is not provided: set it as a list of standard gaussians
        from UQpy.distributions import Normal

        self.proposal = proposal
        self.proposal_is_symmetric = proposal_is_symmetric

        # set default proposal
        if self.proposal is None:
            self.proposal = [Normal(), ] * self.dimension
            self.proposal_is_symmetric = [True, ] * self.dimension
        # Proposal is provided, check it
        else:
            # only one Distribution is provided, check it and transform it to a list
            if isinstance(self.proposal, JointIndependent):
                self.proposal = [m for m in self.proposal.marginals]
                if len(self.proposal) != self.dimension:
                    raise ValueError("UQpy: Proposal given as a list should be of length dimension")
                [self._check_methods_proposal(p) for p in self.proposal]
            elif not isinstance(self.proposal, list):
                self._check_methods_proposal(self.proposal)
                self.proposal = [self.proposal] * self.dimension
            else:  # a list of proposals is provided
                if len(self.proposal) != self.dimension:
                    raise ValueError("UQpy: Proposal given as a list should be of length dimension")
                [self._check_methods_proposal(p) for p in self.proposal]

        # check the symmetry of proposal, assign False as default
        if isinstance(self.proposal_is_symmetric, bool):
            self.proposal_is_symmetric = [self.proposal_is_symmetric, ] * self.dimension
        elif not (isinstance(self.proposal_is_symmetric, list)
                  and all(isinstance(b_, bool) for b_ in self.proposal_is_symmetric)):
            raise TypeError("UQpy: Proposal_is_symmetric should be a (list of) boolean(s)")

        self.logger.info("\nUQpy: Initialization of " + self.__class__.__name__ + " algorithm complete.")



        if (nsamples is not None) or (nsamples_per_chain is not None):
            self.run(nsamples=nsamples, nsamples_per_chain=nsamples_per_chain)

    def run_one_iteration(self, current_state, current_log_pdf):
        """
        Run one iteration of the mcmc chain for MMH algorithm, starting at current state -
        see :class:`MCMC` class.
        """
        # check with algo type is used
        if self.evaluate_log_target_marginals is not None:
            self.target_type = "marginals"
            self.current_log_pdf_marginals = None
        else:
            self.target_type = "joint"
        # The target pdf is provided via its marginals
        accept_vec = np.zeros((self.n_chains,))
        if self.target_type == "marginals":
            # Evaluate the current log_pdf
            if self.current_log_pdf_marginals is None:
                self.current_log_pdf_marginals = [
                    self.evaluate_log_target_marginals[j](current_state[:, j, np.newaxis])
                    for j in range(self.dimension)]

            # Sample candidate (independently in each dimension)
            for j in range(self.dimension):
                candidate_j = current_state[:, j, np.newaxis] + self.proposal[j].rvs(
                    nsamples=self.n_chains, random_state=self.random_state)

                # Compute log_pdf_target of candidate sample
                log_p_candidate_j = self.evaluate_log_target_marginals[j](candidate_j)

                # Compute acceptance ratio
                if self.proposal_is_symmetric[j]:  # proposal is symmetric
                    log_ratios = log_p_candidate_j - self.current_log_pdf_marginals[j]
                else:  # If the proposal is non-symmetric, one needs to account for it in computing acceptance ratio
                    log_prop_j = self.proposal[j].log_pdf
                    log_proposal_ratio = log_prop_j(
                        candidate_j - current_state[:, j, np.newaxis]
                    ) - log_prop_j(current_state[:, j, np.newaxis] - candidate_j)
                    log_ratios = (log_p_candidate_j - self.current_log_pdf_marginals[j] - log_proposal_ratio)

                # Compare candidate with current sample and decide or not to keep the candidate
                unif_rvs = Uniform().rvs(nsamples=self.n_chains, random_state=self.random_state).reshape((-1,))
                for nc, (cand, log_p_cand, r_) in enumerate(zip(candidate_j, log_p_candidate_j, log_ratios)):
                    accept = np.log(unif_rvs[nc]) < r_
                    if accept:
                        current_state[nc, j] = cand
                        self.current_log_pdf_marginals[j][nc] = log_p_cand
                        current_log_pdf = np.sum(self.current_log_pdf_marginals)
                        accept_vec[nc] += 1.0 / self.dimension

        # The target pdf is provided as a joint pdf
        else:
            candidate = np.copy(current_state)
            for j in range(self.dimension):
                candidate_j = current_state[:, j, np.newaxis] + self.proposal[j].rvs(
                    nsamples=self.n_chains, random_state=self.random_state)
                candidate[:, j] = candidate_j[:, 0]

                # Compute log_pdf_target of candidate sample
                log_p_candidate = self.evaluate_log_target(candidate)

                # Compare candidate with current sample and decide or not to keep the candidate
                if self.proposal_is_symmetric[j]:  # proposal is symmetric
                    log_ratios = log_p_candidate - current_log_pdf
                else:  # If the proposal is non-symmetric, one needs to account for it in computing acceptance ratio
                    log_prop_j = self.proposal[j].log_pdf
                    log_proposal_ratio = log_prop_j(candidate_j - current_state[:, j, np.newaxis]) -\
                                         log_prop_j(current_state[:, j, np.newaxis] - candidate_j)
                    log_ratios = log_p_candidate - current_log_pdf - log_proposal_ratio
                unif_rvs = Uniform().rvs(nsamples=self.n_chains, random_state=self.random_state).reshape((-1,))
                for nc, (cand, log_p_cand, r_) in enumerate(zip(candidate_j, log_p_candidate, log_ratios)):
                    accept = np.log(unif_rvs[nc]) < r_
                    if accept:
                        current_state[nc, j] = cand
                        current_log_pdf[nc] = float(log_p_cand)
                        accept_vec[nc] += 1.0 / self.dimension
                    else:
                        candidate[:, j] = current_state[:, j]
        # Update the acceptance rate
        self._update_acceptance_rate(accept_vec)
        return current_state, current_log_pdf


# src/UQpy/sampling/mcmc/Stretch.py

import logging
from typing import Callable
import warnings

warnings.filterwarnings('ignore')

from beartype import beartype
from UQpy.sampling.mcmc.baseclass.MCMC import MCMC
from UQpy.distributions import *
from UQpy.utilities.ValidationTypes import *


class Stretch(MCMC):

    @beartype
    def __init__(
            self,
            pdf_target: Union[Callable, list[Callable]] = None,
            log_pdf_target: Union[Callable, list[Callable]] = None,
            args_target: tuple = None,
            burn_length: Annotated[int, Is[lambda x: x >= 0]] = 0,
            jump: PositiveInteger = 1,
            dimension: int = None,
            seed: list = None,
            save_log_pdf: bool = False,
            concatenate_chains: bool = True,
            scale: float = 2.0,
            random_state: RandomStateType = None,
            n_chains: int = None,
            nsamples: PositiveInteger = None,
            nsamples_per_chain: PositiveInteger = None,
    ):
        """
        Affine-invariant sampler with Stretch moves, parallel implementation. :cite:`Stretch1` :cite:`Stretch2`

        :param pdf_target: Target density function from which to draw random samples. Either `pdf_target` or
         `log_pdf_target` must be provided (the latter should be preferred for better numerical stability).

         If `pdf_target` is a callable, it refers to the joint pdf to sample from, it must take at least one input
         **x**, which are the point(s) at which to evaluate the pdf. Within :class:`.MCMC` the `pdf_target` is evaluated
         as: :code:`p(x) = pdf_target(x, \*args_target)`

         where **x** is a :class:`numpy.ndarray` of shape :code:`(nsamples, dimension)` and `args_target` are additional
         positional arguments that are provided to :class:`.MCMC` via its `args_target` input.

         If `pdf_target` is a list of callables, it refers to independent marginals to sample from. The marginal in
         dimension :code:`j` is evaluated as: :code:`p_j(xj) = pdf_target[j](xj, \*args_target[j])` where **x** is a
         :class:`numpy.ndarray` of shape :code:`(nsamples, dimension)`
        :param log_pdf_target: Logarithm of the target density function from which to draw random samples.
         Either :code:`pdf_target` or :code:`log_pdf_target` must be provided (the latter should be preferred for better
         numerical stability).

         Same comments as for input `pdf_target`.
        :param args_target: Positional arguments of the pdf / log-pdf target function. See `pdf_target`
        :param burn_length: Length of burn-in - i.e., number of samples at the beginning of the chain to discard (note:
         no thinning during burn-in). Default is :math:`0`, no burn-in.
        :param jump: Thinning parameter, used to reduce correlation between samples. Setting :code:`jump=n` corresponds
         to skipping `n-1` states between accepted states of the chain. Default is :math:`1` (no thinning).
        :param dimension: A scalar value defining the dimension of target density function. Either `dimension` and
         `n_chains` or `seed` must be provided.
        :param seed: Seed of the Markov chain(s), shape :code:`(n_chains, dimension)`.
         Default: :code:`zeros(n_chains x dimension)`.

         If `seed` is not provided, both `n_chains` and `dimension` must be provided.
        :param save_log_pdf: Boolean that indicates whether to save log-pdf values along with the samples.
         Default: :any:`False`
        :param concatenate_chains: Boolean that indicates whether to concatenate the chains after a run, i.e., samples
         are stored as an :class:`numpy.ndarray` of shape :code:`(nsamples * n_chains, dimension)` if :any:`True`,
         :code:`(nsamples, n_chains, dimension)` if :any:`False`.
         Default: :any:`True`
        :param n_chains: The number of Markov chains to generate. Either `dimension` and `n_chains` or `seed` must be
         provided.
        :param scale: Scale parameter. Default: :math:`2`.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is
         :any:`None`.
        :param nsamples: Number of samples to generate.
        :param nsamples_per_chain: Number of samples to generate per chain.
        """
        flag_seed = False
        if seed is None:
            if dimension is None or n_chains is None:
                raise ValueError("UQpy: Either `seed` or `dimension` and `n_chains` must be provided.")
            flag_seed = True

        self.nsamples = nsamples
        self.nsamples_per_chain = nsamples_per_chain

        super().__init__(
            pdf_target=pdf_target,
            log_pdf_target=log_pdf_target,
            args_target=args_target,
            dimension=dimension,
            seed=seed,
            burn_length=burn_length,
            jump=jump,
            save_log_pdf=save_log_pdf,
            concatenate_chains=concatenate_chains,
            random_state=random_state,
            n_chains=n_chains, )

        self.logger = logging.getLogger(__name__)
        # Check nchains = ensemble size for the Stretch algorithm
        if flag_seed:
            self.seed = (Uniform().rvs(nsamples=self.dimension * self.n_chains,
                                       random_state=self.random_state, )
                         .reshape((self.n_chains, self.dimension)))
        if self.n_chains < 2:
            raise ValueError("UQpy: For the Stretch algorithm, a seed must be provided with at least two samples.")

        # Check Stretch algorithm inputs: proposal_type and proposal_scale
        self.scale = scale
        if not isinstance(self.scale, float):
            raise TypeError("UQpy: Input scale must be of type float.")

        self.logger.info("\nUQpy: Initialization of " + self.__class__.__name__ + " algorithm complete.")

        if (nsamples is not None) or (nsamples_per_chain is not None):
            self.run(nsamples=nsamples, nsamples_per_chain=nsamples_per_chain, )

    def run_one_iteration(self, current_state, current_log_pdf):
        """
        Run one iteration of the mcmc chain for Stretch algorithm, starting at current state -
        see :class:`.MCMC` class.
        """
        # Start the loop over nsamples - this code uses the parallel version of the stretch algorithm
        all_inds = np.arange(self.n_chains)
        inds = all_inds % 2
        accept_vec = np.zeros((self.n_chains,))
        # Separate the full ensemble into two sets, use one as a complementary ensemble to the other and vice-versa
        for split in range(2):
            set1 = inds == split

            # Get current and complementary sets
            sets = [current_state[inds == j01, :] for j01 in range(2)]
            curr_set, comp_set = (sets[split], sets[1 - split],)  # current and complementary sets respectively
            ns, nc = len(curr_set), len(comp_set)

            # Sample new state for S1 based on S0
            unif_rvs = Uniform().rvs(nsamples=ns, random_state=self.random_state)
            zz = ((self.scale - 1.0) * unif_rvs + 1.0) ** 2.0 / self.scale  # sample Z
            factors = (self.dimension - 1.0) * np.log(zz)  # compute log(Z ** (d - 1))
            multi_rvs = Multinomial(n=1, p=[1.0 / nc, ] * nc).rvs(
                nsamples=ns, random_state=self.random_state)
            rint = np.nonzero(multi_rvs)[1]  # sample X_{j} from complementary set
            candidates = comp_set[rint, :] - (comp_set[rint, :] - curr_set) * np.tile(
                zz, [1, self.dimension])  # new candidates

            # Compute new likelihood, can be done in parallel :)
            logp_candidates = self.evaluate_log_target(candidates)

            # Compute acceptance rate
            unif_rvs = (Uniform().rvs(nsamples=len(all_inds[set1]), random_state=self.random_state).reshape((-1,)))
            for j, f, lpc, candidate, u_rv in zip(all_inds[set1], factors, logp_candidates, candidates, unif_rvs):
                accept = np.log(u_rv) < f + lpc - current_log_pdf[j]
                if accept:
                    current_state[j] = candidate
                    current_log_pdf[j] = lpc
                    accept_vec[j] += 1.0

        # Update the acceptance rate
        self._update_acceptance_rate(accept_vec)
        return current_state, current_log_pdf


# src/UQpy/sampling/mcmc/__init__.py

from UQpy.sampling.mcmc.MetropolisHastings import MetropolisHastings
from UQpy.sampling.mcmc.ModifiedMetropolisHastings import ModifiedMetropolisHastings
from UQpy.sampling.mcmc.Stretch import Stretch
from UQpy.sampling.mcmc.DRAM import DRAM
from UQpy.sampling.mcmc.DREAM import DREAM
from UQpy.sampling.mcmc.baseclass.MCMC import MCMC


# src/UQpy/sampling/stratified_sampling/LatinHypercubeSampling.py

import logging
from typing import Union

from beartype import beartype

from UQpy.sampling.stratified_sampling.baseclass.StratifiedSampling import StratifiedSampling
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria.baseclass import Criterion
from UQpy.utilities.Utilities import process_random_state
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria import Random
from UQpy.utilities.ValidationTypes import PositiveInteger, NumpyFloatArray, RandomStateType
from UQpy.distributions import *
import numpy as np
from UQpy.distributions import DistributionContinuous1D, JointIndependent


class LatinHypercubeSampling(StratifiedSampling):
    @beartype
    def __init__(
        self,
        distributions: Union[Distribution, list[Distribution]],
        nsamples: PositiveInteger,
        criterion: Criterion = Random(),
        random_state: RandomStateType = None
    ):
        """
        Perform Latin hypercube sampling (LHS) of random variables.

        All distributions in :class:`LatinHypercubeSampling` must be independent. :class:`.LatinHypercubeSampling` does
        not generate correlated random variables. Therefore, for multi-variate designs the `distributions` must be a
        list of :class:`.DistributionContinuous1D` objects or an object of the :class:`.JointIndependent` class.


        :param distributions: List of :class:`.Distribution` objects
         corresponding to each random variable.
        :param nsamples: Number of samples to be drawn from each distribution.
        :param random_state: Random seed used to initialize the pseudo-random number generator. If an :any:`int` is
         provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise, the
         object itself can be passed directly.
        :param criterion: The criterion for pairing the generating sample points. This parameter must be of
         type :class:`.Criterion`.

         Options:

         1. :class:`.Random` - completely random. \n
         2. :class:`.Centered` - points only at the centre. \n
         3. :class:`.MaxiMin` - maximizing the minimum distance between points. \n
         4. :class:`.MinCorrelation` - minimizing the correlation between the points. \n
         5. User-defined criterion class, by providing an implementation of the abstract class :class:`Criterion`
        """
        self.random_state = process_random_state(random_state)
        self.distributions = distributions
        self.criterion = criterion
        self.nsamples = nsamples
        self.logger = logging.getLogger(__name__)
        self._samples: NumpyFloatArray = None
        if isinstance(self.distributions, list):
            self._samples = np.zeros([self.nsamples, len(self.distributions)])
        elif isinstance(self.distributions, DistributionContinuous1D):
            self._samples = np.zeros([self.nsamples, 1])
        elif isinstance(self.distributions, JointIndependent):
            self._samples = np.zeros([self.nsamples, len(self.distributions.marginals)])

        self.samplesU01: NumpyFloatArray = np.zeros_like(self._samples)
        """The generated LHS samples on the unit hypercube."""

        if self.nsamples is not None:
            self.run(self.nsamples)

    @property
    def samples(self):
        """ The generated LHS samples."""
        return np.atleast_2d(self._samples)

    @beartype
    def run(self, nsamples: PositiveInteger):
        """
        Execute the random sampling in the :class:`.LatinHypercubeSampling` class.

        :param nsamples: If the :meth:`run` method is invoked multiple times, the newly generated samples will
         overwrite the existing samples.

        The :meth:`run` method is the function that performs random sampling in the :class:`.LatinHypercubeSampling`
        class. If `nsamples` is provided, the :meth:`run` method is automatically called when the
        :class:`.LatinHypercubeSampling` object is defined. The user may also call the :meth:`run` method directly to
        generate samples. The :meth:`run` method of the :class:`.LatinHypercubeSampling` class cannot be invoked
        multiple times for sample size extension.

        The :meth:`run` method has no returns, although it creates and/or appends the :py:attr:`samples` and
        :py:attr:`samplesU01` attributes of the :class:`.LatinHypercubeSampling` object.
        """
        self.nsamples = nsamples
        self.logger.info("UQpy: Running Latin Hypercube sampling...")
        self.criterion.create_bins(self._samples, self.random_state)

        u_lhs = self.criterion.generate_samples(self.random_state)
        self.samplesU01 = u_lhs

        if isinstance(self.distributions, list):
            for j in range(len(self.distributions)):
                if hasattr(self.distributions[j], "icdf"):
                    self._samples[:, j] = self.distributions[j].icdf(u_lhs[:, j])

        elif isinstance(self.distributions, JointIndependent):
            if all(hasattr(m, "icdf") for m in self.distributions.marginals):
                for j in range(len(self.distributions.marginals)):
                    self._samples[:, j] = self.distributions.marginals[j].icdf(u_lhs[:, j])

        elif isinstance(self.distributions, DistributionContinuous1D):
            if hasattr(self.distributions, "icdf"):
                self._samples = self.distributions.icdf(u_lhs)

        self.logger.info("Successful execution of LHS design.")


# src/UQpy/sampling/stratified_sampling/RefinedStratifiedSampling.py


from UQpy.sampling.stratified_sampling.TrueStratifiedSampling import *
from UQpy.sampling.stratified_sampling.refinement.baseclass import Refinement
from UQpy.utilities.ValidationTypes import RandomStateType, PositiveInteger


class RefinedStratifiedSampling(StratifiedSampling):
    @beartype
    def __init__(
        self,
        stratified_sampling: Union[TrueStratifiedSampling],
        refinement_algorithm: Refinement,
        nsamples: PositiveInteger = None,
        samples_per_iteration: int = 1,
        random_state: RandomStateType = None,
    ):
        """

        :param stratified_sampling: Generally, this must be an object of a :class:`.TrueStratifiedSampling` class.
        :param refinement_algorithm: Algorithm used for the refinement of the strata. Two methods exist
         :class:`.RandomRefinement` and :class:`.GradientEnhancedRefinement` .
        :param nsamples: Total number of samples to be drawn (including the initial samples).
         If `nsamples` is provided when instantiating the class, the :meth:`run` method will automatically be
         called. If `nsamples` is not provided, :class:`.RefinedStratifiedSampling` can be executed
         by invoking the :meth:`run` method and passing `nsamples`
        :param samples_per_iteration: Number of samples to be added per refinement iteration.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
         If an :any:`int` is provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise,
         the object itself can be passed directly.
        """
        self.stratified_sampling = stratified_sampling
        self.samples_per_iteration = samples_per_iteration
        self.refinement_algorithm = refinement_algorithm
        self.refinement_algorithm.update_strata(stratified_sampling.samplesU01)
        self.training_points = self.stratified_sampling.samplesU01
        self.samplesU01: Numpy2DFloatArray = self.stratified_sampling.samplesU01
        """The generated samples on the unit hypercube."""
        self.samples: Numpy2DFloatArray = self.stratified_sampling.samples
        """The generated stratified samples following the prescribed distribution."""
        self.dimension = self.samples.shape[1]
        self.random_state = random_state
        self.logger = logging.getLogger(__name__)

        self.nsamples = nsamples

        self.random_state = random_state
        if isinstance(self.random_state, int):
            self.random_state = np.random.RandomState(self.random_state)
        elif not isinstance(self.random_state, (type(None), np.random.RandomState)):
            raise TypeError('UQpy: random_state must be None, an int or an np.random.Generator object.')
        if self.random_state is None:
            self.random_state = self.stratified_sampling.random_state

        if self.nsamples is not None:
            self.run(nsamples=self.nsamples)

    @beartype
    def run(self, nsamples: PositiveInteger):
        """
        Executes refined stratified sampling.

        :param nsamples:  Total number of samples to be drawn (including the initial samples).
         If `nsamples` is provided when instantiating the class, the :meth:`run` method will automatically be
         called. If `nsamples` is not provided, an :class:`.RefinedStratifiedSampling` subclass can be executed
         by invoking the :meth:`run` method and passing `nsamples`
        """
        self.nsamples = nsamples

        if self.nsamples <= self.samples.shape[0]:
            raise ValueError("UQpy Error: The number of requested samples must be larger than the existing "
                             "sample set.")

        initial_number = self.samples.shape[0]

        self.refinement_algorithm.initialize(self.nsamples, self.training_points, self.samples)

        for i in range(initial_number, nsamples, self.samples_per_iteration):
            new_points = self.refinement_algorithm.update_samples(
                self.nsamples, self.samples_per_iteration,
                self.random_state, i, self.dimension,
                self.samplesU01, self.training_points)
            self.append_samples(new_points)

            self.refinement_algorithm.finalize(self.samples, self.samples_per_iteration)

    def append_samples(self, new_points):
        # Adding new sample to training points, samplesU01 and samples attributes
        self.training_points = np.vstack([self.training_points, new_points])
        self.samplesU01 = np.vstack([self.samplesU01, new_points])
        new_point_ = np.zeros_like(new_points)
        for k in range(self.dimension):
            new_point_[:, k] = self.stratified_sampling.distributions[k].icdf(new_points[:, k])
        self.samples = np.vstack([self.samples, new_point_])


# src/UQpy/sampling/stratified_sampling/TrueStratifiedSampling.py

import logging
from beartype import beartype
from numpy.random import RandomState

from UQpy.sampling.stratified_sampling.baseclass.StratifiedSampling import StratifiedSampling
from UQpy.distributions import DistributionContinuous1D, JointIndependent
from UQpy.sampling.stratified_sampling.strata import RectangularStrata
from UQpy.sampling.stratified_sampling.strata.baseclass.Strata import Strata
from UQpy.utilities.ValidationTypes import *


class TrueStratifiedSampling(StratifiedSampling):
    @beartype
    def __init__(
        self,
        distributions: Union[DistributionContinuous1D, JointIndependent, list[DistributionContinuous1D]],
        strata_object: Strata,
        nsamples_per_stratum: Union[int, list[int]] = None,
        nsamples: int = None,
        random_state: RandomStateType = None,
    ):
        """
        Class for Stratified Sampling (:cite:`StratifiedSampling1`).

        :param distributions: List of :class:`.Distribution` objects corresponding to each random variable.
        :param strata_object: Defines the stratification of the unit hypercube. This must be provided and must be an
         object of a :class:`.Strata` child class: :class:`.Rectangular`, :class:`.Voronoi`, or :class:`.Delaunay`.
        :param nsamples_per_stratum: Specifies the number of samples in each stratum. This must be either an
         integer, in which case an equal number of samples are drawn from each stratum, or a list. If it is provided as
         a list, the length of the list must be equal to the number of strata.
         If `nsamples_per_stratum` is provided when the class is defined, the :meth:`run` method will be executed
         automatically.  If neither `nsamples_per_stratum` or `nsamples` are provided when the class is
         defined, the user must call the :meth:`run` method to perform stratified sampling.
        :param nsamples: Specify the total number of samples. If `nsamples` is specified, the samples will
         be drawn in proportion to the volume of the strata. Thus, each stratum will contain
         :code:`round(V_i* nsamples)` samples.
         If `nsamples` is provided when the class is defined, the :meth:`run` method will be executed
         automatically.  If neither `nsamples_per_stratum` or `nsamples` are provided when the class is
         defined, the user must call the :meth:`run` method to perform stratified sampling.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
         If an :any:`int` is provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise,
         the object itself can be passed directly.
        """
        self.logger = logging.getLogger(__name__)
        self.weights: NumpyFloatArray = None
        """Individual sample weights."""
        self.strata_object = strata_object

        self.nsamples_per_stratum = nsamples_per_stratum
        self.nsamples = nsamples

        self.samples:NumpyFloatArray = None
        """The generated samples following the prescribed distribution."""
        self.samplesU01:NumpyFloatArray = None
        """The generated samples on the unit hypercube."""

        self.distributions = distributions
        self.random_state = random_state
        if isinstance(self.random_state, int):
            self.random_state = RandomState(self.random_state)
        elif not isinstance(self.random_state, (type(None), RandomState)):
            raise TypeError('UQpy: random_state must be None, an int or an np.random.RandomState object.')
        if self.random_state is None:
            self.random_state = self.strata_object.random_state

        if isinstance(self.strata_object, RectangularStrata):
            self.strata_object.check_centered(nsamples)
        self.logger.info("UQpy: Stratified_sampling object is created")

        if self.nsamples_per_stratum is not None or self.nsamples is not None:
            self.run(nsamples_per_stratum=self.nsamples_per_stratum,
                     nsamples=self.nsamples)

    def transform_samples(self, samples01):
        """
        Transform samples in the unit hypercube :math:`[0, 1]^n` to the prescribed distribution using the inverse CDF.

        :param samples01: :class:`numpy.ndarray` containing the generated samples on :math:`[0, 1]^n`.

        :return: :class:`numpy.ndarray` containing the generated samples following the prescribed distribution.
        """
        samples_u_to_x = np.zeros_like(samples01)
        for j in range(samples01.shape[1]):
            samples_u_to_x[:, j] = self.distributions[j].icdf(samples01[:, j])

        self.samples = samples_u_to_x

    @beartype
    def run(
        self,
        nsamples_per_stratum: Union[None, int, list[int]] = None,
        nsamples: Union[None, PositiveInteger] = None,
    ):
        """
        Executes stratified sampling.

        This method performs the sampling for each of the child classes by running two methods:
        :meth:`create_samplesu01`, and :meth:`transform_samples`. The :meth:`create_samplesu01` method is
        unique to each child class and therefore must be overwritten when a new child class is defined. The
        :meth:`transform_samples` method is common to all stratified sampling classes and is therefore defined by the
        parent class. It does not need to be modified.

        If `nsamples` or `nsamples_per_stratum` is provided when the class is defined, the :meth:`run`
        method will be executed automatically.  If neither `nsamples_per_stratum` or `nsamples` are provided
        when the class is defined, the user must call the :meth:`run` method to perform stratified sampling.

        :param nsamples_per_stratum: Specifies the number of samples in each stratum. This must be either an
         integer, in which case an equal number of samples are drawn from each stratum, or a list. If it is provided as
         a list, the length of the list must be equal to the number of strata.
         If `nsamples_per_stratum` is provided when the class is defined, the :meth:`run` method will be executed
         automatically.  If neither `nsamples_per_stratum` or `nsamples` are provided when the class is
         defined, the user must call the :meth:`run` method to perform stratified sampling.
        :param nsamples: Specify the total number of samples. If `nsamples` is specified, the samples will
         be drawn in proportion to the volume of the strata. Thus, each stratum will contain
         :code:`round(V_i*nsamples)` samples where :math:`V_i \le 1` is the volume of stratum `i` in the unit
         hypercube.
         If `nsamples` is provided when the class is defined, the :meth:`run` method will be executed
         automatically.  If neither `nsamples_per_stratum` or `nsamples` are provided when the class is
         defined, the user must call the :meth:`run` method to perform stratified sampling.
        """

        self.nsamples_per_stratum = nsamples_per_stratum
        self.nsamples = nsamples
        self._run_checks()

        self.logger.info("UQpy: Performing Stratified Sampling")

        self.create_unit_hypercube_samples()

        self.transform_samples(self.samplesU01)

        self.logger.info("UQpy: Stratified Sampling is completed")

    def _run_checks(self):
        if self.nsamples is not None:
            self.nsamples_per_stratum = (self.strata_object.volume * self.nsamples).round()

        if self.nsamples_per_stratum is not None:
            if isinstance(self.nsamples_per_stratum, int):
                self.nsamples_per_stratum = [self.nsamples_per_stratum] * \
                                                  self.strata_object.volume.shape[0]
            elif isinstance(self.nsamples_per_stratum, list):
                if len(self.nsamples_per_stratum) != self.strata_object.volume.shape[0]:
                    raise ValueError("UQpy: Length of 'nsamples_per_stratum' must match the number of strata.")
            elif self.nsamples is None:
                raise ValueError("UQpy: 'nsamples_per_stratum' must be an integer or a list.")
        else:
            self.nsamples_per_stratum = [1] * self.strata_object.volume.shape[0]

    def create_unit_hypercube_samples(self):
        samples_in_strata, weights = self.strata_object.sample_strata(
            self.nsamples_per_stratum, self.random_state)
        self.weights = np.array(weights)
        self.samplesU01 = np.concatenate(samples_in_strata, axis=0)


# src/UQpy/sampling/stratified_sampling/__init__.py

from UQpy.sampling.stratified_sampling.baseclass import *
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria import *
from UQpy.sampling.stratified_sampling.strata import *
from UQpy.sampling.stratified_sampling.refinement import *


from UQpy.sampling.stratified_sampling.LatinHypercubeSampling import LatinHypercubeSampling
from UQpy.sampling.stratified_sampling.TrueStratifiedSampling import TrueStratifiedSampling
from UQpy.sampling.stratified_sampling.RefinedStratifiedSampling import RefinedStratifiedSampling



# src/UQpy/sensitivity/baseclass/__init__.py

from UQpy.sensitivity.baseclass.sensitivity import *
from UQpy.sensitivity.baseclass.pickfreeze import *


# src/UQpy/sensitivity/baseclass/pickfreeze.py

import copy


def generate_pick_freeze_samples(dist_obj, n_samples, random_state=None):

    """
    Generate samples to be used in the Pick-and-Freeze algorithm.

    **Outputs:**

    * **A_samples** (`ndarray`):
        Sample set A.
        Shape: `(n_samples, num_vars)`.

    * **B_samples** (`ndarray`):
        Sample set B.
        Shape: `(n_samples, num_vars)`.

    * **C_i_generator** (`generator`):
        Generator for the sample set C_i.
        Generator is used so that samples
        do not have to be stored in memory.
        C_i is a 2D array with all columns
        from B_samples, except column `i`,
        which is from A_samples.
        Shape: `(n_samples, num_vars)`.

    * **D_i_generator** (`generator`):
        Generator for the sample set C_i.
        Generator is used so that samples
        do not have to be stored in memory.
        C_i is a 2D array with all columns
        from A_samples, except column `i`,
        which is from B_samples.
        Shape: `(n_samples, num_vars)`.

    """

    # Generate samples for A and B
    samples = dist_obj.rvs(n_samples * 2, random_state=random_state)

    num_vars = samples.shape[1]

    # Split samples into two sets A and B
    A_samples = samples[:n_samples, :]
    B_samples = samples[n_samples:, :]

    # Iterator for generating C_i
    def C_i_generator():
        """Generate C_i for each i."""
        for i in range(num_vars):
            C_i = copy.deepcopy(B_samples)  #! Deepcopy so B is unchanged
            C_i[:, i] = A_samples[:, i]
            yield C_i

    # Iterator for generating D_i
    def D_i_generator():
        """Generate D_i for each i."""
        for i in range(num_vars):
            D_i = copy.deepcopy(A_samples)  #! Deepcopy so A is unchanged
            D_i[:, i] = B_samples[:, i]
            yield D_i

    return A_samples, B_samples, C_i_generator(), D_i_generator()


# src/UQpy/sensitivity/baseclass/sensitivity.py

"""

This module contains the abstract Sensitivity class used by other 
sensitivity classes: 
1. Chatterjee indices
2. Cramer-von Mises indices
3. Generalised Sobol indices
4. Sobol indices

"""

import copy
import numpy as np
import scipy.stats

from UQpy.run_model import RunModel
from UQpy.distributions.baseclass import DistributionContinuous1D
from UQpy.distributions.collection import JointIndependent


class Sensitivity:
    def __init__(
        self, runmodel_object, dist_object, random_state=None, **kwargs
    ) -> None:

        # Check RunModel object
        if not isinstance(runmodel_object, RunModel):
            raise TypeError("UQpy: runmodel_object must be an object of class RunModel")

        self.runmodel_object = runmodel_object

        # Check distributions
        if isinstance(dist_object, list):
            for i in range(len(dist_object)):
                if not isinstance(dist_object[i], (DistributionContinuous1D, JointIndependent)):
                    raise TypeError(
                        "UQpy: A  ``DistributionContinuous1D`` or ``JointInd`` object "
                        "must be provided."
                    )
        else:
            if not isinstance(dist_object, (DistributionContinuous1D, JointIndependent)):
                raise TypeError(
                    "UQpy: A  ``DistributionContinuous1D``  or ``JointInd`` object must be provided."
                )

        self.dist_object = dist_object

        # Check random state
        self.random_state = random_state
        if isinstance(self.random_state, int):
            self.random_state = np.random.RandomState(self.random_state)
        elif not (
            self.random_state is None
            or isinstance(self.random_state, np.random.RandomState)
        ):
            raise TypeError(
                "UQpy: random state should be None, an integer or np.random.RandomState object"
            )

    # wrapper created for convenience to generate model evaluations
    def _run_model(self, samples):
        """Generate model evaluations for a set of samples.

        **Inputs**:

        * **samples** (`numpy.ndarray`):
            A set of samples.
            Shape: `(n_samples, num_vars)`

        **Outputs**:

        * **model_evaluations** (`numpy.ndarray`):
            A set of model evaluations.
            Shape: `(n_samples,)`

            if multioutput: `(n_samples, n_outputs)`

        """

        self.runmodel_object.run(samples=samples, append_samples=False)
        model_evals = copy.deepcopy(np.array(self.runmodel_object.qoi_list))

        return model_evals

    @staticmethod
    def bootstrap_sample_generator_1D(samples):
        """Generate bootstrap samples.

        Generators are used to avoid copying the entire array.

        It will simply pick `N` random rows from the array.

        For example:
        Model evaluations for the samples in A in the pick and freeze estimator.

        **Inputs:**

        * **samples** (`ndarray`):
            Model evaluations for the samples.
            Shape: `(n_samples, 1)`.

        **Outputs:**

        * `generator`:
            Generator for the bootstrap samples.

        """
        n_samples = samples.shape[0]

        while True:
            _indices = np.random.randint(0, high=n_samples, size=n_samples)

            yield samples[_indices]

    @staticmethod
    def bootstrap_sample_generator_2D(samples):
        """Generate bootstrap samples.

        Generators are used to avoid copying the entire array.

        For example:
        Let's say we have '3' random variables
        To pick bootstrap samples from f_C_i, we first
        generate indices to pick values from each column
        num_cols = 3
        cols = [0, 1, 2]
        _indices = [[3, 4, 8]
                    [6, 1, 2]
                    [0, 5, 7]
                    [4, 1, 0]] (4x3)
        elements from f_C_i will be picked column-wise:
        f_C_i[_indices[:, 0], 0]
        f_C_i[_indices[:, 1], 1] etc.

        **Inputs:**

        * **samples** (`ndarray`):
            Model evaluations for the samples.
            Shape: `(n_samples, 1)`.

        **Outputs:**

        * `generator`:
            Generator for the bootstrap samples.

        """
        n_samples = samples.shape[0]

        num_cols = samples.shape[1]
        cols = np.arange(num_cols)

        while True:
            # generate indices to pick N values from f_A, f_B and f_C_i
            _indices = np.random.randint(0, high=n_samples, size=samples.shape)

            yield samples[_indices, cols]

    @staticmethod
    def bootstrap_sample_generator_3D(samples):
        """Generate bootstrap samples.

        Generators are used to avoid copying the entire array.

        For example:
        Let's say we a model with multiple outputs.
        We use the same approach as in the 2D
        case for each slice the 3D array.
        Here, slices refer to the 'depth' of the array,
        given by array.shape[0].

        **Inputs:**

        * **samples** (`ndarray`):
            Model evaluations for the samples.
            Shape: `(n_outputs, n_samples, num_vars)`.

        **Outputs:**

        * `generator`:
            Generator for the bootstrap samples.

        """
        n_samples = samples.shape[1]
        array_shape = samples.shape[1:]
        num_cols = samples.shape[2]
        cols = np.arange(num_cols)

        while True:
            _indices = np.random.randint(0, high=n_samples, size=array_shape)

            yield samples[:, _indices, cols]

    def bootstrapping(
        self,
        estimator,
        estimator_inputs,
        qoi_mean,
        num_bootstrap_samples,
        confidence_level=0.95,
        **kwargs,
    ):

        """An abstract method to implement bootstrapping.

        **Inputs:**

        * **estimator** (`function`):
            A method/func which computes the statistical
            quantities of interest (QoI).
            Example: `compute_first_order_Sobol`
            It must be a method/function that takes several `ndarray`s
            of samples as input and returns a single `ndarray` of estimated value.

        * **estimator_inputs** (`list`):
            Inputs to the estimator concantenated in a list.

        * **qoi_mean** (`ndarray`):
            Mean of the QoI.
            This is the value around which we
            will compute the confidence interval.
            Shape: `(n_qois, n_outputs)`.

        * **num_bootstrap_samples** (`int`):
            Number of bootstrap samples to generate.

        * **confidence_level** (`float`):
            Confidence level for the confidence interval.
            Default: 0.95

        **Outputs:**

        * **confidence_interval_qoi** (`ndarray`):
            Confidence interval for the quantity of interest (QoI).

        """

        n_qois = qoi_mean.shape[0]
        n_outputs = qoi_mean.shape[1]

        ##################### STORAGE #####################

        # store generators of the inputs for bootstrap sampling
        input_generators = []

        # store the qoi computed using bootstrap samples
        bootstrapped_qoi = np.zeros((n_outputs, n_qois, num_bootstrap_samples))

        # store the confidence interval for each qoi
        confidence_interval_qoi = np.zeros((n_outputs, n_qois, 2))

        ##################### CREATE GENERATORS #####################

        for i, input in enumerate(estimator_inputs):

            if isinstance(input, np.ndarray):

                # Example: f_A or f_B of models with single output.
                # Shape: `(n_samples, 1)`.
                if input.ndim == 2 and input.shape[1] == 1:
                    input_generators.append(self.bootstrap_sample_generator_1D(input))

                # Example: f_C_i or f_D_i of models with single output.
                # Shape: `(n_samples, num_vars)`.
                elif input.ndim == 2 and input.shape[1] > 1:
                    input_generators.append(self.bootstrap_sample_generator_2D(input))

                # Example: f_C_i or f_D_i of models with multiple outputs.
                # Shape: `(n_outputs, n_samples, num_vars)`.
                elif input.ndim == 3:
                    input_generators.append(self.bootstrap_sample_generator_3D(input))

            # Example: if models evals is None.
            elif input == None:
                input_generators.append(input)

            else:
                raise ValueError(
                    f"UQpy: estimator_inputs[{i}] should be either None or `ndarray` of dimension 1, 2 or 3"
                )

        ################### BOOTSTRAPPING ##################

        # Compute the qoi for each bootstrap sample
        for j in range(num_bootstrap_samples):

            # inputs to the estimator
            args = []

            # generate samples
            for gen_input in input_generators:
                if gen_input == None:
                    args.append(gen_input)
                else:
                    args.append(gen_input.__next__())

            bootstrapped_qoi[:, :, j] = estimator(*args, **kwargs).T

        ################# CONFIDENCE INTERVAL ################

        # Calculate confidence intervals
        delta = -scipy.stats.norm.ppf((1 - confidence_level) / 2)

        for output_j in range(n_outputs):

            # estimate the standard deviation using the bootstrap indices
            std_qoi = np.std(bootstrapped_qoi[output_j, :, :], axis=1, ddof=1)

            lower_bound = qoi_mean[:, output_j] - delta * std_qoi
            upper_bound = qoi_mean[:, output_j] + delta * std_qoi

            confidence_interval_qoi[output_j, :, 0] = lower_bound
            confidence_interval_qoi[output_j, :, 1] = upper_bound

        # For models with single output, return 2D array.
        if n_outputs == 1:
            confidence_interval_qoi = confidence_interval_qoi[0, :, :]

        return confidence_interval_qoi


# src/UQpy/stochastic_process/supportive/__init__.py

"""Collection of baseclasses"""
from UQpy.stochastic_process.supportive.inverse_wiener_khinchin_transform import (
    inverse_wiener_khinchin_transform,
)
from UQpy.stochastic_process.supportive.wiener_khinchin_transform import (
    wiener_khinchin_transform,
)
from UQpy.stochastic_process.supportive.scaling_correlation_function import (
    scaling_correlation_function,
)


# src/UQpy/stochastic_process/supportive/inverse_wiener_khinchin_transform.py

import numpy as np


def inverse_wiener_khinchin_transform(correlation_function, frequency, time):
    """
    A function to transform the autocorrelation function to a power spectrum by the Inverse Wiener Khinchin
    transformation.

    ** Input:**

    * **correlation_function** (`list or numpy.array`):

        The correlation function of the signal.

    * **frequency** (`list or numpy.array`):

        The frequency discretizations of the power spectrum.

    * **time** (`list or numpy.array`):

        The time discretizations of the signal.

    **Output/Returns:**

    * **power_spectrum** (`list or numpy.array`):

        The power spectrum of the signal.
    """
    time_length = time[1] - time[0]
    fac = np.ones(len(time))
    fac[1 : len(time) - 1 : 2] = 4
    fac[2 : len(time) - 2 : 2] = 2
    fac = fac * time_length / 3
    power_spectrum = np.zeros(len(frequency))
    for i in range(len(frequency)):
        power_spectrum[i] = (
            2
            / (2 * np.pi)
            * np.dot(fac, correlation_function * np.cos(time * frequency[i]))
        )
    power_spectrum[power_spectrum < 0] = 0
    return power_spectrum


# src/UQpy/stochastic_process/supportive/scaling_correlation_function.py

import numpy as np


def scaling_correlation_function(correlation_function):
    """
    A function to scale a correlation function such that correlation at 0 lag is equal to 1

    ** Input:**

    * **correlation_function** (`list or numpy.array`):

        The correlation function of the signal.

    **Output/Returns:**

    * **scaled_correlation_function** (`list or numpy.array`):

        The scaled correlation functions of the signal.
    """
    scaled_correlation_function = correlation_function / np.max(correlation_function)
    return scaled_correlation_function


# src/UQpy/stochastic_process/supportive/wiener_khinchin_transform.py

import numpy as np


def wiener_khinchin_transform(power_spectrum, frequency, time):
    """
    A function to transform the power spectrum to a correlation function by the Wiener Khinchin transformation

    ** Input:**

    * **power_spectrum** (`list or numpy.array`):

        The power spectrum of the signal.

    * **frequency** (`list or numpy.array`):

        The frequency discretizations of the power spectrum.

    * **time** (`list or numpy.array`):

        The time discretizations of the signal.

    **Output/Returns:**

    * **correlation_function** (`list or numpy.array`):

        The correlation function of the signal.

    """
    frequency_interval = frequency[1] - frequency[0]
    fac = np.ones(len(frequency))
    fac[1 : len(frequency) - 1 : 2] = 4
    fac[2 : len(frequency) - 2 : 2] = 2
    fac = fac * frequency_interval / 3
    correlation_function = np.zeros(len(time))
    for i in range(len(time)):
        correlation_function[i] = 2 * np.dot(
            fac, power_spectrum * np.cos(frequency * time[i])
        )
    return correlation_function


# src/UQpy/surrogates/baseclass/Surrogate.py

from abc import ABC, abstractmethod


class Surrogate(ABC):
    @abstractmethod
    def fit(self, samples, values):
        pass

    @abstractmethod
    def predict(self, points, return_std=False):
        pass


# src/UQpy/surrogates/baseclass/__init__.py

from UQpy.surrogates.baseclass.Surrogate import Surrogate


# src/UQpy/surrogates/gaussian_process/GaussianProcessRegression.py

import logging
import numpy as np
from scipy.linalg import cholesky, cho_solve

from beartype import beartype


from UQpy.utilities.Utilities import process_random_state
from UQpy.surrogates.baseclass.Surrogate import Surrogate
from UQpy.utilities.ValidationTypes import RandomStateType
from UQpy.surrogates.gaussian_process.kernels.baseclass.Kernel import Kernel
from UQpy.surrogates.gaussian_process.constraints.baseclass.Constraints import ConstraintsGPR
from UQpy.surrogates.gaussian_process.regression_models.baseclass.Regression import Regression


class GaussianProcessRegression(Surrogate):
    @beartype
    def __init__(
            self,
            kernel: Kernel,
            hyperparameters: list,
            regression_model=None,
            optimizer=None,
            bounds=None,
            optimize_constraints: ConstraintsGPR = None,
            optimizations_number: int = 1,
            normalize: bool = False,
            noise: bool = False,
            random_state: RandomStateType = None,
    ):
        """
        GaussianProcessRegressor an Gaussian process regression-based surrogate model to predict the model output at
        new sample points.

        :param kernel: `kernel` specifies and evaluates the kernel.
         Built-in options: RBF
        :param hyperparameters: List or array of initial values for the kernel hyperparameters/scale parameters. This
         input attribute defines the dimension of input training point, thus its length/shape should be equal to the
         input dimension plus one (d+1), this list/array includes 'd' length scale and process standard deviation. In
         case of noisy observations/output, its length/shape should be equal to the input dimension plus two (d+2), this
         list/array includes 'd' lengthscales, process standar deviation and noise variance.
        :param regression_model: A class object, which computes the basis function at a sample point. If
         regression_model is None, this class will train GP with regression.
         Default: None
        :param bounds: This input attribute defines the bounds of the loguniform distributions, which randomly generates
         new starting point for optimization problem to estimate maximum likelihood estimator. This should be a closed
         bound.
         Default: [10**-3, 10**3] for each hyperparameter and [10**-10, 10**-1] for variance in noise (if applicable).
        :param optimizations_number: Number of times MLE optimization problem is to be solved with a random starting
         point. Default: 1.
        :param normalize: Boolean flag used in case data normalization is required.
        :param optimizer: A class object of 'MinimizeOptimizer' or 'FminCobyla' from UQpy.utilities module. If optimizer
         is not defined, this class will not solve MLE problem and predictions will be based on hyperparameters provided
         as input.
         Default: None.
        :param noise: Boolean flag used in case of noisy training data.
         Default: False
        :param random_state: Random seed used to initialize the pseudo-random number generator. If an integer is
         provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise, the
         object itself can be passed directly.
        """
        self.regression_model = regression_model
        self.kernel = kernel
        self.hyperparameters = np.array(hyperparameters)
        self.bounds = bounds
        self.optimizer = optimizer
        self.optimizations_number = optimizations_number
        self.optimize_constraints = optimize_constraints
        self.normalize = normalize
        self.noise = noise
        self.logger = logging.getLogger(__name__)
        self.random_state = random_state

        # Variables are used outside the __init__
        self.samples = None
        self.values = None
        self.sample_mean, self.sample_std = None, None
        self.value_mean, self.value_std = None, None
        self.rmodel, self.cmodel = None, None
        self.beta = None
        """Regression coefficients."""
        self.gamma = None
        self.err_var = None
        """Variance of the Gaussian random process."""
        self.F_dash = None
        self.C_inv = None
        """Inverse Cholesky decomposition of the correlation matrix."""
        self.G = None
        self.F, self.K = None, None
        self.cc, self.alpha_ = None, None
        self.mu = 0

        if bounds is None:
            if isinstance(self.optimizer, type(None)) or isinstance(self.optimizer._bounds, type(None)):
                self._define_bounds()
            else:
                self.bounds = self.optimizer._bounds

        self.jac = False
        self.random_state = process_random_state(random_state)

    def _define_bounds(self):
        bounds_ = [[10 ** -3, 10 ** 3]] * self.hyperparameters.shape[0]
        if self.noise:
            bounds_[-1] = [10 ** -10, 10 ** -1]
        self.bounds = bounds_

    def fit(
            self,
            samples,
            values,
            optimizations_number=None,
            hyperparameters=None,
    ):
        """
        Fit the surrogate model using the training samples and the corresponding model values.

        The user can run this method multiple time after initiating the :class:`.GaussianProcessRegressor` class object.

        This method updates the samples and parameters of the :class:`.GaussianProcessRegressor` object. This method
        uses `hyperparameters` from previous run as the starting point for MLE problem unless user provides a new
        starting point.

        :param samples: `ndarray` containing the training points.
        :param values: `ndarray` containing the model evaluations at the training points.
        :param optimizations_number: number of optimization iterations
        :param hyperparameters: List or array of initial values for the kernel hyperparameters/scale parameters.

        The :meth:`fit` method has no returns, although it creates the `beta`, `err_var` and `C_inv` attributes of the
        :class:`.GaussianProcessRegressor` class.
        """
        self.logger.info("UQpy: Running gpr.fit")

        if optimizations_number is not None:
            self.optimizations_number = optimizations_number
        if hyperparameters is not None:
            self.hyperparameters = np.array(hyperparameters)
        self.samples = np.array(samples)

        # Number of samples and dimensions of samples and values
        nsamples, input_dim = self.samples.shape
        output_dim = int(np.size(values) / nsamples)

        # Verify the length of hyperparameters and input dimension
        if self.noise:
            if self.hyperparameters.shape[0] != input_dim + 2:
                raise RuntimeError("UQpy: The length/shape of attribute 'hyperparameter' and input dimension are not "
                                   "consistent.")
        elif self.hyperparameters.shape[0] != input_dim + 1:
            raise RuntimeError("UQpy: The length/shape of attribute 'hyperparameter' and input dimension are not "
                               "consistent.")

        self.values = np.array(values).reshape(nsamples, output_dim)

        # Normalizing the data
        if self.normalize:
            self.sample_mean, self.sample_std = np.mean(self.samples, 0), np.std(self.samples, 0)
            self.value_mean, self.value_std = np.mean(self.values, 0), np.std(self.values, 0)
            s_ = (self.samples - self.sample_mean) / self.sample_std
            y_ = (self.values - self.value_mean) / self.value_std
        else:
            s_ = self.samples
            y_ = self.values

        if self.regression_model is not None:
            self.F = self.regression_model.r(s_)

        # Maximum Likelihood Estimation : Solving optimization problem to calculate hyperparameters
        if self.optimizer is not None:
            # To ensure that cholesky of covariance matrix is computed again during optimization
            self.alpha_ = None
            lb = [np.log10(xy[0]) for xy in self.bounds]
            ub = [np.log10(xy[1]) for xy in self.bounds]

            starting_point = np.random.uniform(low=lb, high=ub, size=(self.optimizations_number, len(self.bounds)))
            starting_point[0, :] = np.log10(self.hyperparameters)

            if self.optimize_constraints is not None:
                cons = self.optimize_constraints.define_arguments(self.samples, self.values, self.predict)
                self.optimizer.apply_constraints(constraints=cons)
                self.optimizer.apply_constraints_argument(self.optimize_constraints.constraint_args)
            else:
                log_bounds = [[np.log10(xy[0]), np.log10(xy[1])] for xy in self.bounds]
                self.optimizer.update_bounds(bounds=log_bounds)

            minimizer = np.zeros([self.optimizations_number, len(self.bounds)])
            fun_value = np.zeros([self.optimizations_number, 1])

            for i__ in range(self.optimizations_number):
                p_ = self.optimizer.optimize(function=GaussianProcessRegression.log_likelihood,
                                             initial_guess=starting_point[i__, :],
                                             args=(self.kernel, s_, y_, self.noise, self.F),
                                             jac=self.jac)

                if isinstance(p_, np.ndarray):
                    minimizer[i__, :] = p_
                    fun_value[i__, 0] = GaussianProcessRegression.log_likelihood(p_, self.kernel, s_, y_,
                                                                                 self.noise, self.F)
                else:
                    minimizer[i__, :] = p_.x
                    fun_value[i__, 0] = p_.fun

            if min(fun_value) == np.inf:
                raise NotImplementedError("Maximum likelihood estimator failed: Choose different starting point or "
                                          "increase nopt")
            t = np.argmin(fun_value)
            self.hyperparameters = 10**minimizer[t, :]

        # Updated Correlation matrix corresponding to MLE estimates of hyperparameters
        if self.noise:
            self.K = self.kernel.c(x=s_, s=s_, params=self.hyperparameters[:-1]) + \
                     np.eye(nsamples)*(self.hyperparameters[-1])**2
        else:
            self.K = self.kernel.c(x=s_, s=s_, params=self.hyperparameters)

        self.cc = cholesky(self.K + 1e-10 * np.eye(nsamples), lower=True)
        self.alpha_ = cho_solve((self.cc, True), y_)

        if self.regression_model is not None:
            # Compute the regression coefficient (solving this linear equation: F * beta = Y)
            # Eq: 3.8, DACE
            f_dash = np.linalg.solve(self.cc, self.F)
            y_dash = np.linalg.solve(self.cc, y_)
            q_, g_ = np.linalg.qr(f_dash)  # Eq: 3.11, DACE
            # Check if F is a full rank matrix
            if np.linalg.matrix_rank(g_) != min(np.size(self.F, 0), np.size(self.F, 1)):
                raise NotImplementedError("Chosen regression functions are not sufficiently linearly independent")
            # Design parameters (beta: regression coefficient)
            self.beta = np.linalg.solve(g_, np.matmul(np.transpose(q_), y_dash))
            self.mu = np.einsum("ij,jk->ik", self.F, self.beta)
            self.alpha_ = cho_solve((self.cc, True), y_ - self.mu)

        self.logger.info("UQpy: gpr fit complete.")

    def predict(self, points, return_std: bool = False, hyperparameters: list = None):
        """
        Predict the model response at new points.

        This method evaluates the regression and correlation model at new sample points. Then, it predicts the function
        value and standard deviation.

        :param points: Points at which to predict the model response.
        :param return_std: Indicator to estimate standard deviation.
        :param hyperparameters: Hyperparameters for correlation model.
        :return: Predicted values at the new points, Standard deviation of predicted values at the new points
        """
        x_ = np.atleast_2d(points)
        if self.normalize:
            x_ = (x_ - self.sample_mean) / self.sample_std
            s_ = (self.samples - self.sample_mean) / self.sample_std
            y_ = (self.values - self.value_mean) / self.value_std
        else:
            s_ = self.samples
            y_ = self.values

        if hyperparameters is None:
            hyperparameters = self.hyperparameters

        noise_std = 0
        kernelparameters = hyperparameters
        if self.noise:
            kernelparameters = hyperparameters[:-1]
            noise_std = hyperparameters[-1]

        if hyperparameters is not None:
            # This is used for MLE constraints, if constraints call 'predict' method.
            K = self.kernel.c(x=s_, s=s_, params=kernelparameters) + \
                np.eye(self.samples.shape[0]) * (noise_std ** 2)
            cc = np.linalg.cholesky(K + 1e-10 * np.eye(self.samples.shape[0]))
            mu = 0
            if self.regression_model is not None:
                f_dash = np.linalg.solve(cc, self.F)
                y_dash = np.linalg.solve(cc, y_)
                q_, g_ = np.linalg.qr(f_dash)  # Eq: 3.11, DACE
                # Check if F is a full rank matrix
                if np.linalg.matrix_rank(g_) != min(np.size(self.F, 0), np.size(self.F, 1)):
                    raise NotImplementedError("Chosen regression functions are not sufficiently linearly independent")
                # Design parameters (beta: regression coefficient)
                beta = np.linalg.solve(g_, np.matmul(np.transpose(q_), y_dash))
                mu = np.einsum("ij,jk->ik", self.F, beta)
            alpha_ = cho_solve((cc, True), y_-mu)
        else:
            cc, alpha_ = self.cc, self.alpha_

        mu1 = 0
        if self.regression_model is not None:
            fx = self.regression_model.r(x_)
            if self.beta is None:
                mu1 = np.einsum("ij,jk->ik", fx, beta)
            else:
                mu1 = np.einsum("ij,jk->ik", fx, self.beta)

        k = self.kernel.c(x=x_, s=s_, params=kernelparameters)
        y = mu1 + k @ alpha_
        if self.normalize:
            y = self.value_mean + y * self.value_std
        if x_.shape[1] == 1:
            y = y.flatten()

        if return_std:
            k1 = self.kernel.c(x=x_, s=x_, params=kernelparameters)
            var = (k1 - k @ cho_solve((cc, True), k.T)).diagonal()
            mse = np.sqrt(var)
            if self.normalize:
                mse = self.value_std * mse
            if x_.shape[1] == 1:
                mse = mse.flatten()
            return y, mse
        else:
            return y

    @staticmethod
    def log_likelihood(p0, k_, s, y, ind_noise, fx_):
        """

        :param p0: An 1-D numpy array of hyperparameters, that are identified through MLE. The last two elements are
                   process variance and data noise, rest of the elements are length scale along each input dimension.
        :param k_: Kernel
        :param s: Input training data
        :param y: Output training data
        :param ind_noise: Boolean flag to indicate the noisy output
        :param fx_: Basis function evaluated at training points
        :return:
        """
        # Return the log-likelihood function and it's gradient. Gradient is calculated using Central Difference
        m = s.shape[0]

        if ind_noise:
            k__ = k_.c(x=s, s=s, params=10 ** p0[:-1]) + np.eye(m) * (10 ** p0[-1]) ** 2
        else:
            k__ = k_.c(x=s, s=s, params=10 ** p0)
        cc = cholesky(k__ + 1e-10 * np.eye(m), lower=True)

        mu = 0
        if fx_ is not None:
            f_dash = np.linalg.solve(cc, fx_)
            y_dash = np.linalg.solve(cc, y)
            q_, g_ = np.linalg.qr(f_dash)  # Eq: 3.11, DACE
            # Check if F is a full rank matrix
            if np.linalg.matrix_rank(g_) != min(np.size(fx_, 0), np.size(fx_, 1)):
                raise NotImplementedError("Chosen regression functions are not sufficiently linearly independent")
            # Design parameters (beta: regression coefficient)
            beta = np.linalg.solve(g_, np.matmul(np.transpose(q_), y_dash))
            mu = np.einsum("ij,jk->ik", fx_, beta)

        term1 = (y-mu).T @ (cho_solve((cc, True), y-mu))
        term2 = 2 * np.sum(np.log(np.abs(np.diag(cc))))

        return 0.5 * (term1 + term2 + m * np.log(2 * np.pi))[0, 0]


# src/UQpy/surrogates/gaussian_process/__init__.py

from UQpy.surrogates.gaussian_process.GaussianProcessRegression import GaussianProcessRegression

from UQpy.surrogates.gaussian_process.kernels import *
from UQpy.surrogates.gaussian_process.regression_models import *
from UQpy.surrogates.gaussian_process.constraints import *


# src/UQpy/surrogates/kriging/Kriging.py

import logging
from typing import Callable

import numpy as np
from scipy.linalg import cholesky
import scipy.stats as stats
from beartype import beartype

from UQpy.utilities import MinimizeOptimizer
from UQpy.utilities.Utilities import process_random_state
from UQpy.surrogates.baseclass.Surrogate import Surrogate
from UQpy.utilities.ValidationTypes import RandomStateType
from UQpy.surrogates.kriging.correlation_models.baseclass.Correlation import Correlation
from UQpy.surrogates.kriging.regression_models.baseclass.Regression import Regression


class Kriging(Surrogate):
    @beartype
    def __init__(
        self,
        regression_model: Regression,
        correlation_model: Correlation,
        correlation_model_parameters: list,
        optimizer,
        bounds: list = None,
        optimize: bool = True,
        optimizations_number: int = 1,
        normalize: bool = True,
        random_state: RandomStateType = None,
    ):
        """
        Κriging generates an Gaussian process regression-based surrogate model to predict the model output at new sample
        points.

        :param regression_model: `regression_model` specifies and evaluates the basis functions and their coefficients,
         which defines the trend of the model. Built-in options: :class:`Constant`, :class:`Linear`, :class:`Quadratic`
        :param correlation_model: `correlation_model` specifies and evaluates the correlation function.
         Built-in options: :class:`Exponential`, :class:`Gaussian`, :class:`Linear`, :class:`Spherical`,
         :class:`Cubic`, :class:`Spline`
        :param correlation_model_parameters: List or array of initial values for the correlation model
         hyperparameters/scale parameters.
        :param bounds: Bounds on the hyperparameters used to solve optimization problem to estimate maximum likelihood
         estimator. This should be a closed bound.
         Default: :math:`[0.001, 10^7]` for each hyperparameter.
        :param optimize: Indicator to solve MLE problem or not. If :any:'True' corr_model_params will be used as initial
         solution for optimization problem. Otherwise, correlation_model_parameters will be directly use as the
         hyperparamters.
         Default: :any:`True`.
        :param optimizations_number: Number of times MLE optimization problem is to be solved with a random starting
         point. Default: :math:`1`.
        :param normalize: Boolean flag used in case data normalization is required.
        :param optimizer: Object of the :class:`Optimizer` optimizer used during the Kriging surrogate.
         Default: :class:`.MinimizeOptimizer`.
        :param random_state: Random seed used to initialize the pseudo-random number generator. If an :any:`int` is
         provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise, the
         object itself can be passed directly.
        """
        self.regression_model = regression_model
        self.correlation_model = correlation_model
        self.correlation_model_parameters = np.array(correlation_model_parameters)
        self.bounds = bounds
        self.optimizer = optimizer
        self.optimizations_number = optimizations_number
        self.optimize = optimize
        self.normalize = normalize
        self.logger = logging.getLogger(__name__)
        self.random_state = random_state

        # Variables are used outside the __init__
        self.samples = None
        self.values = None
        self.sample_mean, self.sample_std = None, None
        self.value_mean, self.value_std = None, None
        self.rmodel, self.cmodel = None, None
        self.beta: list = None
        """Regression coefficients."""
        self.gamma = None
        self.err_var: float = None
        """Variance of the Gaussian random process."""
        self.F_dash = None
        self.C_inv = None
        self.G = None
        self.F, self.R = None, None

        if isinstance(self.optimizer, str):
            raise ValueError("The optimization function provided a input parameter cannot be None.")

        if optimizer._bounds is None:
            optimizer.update_bounds([[0.001, 10 ** 7]] * self.correlation_model_parameters.shape[0])

        self.jac = optimizer.supports_jacobian()
        self.random_state = process_random_state(random_state)

    def fit(
        self,
        samples,
        values,
        optimizations_number: int = None,
        correlation_model_parameters: list = None,
    ):
        """
        Fit the surrogate model using the training samples and the corresponding model values.

        The user can run this method multiple time after initiating the :class:`.Kriging` class object.

        This method updates the samples and parameters of the :class:`.Kriging` object. This method uses
        `correlation_model_parameters` from previous run as the starting point for MLE problem unless user provides a
        new starting point.

        :param samples: :class:`numpy.ndarray` containing the training points.
        :param values: :class:`numpy.ndarray` containing the model evaluations at the training points.
        :param optimizations_number: number of optimization iterations
        :param correlation_model_parameters: List or array of initial values for the correlation model
         hyperparameters/scale parameters.

        The :meth:`fit` method has no returns, although it creates the :py:attr:`beta`, :py:attr:`err_var` and
        :py:attr:`C_inv` attributes of the :class:`.Kriging` class.
        """
        self.logger.info("UQpy: Running kriging.fit")

        if optimizations_number is not None:
            self.optimizations_number = optimizations_number
        if correlation_model_parameters is not None:
            self.correlation_model_parameters = np.array(correlation_model_parameters)
        self.samples = np.array(samples)

        # Number of samples and dimensions of samples and values
        nsamples, input_dim = self.samples.shape
        output_dim = int(np.size(values) / nsamples)

        self.values = np.array(values).reshape(nsamples, output_dim)

        # Normalizing the data
        if self.normalize:
            self.sample_mean, self.sample_std = np.mean(self.samples, 0), np.std(self.samples, 0)
            self.value_mean, self.value_std = np.mean(self.values, 0), np.std(self.values, 0)
            s_ = (self.samples - self.sample_mean) / self.sample_std
            y_ = (self.values - self.value_mean) / self.value_std
        else:
            s_ = self.samples
            y_ = self.values

        self.F, jf_ = self.regression_model.r(s_)

        # Maximum Likelihood Estimation : Solving optimization problem to calculate hyperparameters
        if self.optimize:
            starting_point = self.correlation_model_parameters

            minimizer, fun_value = np.zeros([self.optimizations_number, input_dim]),\
                                   np.zeros([self.optimizations_number, 1])
            for i__ in range(self.optimizations_number):
                p_ = self.optimizer.optimize(function=Kriging.log_likelihood,
                                             initial_guess=starting_point,
                                             args=(self.correlation_model, s_, self.F, y_, self.jac),
                                             jac=self.jac)
                print(p_.success)
                # print(self.kwargs_optimizer)
                minimizer[i__, :] = p_.x
                fun_value[i__, 0] = p_.fun
                # Generating new starting points using log-uniform distribution
                if i__ != self.optimizations_number - 1:
                    starting_point = stats.reciprocal.rvs([j[0] for j in self.optimizer._bounds],
                                                          [j[1] for j in self.optimizer._bounds], 1,
                                                          random_state=self.random_state)
                    print(starting_point)

            if min(fun_value) == np.inf:
                raise NotImplementedError("Maximum likelihood estimator failed: Choose different starting point or "
                                          "increase nopt")
            t = np.argmin(fun_value)
            self.correlation_model_parameters = minimizer[t, :]

        # Updated Correlation matrix corresponding to MLE estimates of hyperparameters
        self.R = self.correlation_model.c(x=s_, s=s_, params=self.correlation_model_parameters)

        self.beta, self.gamma, tmp = self._compute_additional_parameters(self.R)
        self.C_inv, self.F_dash, self.G, self.err_var = tmp[1], tmp[3], tmp[2], tmp[5]

        self.logger.info("UQpy: kriging fit complete.")

    def _compute_additional_parameters(self, correlation_matrix):
        if self.normalize:
            y_ = (self.values - self.value_mean) / self.value_std
        else:
            y_ = self.values
        # Compute the regression coefficient (solving this linear equation: F * beta = Y)
        # Eq: 3.8, DACE
        c = cholesky(correlation_matrix + (10 + self.samples.shape[0]) * 2 ** (-52) * np.eye(self.samples.shape[0]),
                     lower=True, check_finite=False)
        c_inv = np.linalg.inv(c)
        f_dash = np.linalg.solve(c, self.F)
        y_dash = np.linalg.solve(c, y_)
        q_, g_ = np.linalg.qr(f_dash)  # Eq: 3.11, DACE
        # Check if F is a full rank matrix
        if np.linalg.matrix_rank(g_) != min(np.size(self.F, 0), np.size(self.F, 1)):
            raise NotImplementedError("Chosen regression functions are not sufficiently linearly independent")
        # Design parameters (beta: regression coefficient)
        beta = np.linalg.solve(g_, np.matmul(np.transpose(q_), y_dash))

        # Design parameter (R * gamma = Y - F * beta = residual)
        gamma = np.linalg.solve(c.T, (y_dash - np.matmul(f_dash, beta)))

        # Computing the process variance (Eq: 3.13, DACE)
        err_var = np.zeros(self.values.shape[1])
        for i in range(self.values.shape[1]):
            err_var[i] = (1 / self.samples.shape[0]) * (np.linalg.norm(y_dash[:, i] -
                                                                       np.matmul(f_dash, beta[:, i])) ** 2)

        return beta, gamma, (c, c_inv, g_, f_dash, y_dash, err_var)

    def predict(self, points: np.ndarray, return_std: bool = False, correlation_model_parameters: list = None):
        """
        Predict the model response at new points.

        This method evaluates the regression and correlation model at new sample points. Then, it predicts the function
        value and standard deviation.

        :param points: Points at which to predict the model response.
        :param  return_std: Indicator to estimate standard deviation.
        :param correlation_model_parameters: Hyperparameters for correlation model.
        :return: Predicted values at the new points, Standard deviation of predicted values at the new points
        """
        x_ = np.atleast_2d(points)
        if self.normalize:
            x_ = (x_ - self.sample_mean) / self.sample_std
            s_ = (self.samples - self.sample_mean) / self.sample_std
        else:
            s_ = self.samples
        fx, jf = self.regression_model.r(x_)
        if correlation_model_parameters is None:
            correlation_model_parameters = self.correlation_model_parameters
        rx = self.correlation_model.c(
            x=x_, s=s_, params=correlation_model_parameters
        )
        if correlation_model_parameters is None:
            beta, gamma = self.beta, self.gamma
            c_inv, f_dash, g_, err_var = self.C_inv, self.F_dash, self.G, self.err_var
        else:
            beta, gamma, tmp = self._compute_additional_parameters(
                self.correlation_model.c(x=s_, s=s_, params=correlation_model_parameters))
            c_inv, f_dash, g_, err_var = tmp[1], tmp[3], tmp[2], tmp[5]
        y = np.einsum("ij,jk->ik", fx, beta) + np.einsum(
            "ij,jk->ik", rx, gamma
        )
        if self.normalize:
            y = self.value_mean + y * self.value_std
        if x_.shape[1] == 1:
            y = y.flatten()
        if return_std:
            r_dash = np.matmul(c_inv, rx.T)
            u = np.matmul(f_dash.T, r_dash) - fx.T
            norm1 = np.linalg.norm(r_dash, 2, 0)
            norm2 = np.linalg.norm(np.linalg.solve(g_, u), 2, 0)
            mse = np.sqrt(err_var * np.atleast_2d(1 + norm2 - norm1).T)
            if self.normalize:
                mse = self.value_std * mse
            if x_.shape[1] == 1:
                mse = mse.flatten()
            return y, mse
        else:
            return y

    def jacobian(self, points: np.ndarray):
        """
        Predict the gradient of the model at new points.

        This method evaluates the regression and correlation model at new sample point. Then, it predicts the gradient
        using the regression coefficients and the training second_order_tensor.

        :param points: Points at which to evaluate the gradient.
        :return: Gradient of the surrogate model evaluated at the new points.
        """
        x_ = np.atleast_2d(points)
        if self.normalize:
            x_ = (x_ - self.sample_mean) / self.sample_std
            s_ = (self.samples - self.sample_mean) / self.sample_std
        else:
            s_ = self.samples

        fx, jf = self.regression_model.r(x_)
        rx, drdx = self.correlation_model.c(
            x=x_, s=s_, params=self.correlation_model_parameters, dx=True
        )
        y_grad = np.einsum("ikj,jm->ik", jf, self.beta) + np.einsum(
            "ijk,jm->ki", drdx.T, self.gamma
        )
        if self.normalize:
            y_grad = y_grad * self.value_std / self.sample_std
        if x_.shape[1] == 1:
            y_grad = y_grad.flatten()
        return y_grad

    @staticmethod
    def log_likelihood(p0, cm, s, f, y, return_grad):
        # Return the log-likelihood function and it's gradient. Gradient is calculate using Central Difference
        m = s.shape[0]
        n = s.shape[1]
        r__, dr_ = cm.c(x=s, s=s, params=p0, dt=True)
        try:
            cc = cholesky(r__ + 2 ** (-52) * np.eye(m), lower=True)
        except np.linalg.LinAlgError:
            return np.inf, np.zeros(n)

        # Product of diagonal terms is negligible sometimes, even when cc exists.
        if np.prod(np.diagonal(cc)) == 0:
            return np.inf, np.zeros(n)

        cc_inv = np.linalg.inv(cc)
        r_inv = np.matmul(cc_inv.T, cc_inv)
        f__ = cc_inv.dot(f)
        y__ = cc_inv.dot(y)

        q__, g__ = np.linalg.qr(f__)  # Eq: 3.11, DACE

        # Check if F is a full rank matrix
        if np.linalg.matrix_rank(g__) != min(np.size(f__, 0), np.size(f__, 1)):
            raise NotImplementedError(
                "Chosen regression functions are not sufficiently linearly independent"
            )

        # Design parameters
        beta_ = np.linalg.solve(g__, np.matmul(np.transpose(q__), y__))

        # Computing the process variance (Eq: 3.13, DACE)
        sigma_ = np.zeros(y.shape[1])

        ll = 0
        for out_dim in range(y.shape[1]):
            sigma_[out_dim] = (1 / m) * (
                    np.linalg.norm(y__[:, out_dim] - np.matmul(f__, beta_[:, out_dim])) ** 2)
            # Objective function:= log(det(sigma**2 * R)) + constant
            ll = (ll + ( np.log(np.linalg.det(sigma_[out_dim] * r__)) + m * (np.log(2 * np.pi) + 1)) / 2)

        # Gradient of loglikelihood
        # Reference: C. E. Rasmussen & C. K. I. Williams, Gaussian Processes for Machine Learning, the MIT Press,
        # 2006, ISBN 026218253X. (Page 114, Eq.(5.9))
        residual = y - np.matmul(f, beta_)
        gamma = np.matmul(r_inv, residual)
        grad_mle = np.zeros(n)
        for in_dim in range(n):
            r_inv_derivative = np.matmul(r_inv, np.matmul(dr_[:, :, in_dim], r_inv))
            tmp = np.matmul(residual.T, np.matmul(r_inv_derivative, residual))
            for out_dim in range(y.shape[1]):
                alpha = gamma / sigma_[out_dim]
                tmp1 = np.matmul(alpha, alpha.T) - r_inv / sigma_[out_dim]
                cov_der = sigma_[out_dim] * dr_[:, :, in_dim] + tmp * r__ / m
                grad_mle[in_dim] = grad_mle[in_dim] - 0.5 * np.trace(
                    np.matmul(tmp1, cov_der)
                )

        if return_grad:
            return ll, grad_mle
        else:
            return ll


# src/UQpy/surrogates/kriging/__init__.py

from UQpy.surrogates.kriging.Kriging import Kriging

from UQpy.surrogates.kriging.regression_models import *
from UQpy.surrogates.kriging.correlation_models import *


# src/UQpy/surrogates/polynomial_chaos/PolynomialChaosExpansion.py

import logging

import numpy as np
from beartype import beartype

from UQpy.utilities.ValidationTypes import NumpyFloatArray
from UQpy.surrogates.baseclass.Surrogate import Surrogate
from UQpy.surrogates.polynomial_chaos.regressions.baseclass.Regression import Regression
from UQpy.surrogates.polynomial_chaos.polynomials.TotalDegreeBasis import PolynomialBasis
from UQpy.distributions import Uniform, Normal
from UQpy.surrogates.polynomial_chaos.polynomials import Legendre,Hermite

class PolynomialChaosExpansion(Surrogate):

    @beartype
    def __init__(self, polynomial_basis: PolynomialBasis, regression_method: Regression):
        """
        Constructs a surrogate model based on the Polynomial Chaos Expansion (polynomial_chaos) method.

        :param regression_method: object for the method used for the calculation of the polynomial_chaos coefficients.
        """
        self.polynomial_basis: PolynomialBasis = polynomial_basis
        """Contains the 1D or ND chaos polynomials that form the PCE basis"""
        self.multi_index_set: NumpyFloatArray = polynomial_basis.multi_index_set
        "Multi-index-set"
        self.regression_method = regression_method
        self.logger = logging.getLogger(__name__)
        self.coefficients: NumpyFloatArray = None
        """Polynomial Chaos Expansion Coefficient"""
        self.bias: NumpyFloatArray = None
        """Bias term in case LASSO or ridge regression are employed to estimate 
        the PCE coefficients"""
        self.outputs_number: int = None
        """Dimensions of the quantities of interest"""
        self.design_matrix: NumpyFloatArray = None
        """Matrix containing the evaluations of the PCE basis on the experimental 
        design that has been used to fit the PCE coefficients"""

        self.experimental_design_input: NumpyFloatArray = None
        """Realizations of the random parameter in the experimental design that 
        has been used to fit the PCE coefficients"""
        self.experimental_design_output: NumpyFloatArray = None
        """Model outputs for the random parameter realizations of the 
        experimental design that has been used to fit the PCE coefficients"""

    @property
    def polynomials_number(self):
        return len(self.polynomial_basis.polynomials)

    @property
    def inputs_number(self):
        return self.polynomial_basis.inputs_number

    def fit(self, x: np.ndarray, y: np.ndarray):
        """
        Fit the surrogate model using the training samples and the corresponding model values. This method calls the
        :py:meth:`run` method of the input method class.

        :param x: containing the training points.
        :param y: containing the model evaluations at the training points.

        The :meth:`fit` method has no returns and it creates an :class:`numpy.ndarray` with the
        polynomial_chaos coefficients.
        """
        self.experimental_design_input = x
        self.experimental_design_output = y
        self.design_matrix = self.polynomial_basis.evaluate_basis(x)
        self.logger.info("UQpy: Running polynomial_chaos.fit")
        self.coefficients, self.bias, self.outputs_number = self.regression_method.run(x, y, self.design_matrix)
        self.logger.info("UQpy: polynomial_chaos fit complete.")

    def predict(self, points: np.ndarray, **kwargs: dict):
        """
        Predict the model response at new points.
        This method evaluates the polynomial_chaos model at new sample points.

        :param points: Points at which to predict the model response.
        :return: Predicted values at the new points.
        """
        a = self.polynomial_basis.evaluate_basis(points)
        y = a.dot(self.coefficients)
        if self.bias is not None:
            y = y + self.bias
        return y

    def leaveoneout_error(self):
        """
        Returns the cross validation error (leave-one-out) based on experimental design.

        The :meth:`.PolynomialChaosExpansion.leaveoneout_error` method can be used to estimate the accuracy of the PCE
        predictor without additional simulations. Leave-one-out error :math:`Q^2` is calculated from differences between the original mathematical model and approximation :math:`\Delta_i= \mathcal{M} (x^{(i)}) - \mathcal{M}^{PCE}(x^{(i)})` as follows:
        
        .. math:: Q^2 = \mathbb{E}[(\Delta_i/(1-h_i))^2]/\sigma_Y^2
        

        where :math:`\sigma_Y^2` is a variance of model response and :math:`h_i` represents the :math:`i` th diagonal term of matrix :math:`\mathbf{H}= \Psi ( \Psi^T \Psi )^{-1} \Psi^T` obtained from design matrix :math:`\Psi` without additional simulations  :cite:`BLATMANLARS`.

        :return: Cross validation error of experimental design.
        """
        
        x=self.experimental_design_input
        y=self.experimental_design_output
        
        if y.ndim == 1 or y.shape[1] == 1:
            y = y.reshape(-1, 1)
            
        n_samples = x.shape[0]
        mu_yval = (1 / n_samples) * np.sum(y, axis=0)
        y_val = self.predict(x, )    
        polynomialbasis= self.design_matrix
        
        H = np.dot(polynomialbasis, np.linalg.pinv(np.dot(polynomialbasis.T, polynomialbasis)))
        H *= polynomialbasis
        Hdiag = np.sum(H, axis=1).reshape(-1,1)


        eps_val=((n_samples - 1) / n_samples *np.sum(((y - y_val)/(1 - Hdiag))**2,axis=0)) / (np.sum((y - mu_yval) ** 2, axis=0))
        if y.ndim == 1 or y.shape[1] == 1:
            eps_val = float(eps_val)

        return np.round(eps_val, 7)
    
    def validation_error(self, x: np.ndarray, y: np.ndarray):
        """
        Returns the validation error.

        The :meth:`.PolynomialChaosExpansion.validation_error` method can be used to estimate the accuracy of the PCE
        predictor.
        Here, we compute the generalization error  in the form of the relative mean squared error normalized by the
        model variance.
        The user must create an independent validation dataset :math:`[x_{val}, y_{val} = M(x_{val})]`
        (i.e. a set of inputs and outputs of the computational model).

        In case where the computational model is very expensive, the use of an alternative error measure is recommended,
        for example the cross-validation error which partitions the existing training dataset into subsets and computes
        the error as the average of the individual errors of each subset.

        :param x: :class:`numpy.ndarray` containing the samples of the validation dataset.
        :param y: :class:`numpy.ndarray` containing model evaluations for the validation dataset.
        :return: Validation error.
        """

        if y.ndim == 1 or y.shape[1] == 1:
            y = y.reshape(-1, 1)

        y_val = self.predict(x, )

        n_samples = x.shape[0]
        mu_yval = (1 / n_samples) * np.sum(y, axis=0)
        eps_val = ((n_samples - 1) / n_samples
                   * ((np.sum((y - y_val) ** 2, axis=0))
                      / (np.sum((y - mu_yval) ** 2, axis=0))))

        if y.ndim == 1 or y.shape[1] == 1:
            eps_val = float(eps_val)

        return np.round(eps_val, 7)

    def get_moments(self, higher: bool = False):
        """
        Returns the first four moments of the polynomial_chaos surrogate which are directly
        estimated from the polynomial_chaos coefficients.

        The :meth:`.PolynomialChaosExpansion.get_moments` method can be used for the calculation of the first two moments of the PCE model directly from the PCE coefficients. This is possible due to the orthonormality of the polynomial basis.

        The first moment (mean value) is calculated as

        .. math:: \mu_{PCE} = \mathbb{E} [ \mathcal{M}^{PCE}(x)] = y_{0}

        where :math:`y_{0}` is the first PCE coefficient associated with the constant term.

        The second moment (variance) is calculated as

        .. math:: \sigma^{2}_{PCE} = \mathbb{E} [( \mathcal{M}^{PCE}(x) - \mu_{PCE} )^{2} ] = \sum_{i=1}^{p} y_{i}

        where :math:`p` is the number of polynomials (first PCE coefficient is excluded).
        
        The third moment (skewness) and the fourth moment (kurtosis) are generally obtained from third and fourth order products obtained by integration, which are extremely computationally demanding. Therefore here we use analytical solution of standard linearization problem for Hermite (based on normalized Feldheim's formula ) and Legendre polynomials (based on normalized Neumann-Adams formula), though it might be still computationally demanding for high number of input random variables.
        
        :param higher: True corresponds to calculation of skewness and kurtosis (computationaly expensive for large basis set).
        :return: Returns the mean and variance.
        """
        
        if self.bias is not None:
            mean = self.coefficients[0, :] + np.squeeze(self.bias)
        else:
            mean = self.coefficients[0, :]

        variance = np.sum(self.coefficients[1:] ** 2, axis=0)

        if self.coefficients.ndim == 1 or self.coefficients.shape[1] == 1:
            variance = float(variance)
            mean = float(mean)

        if higher==False:
            return np.round(mean, 4), np.round(variance, 4)
        
        else:
            multindex=self.multi_index_set
            P,inputs_number=multindex.shape
            
            if inputs_number==1:
                marginals=[self.polynomial_basis.distributions]
                
            else:
                marginals=self.polynomial_basis.distributions.marginals
            
            
            skewness=np.zeros(self.outputs_number)
            kurtosis=np.zeros(self.outputs_number)
            
            for ii in range (0,self.outputs_number):
                
                Beta=self.coefficients[:, ii]
                third_moment=0
                fourth_moment=0
                
                indices=np.array(np.meshgrid(range(1,P),range(1,P),range(1,P),range(1,P))).T.reshape(-1,4)
                i=0
                for index in indices:
                    tripleproduct_ND=1
                    quadproduct_ND=1
                    
                    
                    for m in range (0,inputs_number):
   
                        if i<(P-1)**3:
        
                            if type(marginals[m])==Normal:
                                tripleproduct_1D=Hermite.hermite_triple_product(multindex[index[0],m],multindex[index[1],m],multindex[index[2],m])
                            
                            if type(marginals[m])==Uniform:   
                                tripleproduct_1D=Legendre.legendre_triple_product(multindex[index[0],m],multindex[index[1],m],multindex[index[2],m])
                            
                            tripleproduct_ND=tripleproduct_ND*tripleproduct_1D
                        
                        else:
                            tripleproduct_ND=0
                        
                        quadproduct_1D=0
                        
                        for n in range (0,multindex[index[0],m]+multindex[index[1],m]+1):
                            
                            if type(marginals[m])==Normal:
                                tripproduct1=Hermite.hermite_triple_product(multindex[index[0],m],multindex[index[1],m],n)
                                tripproduct2=Hermite.hermite_triple_product(multindex[index[2],m],multindex[index[3],m],n)
                           
                            if type(marginals[m])==Uniform: 
                                tripproduct1=Legendre.legendre_triple_product(multindex[index[0],m],multindex[index[1],m],n)
                                tripproduct2=Legendre.legendre_triple_product(multindex[index[2],m],multindex[index[3],m],n)

                            quadproduct_1D=quadproduct_1D+tripproduct1*tripproduct2

                        quadproduct_ND=quadproduct_ND*quadproduct_1D
                    
                    third_moment+=tripleproduct_ND*Beta[index[0]]*Beta[index[1]]*Beta[index[2]]
                    fourth_moment+=quadproduct_ND*Beta[index[0]]*Beta[index[1]]*Beta[index[2]]*Beta[index[3]]

                    i+=1

                skewness[ii]=1/(np.sqrt(variance)**3)*third_moment
                kurtosis[ii]=1/(variance**2)*fourth_moment
                
                if self.coefficients.ndim == 1 or self.coefficients.shape[1] == 1:
                    skewness = float(skewness[0])
                    kurtosis = float(kurtosis[0])

            return mean,variance,skewness,kurtosis


# src/UQpy/surrogates/polynomial_chaos/__init__.py

from UQpy.surrogates.polynomial_chaos.polynomials import *
from UQpy.surrogates.polynomial_chaos.regressions import *

from UQpy.surrogates.polynomial_chaos.PolynomialChaosExpansion import PolynomialChaosExpansion
from UQpy.surrogates.polynomial_chaos.polynomials.baseclass.Polynomials import Polynomials
from UQpy.surrogates.polynomial_chaos.regressions.LassoRegression import LassoRegression
from UQpy.surrogates.polynomial_chaos.regressions.LeastSquareRegression import LeastSquareRegression
from UQpy.surrogates.polynomial_chaos.regressions.RidgeRegression import RidgeRegression


# src/UQpy/surrogates/stochastic_reduced_order_models/SROM.py

import logging
from typing import Union
import numpy as np

from UQpy.distributions.baseclass import Distribution
from UQpy.utilities.ValidationTypes import NumpyFloatArray
from UQpy.distributions import DistributionContinuous1D


class SROM:
    def __init__(
        self,
        samples: Union[list, np.ndarray],
        target_distributions: list[Distribution],
        moments:list = None,
        weights_errors: list = None,
        weights_distribution: Union[list, np.ndarray] = None,
        weights_moments: list = None,
        weights_correlation: np.ndarray = None,
        properties: list = None,
        correlation: np.ndarray = None,
    ):
        """
        Stochastic Reduced Order Model(stochastic_reduced_order_models) provide a low-dimensional, discrete
        approximation of a given random quantity.

        :param samples: An array/list of samples corresponding to the points at which the
         stochastic_reduced_order_models is defined.
        :param target_distributions: A list of distribution objects for each random variable.
        :param moments: A list containing first and second order moment about origin of all random variables.
        :param weights_errors: A list of weights associated with the error in distribution, moments and correlation.
         This corresponds to a list of the values :math:`a_{u}` in the objective function above.
         Default: weights_errors = :math:`[1, 0.2, 0]`
        :param weights_distribution: A list or array containing weights associated with matching the distribution at
         each sample value.
         `weights_distribution` is an array or list of shape :code:`(m, d)` where each weight corresponds to the weight
         :math:`w_F(x_{k,i}; i)` assigned for matching the distribution of component :code:`i` at sample point
         :math:`x_{k,i}`.
         If `weights_distribution` is :code:`(1, d)`, it is assumed that each sample is equally weighted according
         to the corresponding weight for its distribution.
         Default: An array of shape :code:`(m, d)` with all elements equal to :math:`1`.
        :param weights_moments: An list or array containing weights associated with matching the moments about the
         origin for each component.
         `weights_moments` is a list or array of shape `(2, d), where each weight corresponds to the weight
         :math:`w_{\mu}(r; i)` assigned for matching the moment of order :math:`r = 1, 2` for component :code:`i`.
         If `weights_moments` is :code:`(1, d)`, it is assumed that moments of all order are equally weighted.
         Default: :code:`weights_moments = [[1/(moment[0][i]^2)], [1/(moment[1][i]^2)]] for i = 1, 2, ..., d`.
        :param weights_correlation: A list or array containing weights associated with matching the correlation of the
         random variables. `weights_correlation` is a list or array of shape :code:`(d, d)` where each weight
         corresponds to the weight :math:`w_R(i, j)` assigned for matching the correlation between component :code:`i`
         and component :code:`j`
         Default: :code:`weights_correlation = (d, d)` array with all elements equal to :math:`1`.
        :param properties: A list of booleans declaring the properties to be matched in the reduced order model.

         - :code:`properties[0] = True` matches the marginal distributions

         - :code:`properties[1] = True` matches the mean values

         - :code:`properties[2] = True` matches the mean square

         - :code:`properties[3] = True` matches the correlation

        :param correlation: Correlation matrix between random variables.

        """
        self.target_distributions = target_distributions
        self.correlation = correlation
        self.moments = moments

        self.weights_distribution = weights_distribution
        self.weights_moments = weights_moments
        self.weights_correlation = weights_correlation
        self.weights_errors = weights_errors

        self.properties = properties
        self.logger = logging.getLogger(__name__)
        self.sample_weights: NumpyFloatArray = None
        """The probability weights defining discrete approximation of continuous random variables."""

        if isinstance(samples, list):
            self.samples = np.array(samples)
            self.samples_number = self.samples.shape[0]
            self.dimension = self.samples.shape[1]
        elif isinstance(samples, np.ndarray):
            self.dimension = samples.shape[1]
            self.samples = samples
            self.samples_number = samples.shape[0]
        else:
            raise NotImplementedError("UQpy: 'samples' should be a list or numpy array")

        if self.target_distributions is None:
            raise NotImplementedError("UQpy: Target Distribution is not defined.")

        if isinstance(self.target_distributions, list):
            for i in range(len(self.target_distributions)):
                if not isinstance(
                    self.target_distributions[i], DistributionContinuous1D
                ):
                    raise TypeError(
                        "UQpy: A DistributionContinuous1D object must be provided."
                    )

        if self.properties is not None:
            self.run()
        else:
            self.logger.info(
                "UQpy: No properties list provided, execute the stochastic_reduced_order_models by calling"
                " run method and specifying a properties list"
            )

    def run(
        self,
        weights_errors: list = None,
        weights_distribution: list = None,
        weights_moments: list = None,
        weights_correlation: list = None,
        properties: list = None,
    ):
        """
        Execute the stochastic reduced order model in the :class:`.SROM` class.

        The :meth:`run` method is the function that computes the probability weights corresponding to the sample. If
        `properties` is provided, the :meth:`run` method is automatically called when the :class:`.SROM`
        object is defined. The user may also call the :meth:`run` method directly to generate samples. The :meth:`run`
        method of the :class:`.SROM` class can be invoked many times with different weights parameters and
        each time computed probability weights are overwritten.

        :param weights_errors: A list of weights associated with the error in distribution, moments and correlation.
         This corresponds to a list of the values :math:`a_{u}` in the objective function above.
         Default: weights_errors = :math:`[1, 0.2, 0]`
        :param weights_distribution: A list or array containing weights associated with matching the distribution at
         each sample value.
         `weights_distribution` is an array or list of shape :code:`(m, d)` where each weight corresponds to the weight
         :math:`w_F(x_{k,i}; i)` assigned for matching the distribution of component :code:`i` at sample point
         :math:`x_{k,i}`.
         If `weights_distribution` is :code:`(1, d)`, it is assumed that each sample is equally weighted according
         to the corresponding weight for its distribution.
         Default: An array of shape :code:`(m, d)` with all elements equal to :math:`1`.
        :param weights_moments: An list or array containing weights associated with matching the moments about the
         origin for each component.
         `weights_moments` is a list or array of shape :code:`(2, d)`, where each weight corresponds to the weight
         :math:`w_{\mu}(r; i)` assigned for matching the moment of order :math:`r = 1, 2` for component :code:`i`.
         If `weights_moments` is :code:`(1, d)`, it is assumed that moments of all order are equally weighted.
         Default: :code:`weights_moments = [[1/(moment[0][i]^2)], [1/(moment[1][i]^2)]] for i = 1, 2, ..., d`.
        :param weights_correlation: A list or array containing weights associated with matching the correlation of the
         random variables.
         `weights_correlation` is a list or array of shape :code:`(d, d)` where each weight corresponds to the weight
         :math:`w_R(i, j)` assigned for matching the correlation between component :code:`i` and component :code:`j`
         Default: :code:`weights_correlation = (d, d)` array with all elements equal to :math:`1`.
        :param properties: A list of booleans declaring the properties to be matched in the reduced order model.

         - :code:`properties[0] = True` matches the marginal distributions
         - :code:`properties[1] = True` matches the mean values
         - :code:`properties[2] = True` matches the mean square
         - :code:`properties[3] = True` matches the correlation
        """
        from scipy import optimize

        self.weights_distribution = weights_distribution
        self.weights_moments = weights_moments
        self.weights_correlation = weights_correlation
        self.weights_errors = weights_errors
        self.properties = properties

        # Check properties to match
        if self.properties is None:
            self.properties = [True, True, True, False]

        self._init_srom()

        self.logger.info("UQpy: Performing stochastic_reduced_order_models...")

        def f(p0, samples, wd, wm, wc, mar, n, d, m, alpha, prop, correlation):
            e1 = 0.0
            e2 = 0.0
            e22 = 0.0
            e3 = 0.0
            com = np.append(samples, np.atleast_2d(p0).T, 1)
            for j in range(d):
                srt = com[np.argsort(com[:, j].flatten())]
                s = srt[:, j]
                a = srt[:, d]
                a0 = np.cumsum(a)
                marginal = mar[j].cdf

                if prop[0] is True:
                    for i in range(n):
                        e1 += wd[i, j] * (a0[i] - marginal(s[i])) ** 2

                if prop[1] is True:
                    e2 += wm[0, j] * (np.sum(p0 * samples[:, j]) - m[0, j]) ** 2

                if prop[2] is True:
                    e22 += (
                        wm[1, j]
                        * (
                            np.sum(np.array(p0) * (samples[:, j] * samples[:, j]))
                            - m[1, j]
                        )
                        ** 2
                    )

                if prop[3] is True:
                    for k in range(d):
                        if k > j:
                            r = (
                                correlation[j, k]
                                * np.sqrt(
                                    (m[1, j] - m[0, j] ** 2) * (m[1, k] - m[0, k] ** 2)
                                )
                                + m[0, j] * m[0, k]
                            )
                            e3 += (
                                wc[k, j]
                                * (np.sum(p0 * (samples[:, j] * samples[:, k])) - r)
                                ** 2
                            )

            return alpha[0] * e1 + alpha[1] * (e2 + e22) + alpha[2] * e3

        def constraint(x):
            return np.sum(x) - 1

        cons = {"type": "eq", "fun": constraint}

        p_ = optimize.minimize(
            f,
            np.zeros(self.samples_number),
            args=(
                self.samples,
                self.weights_distribution,
                self.weights_moments,
                self.weights_correlation,
                self.target_distributions,
                self.samples_number,
                self.dimension,
                self.moments,
                self.weights_errors,
                self.properties,
                self.correlation,
            ),
            constraints=cons,
            method="SLSQP",
            bounds=[[0, 1]] * self.samples_number,
        )

        self.sample_weights = p_.x
        """The probability weights defining discrete approximation of continuous random variables."""
        self.logger.info("UQpy: stochastic_reduced_order_models completed!")

    def _init_srom(self):
        if isinstance(self.moments, list):
            self.moments = np.array(self.moments)

        if isinstance(self.correlation, list):
            self.correlation = np.array(self.correlation)

        # Check moments and correlation
        if (
            self.properties[1] is True
            or self.properties[2] is True
            or self.properties[3] is True
        ):
            if self.moments is None:
                raise NotImplementedError("UQpy: 'moments' are required")
        # Both moments are required, if correlation property is required to be match
        if self.properties[3] is True:
            if self.moments.shape != (2, self.dimension):
                raise NotImplementedError("UQpy: Shape of 'moments' is not correct")
            if self.correlation is None:
                self.correlation = np.identity(self.dimension)
        # moments.shape[0] should be 1 or 2
        if self.moments.shape != (1, self.dimension) and self.moments.shape != (
            2,
            self.dimension,
        ):
            raise NotImplementedError("UQpy: Shape of 'moments' is not correct")
        # If both the moments are to be included in objective function, then moments.shape[0] should be 2
        if self.properties[1] is True and self.properties[2] is True:
            if self.moments.shape != (2, self.dimension):
                raise NotImplementedError("UQpy: Shape of 'moments' is not correct")
        # If only second order moment is to be included in objective function and moments.shape[0] is 1. Then
        # self.moments is converted shape = (2, self.dimension) where is second row contain second order moments.
        if self.properties[1] is False and self.properties[2] is True:
            if self.moments.shape == (1, self.dimension):
                temp = np.ones(shape=(1, self.dimension))
                self.moments = np.concatenate((temp, self.moments))

        # Check weights corresponding to errors
        if self.weights_errors is None:
            self.weights_errors = [1, 0.2, 0]
        elif isinstance(self.weights_errors, list):
            self.weights_errors = np.array(self.weights_errors).astype(np.float64)
        elif not isinstance(self.weights_errors, np.ndarray):
            raise NotImplementedError(
                "UQpy: weights_errors attribute should be a list or numpy array"
            )

        # Check weights corresponding to distribution
        if self.weights_distribution is None:
            self.weights_distribution = np.ones(
                shape=(self.samples.shape[0], self.dimension)
            )
        elif isinstance(self.weights_distribution, list):
            self.weights_distribution = np.array(self.weights_distribution)
        elif not isinstance(self.weights_distribution, np.ndarray):
            raise NotImplementedError(
                "UQpy: weights_distribution attribute should be a list or numpy array"
            )

        if self.weights_distribution.shape == (1, self.dimension):
            self.weights_distribution = self.weights_distribution * np.ones(
                shape=(self.samples.shape[0], self.dimension)
            )
        elif self.weights_distribution.shape != (self.samples.shape[0], self.dimension):
            raise NotImplementedError(
                "UQpy: Size of 'weights for distribution' is not correct"
            )

        # Check weights corresponding to moments and it's default list
        if self.weights_moments is None:
            self.weights_moments = np.reciprocal(np.square(self.moments))
        elif isinstance(self.weights_moments, list):
            self.weights_moments = np.array(self.weights_moments)
        elif not isinstance(self.weights_moments, np.ndarray):
            raise NotImplementedError(
                "UQpy: weights_moments attribute should be a list or numpy array"
            )

        if self.weights_moments.shape == (1, self.dimension):
            self.weights_moments = self.weights_moments * np.ones(
                shape=(2, self.dimension)
            )
        elif self.weights_moments.shape != (2, self.dimension):
            raise NotImplementedError(
                "UQpy: Size of 'weights for moments' is not correct"
            )

        # Check weights corresponding to correlation and it's default list
        if self.weights_correlation is None:
            self.weights_correlation = np.ones(shape=(self.dimension, self.dimension))
        elif isinstance(self.weights_correlation, list):
            self.weights_correlation = np.array(self.weights_correlation)
        elif not isinstance(self.weights_correlation, np.ndarray):
            raise NotImplementedError(
                "UQpy: weights_correlation attribute should be a list or numpy array"
            )

        if self.weights_correlation.shape != (self.dimension, self.dimension):
            raise NotImplementedError(
                "UQpy: Size of 'weights for correlation' is not correct"
            )


# src/UQpy/surrogates/stochastic_reduced_order_models/__init__.py

from .SROM import SROM


# src/UQpy/utilities/distances/__init__.py

from UQpy.utilities.distances.baseclass import *
from UQpy.utilities.distances.euclidean_distances import *
from UQpy.utilities.distances.grassmannian_distances import *


# src/UQpy/utilities/kernels/BinetCauchyKernel.py

import numpy as np

from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.utilities.kernels import GrassmannianKernel


class BinetCauchyKernel(GrassmannianKernel):
    """
    A class to calculate the Binet-Cauchy kernel.

    """
    def apply_method(self, points):
        points.evaluate_matrix(self, self.calculate_kernel_matrix)

    def kernel_entry(self, xi: GrassmannPoint, xj: GrassmannPoint) -> float:
        """
        Compute the Binet-Cauchy kernel entry for two points on the Grassmann manifold.

        :param xi: Orthonormal matrix representing the first point.
        :param xj: Orthonormal matrix representing the second point.

        """
        r = np.dot(xi.data.T, xj.data)
        det = np.linalg.det(r)
        return det * det


# src/UQpy/utilities/kernels/GaussianKernel.py

import numpy as np
import scipy
import scipy.spatial.distance as sd

from UQpy.utilities.ValidationTypes import RandomStateType, Numpy2DFloatArray
from UQpy.utilities.kernels import EuclideanKernel
from scipy.spatial.distance import pdist


class GaussianKernel(EuclideanKernel):
    """
    A class to calculate the Gaussian kernel defined as:

    .. math::
        k(x_j, x_i) = \exp[-(x_j - xj)^2/4\epsilon]

    """
    def __init__(self, epsilon: float = 1.0):
        """
        :param epsilon: Scale parameter of the Gaussian kernel
        """
        super().__init__()
        self.epsilon = epsilon

    def kernel_entry(self, xi: Numpy2DFloatArray, xj: Numpy2DFloatArray):
        """
        Given two points, this method computes the Gaussian kernel value between those two points

        :param xi: First point.
        :param xj: Second point.
        :return: Float representing the kernel entry.
        """
        if len(xi.shape) == 1:
            d = pdist(np.array([xi, xj]), "sqeuclidean")
        else:
            d = np.linalg.norm(xi-xj, 'fro') ** 2
        return np.exp(-d / (2*self.epsilon**2))

    def optimize_parameters(self, data: np.ndarray, tolerance: float,
                            n_nearest_neighbors: int,
                            n_cutoff_samples: int,
                            random_state: RandomStateType = None):
        """

        :param data: Set of data points.
        :param tolerance: Tolerance below which the Gaussian kernel is assumed to be zero.
        :param n_nearest_neighbors: Number of neighbors to use for cut-off estimation.
        :param n_cutoff_samples: Number of samples to use for cut-off estimation.
        :param random_state: Random seed used to initialize the pseudo-random number generator. If an :any:`int` is
         provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise, the
         object itself can be passed directly.
        """

        cut_off = self._estimate_cut_off(data, n_nearest_neighbors, n_cutoff_samples, random_state)
        self.epsilon = cut_off ** 2 / (-np.log(tolerance))

    def _estimate_cut_off(self, data, n_nearest_neighbors, n_partition, random_state):
        data = np.atleast_2d(data)
        n_points = data.shape[0]
        if n_points < 10:
            d = scipy.spatial.distance.pdist(data)
            return np.max(d)

        if n_partition is not None:
            random_indices = np.random.default_rng(random_state).permutation(n_points)
            distance_matrix = sd.cdist(data[random_indices[:n_partition]], data, metric='euclidean')
        else:
            distance_matrix = sd.squareform(sd.pdist(data, metric='euclidean'))
        k = np.min([n_nearest_neighbors, distance_matrix.shape[1]])
        k_smallest_values = np.partition(distance_matrix, k - 1, axis=1)[:, k - 1]

        est_cutoff = np.max(k_smallest_values)
        return float(est_cutoff)


# src/UQpy/utilities/kernels/ProjectionKernel.py

import numpy as np

from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.utilities.kernels import GrassmannianKernel


class ProjectionKernel(GrassmannianKernel):
    """
    A class to calculate the Projection kernel

    """
    def apply_method(self, points):
        points.evaluate_matrix(self, self.calculate_kernel_matrix)

    def kernel_entry(self, xi: GrassmannPoint, xj: GrassmannPoint) -> float:
        """
        Compute the Projection kernel entry for two points on the Grassmann manifold.

        :param xi: Orthonormal matrix representing the first point.
        :param xj: Orthonormal matrix representing the second point.
        """
        r = np.dot(xi.data.T, xj.data)
        n = np.linalg.norm(r, "fro")
        kij = n * n
        return kij


# src/UQpy/utilities/kernels/__init__.py

from UQpy.utilities.kernels.baseclass import *

from UQpy.utilities.kernels.BinetCauchyKernel import BinetCauchyKernel
from UQpy.utilities.kernels.GaussianKernel import GaussianKernel
from UQpy.utilities.kernels.ProjectionKernel import ProjectionKernel



# docs/code/RunModel/abaqus_example.py

"""

Third-party - Abaqus
==================================
"""


# %% md
#
# Import the necessary libraries.

# %%

import pickle
import time

from UQpy.distributions import Normal, Uniform
from UQpy.run_model.RunModel import *
from UQpy.sampling import MonteCarloSampling

calling_directory = os.getcwd()
t = time.time()

# %% md
#
# Building the model
# There are two probabilistic input variables, the fire load density and the yield strength.

# %%

var_names = ['qtd', 'fy']

# %% md
#
# Create the model object

# %%

abaqus_sfe_model = RunModel(model_script='abaqus_fire_analysis.py', input_template='abaqus_input.py',
                            output_script='extract_abaqus_output.py', var_names=var_names, ntasks=24,
                            model_dir='SFE_MCS', verbose=True, cores_per_task=1)
print('Example: Created the model object.')


# %% md
#
# Towards defining the sampling scheme
# The fire load density is assumed to be uniformly distributed between 50 :math:`MJ/m^2` and 450 :math:`MJ/m^2`.
# The yield strength is assumed to be normally distributed, with the parameters
# being: mean = 250 :math:`MPa` and coefficient of variation of :math:`7%`.
#
# Creating samples using MCS.

# %%

d_n = Normal(loc=50, scale=400)
d_u = Uniform(location=2.50e8, scale=1.75e7)
x_mcs = MonteCarloSampling(distributions=[d_n, d_u], samples_number=100, random_state=987979)

# %% md
#
# Running simulations using the previously defined model object and samples

# %%

sample_points = x_mcs.samples
abaqus_sfe_model.run(samples=sample_points)

# %% md
#
# The outputs from the analysis are the values of the performance function.

# %%


qois = abaqus_sfe_model.qoi_list

# %% md
#
# Save the samples and the qois in a dictionary called results with keys 'inputs' and 'outputs'.

# %%
results = {'inputs': sample_points, 'outputs': qois}

# %% md
#
# Pickle the results dictionary in the current directory. The basename and extension of the desired pickle file:

# %%
res_basename = 'MCS_results'
res_extension = '.pkl'

# %% md
#
# Create a new results file with a larger index than any existing results files with the same name in the current
# directory.

# %%
res_file_list = glob.glob(res_basename + '_???' + res_extension)
if len(res_file_list) == 0:
    res_file_name = res_basename + '_000' + res_extension
else:
    max_number = max(res_file_list).split('.')[0].split('_')[-1]
    res_file_name = res_basename + '_%03d' % (int(max_number) + 1) + res_extension

res_file_name = os.path.join(calling_directory, res_file_name)
# %% md
#
# Create a new results file with a larger index than any existing results files with the same name in the current
# directory.

# %%

with open(res_file_name, 'wb') as f:
    pickle.dump(results, f)
print('Saved the results to ' + res_file_name)

print('Example: Done!')
print('Time elapsed: %.2f minutes' % float((time.time() - t) / 60.0))


# docs/code/RunModel/ls_dyna_example_multijob.py

"""

Third-party - Multi-job - LS Dyna
==================================
"""


# %% md
#
# Import the necessary libraries.

# %%
from UQpy.distributions import Uniform
from UQpy.run_model.RunModel import RunModel
from UQpy.sampling import MonteCarloSampling

# %% md
#
# Define the distribution objects.

# %%

d1 = Uniform(location=0.02, scale=0.06)
d2 = Uniform(location=0.02, scale=0.01)
d3 = Uniform(location=0.02, scale=0.01)
d4 = Uniform(location=0.0025, scale=0.0075)
d5 = Uniform(location=0.02, scale=0.06)
d6 = Uniform(location=0.02, scale=0.01)
d7 = Uniform(location=0.02, scale=0.01)
d8 = Uniform(location=0.0025, scale=0.0075)

# %% md
#
# Draw the samples using MCS.

# %%

x = MonteCarloSampling(distributions=[d1, d2, d3, d4, d5, d6, d7, d8], samples_number=12, random_state=349875)

# %% md
#
# Run the model.

# %%

run_ = RunModel(samples=x.samples, ntasks=6, model_script='dyna_script.py', input_template='dyna_input.k',
                var_names=['x0', 'y0', 'z0', 'R0', 'x1', 'y1', 'z1', 'R1'],  model_dir='dyna_test', cluster=True,
                verbose=False, fmt='{:>10.4f}', cores_per_task=12)








# docs/code/RunModel/ls_dyna_example_singlejob.py

"""

Third-party - Single job - LS Dyna
==================================
"""


# %% md
#
# Import the necessary libraries.

# %%
from UQpy.distributions import Uniform
from UQpy.run_model.RunModel import RunModel
from UQpy.sampling import MonteCarloSampling

# %% md
#
# Define the distribution objects.

# %%
d1 = Uniform(location=0.02, scale=0.06)
d2 = Uniform(location=0.02, scale=0.01)
d3 = Uniform(location=0.02, scale=0.01)
d4 = Uniform(location=0.0025, scale=0.0075)
d5 = Uniform(location=0.02, scale=0.06)
d6 = Uniform(location=0.02, scale=0.01)
d7 = Uniform(location=0.02, scale=0.01)
d8 = Uniform(location=0.0025, scale=0.0075)

# %% md
#
# Draw the samples using MCS.

# %%
x = MonteCarloSampling(distributions=[d1, d2, d3, d4, d5, d6, d7, d8], samples_number=1, random_state=349857)

# %% md
#
# Run the model.

# %%
run_ = RunModel(samples=x.samples, ntasks=1, model_script='dyna_script.py', input_template='dyna_input.k',
                var_names=['x0', 'y0', 'z0', 'R0', 'x1', 'y1', 'z1', 'R1'], model_dir='dyna_test', cluster=True,
                verbose=False, fmt='{:>10.4f}', cores_per_task=48)








# docs/code/RunModel/matlab_example.py

"""

Third-party - Matlab
==================================
"""


# %% md
#
# The RunModel class is capable of passing input in different formats into a single computational model. This means that
# the samples passed into a model can be passed as:
# - floating point values
# - numpy arrays
# - lists
# - tuples
# - lists of other iterables
# - numpy arrays of other iterables
# - or any combination of the above
#
# In the examples below, we demonstrate the use of a third-party computational model (in this case, Matlab) with inputs
# that are combinations of the above.
#
# Some notes on their use:
#
# 1. UQpy converts all sample input to a numpy array with at least two dimensions. The first dimension,
# i.e. len(samples) must correspond to the number of samples being passed for model execution. The second dimension,
# i.e. len(samples[0]) must correspond to the number of variables that each sample possesses.
#
# 2. Each individual sample, i.e. sample[j], may be composed of multiple data types -- with each variable having a
# different data type. For example, sample[j][k] may be a floating point value and sample[j][l] may be an array of
# arbitrary dimension.
#
# 3. If a specific variable has multiple dimensions, the user may specify the index to be return in the input file.
# For example, the place holder for a variable x1 corresponding to sample[j][l] that is an array of shape (1,4) can be
# read as <x1[0, 3]>, which will return the final (0,3) component of samples[j][l].
#
# 4. If the user does not specify the index for a multidimensional variable, then the entire multidimensional variable
# is flattened and written with comma delimiters.
#
# All examples are run using Matlab execution through the command line. The user will need to modify the model_script
# in order to provide the correct path to the Matlab application on his/her computer.

# %%

# %% md
#
# Examples 1-2:
# The provided Matlab models take the sum of three random variables:
#
# .. math:: s = \sum_{i=1}^3 x_i
#
# .. math:: x_i \sim N(0,1)
#
# Example 3:
# The provided Matlab model takes the product of a random variable and the determinant of a random matrix:
#
# .. math::z = x \det(Y)
#
# .. math:: x \sim N(0,1)
#
# :math:`y` is a :math:`3x3` matrix of standard normal random variables.

# %%

from UQpy.sampling import MonteCarloSampling
from UQpy.run_model.RunModel import RunModel
from UQpy.distributions import Normal
import time
import numpy as np

# %% md
#
# Pick which model to run
# -----------------------
#
# Options:
# - 'all'
# - 'scalar'
# - 'vector'
# - 'mixed'

# %%

pick_model = 'all'

# %% md
#
# Example 1: Three scalar random variables
# ------------------------------------------
# In this example, we pass three scalar random variables. Note that this is different from assigning a single variable
# with three components, which will be handled in the following example.
#
# Here we will pass the samples both as an ndarray and as a list. Recall that UQpy converts all samples into an ndarray
# of at least two dimensions internally.

# %%

if pick_model == 'scalar' or pick_model == 'vector' or pick_model == 'all':
    # Call MCS to generate samples
    d = Normal(loc=0, scale=1)
    x_mcs = MonteCarloSampling(distributions=[d, d, d], nsamples=5, random_state=987979)
    names = ['var1', 'var11', 'var111']

    # UQpy returns samples as an ndarray. Convert them to a list for part 1.2
    x_mcs_list = list(x_mcs.samples)
    print("Monte Carlo samples of three random variables from a standard normal distribution.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs.samples))
    print('Number of samples:', len(x_mcs.samples))
    print('Dimensions of samples:', np.shape(x_mcs.samples))
    print('Samples')
    print(x_mcs.samples)
    print()
    print('Samples stored as a list:')
    print('Data type:', type(x_mcs_list))
    print('Number of samples:', len(x_mcs_list))
    print('Dimensions of samples:', np.shape(x_mcs_list))
    print('Samples:')
    print(x_mcs_list)

# %% md
#
# 1.1 Pass sampled as ndarray, specify format in generated input file, serial execution
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
#
# - model_script = matlab_model_sum_scalar.py
# - input_template = sum_scalar.m
# - output_script = process_matlab_output.py

# %%

if pick_model == 'scalar' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m = RunModel(ntasks=1, model_script='matlab_model_sum_scalar.py',
                 input_template='sum_scalar.m', var_names=names, model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model', fmt="{:>10.4f}", verbose=True)
    m.run(x_mcs.samples)
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)


# %% md
#
# 1.2 Samples passed as list, no format specification, parallel execution
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
#
# - model_script = matlab_model_sum_scalar.py
# - input_template = sum_scalar.m
# - output_script = process_matlab_output.py

# %%

if pick_model == 'scalar' or pick_model == 'all':
    # Call to RunModel with samples as a list - Again we run the model while instantiating the RunModel object.
    t = time.time()
    m = RunModel(samples=x_mcs_list, ntasks=2, model_script='matlab_model_sum_scalar.py',
                 input_template='sum_scalar.m', var_names=names, model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output', resume=False,
                 model_dir='Matlab_Model', verbose=True)
    t_par_matlab = time.time() - t
    print("\nTime for parallel execution:")
    print(t_par_matlab)
    print()
    print("The values retured from the Matlab simulation:")
    print(m.qoi_list)


# %% md
#
# Example 2: Single tri-variate random variable
# -----------------------------------------------
# In this example, we pass three random variables in as a trivariate random variable. Note that this is different from
# assigning three scalar random variables, which was be handled in Example 1.
#
# Again, we will pass the samples both as an ndarray and as a list. Recall that UQpy converts all samples into an
# ndarray of at least two dimensions internally.

# %%

# %% md
#
# Restructure the samples
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# To pass the samples in as a single tri-variate variable, we need reshape the samples from shape (5, 3) to
# shape (5, 1, 3)

# %%

if pick_model == 'vector' or pick_model == 'all':
    x_mcs_tri = x_mcs.samples.reshape(5, 1, 3)
    x_mcs_tri_list = list(x_mcs_tri)

    print("Monte Carlo samples of three random variables from a standard normal distribution.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs_tri))
    print('Number of samples:', len(x_mcs_tri))
    print('Dimensions of samples:', np.shape(x_mcs_tri))
    print('Samples')
    print(x_mcs_tri)
    print()
    print('Samples stored as a list:')
    print('Data type:', type(x_mcs_tri_list))
    print('Number of samples:', len(x_mcs_tri_list))
    print('Dimensions of samples:', np.shape(x_mcs_tri_list))
    print('Samples:')
    print(x_mcs_tri_list)

# %% md
#
# 2.1 Pass samples as ndarray, specify format in generated input file, serial execution, index samples in input_template
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
#
# - model_script = matlab_model_sum_vector_indexed.py
# - input_template = sum_vector_indexed.m
# - output_script = process_matlab_output.py

# %%


if pick_model == 'vector' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to a single variable with name x0. In this case,
    # we will read them in by indexing in the input_template.
    t = time.time()
    m = RunModel(samples=x_mcs_tri, ntasks=1, model_script='matlab_model_sum_vector_indexed.py',
                 input_template='sum_vector_indexed.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model', fmt="{:>10.4f}")
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)

# %% md
#
# 2.2 Samples passed as list, no format specification, parallel execution, index samples
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = matlab_model_sum_vector_indexed.py
# - input_template = sum_vector_indexed.m
# - output_script = process_matlab_output.py

# %%

if pick_model == 'vector' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to a single variable with name x0. In this case,
    # we will read them in by indexing in the input_template.
    t = time.time()
    m = RunModel(samples=x_mcs_tri_list, ntasks=2, model_script='matlab_model_sum_vector_indexed.py',
                 input_template='sum_vector_indexed.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model')
    t_ser_matlab = time.time() - t
    print("\nTime for parallel execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)

# %% md
#
# 2.3 Samples passed as a ndarray, specify format in generated input file, serial execution, do not index samples
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = matlab_model_sum_vector_indexed.py
# - input_template = sum_vector_indexed.m
# - output_script = process_matlab_output.py

# %%

if pick_model == 'vector' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to a single variable with name x0. In this case,
    # we will read them in by indexing in the input_template.
    t = time.time()
    m = RunModel(samples=x_mcs_tri, ntasks=1, model_script='matlab_model_sum_vector.py',
                 input_template='sum_vector.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model', fmt="{:>10.4f}")
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)


# %% md
#
# Example 3: Passing a scalar and an array to RunModel
# -----------------------------------------------------
# In this example, we pass a single scalar random variable as well as an array into a Matlab model.
#
# Again, we will pass the samples both as an ndarray and as a list. Recall that UQpy converts all samples into an
# ndarray of at least two dimensions internally.

# %%

# %% md
#
# Create the input samples to be passed

# %%

if pick_model == 'mixed' or pick_model == 'vector' or pick_model == 'all':

    # Call MCS to generate samples
    # First generate the scalar random variable
    x_mcs1 = MonteCarloSampling(distributions=d, nsamples=5, random_state=843765)
    # Next generate a 3x3 random matrix
    x_mcs2 = MonteCarloSampling(distributions=[d, d, d], nsamples=15, random_state=438975)
    x_mcs_array = x_mcs2.samples.reshape((5, 3, 3))
    # -------------------------------------------------------------------------------------------------------------

    print("Monte Carlo samples of a single random variable from a standard normal distribution.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs1.samples))
    print('Number of samples:', len(x_mcs1.samples))
    print('Dimensions of samples:', np.shape(x_mcs1.samples))
    print('Samples')
    print(x_mcs1.samples)
    print()
    print("Monte Carlo samples of a 3x3 matrix of standard normal random variables.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs_array))
    print('Number of samples:', len(x_mcs_array))
    print('Dimensions of samples:', np.shape(x_mcs_array))
    print('Samples')
    print(x_mcs_array)
    print()

    # Create a set of samples to be passed into RunModel
    # Here we need to create the mixed samples such that each sample has a single scalar and a single 3x3 matrix.
    # This data structure is essential to passing the input to UQpy correctly.
    x_mixed = []
    for i in range(5):
        x_mixed.append([x_mcs1.samples[i], x_mcs_array[i]])

    print("Combined samples with a scalar and a 3x3 matrix of standard normal random variables.")
    print('Samples stored as a list:')
    print('Data type:', type(x_mixed))
    print('Number of samples:', len(x_mixed))
    print('Dimensions of samples:', np.shape(x_mixed))
    print('Samples')
    print(x_mixed)
    print()

    x_mixed_array = np.atleast_2d(np.asarray(x_mixed))
    print("Combined samples with a scalar and a 3x3 matrix of standard normal random variables.")
    print('Samples stored as ndarray:')
    print('Data type:', type(x_mixed_array))
    print('Number of samples:', len(x_mixed_array))
    print('Dimensions of samples:', np.shape(x_mixed_array))
    print('Samples')
    print(x_mixed_array)
    print()

    # Notice that, in both the ndarray case and the list case, the samples have dimension (5,2). That is, there
    # are five samples of two variables. The first variable is a scalar. The second variable is a 3x3 matrix.

# %% md
#
# 3.1 Pass samples as ndarray, specify format in generated input file, serial execution, do not index samples in
# input_template
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = matlab_model_det.py
# - input_template = prod_determinant.m
# - output_script = process_matlab_output.py

# %%

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to two variables with names x0 and x1. In this
    # case, x0 is a scalar and x1 is a 3x3 matrix. We will read the matrix in without indexing in the
    # input_template.
    t = time.time()
    m = RunModel(samples=x_mixed_array, ntasks=1, model_script='matlab_model_det.py',
                 input_template='prod_determinant.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model', fmt="{:>10.4f}")
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)

# %% md
#
# 3.2 Pass samples as ndarray, specify format in generated input file, serial execution, index samples in input_template
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = matlab_model_det_index.py
# - input_template = prod_determinant_index.m
# - output_script = process_matlab_output.py

# %%

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to two variables with names x0 and x1. In this
    # case, x0 is a scalar and x1 is a 3x3 matrix. We will read the matrix in with indexing in the
    # input_template.
    t = time.time()
    m = RunModel(samples=x_mixed_array, ntasks=1, model_script='matlab_model_det_index.py',
                 input_template='prod_determinant_index.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model', fmt="{:>10.4f}")
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)

# %% md
#
# 3.3 Pass samples as list, do not specify format in generated input file, parallel execution, do not index samples in
# input_template
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = matlab_model_det.py
# - input_template = prod_determinant.m
# - output_script = process_matlab_output.py

# %%


if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to two variables with names x0 and x1. In this
    # case, x0 is a scalar and x1 is a 3x3 matrix. We will read the matrix in without indexing in the
    # input_template.
    t = time.time()
    m = RunModel(samples=x_mixed, ntasks=2, model_script='matlab_model_det.py',
                 input_template='prod_determinant.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model')
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)

    # Notice that the solution changes slightly due to the increased precision by not specifying a fmt.

# %% md
#
# 3.4 Pass samples as list, do not specify format in generated input file, parallel execution, index samples in
# input_template
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = matlab_model_det_index.py
# - input_template = prod_determinant_index.m
# - output_script = process_matlab_output.py

# %%

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to two variables with names x0 and x1. In this
    # case, x0 is a scalar and x1 is a 3x3 matrix. We will read the matrix in with indexing in the
    # input_template.
    t = time.time()
    m = RunModel(samples=x_mixed, ntasks=2, model_script='matlab_model_det_index.py',
                 input_template='prod_determinant_index.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model')
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Matlab simulation:")
    print(m.qoi_list)

    # Notice that the solution changes slightly due to the increased precision by not specifying a fmt.

# %% md
#
# 3.5 Pass samples as ndarray, specify format in generated input file, serial execution, partially index samples in
# input_template
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = matlab_model_det_partial.py
# - input_template = prod_determinant_partial.m
# - output_script = process_matlab_output.py

# %%s

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Notice that we do not specify var_names. This will default to two variables with names x0 and x1. In this
    # case, x0 is a scalar and x1 is a 3x3 matrix. We will read the matrix in with indexing in the
    # input_template.
    t = time.time()
    m = RunModel(samples=x_mixed_array, ntasks=1, model_script='matlab_model_det_partial.py',
                 input_template='prod_determinant_partial.m', model_object_name="matlab",
                 output_script='process_matlab_output.py', output_object_name='read_output',
                 resume=False, model_dir='Matlab_Model', fmt="{:>10.4f}")
    t_ser_matlab = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_matlab)
    print()
    print("The values returned from the Python simulation:")
    print(m.qoi_list)

# docs/code/RunModel/opensees_example.py

"""

Third-party - OpenSees
==================================
"""


# %% md
#
# Import the necessary libraries.

# %%
import numpy as np

from UQpy.distributions import Uniform
from UQpy.run_model.RunModel import RunModel
from UQpy.sampling import MonteCarloSampling

# %% md
#
# Define the distribution objects.

# %%

dist1 = Uniform(location=15000, scale=10000)
dist2 = Uniform(location=450000, scale=80000)
dist3 = Uniform(location=2.0e8, scale=0.5e8)

# %% md
#
# Draw the samples using MCS.

# %%

x = MonteCarloSampling(distributions=[dist1, dist2, dist3] * 6, samples_number=5, random_state=938475)
samples = np.array(x.samples).round(2)

# %% md
#
# Run the model.

# %%

names_ = ['fc1', 'fy1', 'Es1',  'fc2', 'fy2', 'Es2', 'fc3', 'fy3', 'Es3', 'fc4', 'fy4', 'Es4', 'fc5', 'fy5', 'Es5',
          'fc6', 'fy6', 'Es6']

opensees_rc6_model = RunModel(samples=samples, ntasks=5, model_script='opensees_model.py',
                              input_template='import_variables.tcl', var_names=names_, model_object_name="opensees_run",
                              output_script='process_opensees_output.py', output_object_name='read_output')

outputs = opensees_rc6_model.qoi_list
print(outputs)


# docs/code/RunModel/python_example.py

"""

Python
==================================
"""






# %% md
#
# The RunModel class is capable of passing input in different formats into a single computational model. This means that
# the samples passed into a model can be passed as:
#
# - floating point values
# - numpy arrays
# - lists
# - tuples
# - lists of other iterables
# - numpy arrays of other iterables
# - or any combination of the above
#
# In the examples below, we demonstrate the use of a Python computational model with inputs that are combinations of the
# above.
#
# Some notes on their use:
#
# 1. UQpy converts all sample input to a numpy array with at least two dimensions. The first dimension,
# i.e. len(samples) must correspond to the number of samples being passed for model execution. The second dimension,
# i.e. len(samples[0]) must correspond to the number of variables that each sample possesses.
#
# 2. Each individual sample, i.e. sample[j], may be composed of multiple data types -- with each variable having a
# different data type. For example, sample[j][k] may be a floating point value and sample[j][l] may be an array of
# arbitrary dimension.
#
# 3. If a specific variable has multiple dimensions, the user may specify the index to be return in the input file.
# For example, the placeholder for a variable x1 corresponding to sample[j][l] that is an array of shape (1,4) can be
# read as <x1[0, 3]>, which will return the final (0,3) component of samples[j][l].
#
# 4. If the user does not specify the index for a multidimensional variable, then the entire multidimensional variable
# is flattened and written with comma delimiters.

# %%

# %% md
#
# Python Model Summary
# --------------------
# Examples 1-2:
# The provided Matlab models take the sum of three random variables:
#
# .. math:: s = \sum_{i=1}^3 x_i
#
# .. math:: x_i \sim N(0,1)
#
# Example 3:
# The provided Matlab model takes the product of a random variable and the determinant of a random matrix:
#
# .. math:: z = x \det(Y)
#
# .. math:: x \sim N(0,1)
#
# :math:`y` is a 3x3 matrix of standard normal random variables.
#
# The Python model may be provided as either a class or a function. The examples below explore both cases.

# %%

from UQpy.sampling import MonteCarloSampling
from UQpy.run_model.RunModel import RunModel
from UQpy.distributions import Normal
import time
import numpy as np


# %%

# %% md
#
# Pick which model to run
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Options:
# - 'all'
# - 'scalar'
# - 'vector'
# - 'mixed'
# - 'fixed'

# %%

pick_model = 'all'


# %% md
#
# Example 1: Three scalar random variables
# ----------------------------------------------------
# In this example, we pass three scalar random variables. Note that this is different from assigning a single variable
# with three components, which will be handled in the following example.
#
# Here we will pass the samples both as an ndarray and as a list. Recall that UQpy converts all samples into an ndarray
# of at least two dimensions internally.

# %%

if pick_model in {'scalar', 'vector', 'all'}:
    d = Normal(loc=0, scale=1)
    x_mcs = MonteCarloSampling(distributions=[d, d, d], nsamples=5, random_state=987979)
    names = ['var1', 'var11', 'var111']

    # UQpy returns samples as an ndarray. Convert them to a list for part 1.2
    x_mcs_list = list(x_mcs.samples)
    print("Monte Carlo samples of three random variables from a standard normal distribution.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs.samples))
    print('Number of samples:', len(x_mcs.samples))
    print('Dimensions of samples:', np.shape(x_mcs.samples))
    print('Samples')
    print(x_mcs.samples)
    print()
    print('Samples stored as a list:')
    print('Data type:', type(x_mcs_list))
    print('Number of samples:', len(x_mcs_list))
    print('Dimensions of samples:', np.shape(x_mcs_list))
    print('Samples:')
    print(x_mcs_list)

# %% md
#
# 1.1 Pass samples as ndarray, Python class called, serial execution
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = python_model.py

# %%

if pick_model in {'scalar', 'all'}:
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m11 = RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs', model_dir='Python_Runs', verbose=True)
    m11.run(samples=x_mcs.samples,)
    t_ser_python = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m11.qoi_list)

# %% md
#
# 1.2 Pass samples as list, Python function called, parallel execution
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = python_model.py

# %%

if pick_model in {'scalar', 'all'}:
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m12 = RunModel(samples=x_mcs_list, ntasks=2, model_script='python_model.py',
                 model_object_name='sum_rvs', model_dir='Python_Runs')
    t_par_python = time.time() - t
    print("\nTime for parallel execution:")
    print(t_par_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m12.qoi_list)


# %% md
#
# Example 2: Single tri-variate random variable
# -----------------------------------------------------
# In this example, we pass three random variables in as a trivariate random variable. Note that this is different from
# assigning three scalar random variables, which was be handled in Example 1.
#
# Again, we will pass the samples both as an ndarray and as a list. Recall that UQpy converts all samples into an
# ndarray of at least two dimensions internally.

# %%

# %% md
#
# Restructure the samples
# -----------------------------------------------------
# To pass the samples in as a single tri-variate variable, we need reshape the samples from shape (5, 3) to
# shape (5, 1, 3)

# %%

if pick_model in {'vector', 'all'}:
    x_mcs_tri = x_mcs.samples.reshape(5, 1, 3)
    x_mcs_tri_list = list(x_mcs_tri)

    print("Monte Carlo samples of three random variables from a standard normal distribution.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs_tri))
    print('Number of samples:', len(x_mcs_tri))
    print('Dimensions of samples:', np.shape(x_mcs_tri))
    print('Samples')
    print(x_mcs_tri)
    print()
    print('Samples stored as a list:')
    print('Data type:', type(x_mcs_tri_list))
    print('Number of samples:', len(x_mcs_tri_list))
    print('Dimensions of samples:', np.shape(x_mcs_tri_list))
    print('Samples:')
    print(x_mcs_tri_list)

# %% md
#
# 2.1 Pass samples as ndarray, Python function called, serial execution
# -----------------------------------------------------------------------
# This example uses the following files:
# - model_script = python_model.py

# %%

if pick_model in {'vector', 'all'}:
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m21 = RunModel(samples=x_mcs_tri, ntasks=1, model_script='python_model.py',
                 model_object_name='sum_rvs_vec', model_dir='Python_Runs')
    t_ser_python = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m21.qoi_list)

# %% md
#
# 2.2 Pass samples as list, Python class called, parallel execution
# --------------------------------------------------------------------
# This example uses the following files:
# - model_script = python_model.py

# %%

if pick_model == 'vector' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m22 = RunModel(samples=x_mcs_tri_list, ntasks=2, model_script='python_model.py',
                 model_object_name='SumRVs', model_dir='Python_Runs')
    t_par_python = time.time() - t
    print("\nTime for parallel execution:")
    print(t_par_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m22.qoi_list)

# %% md
#
# Example 3: Passing a scalar and an array to RunModel
# -----------------------------------------------------
# In this example, we pass a single scalar random variable as well as an array into a Matlab model.
#
# Again, we will pass the samples both as an ndarray and as a list. Recall that UQpy converts all samples into an
# ndarray of at least two dimensions internally.

# %%

if pick_model == 'mixed' or pick_model == 'vector' or pick_model == 'all':
    # Call MCS to generate samples
    # THIS WILL NEED TO BE REWRITTEN WITH DISTRIBUTION AND MCS UPDATES --------------------------------------------
    # First generate the scalar random variable
    #     x_mcs1 = MCS(dist_name=['Normal'], dist_params=[[0,1]], nsamples=5, var_names = ['var1'])
    # Next generate a 3x3 random matrix
    #     x_mcs2 = MCS(dist_name=['Normal','Normal','Normal'], dist_params=[[0,1],[0,1],[0,1]], nsamples=15)
    #     x_mcs_array = x_mcs2.samples.reshape((5,3,3))
    # -------------------------------------------------------------------------------------------------------------

    # Call MCS to generate samples
    d = Normal(loc=0, scale=1)
    x_mcs1 = MonteCarloSampling(distributions=d, nsamples=5, random_state=987979)
    x_mcs2 = MonteCarloSampling(distributions=[d, d, d], nsamples=15, random_state=34876)
    x_mcs_array = x_mcs2.samples.reshape((5, 3, 3))

    print("Monte Carlo samples of a single random variable from a standard normal distribution.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs1.samples))
    print('Number of samples:', len(x_mcs1.samples))
    print('Dimensions of samples:', np.shape(x_mcs1.samples))
    print('Samples')
    print(x_mcs1.samples)
    print()
    print("Monte Carlo samples of a 3x3 matrix of standard normal random variables.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs_array))
    print('Number of samples:', len(x_mcs_array))
    print('Dimensions of samples:', np.shape(x_mcs_array))
    print('Samples')
    print(x_mcs_array)
    print()

    # Create a set of samples to be passed into RunModel
    # Here we need to create the mixed samples such that each sample has a single scalar and a single 3x3 matrix.
    # This data structure is essential to passing the input to UQpy correctly.
    x_mixed = []
    for i in range(5):
        x_mixed.append([x_mcs1.samples[i], x_mcs_array[i]])

    print("Combined samples with a scalar and a 3x3 matrix of standard normal random variables.")
    print('Samples stored as a list:')
    print('Data type:', type(x_mixed))
    print('Number of samples:', len(x_mixed))
    print('Dimensions of samples:', np.shape(x_mixed))
    print('Samples')
    print(x_mixed)
    print()

    x_mixed_array = np.atleast_2d(np.asarray(x_mixed))
    print("Combined samples with a scalar and a 3x3 matrix of standard normal random variables.")
    print('Samples stored as ndarray:')
    print('Data type:', type(x_mixed_array))
    print('Number of samples:', len(x_mixed_array))
    print('Dimensions of samples:', np.shape(x_mixed_array))
    print('Samples')
    print(x_mixed_array)
    print()

    # Notice that, in both the ndarray case and the list case, the samples have dimension (5,2). That is, there
    # are five samples of two variables. The first variable is a scalar. The second variable is a 3x3 matrix.

# %% md
#
# 3.1 Pass samples as ndarray, Python class called, serial execution
# ---------------------------------------------------------------------
# This examples uses the following files:
# - model_script = python_model.py

# %%

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m31 = RunModel(samples=x_mixed_array, ntasks=1, model_script='python_model.py',
                 model_object_name='DetRVs', model_dir='Python_Runs', vec=False)
    t_ser_python = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m31.qoi_list)

# %% md
#
# 3.2 Pass samples as list, Python function called, parallel execution
# ------------------------------------------------------------------------
# This examples uses the following files:
# - model_script = python_model.py

# %%

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    # Note that the parallel model_object handles only one sample at a time.
    t = time.time()
    m32 = RunModel(samples=x_mixed, ntasks=1, model_script='python_model.py',
                 model_object_name='det_rvs_par', model_dir='Python_Runs', vec=False)
    t_par_python = time.time() - t
    print("\nTime for parallel execution:")
    print(t_par_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m32.qoi_list)

# %% md
#
# Example 4: Passing a fixed variable and an array of Random Variables to RunModel
# -----------------------------------------------------------------------------------
# In this example, we pass a fixed value coefficient as well as an array of random variables into a Python model.
#
# Again, we will pass the samples both as an ndarray and as a list. Recall that UQpy converts all samples into an
# ndarray of at least two dimensions internally.

# %%

if pick_model == 'mixed' or pick_model == 'all':
    x = 2.5

    print('Constant Coefficient:')
    print(x)
    print()

    print("Monte Carlo samples of a 3x3 matrix of standard normal random variables.")
    print('Samples stored as an array:')
    print('Data type:', type(x_mcs_array))
    print('Number of samples:', len(x_mcs_array))
    print('Dimensions of samples:', np.shape(x_mcs_array))
    print('Samples')
    print(x_mcs_array)
    print()

    x_mcs_list = list(x_mcs_array)

    print("3x3 matrix of standard normal random variables.")
    print('Samples stored as ndarray:')
    print('Data type:', type(x_mcs_list))
    print('Number of samples:', len(x_mcs_list))
    print('Dimensions of samples:', np.shape(x_mcs_list))
    print('Samples')
    print(x_mcs_list)
    print()

    # Notice that, in both the ndarray case and the list case, the samples have dimension (5,2). That is, there
    # are five samples of two variables. The first variable is a scalar. The second variable is a 3x3 matrix.

# %% md
#
# 4.1 Pass samples as ndarray, Python class called, serial execution
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = python_model.py

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m41 = RunModel(samples=x_mcs_array, ntasks=1, model_script='python_model.py',
                 model_object_name='det_rvs_fixed', model_dir='Python_Runs', vec=False, coeff=x)
    t_ser_python = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m41.qoi_list)

# %% md
#
# 4.2 Pass samples as list, Python class called, serial execution
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This examples uses the following files:
# - model_script = python_model.py

if pick_model == 'mixed' or pick_model == 'all':
    # Call to RunModel - Here we run the model while instantiating the RunModel object.
    t = time.time()
    m42 = RunModel(samples=x_mcs_list, ntasks=1, model_script='python_model.py',
                 model_object_name='det_rvs_fixed', model_dir='Python_Runs', vec=False, coeff=x)
    t_ser_python = time.time() - t
    print("\nTime for serial execution:")
    print(t_ser_python)
    print()
    print("The values returned from the Matlab simulation:")
    print(m42.qoi_list)

# src/UQpy/dimension_reduction/__init__.py

from UQpy.dimension_reduction.grassmann_manifold import *
from UQpy.dimension_reduction.pod import *
from UQpy.dimension_reduction.hosvd import *
from UQpy.dimension_reduction.diffusion_maps import *


# src/UQpy/distributions/__init__.py

# pylint: disable=wildcard-import

from UQpy.distributions.baseclass import *
from UQpy.distributions.copulas import *
from UQpy.distributions.collection import *

from . import baseclass, copulas, collection


# src/UQpy/inference/BayesModelSelection.py

import logging
from typing import Union

import numpy as np
from beartype import beartype

from UQpy.inference.BayesParameterEstimation import BayesParameterEstimation
from UQpy.inference.evidence_methods.HarmonicMean import HarmonicMean
from UQpy.inference.evidence_methods.baseclass import EvidenceMethod
from UQpy.inference.inference_models.baseclass.InferenceModel import InferenceModel
from UQpy.sampling import ImportanceSampling
from UQpy.utilities.ValidationTypes import PositiveInteger


class BayesModelSelection:

    # Authors: Audrey Olivier, Yuchen Zhou
    # Last modified: 01/24/2020 by Audrey Olivier
    @beartype
    def __init__(
            self,
            parameter_estimators: list[BayesParameterEstimation],
            prior_probabilities=None,
            evidence_method: EvidenceMethod = HarmonicMean(),
            nsamples: list[PositiveInteger] = None,
    ):
        """
        Perform model selection via Bayesian inference, i.e., compute model posterior probabilities given data.

        This class leverages the :class:`.BayesParameterEstimation` class to get samples from the parameter posterior
        densities. These samples are then used to compute the model evidence :code:`p(data|model)` for all models and
        the model posterior probabilities.

        :param data: Available data
        :param parameter_estimators: Parameter estimators used during the model selection algorithm.
        :param prior_probabilities: Prior probabilities of each model, default is :code:`[1/nmodels, ] * nmodels`
        :param evidence_method: as of v3, only the harmonic mean method is supported
        :param nsamples: Number of samples used in :class:`.MCMC`/:class:`.ImportanceSampling`, for each model
        """
        self.bayes_estimators: list[BayesParameterEstimation] = parameter_estimators
        """Results of the Bayesian parameter estimation."""
        self.candidate_models: list[InferenceModel] = [x.inference_model for x in self.bayes_estimators]
        """Probabilistic models used during the model selection process."""
        self.models_number = len(self.candidate_models)
        self.evidence_method = evidence_method
        self.logger = logging.getLogger(__name__)

        if prior_probabilities is None:
            self.prior_probabilities = [1.0 / len(self.candidate_models) for _ in self.candidate_models]
        else:
            self.prior_probabilities = prior_probabilities

        # Instantiate the Bayesian parameter estimators (without running them)

        self._update_bayes_estimators()

        # Initialize the outputs
        self.evidences: list = [0.0] * self.models_number
        """Value of the evidence for all models."""
        self.probabilities: list = [0.0] * self.models_number
        """Posterior probability for all models"""

        # Run the model selection procedure
        if nsamples is not None:
            self.run(nsamples=nsamples)

    def _update_bayes_estimators(self):
        for i, estimator in enumerate(self.bayes_estimators):
            if not isinstance(estimator.sampler, ImportanceSampling):
                estimator.sampler.save_log_pdf = True
                estimator.sampler.concatenate_chains = True
                estimator.sampler.dimension = estimator.inference_model.n_parameters

    @beartype
    def run(self, nsamples: Union[None, list[int]]):
        """
        Run the Bayesian model selection procedure, i.e., compute model posterior probabilities.

        This function calls the :py:meth:`run_estimation` method of the :class:`.BayesParameterEstimation` object for
        each model to sample from the parameter posterior probability, then computes the model evidence and model
        posterior probability. This function updates attributes :py:attr:`bayes_estimators`, :py:attr:`evidences` and
        :py:attr:`probabilities`. If `nsamples` are given when creating the object, this method is called
        directly when the object is created. It can also be called separately.

        :param nsamples: Number of samples used in :class:`.MCMC`/:class:`.ImportanceSampling``, for each model
        """
        self.logger.info("UQpy: Running Bayesian Model Selection.")
        # Perform mcmc for all candidate models
        for i, (inference_model, bayes_estimator) in enumerate(zip(self.candidate_models, self.bayes_estimators)):
            self.logger.info("UQpy: Running mcmc for model " + inference_model.name)
            if nsamples[i] == 0:
                continue
            bayes_estimator.run(nsamples=nsamples[i])
            self.evidences[i] = \
                self.evidence_method.estimate_evidence(inference_model=inference_model,
                                                       posterior_samples=bayes_estimator.sampler.samples,
                                                       log_posterior_values=bayes_estimator.sampler.log_pdf_values, )

        # Compute posterior probabilities
        self.probabilities = self._compute_posterior_probabilities(
            prior_probabilities=self.prior_probabilities,
            evidence_values=self.evidences)

        self.logger.info("UQpy: Bayesian Model Selection analysis completed!")

    def sort_models(self):
        """
        Sort models in descending order of model probability (increasing order of criterion value).

        This function sorts - in place - the attribute lists :py:attr:`candidate_models`, :py:attr:`probabilities`
        and :py:attr:`evidences` so that they are sorted from most probable to the least probable model. It is a
        stand-alone function that is provided to help the user to easily visualize which model is the best.

        No inputs/outputs.

        """
        sort_idx = list(np.argsort(np.array(self.probabilities)))[::-1]

        self.candidate_models = [self.candidate_models[i] for i in sort_idx]
        self.prior_probabilities = [self.prior_probabilities[i] for i in sort_idx]
        self.probabilities = [self.probabilities[i] for i in sort_idx]
        self.evidences = [self.evidences[i] for i in sort_idx]

    @staticmethod
    def _compute_posterior_probabilities(prior_probabilities, evidence_values):
        """
        Compute the model probability given prior probabilities P(M) and evidence values p(data|M).

        Model posterior probability P(M|data) is proportional to p(data|M)P(M). Posterior probabilities sum up to 1 over
        all models. This function is a utility function (static method), called within the run_estimation method.

        **Inputs:**

        :param prior_probabilities: Values of prior probabilities for all models.
        :type prior_probabilities: list (length nmodels) of floats

        :param prior_probabilities: Values of evidence for all models.
        :type prior_probabilities: list (length nmodels) of floats

        **Output/Returns:**

        :return probabilities: Values of model posterior probabilities
        :rtype probabilities: list (length nmodels) of floats

        """
        scaled_evidences = [evidence * prior_probability for (evidence, prior_probability)
                            in zip(evidence_values, prior_probabilities)]
        return scaled_evidences / np.sum(scaled_evidences)


# src/UQpy/inference/BayesParameterEstimation.py

import logging
from typing import Union

import numpy as np
from beartype import beartype
from UQpy.utilities.ValidationTypes import PositiveInteger
from UQpy.inference.inference_models.baseclass.InferenceModel import InferenceModel
from UQpy.sampling import MCMC, ImportanceSampling


class BayesParameterEstimation:
    # Authors: Audrey Olivier, Dimitris Giovanis
    # Last Modified: 12/19 by Audrey Olivier
    @beartype
    def __init__(
        self,
        inference_model: InferenceModel,
        data: Union[list, np.ndarray],
        sampling_class: Union[MCMC, ImportanceSampling] = None,
        nsamples: Union[None, int] = None,
    ):
        """
        Estimate the parameter posterior density given some data.

        This class generates samples from the parameter posterior distribution using Markov Chain Monte Carlo or
        Importance Sampling. It leverages the :class:`.MCMC` and :class:`.ImportanceSampling` classes from the
        :py:mod:`.sampling` module.

        :param inference_model: The inference model that defines the likelihood function.
        :param data: Available data, :class:`numpy.ndarray` of shape consistent with log-likelihood function in
         :class:`.InferenceModel`
        :param sampling_class: Class instance, must be a subclass of :class:`.MCMC` or :class:`.ImportanceSampling`.
        :param nsamples: Number of samples used in :class:`.MCMC`/:class:`ImportanceSampling`, see
         :meth:`run` method. If the `nsamples` parameter is provided then :class:`.BayesParameterEstimation` is
         automatically performed. In case an :class:`.ImportanceSampling` method is used to perform the parameter
         estimation, then `nsamples` equal to the total number of samples. In case an :class:`.MCMC` sampler is used,
         and the given `nsamples` is not a multiple of `n_chains`, then `nsamples` is set to the next largest integer
         that is a multiple of  `nchains`.
        """
        self.inference_model = inference_model
        self.data = data
        self.logger = logging.getLogger(__name__)
        self.sampler: Union[MCMC, ImportanceSampling] = sampling_class
        """Sampling method object, contains e.g. the posterior samples.

        This must be created along side the :class:`.BayesParameterEstimation` object, and its run method is called 
        whenever the :py:meth:`run` method of the :class:`.BayesParameterEstimation` is called.
        """
        if isinstance(self.sampler, ImportanceSampling):
            if self.sampler.proposal is None:
                if inference_model.prior is None:
                    raise NotImplementedError(
                        "UQpy: A proposal density of the ImportanceSampling"
                        " or a prior to the Inference model  must be provided.")
                self.sampler.proposal = inference_model.prior
            self.sampler._args_target = (data,)
            self.sampler.log_pdf_target = inference_model.evaluate_log_posterior
        elif isinstance(self.sampler, MCMC):
            if self.sampler._initialization_seed is None:
                if inference_model.prior is None or not hasattr(inference_model.prior, "rvs"):
                    raise ValueError(
                        "UQpy: A prior with a rvs method must be provided for the InferenceModel"
                        " or a seed must be provided for MCMC.")
                else:
                    self.sampler.seed = inference_model.prior.rvs(nsamples=self.sampler.n_chains,
                                                                  random_state=self.sampler.random_state, ).tolist()
            self.sampler.log_pdf_target = inference_model.evaluate_log_posterior
            self.sampler.pdf_target = None
            self.sampler.args_target = (data,)
            (self.sampler.evaluate_log_target,
             self.sampler.evaluate_log_target_marginals,) = \
                self.sampler._preprocess_target(pdf_=None,
                                                log_pdf_=self.sampler.log_pdf_target,
                                                args=self.sampler.args_target)

        if nsamples is not None:
            self.run(nsamples=nsamples)

    @beartype
    def run(self, nsamples: PositiveInteger = None):
        """
        Run the Bayesian inference procedure, i.e., sample from the parameter posterior distribution.

        This function calls the :meth:`run` method of the :py:attr:`.sampler` attribute to generate samples from the
        parameter posterior distribution.

        :param nsamples: Number of samples used in :class:`.MCMC`/:class:`.ImportanceSampling`.  In case an
         :class:`.ImportanceSampling` method is used to perform the parameter estimation, then `nsamples` equal to the
         total number of samples. In case an :class:`.MCMC` sampler is used, and the given `nsamples` is not a multiple
         of `n_chains`, then `nsamples` is set to the next largest integer that is a multiple of  `nchains`.
        """
        self.sampler.run(nsamples=nsamples)

        self.logger.info("UQpy: Parameter estimation with " + self.sampler.__class__.__name__
                         + " completed successfully!")


# src/UQpy/inference/InformationModelSelection.py

import logging
from typing import Union

from UQpy.inference.information_criteria import AIC
from UQpy.inference.information_criteria.baseclass.InformationCriterion import InformationCriterion
from beartype import beartype
from UQpy.inference.MLE import MLE
import numpy as np


class InformationModelSelection:

    # Authors: Audrey Olivier, Dimitris Giovanis
    # Last Modified: 12/19 by Audrey Olivier
    @beartype
    def __init__(
            self,
            parameter_estimators: list[MLE],
            criterion: InformationCriterion = AIC(),
            n_optimizations: list[int] = None,
            initial_parameters: list[np.ndarray] = None
    ):
        """
        Perform model selection using information theoretic criteria.

        Supported criteria are :class:`.BIC`, :class:`.AIC` (default), :class:`.AICc`. This class leverages the
        :class:`.MLE` class for maximum likelihood estimation, thus inputs to :class:`.MLE` can also be provided to
        :class:`InformationModelSelection`, as lists of length equal to the number of models.

        :param parameter_estimators: A list containing a maximum-likelihood estimator (:class:`.MLE`) for each one of the
         models to be compared.
        :param criterion: Criterion to be used (:class:`.AIC`, :class:`.BIC`, :class:`.AICc)`. Default is :class:`.AIC`
        :param initial_parameters: Initial guess(es) for optimization, :class:`numpy.ndarray` of shape
         :code:`(nstarts, n_parameters)` or :code:`(n_parameters, )`, where :code:`nstarts` is the number of times the
         optimizer will be called. Alternatively, the user can provide input `n_optimizations` to randomly sample
         initial guess(es). The identified MLE is the one that yields the maximum log likelihood over all calls of the
         optimizer.
        """
        self.candidate_models = [mle.inference_model for mle in parameter_estimators]
        self.models_number = len(parameter_estimators)
        self.criterion: InformationCriterion = criterion
        self.logger = logging.getLogger(__name__)

        self.n_optimizations = n_optimizations
        self.initial_parameters= initial_parameters

        self.parameter_estimators: list = parameter_estimators
        """:class:`.MLE` results for each model (contains e.g. fitted parameters)"""

        # Initialize the outputs
        self.criterion_values: list = [None, ] * self.models_number
        """Value of the criterion for all models."""
        self.penalty_terms: list = [None, ] * self.models_number
        """Value of the penalty term for all models. Data fit term is then criterion_value - penalty_term."""
        self.probabilities: list = [None, ] * self.models_number
        """Value of the model probabilities, computed as
        
        .. math:: P(M_i|d) = \dfrac{\exp(-\Delta_i/2)}{\sum_i \exp(-\Delta_i/2)}
        
        where :math:`\Delta_i = criterion_i - min_i(criterion)`"""

        # Run the model selection procedure
        if (self.n_optimizations is not None) or (self.initial_parameters is not None):
            self.run(self.n_optimizations, self.initial_parameters)

    def run(self, n_optimizations: list[int], initial_parameters: list[np.ndarray]=None):
        """
        Run the model selection procedure, i.e. compute criterion value for all models.

        This function calls the :meth:`run` method of the :class:`.MLE` object for each model to compute the maximum
        log-likelihood, then computes the criterion value and probability for each model. If `data` are given when
        creating the :class:`.MLE` object, this method is called automatically when the object is created.

        :param n_optimizations: Number of iterations that the optimization is run, starting at random initial
         guesses. It is only used if `initial_parameters` is not provided. Default is :math:`1`.
         The random initial guesses are sampled uniformly between :math:`0` and :math:`1`, or uniformly between
         user-defined bounds if an input bounds is provided as a keyword argument to the `optimizer` input parameter.
        :param initial_parameters: Initial guess(es) for optimization, :class:`numpy.ndarray` of shape
         :code:`(nstarts, n_parameters)` or :code:`(n_parameters, )`, where :code:`nstarts` is the number of times the
         optimizer will be called. Alternatively, the user can provide input `n_optimizations` to randomly sample
         initial guess(es). The identified MLE is the one that yields the maximum log likelihood over all calls of the
         optimizer.
        """
        if (n_optimizations is not None and (len(n_optimizations) != len(self.parameter_estimators))) or \
           (initial_parameters is not None and len(initial_parameters) != len(self.parameter_estimators)):
            raise ValueError("The length of n_optimizations and initial_parameters should be equal to the number of "
                             "parameter estimators")
        # Loop over all the models
        for i, parameter_estimator in enumerate(self.parameter_estimators):
            # First evaluate ML estimate for all models, do several iterations if demanded
            parameters = None
            if initial_parameters is not None:
                parameters = initial_parameters[i]

            optimizations = 0
            if n_optimizations is not None:
                optimizations = n_optimizations[i]

            parameter_estimator.run(n_optimizations=optimizations, initial_parameters=parameters)

            # Then minimize the criterion
            self.criterion_values[i], self.penalty_terms[i] = \
                self.criterion.minimize_criterion(data=parameter_estimator.data,
                                                  parameter_estimator=parameter_estimator,
                                                  return_penalty=True)
        # Compute probabilities from criterion values
        self.probabilities = self._compute_probabilities(self.criterion_values)

    def sort_models(self):
        """
        Sort models in descending order of model probability (increasing order of `criterion` value).

        This function sorts - in place - the attribute lists :py:attr:`.candidate_models`, :py:attr:`.ml_estimators`,
        :py:attr:`criterion_values`, :py:attr:`penalty_terms` and :py:attr:`probabilities` so that they are sorted from
        most probable to least probable model. It is a stand-alone function that is provided to help the user to easily
        visualize which model is the best.

        No inputs/outputs.

        """
        sort_idx = list(np.argsort(np.array(self.criterion_values)))

        self.candidate_models = [self.candidate_models[i] for i in sort_idx]
        self.parameter_estimators = [self.parameter_estimators[i] for i in sort_idx]
        self.criterion_values = [self.criterion_values[i] for i in sort_idx]
        self.penalty_terms = [self.penalty_terms[i] for i in sort_idx]
        self.probabilities = [self.probabilities[i] for i in sort_idx]

    @staticmethod
    def _compute_probabilities(criterion_values):
        delta = np.array(criterion_values) - min(criterion_values)
        prob = np.exp(-delta / 2)
        return prob / np.sum(prob)


# src/UQpy/inference/MLE.py

import logging
from typing import Union

import numpy as np
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from beartype import beartype

from UQpy.inference.inference_models.baseclass.InferenceModel import InferenceModel
from UQpy.utilities.Utilities import process_random_state
from UQpy.utilities.ValidationTypes import NumpyFloatArray, RandomStateType


class MLE:
    # Authors: Audrey Olivier, Dimitris Giovanis
    # Last Modified: 12/19 by Audrey Olivier
    @beartype
    def __init__(
            self,
            inference_model: InferenceModel,
            data: Union[list, np.ndarray],
            n_optimizations: Union[None, int] = 1,
            initial_parameters: Union[list, np.ndarray, None] = None,
            optimizer=MinimizeOptimizer(),
            random_state: RandomStateType = None,
    ):
        """
        Estimate the maximum likelihood parameters of a model given some data.

        :param inference_model: The inference model that defines the likelihood function.
        :param data: Available data, :class:`numpy.ndarray` of shape consistent with log likelihood function in
         :class:`.InferenceModel`. If this parameter is provided at :class:`.InformationModelSelection` object
         initialization, the maximum likelihood estimation algorithm will be automatically performed. Alternatively,
         the user must execute the :meth:`.run` method.
        :param n_optimizations: Number of iterations that the optimization is run, starting at random initial
         guesses. It is only used if `initial_parameters` is not provided. Default is :math:`1`.
         The random initial guesses are sampled uniformly between :math:`0` and :math:`1`, or uniformly between
         user-defined bounds if an input bounds is provided as a keyword argument to the `optimizer` input parameter.
        :param initial_parameters: Initial guess(es) for optimization, :class:`numpy.ndarray` of shape
         :code:`(nstarts, n_parameters)` or :code:`(n_parameters, )`, where :code:`nstarts` is the number of times the
         optimizer will be called. Alternatively, the user can provide input `n_optimizations` to randomly sample
         initial guess(es). The identified MLE is the one that yields the maximum log likelihood over all calls of the
         optimizer.
        :param optimizer: This parameter takes as input an object that implements the :class:`Optimizer` class.
         Default is the :class:`.Minimize` which utilizes the :class:`scipy.optimize.minimize` method.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
        """
        # Initialize variables
        self.inference_model = inference_model
        self.data = data
        self.random_state = process_random_state(random_state)
        self.logger = logging.getLogger(__name__)
        self.optimizer = optimizer
        self.mle: NumpyFloatArray = None
        """Value of parameter vector that maximizes the likelihood function."""
        self.max_log_like: NumpyFloatArray = None
        """Value of the likelihood function at the MLE."""
        self.logger.info("UQpy: Initialization of MLEstimation object completed.")
        self.n_optimizations = n_optimizations
        self.initial_parameters = initial_parameters

        # Run the optimization procedure
        if (n_optimizations is not None) or (initial_parameters is not None):
            self.run(n_optimizations=n_optimizations, initial_parameters=initial_parameters)

    @beartype
    def run(self, n_optimizations: Union[None, int] = 1, initial_parameters: Union[list, np.ndarray, None] = None):
        """
        Run the maximum likelihood estimation procedure.

        This function runs the optimization and updates the :py:attr:`.mle` and :py:attr:`.max_log_like` attributes of
        the class. When learning the parameters of a distribution, if `distributions` possesses an :meth:`mle` method
        this method is used. If `data` are given when creating the :class:`.MLE` object, this method is called
        automatically when the object is created.

        :param n_optimizations: Number of iterations that the optimization is run, starting at random initial
         guesses. It is only used if `initial_parameters` is not provided. Default is :math:`1`.
         The random initial guesses are sampled uniformly between :math:`0` and :math:`1`, or uniformly between
         user-defined bounds if an input bounds is provided as a keyword argument to the `optimizer` input parameter.
        :param initial_parameters: Initial guess(es) for optimization, :class:`numpy.ndarray` of shape
         :code:`(nstarts, n_parameters)` or :code:`(n_parameters, )`, where :code:`nstarts` is the number of times the
         optimizer will be called. Alternatively, the user can provide input `n_optimizations` to randomly sample
         initial guess(es). The identified MLE is the one that yields the maximum log likelihood over all calls of the
         optimizer.
        """
        # Run optimization (use x0 if provided, otherwise sample starting point from [0, 1] or bounds)
        self.logger.info(
            "UQpy: Evaluating maximum likelihood estimate for inference model "
            + self.inference_model.name)

        self.n_optimizations = n_optimizations
        self.initial_parameters = initial_parameters

        use_distribution_fit = (
                hasattr(self.inference_model, "distributions")
                and self.inference_model.distributions is not None
                and hasattr(self.inference_model.distributions, "fit"))

        if use_distribution_fit:
            self._run_distribution_fit(self.n_optimizations)
        else:
            self._run_optimization(self.initial_parameters, self.n_optimizations)

    def _run_distribution_fit(self, n_optimizations):
        for _ in range(n_optimizations):
            self.inference_model.distributions.update_parameters(
                **{key: None for key in self.inference_model.list_params})
            mle_dict = self.inference_model.distributions.fit(data=self.data)
            mle_tmp = np.array([mle_dict[key] for key in self.inference_model.list_params])
            max_log_like_tmp = self.inference_model.evaluate_log_likelihood(
                parameters=mle_tmp[np.newaxis, :], data=self.data)[0]
            # Save result
            if self.mle is None or max_log_like_tmp > self.max_log_like:
                self.mle = mle_tmp
                self.max_log_like = max_log_like_tmp

    def _run_optimization(self, initial_parameters, n_optimizations):
        if initial_parameters is None:
            from UQpy.distributions import Uniform
            initial_parameters = (
                Uniform()
                    .rvs(nsamples=n_optimizations * self.inference_model.n_parameters, random_state=self.random_state, )
                    .reshape((n_optimizations, self.inference_model.n_parameters)))
            if self.optimizer._bounds is not None:
                bounds = np.array(self.optimizer._bounds)
                initial_parameters = (bounds[:, 0].reshape((1, -1))
                                      + (bounds[:, 1] - bounds[:, 0]).reshape((1, -1)) * initial_parameters)
        else:
            initial_parameters = np.atleast_2d(initial_parameters)
            if initial_parameters.shape[1] != self.inference_model.n_parameters:
                raise ValueError("UQpy: Wrong dimensions in x0")
        for x0_ in initial_parameters:
            res = self.optimizer.optimize(self._evaluate_func_to_minimize, x0_)
            mle_tmp = res.x
            max_log_like_tmp = (-1.0) * res.fun
            # Save result
            if self.mle is None or max_log_like_tmp > self.max_log_like:
                self.mle = mle_tmp
                self.max_log_like = max_log_like_tmp
        self.logger.info("UQpy: ML estimation completed.")

    @beartype
    def _evaluate_func_to_minimize(self, one_param: np.ndarray):
        return (-1 * self.inference_model.evaluate_log_likelihood(
            parameters=one_param.reshape((1, -1)), data=self.data)[0])


# src/UQpy/inference/__init__.py

from UQpy.inference.BayesModelSelection import BayesModelSelection
from UQpy.inference.inference_models import *
from UQpy.inference.evidence_methods import *
from UQpy.inference.information_criteria import *
from UQpy.inference.InformationModelSelection import InformationModelSelection
from UQpy.inference.BayesParameterEstimation import BayesParameterEstimation
from UQpy.inference.MLE import MLE

from UQpy.inference.BayesModelSelection import *
from UQpy.inference.InformationModelSelection import *
from UQpy.inference.BayesParameterEstimation import *
from UQpy.inference.MLE import *



# src/UQpy/reliability/SubsetSimulation.py

import copy
import logging

from UQpy.run_model import RunModel
from UQpy.utilities.Utilities import process_random_state
from UQpy.sampling import *


class SubsetSimulation:

    @beartype
    def __init__(
        self,
        runmodel_object: RunModel,
        sampling: MCMC,
        samples_init: np.ndarray = None,
        conditional_probability: Annotated[Union[float, int], Is[lambda x: 0 <= x <= 1]] = 0.1,
        nsamples_per_subset: int = 1000,
        max_level: int = 10,
    ):
        """
        Perform Subset Simulation to estimate probability of failure.

        This class estimates probability of failure for a user-defined model using Subset Simulation. The class can
        use one of several mcmc algorithms to draw conditional samples.

        :param runmodel_object: The computational model. It should be of type :class:`.RunModel`.
        :param sampling: Specifies the :class:`MCMC` algorithm.
        :param samples_init: A set of samples from the specified probability distribution. These are the samples from
         the original distribution. They are not conditional samples. The samples must be an array of size
         :code:`samples_number_per_subset x dimension`.
         If `samples_init` is not specified, the :class:`.SubsetSimulation` class will use the :class:`.MCMC` class
         created with the aid of `sampling` to draw the initial samples.
        :param conditional_probability: Conditional probability for each conditional level.
        :param nsamples_per_subset: Number of samples to draw in each conditional level.
        :param max_level: Maximum number of allowable conditional levels.
        """
        # Initialize other attributes
        self._sampling_class = sampling
        self._random_state = process_random_state(sampling._random_state)
        self.runmodel_object = runmodel_object
        self.samples_init = samples_init
        self.conditional_probability = conditional_probability
        self.nsamples_per_subset = nsamples_per_subset
        self.max_level = max_level
        self.logger = logging.getLogger(__name__)

        self.mcmc_objects = [sampling]

        self.samples: list = list()
        """A list of arrays containing the samples in each conditional level. The size of the list is equal to the 
        number of levels."""
        self.performance_function_per_level: list = []
        """A list of arrays containing the evaluation of the performance function at each sample in each conditional 
        level. The size of the list is equal to the number of levels."""
        self.performance_threshold_per_level: list = []
        """Threshold value of the performance function for each conditional level. The size of the list is equal to the 
        number of levels."""

        self.logger.info("UQpy: Running Subset Simulation with mcmc of type: " + str(type(sampling)))
        self.failure_probability: float = None
        """Probability of failure estimate."""
        self.independent_chains_CoV: float = None
        """Coefficient of variation of the probability of failure estimate assuming independent chains."""
        self.dependent_chains_CoV: float = None
        """Coefficient of variation of the probability of failure estimate with dependent chains."""

        [self.failure_probability, self.independent_chains_CoV, self.dependent_chains_CoV] = self._run()

        self.logger.info("UQpy: Subset Simulation Complete!")

    # The run function executes the chosen subset simulation algorithm
    def _run(self):
        """
        Execute subset simulation

        This is an instance method that runs subset simulation. It is automatically called when the
        :class:`.SubsetSimulation` class is instantiated.
        """
        step = 0
        n_keep = int(self.conditional_probability * self.nsamples_per_subset)
        d12 = []
        d22 = []

        # Generate the initial samples - Level 0
        # Here we need to make sure that we have good initial samples from the target joint density.
        if self.samples_init is None:
            self.logger.warning(
                "UQpy: You have not provided initial samples.\n Subset simulation is highly sensitive "
                "to the initial sample set. It is recommended that the user either:\n"
                "- Provide an initial set of samples (samples_init) known to follow the distribution; "
                "or\n - Provide a robust mcmc object that will draw independent initial samples from "
                "the distribution."
            )
            self.mcmc_objects[0].run(nsamples=self.nsamples_per_subset)
            self.samples.append(self.mcmc_objects[0].samples)
        else:
            self.samples.append(self.samples_init)

        # Run the model for the initial samples, sort them by their performance function, and identify the
        # conditional level
        self.runmodel_object.run(samples=np.atleast_2d(self.samples[step]))
        self.performance_function_per_level.append(np.squeeze(self.runmodel_object.qoi_list))
        g_ind = np.argsort(self.performance_function_per_level[step])
        self.performance_threshold_per_level.append(self.performance_function_per_level[step][g_ind[n_keep - 1]])

        # Estimate coefficient of variation of conditional probability of first level
        d1, d2 = self._compute_coefficient_of_variation(step)
        d12.append(d1 ** 2)
        d22.append(d2 ** 2)

        self.logger.info("UQpy: Subset Simulation, conditional level 0 complete.")

        while self.performance_threshold_per_level[step] > 0 and step < self.max_level:

            # Increment the conditional level
            step += 1

            # Initialize the samples and the performance function at the next conditional level
            self.samples.append(np.zeros_like(self.samples[step - 1]))
            self.samples[step][:n_keep] = self.samples[step - 1][g_ind[0:n_keep], :]
            self.performance_function_per_level.append(np.zeros_like(self.performance_function_per_level[step - 1]))
            self.performance_function_per_level[step][:n_keep] = self.performance_function_per_level[step - 1][g_ind[:n_keep]]

            # Unpack the attributes
            new_sampler = copy.deepcopy(self._sampling_class)
            new_sampler.seed = np.atleast_2d(self.samples[step][:n_keep, :])
            new_sampler.n_chains = len(np.atleast_2d(self.samples[step][:n_keep, :]))
            new_sampler.random_state = process_random_state(self._random_state)
            self.mcmc_objects.append(new_sampler)

            # Set the number of samples to propagate each chain (n_prop) in the conditional level
            n_prop_test = (self.nsamples_per_subset / self.mcmc_objects[step].n_chains)
            if n_prop_test.is_integer():
                n_prop = (self.nsamples_per_subset // self.mcmc_objects[step].n_chains)
            else:
                raise AttributeError(
                    "UQpy: The number of samples per subset (nsamples_per_ss) must be an integer multiple of "
                    "the number of mcmc chains.")

            # Propagate each chain n_prop times and evaluate the model to accept or reject.
            for i in range(n_prop - 1):

                # Propagate each chain
                if i == 0:
                    self.mcmc_objects[step].run(nsamples=2 * self.mcmc_objects[step].n_chains)
                else:
                    self.mcmc_objects[step].run(nsamples=self.mcmc_objects[step].n_chains)

                # Decide whether a new simulation is needed for each proposed state
                a = self.mcmc_objects[step].samples[i * n_keep : (i + 1) * n_keep, :]
                b = self.mcmc_objects[step].samples[(i + 1) * n_keep : (i + 2) * n_keep, :]
                test1 = np.equal(a, b)
                test = np.logical_and(test1[:, 0], test1[:, 1])

                # Pull out the indices of the false values in the test list
                ind_false = [i for i, val in enumerate(test) if not val]
                # Pull out the indices of the true values in the test list
                ind_true = [i for i, val in enumerate(test) if val]

                # Do not run the model for those samples where the mcmc state remains unchanged.
                self.samples[step][[x + (i + 1) * n_keep for x in ind_true], :] = \
                    self.mcmc_objects[step].samples[ind_true, :]
                self.performance_function_per_level[step][[x + (i + 1) * n_keep for x in ind_true]] = \
                    self.performance_function_per_level[step][ind_true]

                # Run the model at each of the new sample points
                x_run = self.mcmc_objects[step].samples[[x + (i + 1) * n_keep for x in ind_false], :]
                if x_run.size != 0:
                    self.runmodel_object.run(samples=x_run)

                    # Temporarily save the latest model runs
                    g_temp = np.asarray(self.runmodel_object.qoi_list[-len(x_run) :])

                    # Accept the states with g <= g_level
                    ind_accept = np.where(g_temp <= self.performance_threshold_per_level[step - 1])[0]
                    for ii in ind_accept:
                        self.samples[step][(i + 1) * n_keep + ind_false[ii]] = x_run[ii]
                        self.performance_function_per_level[step][(i + 1) * n_keep + ind_false[ii]] = g_temp[ii]

                    # Reject the states with g > g_level
                    ind_reject = np.where(g_temp > self.performance_threshold_per_level[step - 1])[0]
                    for ii in ind_reject:
                        self.samples[step][(i + 1) * n_keep + ind_false[ii]] =\
                            self.samples[step][i * n_keep + ind_false[ii]]
                        self.performance_function_per_level[step][(i + 1) * n_keep + ind_false[ii]] = \
                            self.performance_function_per_level[step][i * n_keep + ind_false[ii]]

            g_ind = np.argsort(self.performance_function_per_level[step])
            self.performance_threshold_per_level.append(self.performance_function_per_level[step][g_ind[n_keep]])

            # Estimate coefficient of variation of conditional probability of first level
            d1, d2 = self._compute_coefficient_of_variation(step)
            d12.append(d1 ** 2)
            d22.append(d2 ** 2)

            self.logger.info("UQpy: Subset Simulation, conditional level " + str(step) + " complete.")

        n_fail = len([value for value in self.performance_function_per_level[step] if value < 0])

        failure_probability = (self.conditional_probability ** step * n_fail / self.nsamples_per_subset)
        probability_cov_independent = np.sqrt(np.sum(d12))
        probability_cov_dependent = np.sqrt(np.sum(d22))

        return failure_probability, probability_cov_independent, probability_cov_dependent

    def _compute_coefficient_of_variation(self, step):
        # Here, we assume that the initial samples are drawn to be uncorrelated such that the correction factors do not
        # need to be computed.
        if step == 0:
            independent_chains_cov = np.sqrt((1 - self.conditional_probability)
                                             / (self.conditional_probability * self.nsamples_per_subset))
            dependent_chains_cov = np.sqrt((1 - self.conditional_probability)
                                           / (self.conditional_probability * self.nsamples_per_subset))
        else:
            n_c = int(self.conditional_probability * self.nsamples_per_subset)
            n_s = int(1 / self.conditional_probability)
            indicator = np.reshape(self.performance_function_per_level[step] < self.performance_threshold_per_level[step], (n_s, n_c))
            gamma = self._correlation_factor_gamma(indicator, n_s, n_c)
            g_temp = np.reshape(self.performance_function_per_level[step], (n_s, n_c))
            beta_hat = self._correlation_factor_beta(g_temp, step)

            independent_chains_cov = np.sqrt(((1 - self.conditional_probability)
                                              / (self.conditional_probability * self.nsamples_per_subset))
                                             * (1 + gamma))
            dependent_chains_cov = np.sqrt(((1 - self.conditional_probability)
                                            / (self.conditional_probability * self.nsamples_per_subset))
                                           * (1 + gamma + beta_hat))

        return independent_chains_cov, dependent_chains_cov

    # Computes the conventional correlation factor gamma from Au and Beck
    def _correlation_factor_gamma(self, indicator, n_s, n_c):
        gam = np.zeros(n_s - 1)
        r = np.zeros(n_s)

        ii = indicator * 1
        r_ = ii @ ii.T / n_c - self.conditional_probability ** 2
        for i in range(r_.shape[0]):
            r[i] = np.sum(np.diag(r_, i)) / (r_.shape[0] - i)

        r0 = 0.1 * (1 - 0.1)
        r = r / r0

        for i in range(n_s - 1):
            gam[i] = (1 - ((i + 1) / n_s)) * r[i + 1]
        gam = 2 * np.sum(gam)

        return gam

    # Computes the updated correlation factor beta from Shields et al.
    def _correlation_factor_beta(self, g, step):
        beta = 0
        for i in range(np.shape(g)[1]):
            for j in range(i + 1, np.shape(g)[1]):
                if g[0, i] == g[0, j]:
                    beta += 1
        beta *= 2

        ar = np.asarray(self.mcmc_objects[step].acceptance_rate)
        ar_mean = np.mean(ar)

        factor = sum((1 - (i + 1) * np.shape(g)[0] / np.shape(g)[1]) * (1 - ar_mean) for i in range(np.shape(g)[0] - 1))
        factor = factor * 2 + 1

        beta = beta / np.shape(g)[1] * factor

        return beta


# src/UQpy/reliability/__init__.py

from UQpy.reliability.SubsetSimulation import SubsetSimulation
from UQpy.reliability.taylor_series import *

from . import TaylorSeries


# src/UQpy/run_model/RunModel.py

# UQpy is distributed under the MIT license.
#
# Copyright (C) 2018  -- Michael D. Shields
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
# Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import logging
import os
import pickle
from typing import Union

import numpy as np
from beartype import beartype

from UQpy.utilities.ValidationTypes import NumpyFloatArray


class RunModel:
    # Authors:
    # B.S. Aakash, Lohit Vandanapu, Michael D.Shields
    #
    # Last
    # modified: 5 / 8 / 2020 by Michael D.Shields
    @beartype
    def __init__(
            self,
            model,
            samples: Union[list, NumpyFloatArray] = None,
            ntasks: int = 1,
            cores_per_task: int = 1,
            nodes: int = 1,
            resume: bool = False,
    ):
        """
        Run a computational model at specified sample points.

        This class is the interface between :py:mod:`UQpy` and computational models. The model is called in a Python script whose
        name must be passed as one the arguments to the :class:`.RunModel` call. If the model is in Python, :py:mod:`UQpy` import the
        model and executes it directly. If the model is not in Python, :class:`.RunModel` must be provided the name of a template
        input file, the name of the Python script that runs the model, and an (optional) output Python script.


        :param samples: Samples to be passed as inputs to the model.
         Regardless of data type, the first dimension of ``samples`` must be equal to the number of samples at which
         to execute the model. That is, ``len(samples) = nsamples``.

         Regardless of data type, the second dimension of ``samples`` must be equal to the number of variables to
         to pass for each model evaluation. That is, ``len(samples[0]) = n_vars``. Each variable need not be a scalar.
         Variables may be scalar, vector, matrix, or tensor type (i.e. `float`, `list`, or `ndarray`).

         If `samples` are not passed, a :class:`.RunModel` object will be instantiated that can be used later, with the
         :meth:`run` method, to evaluate the model.

         Used in both python and third-party model execution.

        :param ntasks: Number of tasks to be run in parallel. By default, ``ntasks = 1`` and the models are executed
         serially.

         Setting ntasks equal to a positive integer greater than 1 will trigger the parallel workflow.

         `ntasks` is used for both the Python and third-party model workflows. :class:`.RunModel` uses `GNU parallel` to
         execute third-party models in parallel and the multiprocessing module to execute Python models in parallel.

        :param cores_per_task: Number of cores to be used by each task. In cases where a third-party model runs across
         multiple CPUs, this optional attribute allocates the necessary resources to each model evaluation.

         `cores_per_task` is not used in the Python model workflow.

        :param nodes: Number of nodes across which to distribute individual tasks on an HPC cluster in the third-party
         model workflow. If more than one compute node is necessary to execute individual runs in parallel, `nodes` must
         be specified.
        """
        self.logger = logging.getLogger(__name__)
        self.model = model
        # Save option for resuming parallel execution
        self.resume = resume

        self.nodes = nodes
        self.ntasks = ntasks
        self.cores_per_task = cores_per_task

        self.is_serial = ntasks <= 1 and cores_per_task <= 1 and nodes <= 1

        # Initialize sample related variables
        self.samples: NumpyFloatArray = []
        """Internally, :class:`.RunModel` converts the input `samples` into a numpy `ndarray` with at least two 
        dimension where the first dimension of the :class:`numpy.ndarray` corresponds to a single sample to be executed 
        by the model."""
        self.samples = np.atleast_2d(self.samples)
        self.qoi_list: list = []
        """A list containing the output quantities of interest

        In the third-party model workflow, these output quantities of interest are extracted from the model output files
        by `output_script`.

        In the Python model workflow, the returned quantity of interest from the model evaluations is stored as
        :py:attr:`qoi_list`.

        This attribute is commonly used for adaptive algorithms that employ learning functions based on previous model
        evaluations."""
        self.n_existing_simulations: int = 0
        """Number of pre-existing model evaluations, prior to a new :meth:`run` method call.

        If the :meth:`run` methods has previously been called and model evaluations performed, subsequent calls to the
        :meth:`run` method will be appended to the :class:`RunModel` object. :py:attr:`nexist` stores the number of 
        previously existing model evaluations."""
        self.n_new_simulations: int = 0
        """Number of model evaluations to be performed, ``nsim = len(samples)``."""

        # Check if samples are provided.
        if samples is None:
            self.logger.info("\nUQpy: No samples are provided. Creating the object and building the model directory.\n")
        elif isinstance(samples, (list, np.ndarray)):
            self.run(samples)
        else:
            raise ValueError("\nUQpy: samples must be passed as a list or numpy ndarray\n")

    def run(self, samples=None, append_samples=True):
        """
        Execute a computational model at given sample values.

        If `samples` are passed when defining the :class:`.RunModel` object, the :meth:`run` method is called automatically.

        The :meth:`run` method may also be called directly after defining the :meth:`RunModel` object.

        :param list samples: Samples to be passed as inputs to the model defined by the :class:`.RunModel` object.

         Regardless of data type, the first dimension of ``samples`` must be equal to the number of samples at which
         to execute the model. That is, ``len(samples) = nsamples``.

         Regardless of data type, the second dimension of ``samples`` must be equal to the number of variables to
         to pass for each model evaluation. That is, ``len(samples[0]) = n_vars``. Each variable need not be a
         scalar. Variables may be scalar, vector, matrix, or tensor type (i.e. `float`, `list`, or `ndarray`).

         Used in both python and third-party model execution.
        :param bool append_samples: Append over overwrite existing samples and model evaluations.

         If ``append_samples = False``, all previous samples and the corresponding quantities of interest from their
         model evaluations are deleted.

         If ``append_samples = True``, samples and their resulting quantities of interest are appended to the
         existing ones.
        """
        # Ensure the input samples have the correct structure
        # --> If a list is provided, convert to at least 2d ndarray. dim1 = nsim, dim2 = n_vars
        # --> If 1D array/list is provided, convert it to a 2d array. dim1 = 1, dim2 = n_vars
        # --> If samples cannot be converted to an array, this will fail.
        samples = np.atleast_2d(samples)

        # Number of simulations to be performed
        self.n_new_simulations = len(samples)

        # If append_samples is False, a new set of samples is created, the previous ones are deleted!
        if not append_samples:
            self.samples = []
            self.samples = np.atleast_2d(self.samples)
            self.qoi_list = []

        # Check if samples already exist, if yes append new samples to old ones
        # if not self.samples:  # There are currently no samples
        if self.samples.size == 0:
            self.n_existing_simulations = 0
            self.samples = samples

        else:  # Samples already exist in the RunModel object, append new ones
            self.n_existing_simulations = len(self.samples)
            self.samples = np.vstack((self.samples, samples))

        self.model.initialize(samples)

        self.qoi_list.extend(self.serial_execution() if self.is_serial else self.parallel_execution())

        self.model.finalize()

    def parallel_execution(self):
        # TODO: Check if files with the names used below already exist and raise error
        with open('model.pkl', 'wb') as filehandle:
            pickle.dump(self.model, filehandle)
        with open('samples.pkl', 'wb') as filehandle:
            pickle.dump(self.samples, filehandle)
        os.system(f"mpirun python -m "
                  f"UQpy.run_model.model_execution.ParallelExecution {self.n_existing_simulations} "
                  f"{self.n_new_simulations}")
        with open('qoi.pkl', 'rb') as filehandle:
            results = pickle.load(filehandle)

        os.remove("model.pkl")
        os.remove("samples.pkl")
        os.remove("qoi.pkl")

        self.logger.info("\nUQpy: Parallel execution of the python model complete.\n")
        return results

    def serial_execution(self):
        results = []
        for i in range(self.n_existing_simulations, self.n_existing_simulations + self.n_new_simulations):
            sample = self.model.preprocess_single_sample(i, self.samples[i])

            execution_output = self.model.execute_single_sample(i, sample)

            results.append(self.model.postprocess_single_file(i, execution_output))

        self.logger.info("\nUQpy: Serial execution of the python model complete.\n")
        return results


# src/UQpy/run_model/__init__.py

from UQpy.run_model.RunModel import RunModel

from UQpy.run_model.model_execution import *


# src/UQpy/sampling/AdaptiveKriging.py

import logging
from beartype import beartype
from sklearn.gaussian_process import GaussianProcessRegressor

from UQpy.run_model.RunModel import RunModel
from UQpy.distributions.baseclass import Distribution
from UQpy.sampling.stratified_sampling.LatinHypercubeSampling import LatinHypercubeSampling
from UQpy.sampling.adaptive_kriging_functions.baseclass.LearningFunction import (
    LearningFunction,
)
from UQpy.distributions import DistributionContinuous1D, JointIndependent
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria import Random
from UQpy.surrogates.baseclass import Surrogate
from UQpy.utilities.ValidationTypes import *
from UQpy.utilities.Utilities import process_random_state

SurrogateType = Union[Surrogate, GaussianProcessRegressor,
                      Annotated[object, Is[lambda x: hasattr(x, 'fit') and hasattr(x, 'predict')]]]


class AdaptiveKriging:
    @beartype
    def __init__(
            self,
            distributions: Union[Distribution, list[Distribution]],
            runmodel_object: RunModel,
            surrogate: SurrogateType,
            learning_function: LearningFunction,
            samples: Numpy2DFloatArray = None,
            nsamples: PositiveInteger = None,
            learning_nsamples: PositiveInteger = None,
            qoi_name: str = None,
            n_add: int = 1,
            random_state: RandomStateType = None,
    ):
        """
        Adaptively sample for construction of a kriging surrogate for different objectives including reliability,
        optimization, and global fit.

        :param distributions: List of :class:`.Distribution` objects corresponding to each random variable.
        :param runmodel_object: A :class:`.RunModel` object, which is used to evaluate the model.
        :param surrogate: A kriging surrogate model, this object must have :meth:`fit` and :meth:`predict` methods.
         May be an object of the :py:meth:`UQpy` :class:`Kriging` class or an object of the :py:mod:`scikit-learn`
         :class:`GaussianProcessRegressor`
        :param learning_function: Learning function used as the selection criteria to identify new samples.
        :param samples: The initial samples at which to evaluate the model.
         Either `samples` or `nstart` must be provided.
        :param nsamples: Total number of samples to be drawn (including the initial samples).
         If `nsamples` and `samples` are provided when instantiating the class, the :meth:`run` method will
         automatically be called. If either `nsamples` or `samples` is not provided, :class:`.AdaptiveKriging`
         can be executed by invoking the :meth:`run` method and passing `nsamples`.
        :param learning_nsamples: Number of samples generated for evaluation of the learning function. Samples for
         the learning set are drawn using :class:`.LatinHypercubeSampling`.
        :param qoi_name: Name of the quantity of interest. If the quantity of interest is a dictionary, this is used to
         convert it to a list
        :param n_add: Number of samples to be added per iteration.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
         If an :any:`int` is provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise,
         the object itself can be passed directly.
        """
        # Initialize the internal variables of the class.
        self.runmodel_object = runmodel_object
        self.samples: Numpy2DFloatArray = np.array(samples)
        """contains the samples at which the model is evaluated."""
        self.learning_nsamples = learning_nsamples
        self.initial_nsamples = None
        self.logger = logging.getLogger(__name__)
        self.qoi_name = qoi_name

        self.learning_function = learning_function
        self.learning_set = None
        self.dist_object = distributions
        self.nsamples = nsamples

        self.moments = None
        self.n_add = n_add
        self.indicator = False
        self.pf = []
        self.cov_pf = []
        self.dimension = 0
        self.qoi = None
        self.prediction_model = None

        # Initialize and run preliminary error checks.
        self.dimension = len(distributions)

        if samples is not None and self.dimension != self.samples.shape[1]:
            raise NotImplementedError("UQpy Error: Dimension of samples and distribution are inconsistent.")

        if isinstance(distributions, list):
            for i in range(len(distributions)):
                if not isinstance(distributions[i], DistributionContinuous1D):
                    raise TypeError("UQpy: A DistributionContinuous1D object must be provided.")
        elif not isinstance(distributions, (DistributionContinuous1D, JointIndependent)):
            raise TypeError("UQpy: A DistributionContinuous1D or JointInd object must be provided.")

        self.random_state = process_random_state(random_state)

        self.surrogate = surrogate

        self.logger.info("UQpy: Adaptive Kriging - Running the initial sample set using RunModel.")

        # Evaluate model at the training points
        if len(self.runmodel_object.qoi_list) == 0 and samples is not None:
            self.runmodel_object.run(samples=self.samples, append_samples=False)
        if samples is not None and len(self.runmodel_object.qoi_list) != self.samples.shape[0]:
            raise NotImplementedError("UQpy: There should be no model evaluation or Number of samples and model "
                                      "evaluation in RunModel object should be same.")

        if self.nsamples is not None and samples is not None:
            self.run(nsamples=self.nsamples)

    def run(
            self,
            nsamples: int,
            samples: np.ndarray = None,
            append_samples: bool = True,
            initial_nsamples: int = None,
    ):
        """
        Execute the :class:`.AdaptiveKriging` learning iterations.

        The :meth:`run` method is the function that performs iterations in the :class:`.AdaptiveKriging` class. If
        `nsamples` is provided when defining the :class:`.AdaptiveKriging` object, the :meth:`run` method is
        automatically called. The user may also call the :meth:`run` method directly to generate samples.
        The :meth:`run` method of the :class:`.AdaptiveKriging` class can be invoked many times.

        The :meth:`run` method has no returns, although it creates and/or appends the :py:attr:`samples` attribute of
        the :class:`.AdaptiveKriging` class.

        :param nsamples: Total number of samples to be drawn (including the initial samples).
        :param samples: Samples at which to evaluate the model.
        :param append_samples: Append new samples and model evaluations to the existing samples and model evaluations.
         If ``append_samples = False``, all previous samples and the corresponding quantities of interest from their
         model evaluations are deleted.
         If ``append_samples = True``, samples and their resulting quantities of interest are appended to the
         existing ones.
        :param initial_nsamples: Number of initial samples, randomly generated using
         :class:`.LatinHypercubeSampling` class.
        """

        self.nsamples = nsamples
        self.initial_nsamples = initial_nsamples

        if samples is not None:
            # New samples are appended to existing samples, if append_samples is TRUE
            if append_samples:
                if len(self.samples.shape) == 0:
                    self.samples = np.array(samples)
                else:
                    self.samples = np.vstack([self.samples, np.array(samples)])
            else:
                self.samples = np.array(samples)
                self.runmodel_object.qoi_list = []

            self.logger.info("UQpy: Adaptive Kriging - Evaluating the model at the sample set using RunModel.")

            self.runmodel_object.run(samples=samples, append_samples=append_samples)
        elif len(self.samples.shape) == 0:
            if self.initial_nsamples is None:
                raise NotImplementedError("UQpy: User should provide either 'samples' or 'nstart' value.")
            self.logger.info("UQpy: Adaptive Kriging - Generating the initial sample set using Latin hypercube "
                             "sampling.")

            random_criterion = Random()
            latin_hypercube_sampling = LatinHypercubeSampling(
                distributions=self.dist_object,
                nsamples=2,
                criterion=random_criterion,
                random_state=self.random_state)
            self.samples = latin_hypercube_sampling._samples
            self.runmodel_object.run(samples=self.samples)

        self.logger.info("UQpy: Performing AK-MCS design...")

        # If the quantity of interest is a dictionary, convert it to a list
        self._convert_qoi_tolist()

        # Train the initial kriging model.
        self.surrogate.fit(self.samples, self.qoi)
        self.prediction_model = self.surrogate.predict

        # ---------------------------------------------
        # Primary loop for learning and adding samples.
        # ---------------------------------------------

        for i in range(self.samples.shape[0], self.nsamples):
            # Initialize the population of samples at which to evaluate the learning function and from which to draw
            # in the sampling.
            random_criterion = Random()
            lhs = LatinHypercubeSampling(
                random_state=self.random_state,
                distributions=self.dist_object,
                nsamples=self.learning_nsamples,
                criterion=random_criterion,
            )

            self.learning_set = lhs._samples.copy()

            # Find all of the points in the population that have not already been integrated into the training set
            rest_pop = np.array([x for x in self.learning_set.tolist() if x not in self.samples.tolist()])

            # Apply the learning function to identify the new point to run the model.

            # new_point, lf, ind = self.learning_function(self.krig_model, rest_pop, **kwargs)
            new_point, lf, ind = self.learning_function.evaluate_function(
                distributions=self.dist_object,
                n_add=self.n_add,
                surrogate=self.surrogate,
                population=rest_pop,
                qoi=self.qoi,
                samples=self.samples,
            )

            # Add the new points to the training set and to the sample set.
            self.samples = np.vstack([self.samples, np.atleast_2d(new_point)])

            # Run the model at the new points
            self.runmodel_object.run(samples=new_point, append_samples=True)

            # If the quantity of interest is a dictionary, convert it to a list
            self._convert_qoi_tolist()

            # Retrain the surrogate model
            self.surrogate.fit(self.samples, self.qoi, optimizations_number=1)
            self.prediction_model = self.surrogate.predict

            # Exit the loop, if error criteria is satisfied
            if ind:
                self.logger.info("UQpy: Learning stops at iteration: %(iteration)s" % {"iteration": i})
                break

            self.logger.info("Iteration: %(iteration)s" % {"iteration": i})

        self.logger.info("UQpy: Adaptive Kriging complete")

    def _convert_qoi_tolist(self):
        self.qoi = [None] * len(self.runmodel_object.qoi_list)
        if type(self.runmodel_object.qoi_list[0]) is dict:
            for j in range(len(self.runmodel_object.qoi_list)):
                self.qoi[j] = self.runmodel_object.qoi_list[j][self.qoi_name]
        else:
            self.qoi = self.runmodel_object.qoi_list


# src/UQpy/sampling/ImportanceSampling.py

import logging
from typing import Union, Callable

from beartype import beartype

from UQpy.utilities.ValidationTypes import PositiveInteger, RandomStateType, NumpyFloatArray
from UQpy.utilities.Utilities import process_random_state
from UQpy.distributions import Distribution
import numpy as np


class ImportanceSampling:

    # Last Modified: 10/05/2020 by Audrey Olivier
    @beartype
    def __init__(self,
                 pdf_target: Callable = None,
                 log_pdf_target: Callable = None,
                 args_target: tuple = None,
                 proposal: Union[None, Distribution] = None,
                 random_state: RandomStateType = None,
                 nsamples: PositiveInteger = None):
        """
        Sample from a user-defined target density using importance sampling.

        :param pdf_target: Callable that evaluates the pdf of the target distribution. Either `log_pdf_target` or
         `pdf_target` must be specified (the former is preferred).
        :param log_pdf_target: Callable that evaluates the log-pdf of the target distribution. Either `log_pdf_target`
         or `pdf_target` must be specified (the former is preferred).
        :param args_target: Positional arguments of the target log_pdf / pdf callable.
        :param proposal: Proposal to sample from. This :class:`.Distribution` object must have an :py:meth:`rvs` method
         and a `log_pdf` (or pdf) method.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is
         :any:`None`.

         If an :any:`int` is provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise,
         the object itself can be passed directly.
        :param nsamples: Number of samples to generate - see :meth:`run` method. If not :any:`None`, the :py:meth:`run`
         method is called when the object is created. Default is :any:`None`.
        """
        # Initialize proposal: it should have an rvs and log pdf or pdf method
        self.evaluate_log_target = None
        self.proposal = proposal
        self._args_target = args_target
        self.log_pdf_target = log_pdf_target
        self.pdf_target = pdf_target

        self.logger = logging.getLogger(__name__)
        self.random_state = process_random_state(random_state)

        # Initialize the samples and weights
        self.samples: NumpyFloatArray = None
        """Set of samples, :class:`numpy.ndarray` of shape :code:`(nsamples, dimensions)`"""
        self.unnormalized_log_weights: NumpyFloatArray = None
        """Unnormalized log weights, i.e., :code:`log_w(x) = log_target(x) - log_proposal(x)`, :class:`numpy.ndarray` of 
        shape :code:`(nsamples, )`"""
        self.weights: NumpyFloatArray = None
        """Importance weights, weighted so that they sum up to 1, :class:`numpy.ndarray` of shape :code:`(nsamples, )`
        """
        self.unweighted_samples: NumpyFloatArray = None
        """Set of un-weighted samples (useful for instance for plotting), computed by calling the :meth:`resample` 
        method"""

        # Run IS if nsamples is provided
        if nsamples is not None and nsamples != 0:
            self.run(nsamples)

    def _preprocess_proposal(self):
        if not isinstance(self.proposal, Distribution):
            raise TypeError("UQpy: The proposal should be of type Distribution.")
        if not hasattr(self.proposal, "rvs"):
            raise AttributeError("UQpy: The proposal should have an rvs method")
        if not hasattr(self.proposal, "log_pdf"):
            if not hasattr(self.proposal, "pdf"):
                raise AttributeError("UQpy: The proposal should have a log_pdf or pdf method")
            self.proposal.log_pdf = lambda x: np.log(
                np.maximum(self.proposal.pdf(x), 10 ** (-320) * np.ones((x.shape[0],))))

    @beartype
    def run(self, nsamples: PositiveInteger):
        """
        Generate and weight samples.

        This function samples from the proposal and appends samples to existing ones (if any). It then weights the
        samples as :code:`log_w_unnormalized) = log(target)-log(proposal)`.

        :param nsamples: Number of weighted samples to generate.

        This function has no returns, but it updates the output attributes :py:attr:`samples`,
        :py:attr:`unnormalized_log_weights` and :py:attr:`weights` of the :class:`.ImportanceSampling` object.
        """
        if self.evaluate_log_target is None:
            self._preprocess_proposal()
            self.evaluate_log_target = self._preprocess_target(
                log_pdf_=self.log_pdf_target,
                pdf_=self.pdf_target,
                args=self._args_target, )

        self.logger.info("UQpy: Running Importance Sampling...")
        # Sample from proposal
        new_samples = self.proposal.rvs(nsamples=nsamples, random_state=self.random_state)
        # Compute un-scaled weights of new samples
        a = self.evaluate_log_target(x=new_samples)
        new_log_weights = self.evaluate_log_target(x=new_samples) - self.proposal.log_pdf(x=new_samples)

        # Save samples and weights (append to existing if necessary)
        if self.samples is None:
            self.samples = new_samples
            self.unnormalized_log_weights = new_log_weights
        else:
            self.samples = np.concatenate([self.samples, new_samples], axis=0)
            self.unnormalized_log_weights = np.concatenate(
                [self.unnormalized_log_weights, new_log_weights], axis=0)

        # Take the exponential and normalize the weights
        weights = np.exp(self.unnormalized_log_weights - max(self.unnormalized_log_weights))
        # note: scaling with max avoids having NaN of Inf when taking the exp
        sum_w = np.sum(weights, axis=0)
        self.weights = weights / sum_w
        self.logger.info("UQpy: Importance Sampling performed successfully")

        # If a set of unweighted samples exist, delete them as they are not representative of the distribution anymore
        if self.unweighted_samples is not None:
            self.logger.info("UQpy: unweighted samples are being deleted, call the resample method to regenerate them")
            self.unweighted_samples = None

    def resample(self, method: str = "multinomial", nsamples: int = None):
        """
        Resample to get a set of un-weighted samples that represent the target pdf.

        Utility function that creates a set of un-weighted samples from a set of weighted samples. Can be useful for
        plotting for instance.

        The :meth:`resample` method is not called automatically when instantiating the :class:`.ImportanceSampling`
        class or when invoking its :meth:`run` method.

        :param method: Resampling method, as of V4 only multinomial resampling is supported. Default: 'multinomial'.
        :param nsamples: Number of un-weighted samples to generate. Default: None (sets `nsamples` equal to the
         number of existing weighted samples).
        """

        if nsamples is None:
            nsamples = self.samples.shape[0]
        if method != "multinomial":
            raise ValueError("Exit code: Current available method: multinomial")
        multinomial_run = self.random_state.multinomial(nsamples, self.weights, size=1)[0]
        idx = []
        for j in range(self.samples.shape[0]):
            if multinomial_run[j] > 0:
                idx.extend([j for _ in range(multinomial_run[j])])
        self.unweighted_samples = self.samples[idx, :]

    @staticmethod
    def _preprocess_target(log_pdf_, pdf_, args):
        # log_pdf is provided
        if log_pdf_ is not None:
            if not callable(log_pdf_):
                raise TypeError("UQpy: log_pdf_target must be a callable")
            if args is None:
                args = ()
            evaluate_log_pdf = lambda x: log_pdf_(x, *args)
        elif pdf_ is not None:
            if not callable(pdf_):
                raise TypeError("UQpy: pdf_target must be a callable")
            if args is None:
                args = ()
            evaluate_log_pdf = lambda x: np.log(np.maximum(pdf_(x, *args), 10 ** (-320) * np.ones((x.shape[0],))))
        else:
            raise ValueError("UQpy: log_pdf_target or pdf_target should be provided.")
        return evaluate_log_pdf


# src/UQpy/sampling/MonteCarloSampling.py

from typing import Union
from typing import Optional
from beartype import beartype

from UQpy.utilities.ValidationTypes import RandomStateType, PositiveInteger, NumpyFloatArray
from UQpy.distributions import *
from UQpy.utilities.Utilities import process_random_state
import numpy as np
import logging


class MonteCarloSampling:

    @beartype
    def __init__(
        self,
        distributions: Union[Distribution, list[Distribution]],
        nsamples: Optional[int] = None,
        random_state: RandomStateType = None,
    ):
        """
        Perform Monte Carlo sampling (MCS) of random variables.

        :param distributions: Probability distribution of each random variable.
        :param nsamples: Number of samples to be drawn from each distribution. The :meth:`run` method is
         automatically called if `nsamples` is provided. If `nsamples` is not provided,
         then the :class:`.MonteCarloSampling` object is created but samples are not generated.
        :param random_state: Random seed used to initialize the pseudo-random number generator. If an :any:`int` is
         provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise, the
         object itself can be passed directly.
        """
        self.logger = logging.getLogger(__name__)
        self.random_state = process_random_state(random_state)

        self.list = False
        self.array = False
        self._process_distributions(distributions)

        self.samples: NumpyFloatArray = None
        """Generated samples.
        
        If a list of :class:`.DistributionContinuous1D` objects is provided for `distributions`, then `samples` is an
        :class:`numpy.ndarray` with ``samples.shape=(nsamples, len(distributions))``.
        
        If a :class:`.DistributionContinuous1D` object is provided for `distributions` then `samples` is an array with
        ``samples.shape=(nsamples, 1)``.
        
        If a :class:`.DistributionContinuousND` object is provided for `distributions` then `samples` is an array with
        ``samples.shape=(nsamples, ND)``.
        
        If a list of mixed :class:`.DistributionContinuous1D` and :class:`.DistributionContinuousND` objects is provided
        then `samples` is a list with ``len(samples)=nsamples`` and ``len(samples[i]) = len(distributions)``.
        """
        self.x = None
        self.samplesU01: NumpyFloatArray = None
        """
        Generated samples transformed to the unit hypercube.
        
        This attribute exists only if the :meth:`transform_u01` method is invoked by the user.
        
        If a list of :class:`.DistributionContinuous1D` objects is provided for `distributions`, then `samplesU01` is an
        :class:`numpy.ndarray` with ``samples.shape=(nsamples, len(distributions))``.
        
        If a :class:`.DistributionContinuous1D` object is provided for `distributions` then `samplesU01` is an array 
        with ``samples.shape=(nsamples, 1)``.
        
        If a :class:`.DistributionContinuousND` object is provided for `distributions` then `samplesU01` is an array 
        with ``samples.shape=(nsamples, ND)``.
        
        If a list of mixed :class:`.DistributionContinuous1D` and :class:`.DistributionContinuousND` objects is provided
        then `samplesU01` is a list with ``len(samples)=nsamples`` and ``len(samples[i]) = len(distributions)``.
        """
        self.nsamples = nsamples

        # Run Monte Carlo sampling
        if nsamples is not None:
            self.run(nsamples=self.nsamples, random_state=self.random_state)

    def _process_distributions(self, distributions):
        if isinstance(distributions, list):
            add_continuous_1d = 0
            add_continuous_nd = 0
            for i in range(len(distributions)):
                if not isinstance(distributions[i], Distribution):
                    raise TypeError("UQpy: A UQpy.Distribution object must be provided.")
                if isinstance(distributions[i], DistributionContinuous1D):
                    add_continuous_1d += 1
                elif isinstance(distributions[i], DistributionND):
                    add_continuous_nd += 1
            if add_continuous_1d == len(distributions):
                self.list = False
                self.array = True
            else:
                self.list = True
                self.array = False
            self.dist_object = distributions
        else:
            self.dist_object = distributions
            self.list = False
            self.array = True

    @beartype
    def run(self, nsamples: PositiveInteger, random_state: RandomStateType = None):
        """
        Execute the random sampling in the :class:`.MonteCarloSampling` class.

        The :meth:`run` method is the function that performs random sampling in the :class:`.MonteCarloSampling` class.
        If `nsamples` is provided, the :meth:`run` method is automatically called when the
        :class:`MonteCarloSampling` object is defined. The user may also call the :meth:`run` method directly to
        generate samples. The :meth:`run` method of the :class:`.MonteCarloSampling` class can be  invoked many times
        and each time the generated samples are appended to the existing samples.

        :param nsamples: Number of samples to be drawn from each distribution.

         If the :meth:`run` method is invoked multiple times, the newly generated samples will be appended to the
         existing samples.
        :param random_state: Random seed used to initialize the pseudo-random number generator.
        """
        # Check if a random_state is provided.
        self.random_state = (process_random_state(random_state) if random_state is not None else self.random_state)

        self.logger.info("UQpy: Running Monte Carlo Sampling.")

        if isinstance(self.dist_object, list):
            temp_samples = []
            for i in range(len(self.dist_object)):
                if hasattr(self.dist_object[i], "rvs"):
                    temp_samples.append(self.dist_object[i].rvs(nsamples=nsamples, random_state=self.random_state))
                else:
                    raise ValueError("UQpy: rvs method is missing.")
            self.x = []
            for j in range(nsamples):
                y = [temp_samples[k][j] for k in range(len(self.dist_object))]
                self.x.append(np.array(y))
        elif hasattr(self.dist_object, "rvs"):
            temp_samples = self.dist_object.rvs(nsamples=nsamples, random_state=self.random_state)
            self.x = temp_samples

        if self.samples is None:
            if isinstance(self.dist_object, list) and self.array is True:
                self.samples = np.hstack(np.array(self.x)).T
            else:
                self.samples = np.array(self.x)
        elif isinstance(self.dist_object, list) and self.array is True:
            self.samples = np.concatenate([self.samples, np.hstack(np.array(self.x)).T], axis=0)
        elif isinstance(self.dist_object, Distribution):
            self.samples = np.vstack([self.samples, self.x])
        else:
            self.samples = np.vstack([self.samples, self.x])
        self.nsamples = len(self.samples)

        self.logger.info("UQpy: Monte Carlo Sampling Complete.")

    def transform_u01(self):
        """
        Transform random samples to uniform on the unit hypercube.

        The :meth:`transform_u01` method has no returns, although it creates and/or appends the :py:attr:`samplesU01`
        attribute of the :meth:`.MonteCarloSampling` class.
        """
        if isinstance(self.dist_object, list) and self.array is True:
            zi = np.zeros_like(self.samples)
            for i in range(self.nsamples):
                z = self.samples[i, :]
                for j in range(len(self.dist_object)):
                    if hasattr(self.dist_object[j], "cdf"):
                        zi[i, j] = self.dist_object[j].cdf(z[j])
                    else:
                        raise ValueError("UQpy: All distributions must have a cdf method.")
            self.samplesU01 = zi

        elif isinstance(self.dist_object, Distribution):
            if not hasattr(self.dist_object, "cdf"):
                raise ValueError("UQpy: All distributions must have a cdf method.")

            zi = np.zeros_like(self.samples)
            for i in range(self.nsamples):
                z = self.samples[i, :]
                zi[i, :] = self.dist_object.cdf(z)
            self.samplesU01 = zi
        elif isinstance(self.dist_object, list) and self.list is True:
            temp_samples_u01 = []
            for i in range(self.nsamples):
                z = self.samples[i][:]
                y = [None] * len(self.dist_object)
                for j in range(len(self.dist_object)):
                    if hasattr(self.dist_object[j], "cdf"):
                        zi = self.dist_object[j].cdf(z[j])
                    else:
                        raise ValueError("UQpy: All distributions must have a cdf method.")
                    y[j] = zi
                temp_samples_u01.append(np.array(y))
            self.samplesU01 = temp_samples_u01


# src/UQpy/sampling/SimplexSampling.py

import scipy.stats as stats
from beartype import beartype
from UQpy.utilities.ValidationTypes import *
from UQpy.utilities.Utilities import process_random_state


class SimplexSampling:
    @beartype
    def __init__(
        self,
        nodes: Union[list, Numpy2DFloatArray] = None,
        nsamples: PositiveInteger = None,
        random_state: RandomStateType = None,
    ):
        """
        Generate uniform random samples inside an n-dimensional simplex.

        :param nodes: The vertices of the simplex.
        :param nsamples: The number of samples to be generated inside the simplex.
         If `nsamples` is provided when the object is defined, the :meth:`run` method will be called
         automatically. If `nsamples` is not provided when the object is defined, the user must invoke the
         :meth:`run` method and specify `nsamples`.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
         If an :any:`int` is provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise,
         the object itself can be passed directly.
        """
        self.samples: NumpyFloatArray = None
        self.nodes = np.atleast_2d(nodes)
        self.nsamples = nsamples

        if self.nodes.shape[0] != self.nodes.shape[1] + 1:
            raise NotImplementedError("UQpy: Size of simplex (nodes) is not consistent.")

        self.random_state = process_random_state(random_state)

        if nsamples is not None:
            self.run(nsamples=nsamples)
            """New random samples distributed uniformly inside the simplex."""

    @beartype
    def run(self, nsamples: PositiveInteger):
        """
        Execute the random sampling in the :class:`.SimplexSampling` class.
        The :meth:`run` method is the function that performs random sampling in the :class:`.SimplexSampling` class.
        If `nsamples` is provided called when the :class:`.SimplexSampling` object is defined, the
        :meth:`run` method is automatically. The user may also call the :meth:`run` method directly to generate samples.
        The :meth:`run` method of the :class:`.SimplexSampling` class can be
        invoked many times and each time the generated samples are appended to the existing samples.

        :param nsamples: Number of samples to be generated inside the simplex.
         If the :meth:`run` method is invoked multiple times, the newly generated samples will be appended to the
         existing samples.
        :return: The :meth:`run` method has no returns, although it creates and/or appends the :py:attr:`samples`
         attribute of the :class:`.SimplexSampling` class.
        """
        self.nsamples = nsamples
        dimension = self.nodes.shape[1]
        if dimension > 1:
            sample = np.zeros([self.nsamples, dimension])
            for i in range(self.nsamples):
                r = np.zeros([dimension])
                ad = np.zeros(shape=(dimension, len(self.nodes)))
                for j in range(dimension):
                    b_ = []
                    for k in range(1, len(self.nodes)):
                        ai = self.nodes[k, j] - self.nodes[k - 1, j]
                        b_.append(ai)
                    ad[j] = np.hstack((self.nodes[0, j], b_))
                    r[j] = stats.uniform.rvs(loc=0, scale=1, random_state=self.random_state) ** (1 / (dimension - j))
                d = np.cumprod(r)
                r_ = np.hstack((1, d))
                sample[i, :] = np.dot(ad, r_)
        else:
            a = min(self.nodes)
            b = max(self.nodes)
            sample = a + (b - a) * stats.uniform.rvs(size=[self.nsamples, dimension], random_state=self.random_state)
        self.samples: NumpyFloatArray = sample



# src/UQpy/sampling/__init__.py

from UQpy.sampling.mcmc import *
from UQpy.sampling.adaptive_kriging_functions import *
from UQpy.sampling.stratified_sampling import *

from UQpy.sampling.AdaptiveKriging import AdaptiveKriging
from UQpy.sampling.ImportanceSampling import ImportanceSampling

from UQpy.sampling.MonteCarloSampling import MonteCarloSampling
from UQpy.sampling.SimplexSampling import SimplexSampling


# src/UQpy/sensitivity/MorrisSensitivity.py

"""This module contains functionality for sensitivity analysis in ``UQpy``.

The module currently contains the following classes:

- ``Morris``: Class to compute sensitivity indices based on the Morris method.
"""
from typing import Union, Annotated

from beartype import beartype
from beartype.vale import Is

from UQpy.utilities.Utilities import process_random_state
from UQpy.utilities.ValidationTypes import RandomStateType, PositiveInteger, NumpyFloatArray
from UQpy.distributions import *
from UQpy.run_model.RunModel import RunModel
import numpy as np
from scipy.stats import randint


class MorrisSensitivity:
    @beartype
    def __init__(
            self,
            runmodel_object: RunModel,
            distributions: Union[JointIndependent, Union[list, tuple]],
            n_levels: Annotated[int, Is[lambda x: x >= 3]],
            delta: Union[float, int] = None,
            random_state: RandomStateType = None,
            n_trajectories: PositiveInteger = None,
            maximize_dispersion: bool = False,
    ):
        """
        Compute sensitivity indices based on the Morris screening method.

        :param runmodel_object: The computational model. It should be of type :class:`.RunModel`. The
         output QoI can be a scalar or vector of length :code:`ny`, then the sensitivity indices of all :code:`ny` outputs are
         computed independently.
        :param distributions: List of :class:`.Distribution` objects corresponding to each random variable, or
         :class:`.JointIndependent` object (multivariate RV with independent marginals).
        :param n_levels: Number of levels that define the grid over the hypercube where evaluation points are
         sampled. Must be an integer :math:`\ge 3`.
        :param delta: Size of the jump between two consecutive evaluation points, must be a multiple of delta should be
         in :code:`{1/(n_levels-1), ..., 1-1/(n_levels-1)}`.
         Default: :math:`delta=\\frac{levels\_number}{2 * (levels\_number-1)}` if `n_levels` is even,
         :math:`delta=0.5` if n_levels is odd.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
        :param n_trajectories: Number of random trajectories, usually chosen between :math:`5` and :math:`10`.
         The number of model evaluations is :code:`n_trajectories * (d+1)`. If None, the `Morris` object is created
         but not run (see :py:meth:`run` method)
        :param maximize_dispersion: If :any:`True`, generate a large number of design trajectories and keep the ones
         that maximize dispersion between all trajectories, allows for a better coverage of the input space.

         Default :any:`False`.
        """
        # Check RunModel object and distributions
        self.runmodel_object = runmodel_object
        marginals = (distributions.marginals if isinstance(distributions, JointIndependent) else distributions)
        self.icdfs = [getattr(dist, "icdf", None) for dist in marginals]
        if any(icdf is None for icdf in self.icdfs):
            raise ValueError("At least one of the distributions provided has a None icdf")
        self.dimension = len(self.icdfs)
        if self.dimension != len(self.runmodel_object.model.var_names):
            raise ValueError("The number of distributions provided does not match the number of RunModel variables")

        self.n_levels = n_levels
        self.delta = delta
        self.check_levels_delta()
        self.random_state = process_random_state(random_state)
        self.maximize_dispersion = maximize_dispersion

        self.trajectories_unit_hypercube: NumpyFloatArray = None
        """Trajectories in the unit hypercube, :class:`numpy.ndarray` of shape :code:`(n_trajectories, d+1, d)`"""
        self.trajectories_physical_space: NumpyFloatArray = None
        """Trajectories in the physical space, :class:`numpy.ndarray` of shape :code:`(n_trajectories, d+1, d)`"""
        self.elementary_effects: NumpyFloatArray = None
        """Elementary effects :math:`EE_{k}`, :class:`numpy.ndarray` of shape :code:`(n_trajectories, d, ny)`."""
        self.mustar_indices: NumpyFloatArray = None
        """First Morris sensitivity index :math:`\mu_{k}^{\star}`, :class:`numpy.ndarray` of shape :code:`(d, ny)`"""
        self.sigma_indices: NumpyFloatArray = None
        """Second Morris sensitivity index :math:`\sigma_{k}`, :class:`numpy.ndarray` of shape :code:`(d, ny)`"""

        if n_trajectories is not None:
            self.run(n_trajectories)

    def check_levels_delta(self):
        # delta should be in {1/(nlevels-1), ..., 1-1/(nlevels-1)}
        if (self.delta is None) and (self.n_levels % 2) == 0:
            # delta = trial_probability / (2 * (trial_probability-1))
            self.delta = self.n_levels / (2 * (self.n_levels - 1))
        elif (self.delta is None) and (self.n_levels % 2) == 1:
            self.delta = (1 / 2)  # delta = (trial_probability-1) / (2 * (trial_probability-1))
        elif not (isinstance(self.delta, (int, float)) and float(self.delta)
                  in [float(j / (self.n_levels - 1)) for j in range(1, self.n_levels - 1)]):
            raise ValueError("UQpy: delta should be in {1/(nlevels-1), ..., 1-1/(nlevels-1)}")

    @beartype
    def run(self, n_trajectories: PositiveInteger):
        """
        Run the Morris indices evaluation.

        The code first sample trajectories in the unit hypercube and transform them to the physical space (see method
        :py:meth:`sample_trajectories`), then runs the forward model to compute the elementary effects,
        and finally computes the sensitivity indices.

        :param n_trajectories: Number of random trajectories. Usually chosen between :math:`5` and :math:`10`.
         The number of model evaluations is :code:`n_trajectories * (d+1)`.
        """
        # Compute trajectories and elementary effects - append if any already exist
        (trajectories_unit_hypercube, trajectories_physical_space,) = \
            self.sample_trajectories(n_trajectories=n_trajectories, maximize_dispersion=self.maximize_dispersion,)
        elementary_effects = self._compute_elementary_effects(trajectories_physical_space)
        self.store_data(elementary_effects, trajectories_physical_space, trajectories_unit_hypercube)
        self.mustar_indices, self.sigma_indices = self._compute_indices(self.elementary_effects)

    def store_data(
            self,
            elementary_effects,
            trajectories_physical_space,
            trajectories_unit_hypercube,
    ):
        if self.elementary_effects is None:
            self.elementary_effects = elementary_effects
            self.trajectories_unit_hypercube = trajectories_unit_hypercube
            self.trajectories_physical_space = trajectories_physical_space
        else:
            self.elementary_effects = np.concatenate([self.elementary_effects, elementary_effects], axis=0)
            self.trajectories_unit_hypercube = np.concatenate([self.trajectories_unit_hypercube,
                                                               trajectories_unit_hypercube], axis=0)
            self.trajectories_physical_space = np.concatenate([self.trajectories_physical_space,
                                                               trajectories_physical_space], axis=0)

    @beartype
    def sample_trajectories(self, n_trajectories: PositiveInteger, maximize_dispersion: bool = False):
        """
        Create the trajectories, first in the unit hypercube then transform them in the physical space.

        :param n_trajectories: Number of random trajectories. Usually chosen between :math:`5` and :math:`10`.
         The number of model evaluations is :code:`n_trajectories * (d+1)`.
        :param maximize_dispersion: If :any:`True`, generate a large number of design trajectories and keep the ones
         that maximize dispersion between all trajectories, allows for a better coverage of the input space.
         Default :any:`False`.
        """

        trajectories_unit_hypercube = []
        perms_indices = []
        ntrajectories_all = (10 * n_trajectories if maximize_dispersion else 1 * n_trajectories)
        for r in range(ntrajectories_all):
            if self.random_state is None:
                perms = np.random.permutation(self.dimension)
            else:
                perms = self.random_state.permutation(self.dimension)
            initial_state = (1.0 / (self.n_levels - 1) *
                             randint(low=0, high=int((self.n_levels - 1) * (1 - self.delta) + 1))
                             .rvs(size=(1, self.dimension), random_state=self.random_state))
            trajectory_uh = np.tile(initial_state, [self.dimension + 1, 1])
            for count_d, d in enumerate(perms):
                trajectory_uh[count_d + 1:, d] = initial_state[0, d] + self.delta
            trajectories_unit_hypercube.append(trajectory_uh)
            perms_indices.append(perms)
        trajectories_unit_hypercube = np.array(trajectories_unit_hypercube)  # ndarray (r, d+1, d)

        # if maximize_dispersion, compute the 'best' trajectories
        if maximize_dispersion:
            from itertools import combinations

            distances = np.zeros((ntrajectories_all, ntrajectories_all))
            for r in range(ntrajectories_all):
                des_r = np.tile(trajectories_unit_hypercube[r, :, :][np.newaxis, :, :],[self.dimension + 1, 1, 1],)
                for r2 in range(r + 1, ntrajectories_all):
                    des_r2 = np.tile(trajectories_unit_hypercube[r2, :, :][:, np.newaxis, :],
                                     [1, self.dimension + 1, 1],)
                    distances[r, r2] = np.sum(np.sqrt(np.sum((des_r - des_r2) ** 2, axis=-1)))

            # try 20000 combinations of ntrajectories trajectories, keep the one that maximizes the distance
            def compute_combi_and_dist():
                if self.random_state is None:
                    combi = np.random.choice(ntrajectories_all, replace=False, size=n_trajectories)
                else:
                    combi = self.random_state.choice(ntrajectories_all, replace=False, size=n_trajectories)
                dist_combi = 0.0
                for pairs in list(combinations(combi, 2)):
                    dist_combi += distances[min(pairs), max(pairs)] ** 2
                return combi, np.sqrt(dist_combi)

            comb_to_keep, dist_comb = compute_combi_and_dist()
            for _ in range(1, 20000):
                comb, new_dist_comb = compute_combi_and_dist()
                if new_dist_comb > dist_comb:
                    comb_to_keep, dist_comb = comb, new_dist_comb
            trajectories_unit_hypercube = np.array([trajectories_unit_hypercube[j] for j in comb_to_keep])

        # Avoid 0 and 1 cdf values
        trajectories_unit_hypercube[trajectories_unit_hypercube < 0.01] = 0.01
        trajectories_unit_hypercube[trajectories_unit_hypercube > 0.99] = 0.99

        # Transform to physical space via icdf
        trajectories_physical_space = []
        for trajectory_uh in trajectories_unit_hypercube:
            trajectory_ps = np.zeros_like(trajectory_uh)
            for count_d, (design_d, icdf_d) in enumerate(zip(trajectory_uh.T, self.icdfs)):
                trajectory_ps[:, count_d] = icdf_d(x=design_d)
            trajectories_physical_space.append(trajectory_ps)
        trajectories_physical_space = np.array(trajectories_physical_space)
        return trajectories_unit_hypercube, trajectories_physical_space

    def _compute_elementary_effects(self, trajectories_physical_space):
        r, _, d = trajectories_physical_space.shape
        # Run the model for all replicates
        elementary_effects = []
        for samples in trajectories_physical_space:
            self.runmodel_object.run(samples=samples, append_samples=False)
            qoi = np.array(self.runmodel_object.qoi_list)
            el_effect = np.zeros((self.dimension,))
            perms = [np.argwhere(bi != 0.0)[0, 0] for bi in (samples[1:] - samples[:-1])]
            for count_d, d in enumerate(perms):
                el_effect[d] = (qoi[count_d + 1] - qoi[count_d]) / self.delta
            elementary_effects.append(el_effect)
        return np.array(elementary_effects)

    @staticmethod
    def _compute_indices(elementary_effects):
        # elementary_effects is an array of shape (r, d, ny) or (r, d)
        mu_star = np.mean(np.abs(elementary_effects), axis=0)
        mu = np.mean(elementary_effects, axis=0)
        sigma = np.sqrt(np.mean((elementary_effects - mu) ** 2, axis=0))
        return mu_star, sigma


# src/UQpy/sensitivity/PceSensitivity.py

from typing import Annotated

import numpy as np
from beartype.vale import Is

from UQpy.surrogates.polynomial_chaos import PolynomialChaosExpansion

FittedPce = Annotated[PolynomialChaosExpansion, Is[lambda pce: pce.coefficients is not None]]


class PceSensitivity:

    def __init__(self, pce_object: FittedPce):
        """
        Compute Sobol sensitivity indices based on a PCE surrogate approximation of the QoI.

        :param pce_object: Polynomial Chaos Expansion surrogate.
        """
        self.pce_object = pce_object
        self.first_order_indices = None
        """PCE estimates for the first order Sobol indices"""
        self.total_order_indices = None
        """PCE estimates for the total order Sobol indices"""
        self.generalized_first_order_indices = None
        """PCE estimates of generalized first order Sobol indices, which characterize
        the sensitivity of a vector-valued quantity of interest on the random
        inputs."""
        self.generalized_total_order_indices = None
        """PCE estimates of generalized total order Sobol indices, which characterize
        the sensitivity of a vector-valued quantity of interest on the random
        inputs."""

    def run(self):
        """
        This method calculates and saves as attributes first and total order as well as generalized sensitivity indices.
        """
        self.calculate_first_order_indices()
        self.calculate_total_order_indices()
        self.calculate_generalized_first_order_indices()
        self.calculate_generalized_total_order_indices()

    def calculate_first_order_indices(self) -> np.ndarray:
        """
        PCE estimates for the first order Sobol indices.

        :return: First order Sobol indices.
        """
        outputs_number = np.shape(self.pce_object.coefficients)[1]
        variance = self.pce_object.get_moments()[1]
        inputs_number = self.pce_object.inputs_number
        multi_index_set = self.pce_object.multi_index_set

        first_order_indices = np.zeros([inputs_number, outputs_number])
        # take all multi-indices except 0-index
        idx_no_0 = np.delete(multi_index_set, 0, axis=0)
        for nn in range(inputs_number):
            # remove nn-th column
            idx_no_0_nn = np.delete(idx_no_0, nn, axis=1)
            # we want the rows with all indices (except nn) equal to zero
            sum_idx_rows = np.sum(idx_no_0_nn, axis=1)
            zero_rows = np.asarray(np.where(sum_idx_rows == 0)).flatten() + 1
            variance_contribution = np.sum(self.pce_object.coefficients[zero_rows, :] ** 2, axis=0)
            first_order_indices[nn, :] = variance_contribution / variance
        self.first_order_indices = first_order_indices
        return first_order_indices

    def calculate_total_order_indices(self) -> np.ndarray:
        """
        PCE estimates for the total order Sobol indices.

        :return: Total order Sobol indices.
        """
        outputs_number = np.shape(self.pce_object.coefficients)[1]
        variance = self.pce_object.get_moments()[1]
        inputs_number = self.pce_object.inputs_number
        multi_index_set = self.pce_object.multi_index_set

        total_order_indices = np.zeros([inputs_number, outputs_number])
        for nn in range(inputs_number):
            # we want all multi-indices where the nn-th index is NOT zero
            idx_column_nn = np.array(multi_index_set)[:, nn]
            nn_rows = np.asarray(np.where(idx_column_nn != 0)).flatten()
            variance_contribution = np.sum(self.pce_object.coefficients[nn_rows, :] ** 2, axis=0)
            total_order_indices[nn, :] = variance_contribution / variance
        self.total_order_indices = total_order_indices
        return total_order_indices

    def calculate_generalized_first_order_indices(self) -> np.ndarray:
        """
        PCE estimates of generalized first order Sobol indices, which characterize
        the sensitivity of a vector-valued quantity of interest on the random
        inputs.

        :return: Generalized first order Sobol indices.
        """
        inputs_number = self.pce_object.inputs_number
        if inputs_number == 1:
            raise ValueError('Not applicable for scalar model outputs.')

        variance = self.pce_object.get_moments()[1]
        first_order_indices = self.calculate_first_order_indices()
        variance_contributions = first_order_indices * variance
        total_variance = np.sum(variance)
        total_variance_contribution_per_input = np.sum(variance_contributions, axis=1)
        generalized_first_order_indices = total_variance_contribution_per_input / total_variance
        self.generalized_first_order_indices = generalized_first_order_indices
        return generalized_first_order_indices

    def calculate_generalized_total_order_indices(self):
        """
        PCE estimates of generalized total order Sobol indices, which characterize
        the sensitivity of a vector-valued quantity of interest on the random
        inputs.

        :return: Generalized total order Sobol indices.
        """
        inputs_number = self.pce_object.inputs_number

        if inputs_number == 1:
            raise ValueError('Not applicable for scalar model outputs.')

        variance = self.pce_object.get_moments()[1]
        total_order_indices = self.calculate_total_order_indices()
        variance_contributions = total_order_indices * variance
        total_variance = np.sum(variance)
        total_variance_contribution_per_input = np.sum(variance_contributions, axis=1)
        generalized_total_order_indices = total_variance_contribution_per_input / total_variance
        self.generalized_total_order_indices = generalized_total_order_indices
        return generalized_total_order_indices


# src/UQpy/sensitivity/__init__.py

from UQpy.sensitivity.MorrisSensitivity import MorrisSensitivity
from UQpy.sensitivity.PceSensitivity import PceSensitivity
from UQpy.sensitivity.sobol import Sobol
from UQpy.sensitivity.cramer_von_mises import CramervonMises
from UQpy.sensitivity.chatterjee import Chatterjee
from UQpy.sensitivity.generalised_sobol import GeneralisedSobol

from . import MorrisSensitivity
from . import PceSensitivity
from . import Sobol
from . import CramervonMises
from . import Chatterjee
from . import GeneralisedSobol


# src/UQpy/sensitivity/chatterjee.py

"""
This module contains the Chatterjee coefficient of correlation proposed 
in [1]_. 

Using the rank statistics, we can also estimate the Sobol indices proposed by 
Gamboa et al. [2]_.

References
----------

.. [1] Sourav Chatterjee (2021) A New Coefficient of Correlation, Journal of the
        American Statistical Association, 116:536, 2009-2022, 
        DOI: 10.1080/01621459.2020.1758115

.. [2] Fabrice Gamboa, Pierre Gremaud, Thierry Klein, and Agnès Lagnoux. (2020). 
        Global Sensitivity Analysis: a new generation of mighty estimators 
        based on rank statistics.

"""

import logging

import numpy as np
import scipy.stats

from UQpy.sensitivity.baseclass.sensitivity import Sensitivity
from UQpy.sensitivity.sobol import compute_first_order as compute_first_order_sobol
from UQpy.utilities.UQpyLoggingFormatter import UQpyLoggingFormatter


class Chatterjee(Sensitivity):
    """
    Compute sensitivity indices using the Chatterjee correlation coefficient.

    Using the same model evaluations, we can also estimate the Sobol indices.

    :param runmodel_object: The computational model. It should be of type :class:`.RunModel`. \
        The output QoI can be a scalar or vector of length :code:`ny`, then the sensitivity \
        indices of all :code:`ny` outputs are computed independently.

    :param distributions: List of :class:`.Distribution` objects corresponding to each \
        random variable, or :class:`.JointIndependent` object \
        (multivariate RV with independent marginals).

    :param random_state: Random seed used to initialize the pseudo-random number \
        generator. Default is :any:`None`.

    **Methods:**
    """

    def __init__(self, runmodel_object, dist_object, random_state=None, **kwargs):
        super().__init__(
            runmodel_object, dist_object, random_state=random_state, **kwargs
        )

        # Create logger with the same name as the class
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.ERROR)
        frmt = UQpyLoggingFormatter()

        # create console handler with a higher log level
        ch = logging.StreamHandler()
        ch.setFormatter(frmt)

        # add the handler to the logger
        self.logger.addHandler(ch)

        self.chatterjee_i = None
        "Chatterjee sensitivity indices (First order), :class:`numpy.ndarray` of shape :code:`(num_vars, 1)`"

        self.sobol_i = None
        "Sobol indices computed using the rank statistics, :class:`numpy.ndarray` of shape :code:`(num_vars, 1)`"

        self.CI_chatterjee_i = None
        "Confidence intervals for the Chatterjee sensitivity indices, :class:`numpy.ndarray` of shape :code:`(num_vars, 2)`"

        self.num_vars = None
        "Number of input random variables, :class:`int`"

        self.n_samples = None
        "Number of samples used to estimate the sensitivity indices, :class:`int`"

    def run(
        self,
        n_samples=1_000,
        estimate_sobol_indices=False,
        num_bootstrap_samples=None,
        confidence_level=0.95,
    ):
        """
        Compute the sensitivity indices using the Chatterjee method.

        :param n_samples: Number of samples used to compute the Cramér-von Mises indices. \
            Default is 1,000.

        :param estimate_sobol_indices: If :code:`True`, the Sobol indices are estimated \
            using the pick-and-freeze samples.

        :param num_bootstrap_samples: Number of bootstrap samples used to estimate the \
            Sobol indices. Default is :any:`None`.

        :param confidence_level: Confidence level used to compute the confidence \
            intervals of the Cramér-von Mises indices.

        :return: A :class:`dict` with the following keys: \
            :code:`'chatterjee_i'` of shape :code:`(num_vars, 1)`, \
            :code:`'CI_chatterjee_i'` of shape :code:`(num_vars, 2)`, \
            :code:`'sobol_i'` of shape :code:`(num_vars, 1)`.

        """

        # Check nsamples
        self.n_samples = n_samples
        if not isinstance(self.n_samples, int):
            raise TypeError("UQpy: nsamples should be an integer")

        # Check num_bootstrap_samples data type
        if num_bootstrap_samples is not None:
            if not isinstance(num_bootstrap_samples, int):
                raise TypeError("UQpy: num_bootstrap_samples should be an integer.\n")
        elif num_bootstrap_samples is None:
            self.logger.info(
                "UQpy: num_bootstrap_samples is set to None, confidence intervals will not be computed.\n"
            )

        ################## GENERATE SAMPLES ##################

        A_samples = self.dist_object.rvs(self.n_samples, random_state=self.random_state)

        self.logger.info("UQpy: Generated samples successfully.\n")

        self.num_vars = A_samples.shape[1]  # number of variables

        ################# MODEL EVALUATIONS ####################

        A_model_evals = self._run_model(A_samples).reshape(-1, 1)

        self.logger.info("UQpy: Model evaluations completed.\n")

        ######################### STORAGE ########################
        # Create dictionary to store the sensitivity indices
        computed_indices = {}

        ################## COMPUTE CHATTERJEE INDICES ##################

        self.chatterjee_i = self.compute_chatterjee_indices(A_samples, A_model_evals)

        self.logger.info("UQpy: Chatterjee indices computed successfully.\n")

        # Store the indices in the dictionary
        computed_indices["chatterjee_i"] = self.chatterjee_i

        ################## COMPUTE SOBOL INDICES ##################

        self.logger.info("UQpy: Computing First order Sobol indices ...\n")

        if estimate_sobol_indices:
            f_C_i_model_evals = self.compute_rank_analog_of_f_C_i(
                A_samples, A_model_evals
            )

            self.sobol_i = self.compute_Sobol_indices(A_model_evals, f_C_i_model_evals)

            self.logger.info("UQpy: First order Sobol indices computed successfully.\n")

            # Store the indices in the dictionary
            computed_indices["sobol_i"] = self.sobol_i

        ################## CONFIDENCE INTERVALS ####################

        if num_bootstrap_samples is not None:

            self.logger.info("UQpy: Computing confidence intervals ...\n")

            estimator_inputs = [A_samples, A_model_evals]

            self.CI_chatterjee_i = self.bootstrapping(
                self.compute_chatterjee_indices,
                estimator_inputs,
                computed_indices["chatterjee_i"],
                num_bootstrap_samples,
                confidence_level,
            )

            self.logger.info(
                "UQpy: Confidence intervals for Chatterjee indices computed successfully.\n"
            )

            computed_indices["CI_chatterjee_i"] = self.CI_chatterjee_i

        return computed_indices

    @staticmethod
    def compute_chatterjee_indices(X, Y, seed=None):
        r"""

        Compute the Chatterjee sensitivity indices
        between the input random vectors :math:`X=\left[ X_{1}, X_{2},…,X_{d} \right]`
        and output random vector Y.

        :param X: Input random vectors, :class:`numpy.ndarray` of shape :code:`(n_samples, num_vars)`

        :param Y: Output random vector, :class:`numpy.ndarray` of shape :code:`(n_samples, 1)`

        :param seed: Seed for the random number generator.

        :return: Chatterjee sensitivity indices, :class:`numpy.ndarray` of shape :code:`(num_vars, 1)`

        """

        if seed is not None:
            # set seed for reproducibility
            np.random.seed(seed)

        N = X.shape[0]  # number of samples
        m = X.shape[1]  # number of variables

        chatterjee_indices = np.zeros((m, 1))

        for i in range(m):

            # Samples of random variable X_i
            X_i = X[:, i].reshape(-1, 1)

            #! For ties in X_i
            # we break ties uniformly at random
            # Shuffle X_i and Y
            _ix = np.arange(N)  # indices of X_i
            np.random.shuffle(_ix)  # shuffle indices
            X_i_shuffled = X_i[_ix]  # shuffle X_i
            Y_shuffled = Y[_ix]  # shuffle Y

            Z = np.hstack((X_i_shuffled, Y_shuffled))

            # Sort the columns of Z by X_i
            # such that the tuple (X_i, Y_i) is unchanged
            Z_sorted = Z[Z[:, 0].argsort()]

            # Find rank of y_i in the sorted columns of Y
            # r[i] is number of j s.t. y[j] <= y[i],
            # This is accomplished using rankdata with method='max'
            # Example: Y = [1, 2, 3, 3, 4, 5], rank = [1, 2, 4, 4, 5, 6]
            rank = scipy.stats.rankdata(Z_sorted[:, 1], method="max")

            #! For ties in Y
            # l[i] is number of j s.t. y[i] <= y[j],
            # This is accomplished using rankdata with method='max'
            # Example: Y = [1, 2, 3, 3, 4, 5], l = [6, 5, 4, 4, 2, 1]
            # One could also use the Y_shuffled array, since sum2 only
            # multiplies terms of same index, i.e l_i*(n - l_i)
            L = scipy.stats.rankdata(-Z_sorted[:, 1], method="max")

            sum1 = np.abs(rank[1:] - rank[:-1]).sum()

            sum2 = np.sum(L * (N - L))

            chatterjee_indices[i] = 1 - N * sum1 / (2 * sum2)

        return chatterjee_indices

    @staticmethod
    def rank_analog_to_pickfreeze(X, j):
        r"""
        Computing the :math:`N(j)` for each :math:`j \in \{1, \ldots, n\}`
        as in eq.(8) in [6]_, where :math:`n` is the size of :math:`X`.

        .. math::
            :nowrap:

                \begin{equation}
                    N(j):=
                    \begin{cases}
                        \pi^{-1}(\pi(j)+1) &\text { if } \pi(j)+1 \leqslant n \\
                        \pi^{-1}(1) &\text { if } \pi(j)=n
                    \end{cases}
                \end{equation}

        where, :math:`\pi(j) := \mathrm{rank}(x_j)`

        :param X: Input random vector, :class:`numpy.ndarray` of shape :code:`(n_samples, 1)`

        :param j: Index of the sample :math:`j \in \{1, \ldots, n\}`

        :return: :math:`N(j)` :class:`int`

        """

        N = X.shape[0]  # number of samples

        # Ranks of elements of X_i
        # -1 so that the ranks are 0-based
        # for convenience in indexing
        rank_X = scipy.stats.rankdata(X) - 1
        rank_X = rank_X.astype(int)

        # Find rank of element j
        rank_j = rank_X[j]

        if rank_j + 1 <= N - 1:
            # Get index of element: rank_j + 1
            return np.where(rank_X == rank_j + 1)[0][0]

        if rank_j == N - 1:
            return np.where(rank_X == 0)[0][0]

    @staticmethod
    def rank_analog_to_pickfreeze_vec(X):
        r"""
        Computing the :math:`N(j)` for each :math:`j \in \{1, \ldots, n\}`
        in a vectorized manner., where :math:`n` is the size of :math:`X`.

        This method is significantly faster than the looping version
        ``rank_analog_to_pickfreeze`` but is also more complicated.

        .. math::
            :nowrap:

                \begin{equation}
                    N(j):=
                    \begin{cases}
                        \pi^{-1}(\pi(j)+1) &\text { if } \pi(j)+1 \leqslant n \\
                        \pi^{-1}(1) &\text { if } \pi(j)=n
                    \end{cases}
                \end{equation}
        
        where, :math:`\pi(j) := \mathrm{rank}(x_j)`

        Key idea: :math:`\pi^{-1}` is rank_X.argsort() (
        `see also <https://rdrr.io/cran/sensitivity/src/R/sobolrank.R>`_)

        Example:
        X = [22, 74, 44, 11, 1]

        N_J = [3, 5, 2, 1, 4] (1-based indexing)

        N_J = [2, 4, 1, 0, 3] (0-based indexing)

        :param X: Input random vector, :class:`numpy.ndarray` of shape :code:`(n_samples, 1)`

        :return: :math:`N(j)`, :class:`numpy.ndarray` of shape :code:`(n_samples, 1)`

        """

        N = X.shape[0]  # number of samples
        N_func = np.zeros((N, 1))

        # Ranks of elements of X_i
        # -1 since ranks are 0-based
        rank_X = scipy.stats.rankdata(X, method="ordinal") - 1
        rank_X = rank_X.astype(int)

        # Inverse of pi(j): j = pi^-1(rank_X(j))
        #! This is non-trivial
        pi_inverse = rank_X.argsort()  # complexity: N*log(N)

        # CONDITION 2
        # Find j with rank_j == N-1
        j_meets_condition_2 = pi_inverse[N - 1]
        N_func[j_meets_condition_2] = pi_inverse[0]

        # CONDITION 1
        # Find j's with rank_j + 1 <= N-1
        # term_1 = pi(j) + 1
        j_remaining = np.delete(np.arange(N), j_meets_condition_2)
        term_1 = rank_X[j_remaining] + 1

        j_remaining_meet_condition_1 = pi_inverse[term_1]

        # j_remaining_meet_condition_1 = np.where(rank_X_i == condition)
        N_func[j_remaining, 0] = j_remaining_meet_condition_1

        return N_func.astype(int)

    @staticmethod
    def compute_Sobol_indices(A_model_evals, C_i_model_evals):
        r"""
        A method to estimate the first order Sobol indices using
        the Chatterjee method.

        .. math::
            :nowrap:

            \begin{equation}
                \xi_{n}^{\mathrm{Sobol}}\left(X_{1}, Y\right):=
                \frac{\frac{1}{n} \sum_{j=1}^{n} Y_{j} Y_{N(j)}-\left(\frac{1}{n} \sum_{j=1}^{n} Y_{j}\right)^{2}}
                {\frac{1}{n} \sum_{j=1}^{n}\left(Y_{j}\right)^{2}-\left(\frac{1}{n} \sum_{j=1}^{n} Y_{j}\right)^{2}}
            \end{equation}

        where the term :math:`Y_{N(j)}` is computed using the method:``rank_analog_to_pickfreeze_vec``.

        :param A_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_samples, 1)`

        :param C_i_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_samples, num_vars)`

        :return: First order Sobol indices, :class:`numpy.ndarray` of shape :code:`(num_vars, 1)`

        """

        # extract shape
        _shape = C_i_model_evals.shape

        # convert C_i_model_evals to 3D array
        # with n_outputs=1 in first dimension
        n_outputs = 1
        C_i_model_evals = C_i_model_evals.reshape((n_outputs, *_shape))

        first_order_sobol = compute_first_order_sobol(
            A_model_evals, None, C_i_model_evals, scheme="Sobol1993"
        )

        return first_order_sobol

    def compute_rank_analog_of_f_C_i(self, A_samples, A_model_evals):
        r"""
        In the Pick and Freeze method, we use model evaluations
        :math:`f_A`, :math:`f_B`, :math:`f_{C_{i}}`
        to compute the Sobol indices.

        Gamboa et al. provide a rank analog to :math:`f_{C_{i}}` in eq. (6) in [6]_.

        **Inputs:**

        * **A_samples** (`ndarray`):
            Shape: `(n_samples, num_vars)`.

        * **A_model_evals** (`ndarray`):
            Shape: `(n_samples, 1)`.

        **Outputs:**

        * **A_i_model_evals** (`ndarray`):
            Shape: `(n_samples, num_vars)`.

        """

        f_A = A_model_evals
        N = f_A.shape[0]
        m = self.num_vars

        A_i_model_evals = np.zeros((N, m))

        for i in range(m):

            K = self.rank_analog_to_pickfreeze_vec(A_samples[:, i])

            A_i_model_evals[:, i] = f_A[K].ravel()

        return A_i_model_evals


# src/UQpy/sensitivity/cramer_von_mises.py

"""
Computing the Cramér-von Mises sensitivity indices.

References
----------

.. [1] Gamboa, F., Klein, T., & Lagnoux, A. (2018). Sensitivity Analysis
       Based on Cramér-von Mises Distance. SIAM/ASA Journal on Uncertainty
       Quantification, 6(2), 522-548. doi:10.1137/15M1025621

.. [2] Gamboa, F., Gremaud, P., Klein, T., & Lagnoux, A. (2020). Global
       Sensitivity Analysis: a new generation of mighty estimators based on
       rank statistics. arXiv [math.ST]. http://arxiv.org/abs/2003.01772

"""

import logging

import numpy as np

from UQpy.sensitivity.baseclass.sensitivity import Sensitivity
from UQpy.sensitivity.baseclass.pickfreeze import generate_pick_freeze_samples
from UQpy.sensitivity.sobol import compute_first_order as compute_first_order_sobol
from UQpy.sensitivity.sobol import compute_total_order as compute_total_order_sobol
from UQpy.utilities.UQpyLoggingFormatter import UQpyLoggingFormatter

# TODO: Sampling strategies


class CramervonMises(Sensitivity):
    """
    Compute the Cramér-von Mises indices.

    Currently only available for models with scalar output.

    :param runmodel_object: The computational model. It should be of type :class:`.RunModel`. \
        The output QoI can be a scalar or vector of length :code:`ny`, then the sensitivity \
        indices of all :code:`ny` outputs are computed independently.

    :param distributions: List of :class:`.Distribution` objects corresponding to each \
        random variable, or :class:`.JointIndependent` object \
        (multivariate RV with independent marginals).

    :param random_state: Random seed used to initialize the pseudo-random number \
        generator. Default is :any:`None`.

    **Methods:**
    """

    def __init__(
        self, runmodel_object, dist_object, random_state=None, **kwargs
    ) -> None:

        super().__init__(
            runmodel_object, dist_object, random_state=random_state, **kwargs
        )

        # Create logger with the same name as the class
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.ERROR)
        frmt = UQpyLoggingFormatter()

        # create console handler with a higher log level
        ch = logging.StreamHandler()
        ch.setFormatter(frmt)

        # add the handler to the logger
        self.logger.addHandler(ch)

        self.CVM_i = None
        "First order Cramér-von Mises indices, :class:`numpy.ndarray` of shape :code:`(num_vars, 1)`"

        self.CI_CVM_i = None
        "Confidence intervals of the first order Cramér-von Mises indices, :class:`numpy.ndarray` of shape :code:`(num_vars, 2)`"

        self.sobol_i = None
        "First order Sobol indices computed using the pick-and-freeze samples, :class:`numpy.ndarray` of shape :code:`(num_vars, 1)`"

        self.sobol_total_i = None
        "Total order Sobol indices computed using the pick-and-freeze samples, :class:`numpy.ndarray` of shape :code:`(num_vars, 1)`"

        self.n_samples = None
        "Number of samples used to compute the Cramér-von Mises indices, :class:`int`"

        self.num_vars = None
        "Number of input random variables, :class:`int`"

    def run(
        self,
        n_samples=1_000,
        estimate_sobol_indices=False,
        num_bootstrap_samples=None,
        confidence_level=0.95,
        disable_CVM_indices=False,
    ):

        """
        Compute the Cramér-von Mises indices.

        :param n_samples: Number of samples used to compute the Cramér-von Mises indices. \
            Default is 1,000.

        :param estimate_sobol_indices: If :code:`True`, the Sobol indices are estimated \
            using the pick-and-freeze samples.

        :param num_bootstrap_samples: Number of bootstrap samples used to estimate the \
            Sobol indices. Default is :any:`None`.

        :param confidence_level: Confidence level used to compute the confidence \
            intervals of the Cramér-von Mises indices.

        :param disable_CVM_indices: If :code:`True`, the Cramér-von Mises indices \
            are not computed.

        :return: A :class:`dict` with the following keys: \
            :code:`CVM_i` of shape :code:`(num_vars, 1)`, \
            :code:`CI_CVM_i` of shape :code:`(num_vars, 2)`, \
            :code:`sobol_i` of shape :code:`(num_vars, 1)`, \
            :code:`sobol_total_i` of shape :code:`(num_vars, 1)`.

        """

        # Check nsamples
        self.n_samples = n_samples
        if not isinstance(self.n_samples, int):
            raise TypeError("UQpy: nsamples should be an integer")

        # Check num_bootstrap_samples data type
        if num_bootstrap_samples is not None:
            if not isinstance(num_bootstrap_samples, int):
                raise TypeError("UQpy: num_bootstrap_samples should be an integer.\n")
        elif num_bootstrap_samples is None:
            self.logger.info(
                "UQpy: num_bootstrap_samples is set to None, confidence intervals will not be computed.\n"
            )

        ################## GENERATE SAMPLES ##################

        A_samples, W_samples, C_i_generator, _ = generate_pick_freeze_samples(
            self.dist_object, self.n_samples, self.random_state
        )

        self.logger.info("UQpy: Generated samples using the pick-freeze scheme.\n")

        ################# MODEL EVALUATIONS ####################

        A_model_evals = self._run_model(A_samples).reshape(-1, 1)

        self.logger.info("UQpy: Model evaluations A completed.\n")

        W_model_evals = self._run_model(W_samples).reshape(-1, 1)

        self.logger.info("UQpy: Model evaluations W completed.\n")

        self.num_vars = A_samples.shape[1]

        C_i_model_evals = np.zeros((self.n_samples, self.num_vars))

        for i, C_i in enumerate(C_i_generator):
            C_i_model_evals[:, i] = self._run_model(C_i).ravel()

        self.logger.info("UQpy: Model evaluations C completed.\n")

        self.logger.info("UQpy: All model evaluations computed successfully.\n")

        ######################### STORAGE ########################

        # Create dictionary to store the sensitivity indices
        computed_indices = {}

        ################## COMPUTE CVM INDICES ##################

        # flag is used to disable computation of
        # CVM indices during testing
        if not disable_CVM_indices:
            # Compute the Cramér-von Mises indices
            self.CVM_i = self.pick_and_freeze_estimator(
                A_model_evals, W_model_evals, C_i_model_evals
            )

            self.logger.info("UQpy: Cramér-von Mises indices computed successfully.\n")

            # Store the indices in the dictionary
            computed_indices["CVM_i"] = self.CVM_i

        ################# COMPUTE CONFIDENCE INTERVALS ##################

        if num_bootstrap_samples is not None:

            self.logger.info("UQpy: Computing confidence intervals ...\n")

            estimator_inputs = [
                A_model_evals,
                W_model_evals,
                C_i_model_evals,
            ]

            self.CI_CVM_i = self.bootstrapping(
                self.pick_and_freeze_estimator,
                estimator_inputs,
                computed_indices["CVM_i"],
                num_bootstrap_samples,
                confidence_level,
            )

            self.logger.info(
                "UQpy: Confidence intervals for Cramér-von Mises indices computed successfully.\n"
            )

            # Store the indices in the dictionary
            computed_indices["CI_CVM_i"] = self.CI_CVM_i

        ################## COMPUTE SOBOL INDICES ##################

        if estimate_sobol_indices:

            self.logger.info("UQpy: Computing First order Sobol indices ...\n")

            # extract shape
            _shape = C_i_model_evals.shape

            # convert C_i_model_evals to 3D array
            # with n_outputs=1 in first dimension
            n_outputs = 1
            C_i_model_evals = C_i_model_evals.reshape((n_outputs, *_shape))

            self.sobol_i = compute_first_order_sobol(
                A_model_evals, W_model_evals, C_i_model_evals
            )

            self.logger.info("UQpy: First order Sobol indices computed successfully.\n")

            self.sobol_total_i = compute_total_order_sobol(
                A_model_evals, W_model_evals, C_i_model_evals
            )

            self.logger.info("UQpy: Total order Sobol indices computed successfully.\n")

            # Store the indices in the dictionary
            computed_indices["sobol_i"] = self.sobol_i
            computed_indices["sobol_total_i"] = self.sobol_total_i

        return computed_indices

    @staticmethod
    def indicator_function(Y, W):
        """
        Vectorized version of the indicator function.

        .. math::
           \mathbb{I}(Y,W) = \mathbf{1}_{Y \leq W}

        **Inputs:**

        * **Y** (`ndarray`):
            Vector of values of the random variable.
            Shape: `(N, 1)`

        * **W** (`ndarray`):
            Vector of values of the random variable.
            Shape: `(N, 1)`

        **Outputs:**

        * **indicator** (`ndarray`):
            Shape: `(N, 1)`

        """
        return (Y <= W.T).astype(int)

    def pick_and_freeze_estimator(self, A_model_evals, W_model_evals, C_i_model_evals):

        """
        Compute the first order Cramér-von Mises indices
        using the Pick-and-Freeze estimator.

        **Inputs**

        * **A_model_evals** (`np.array`):
            Shape: `(n_samples, 1)`

        * **W_model_evals** (`np.array`):
            Shape: `(n_samples, 1)`

        * **C_i_model_evals** (`np.array`):
            Shape: `(n_samples, num_vars)`

        **Outputs**

        * **First_order_CVM** (`np.array`):
            Shape: `(num_vars)`

        """

        ## **Notes**

        # Implementation using 2 `for` loops. This is however
        # faster than the vectorized version which has only 1 `for` loop.

        # For N = 50_000 runs
        # With 2 `for` loops: 26.75 seconds (this implementation)
        # With 1 `for` loops: 62.42 seconds (vectorized implementation)

        # Possible improvements:
        # Check indicator function run time using a profiler
        # as it results in an `N` x `N` array.
        # Q. Does it use a for loop under the hood?
        # Computations such as `np.sum` and `np.mean`
        # are handled by numpy so they are fast.
        # (This should however be faster for small `N`, e.g. N=10_000)

        N = self.n_samples
        m = self.num_vars

        # Model evaluations
        f_A = A_model_evals.ravel()
        f_W = W_model_evals.ravel()
        f_C_i = C_i_model_evals

        # Store CramérvonMises indices
        First_order_indices = np.zeros((m, 1))

        # Compute Cramér-von Mises indices
        for i in range(m):
            sum_numerator = 0
            sum_denominator = 0

            for k in range(N):

                term_1 = self.indicator_function(f_A, f_W[k])
                term_2 = self.indicator_function(f_C_i[:, i], f_W[k])

                mean_sum = (1 / (2 * N)) * np.sum(term_1 + term_2)
                mean_product = (1 / N) * np.sum(term_1 * term_2)

                sum_numerator += mean_product - mean_sum**2
                sum_denominator += mean_sum - mean_sum**2

            First_order_indices[i] = sum_numerator / sum_denominator

        return First_order_indices


# src/UQpy/sensitivity/generalised_sobol.py

"""

The GeneralisedSobol class computes the generalised Sobol indices for a given
multi-ouput model. The class is based on the work of [1]_ and [2]_.

Additionally, we can compute the confidence intervals for the Sobol indices 
using bootstrapping [3]_.

References
----------

 .. [1] Gamboa F, Janon A, Klein T, Lagnoux A, others.
        Sensitivity analysis for multidimensional and functional outputs.
        Electronic journal of statistics 2014; 8(1): 575-603.

 .. [2] Alexanderian A, Gremaud PA, Smith RC. Variance-based sensitivity 
        analysis for time-dependent processes. Reliability engineering 
        & system safety 2020; 196: 106722.

.. [3] Jeremy Orloff and Jonathan Bloom (2014), Bootstrap confidence intervals, 
       Introduction to Probability and Statistics, MIT OCW. 

"""

import logging

import numpy as np

from UQpy.sensitivity.baseclass.sensitivity import Sensitivity
from UQpy.sensitivity.baseclass.pickfreeze import generate_pick_freeze_samples
from UQpy.utilities.UQpyLoggingFormatter import UQpyLoggingFormatter


class GeneralisedSobol(Sensitivity):
    """
    Compute the generalised Sobol indices for models with multiple outputs
    (vector-valued response) using the Pick-and-Freeze method.

    :param runmodel_object: The computational model. It should be of type :class:`.RunModel`. \
        The output QoI can be a scalar or vector of length :code:`ny`, then the sensitivity \
        indices of all :code:`ny` outputs are computed independently.

    :param distributions: List of :class:`.Distribution` objects corresponding to each \
        random variable, or :class:`.JointIndependent` object \
        (multivariate RV with independent marginals).

    :param random_state: Random seed used to initialize the pseudo-random number \
        generator. Default is :any:`None`.

    **Methods:**
    """

    def __init__(
        self, runmodel_object, dist_object, random_state=None, **kwargs
    ) -> None:

        super().__init__(runmodel_object, dist_object, random_state, **kwargs)

        # Create logger with the same name as the class
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.ERROR)
        frmt = UQpyLoggingFormatter()

        # create console handler with a higher log level
        ch = logging.StreamHandler()
        ch.setFormatter(frmt)

        # add the handler to the logger
        self.logger.addHandler(ch)

        self.gen_sobol_i = None
        "Generalised first order Sobol indices, :class:`ndarray` of shape (num_vars, 1)"

        self.gen_sobol_total_i = None
        "Generalised total order Sobol indices, :class:`ndarray` of shape (num_vars, 1)"

        self.n_samples = None
        "Number of samples used to compute the sensitivity indices, :class:`int`"

        self.num_vars = None
        "Number of model input variables, :class:`int`"

    def run(
        self,
        n_samples=1_000,
        num_bootstrap_samples=None,
        confidence_level=0.95,
    ):

        """
        Compute the generalised Sobol indices for models with multiple outputs
        (vector-valued response) using the Pick-and-Freeze method.

        :param n_samples: Number of samples used to compute the sensitivity indices. \
            Default is 1,000.

        :param num_bootstrap_samples: Number of bootstrap samples used to compute the \
            confidence intervals. Default is :any:`None`.

        :param confidence_level: Confidence level used to compute the confidence \
            intervals. Default is 0.95.

        :return: A :class:`dict` with the following keys: \
            :code:`gen_sobol_i` of shape :code:`(num_vars, 1)`, \
            :code:`gen_sobol_total_i` of shape :code:`(num_vars, 1)`, \
            :code:`CI_gen_sobol_i` of shape :code:`(num_vars, 2)`, \
            :code:`CI_gen_sobol_total_i` of shape :code:`(num_vars, 2)`.

        """

        # Check n_samples data type
        self.n_samples = n_samples
        if not isinstance(self.n_samples, int):
            raise TypeError("UQpy: n_samples should be an integer")

        # Check num_bootstrap_samples data type
        if num_bootstrap_samples is not None:
            if not isinstance(num_bootstrap_samples, int):
                raise TypeError("UQpy: num_bootstrap_samples should be an integer.\n")
        elif num_bootstrap_samples is None:
            self.logger.info(
                "UQpy: num_bootstrap_samples is set to None, confidence intervals will not be computed.\n"
            )

        ################## GENERATE SAMPLES ##################

        (A_samples, B_samples, C_i_generator, _,) = generate_pick_freeze_samples(
            self.dist_object, self.n_samples, self.random_state
        )

        self.logger.info("UQpy: Generated samples using the pick-freeze scheme.\n")

        self.num_vars = A_samples.shape[1]  # Number of variables

        ################# MODEL EVALUATIONS ####################

        A_model_evals = self._run_model(A_samples)  # shape: (n_samples, n_outputs)

        # if model output is vectorised,
        # shape retured by model is (n_samples, n_outputs, 1)
        # we need to reshape it to (n_samples, n_outputs)
        if A_model_evals.ndim == 3:
            A_model_evals = A_model_evals[:, :, 0]  # shape: (n_samples, n_outputs)

        self.logger.info("UQpy: Model evaluations A completed.\n")

        B_model_evals = self._run_model(B_samples)  # shape: (n_samples, n_outputs)

        # if model output is vectorised,
        # shape retured by model is (n_samples, n_outputs, 1)
        # we need to reshape it to (n_samples, n_outputs)
        if B_model_evals.ndim == 3:
            B_model_evals = B_model_evals[:, :, 0]  # shape: (n_samples, n_outputs)

        self.logger.info("UQpy: Model evaluations B completed.\n")

        self.n_outputs = A_model_evals.shape[1]

        # shape: (n_outputs, n_samples, num_vars)
        C_i_model_evals = np.zeros((self.n_outputs, self.n_samples, self.num_vars))

        for i, C_i in enumerate(C_i_generator):

            # if model output is vectorised,
            # shape retured by model is (n_samples, n_outputs, 1)
            # we need to reshape it to (n_samples, n_outputs)
            model_evals = self._run_model(C_i)

            if model_evals.ndim == 3:
                C_i_model_evals[:, :, i] = self._run_model(C_i)[:, :, 0].T
            else:
                C_i_model_evals[:, :, i] = model_evals.T

        self.logger.info("UQpy: Model evaluations C completed.\n")

        self.logger.info("UQpy: All model evaluations computed successfully.\n")

        ######################### STORAGE ########################

        # Create dictionary to store the sensitivity indices
        computed_indices = {}

        ################## COMPUTE GENERALISED SOBOL INDICES ##################

        self.gen_sobol_i = self.compute_first_order_generalised_sobol_indices(
            A_model_evals, B_model_evals, C_i_model_evals
        )

        self.logger.info(
            "UQpy: First order Generalised Sobol indices computed successfully.\n"
        )

        self.gen_sobol_total_i = self.compute_total_order_generalised_sobol_indices(
            A_model_evals, B_model_evals, C_i_model_evals
        )

        self.logger.info(
            "UQpy: Total order Generalised Sobol indices computed successfully.\n"
        )

        # Store the indices in the dictionary
        computed_indices["gen_sobol_i"] = self.gen_sobol_i
        computed_indices["gen_sobol_total_i"] = self.gen_sobol_total_i

        ################## CONFIDENCE INTERVALS ####################

        if num_bootstrap_samples is not None:

            self.logger.info("UQpy: Computing confidence intervals ...\n")

            estimator_inputs = [
                A_model_evals,
                B_model_evals,
                C_i_model_evals,
            ]

            # First order generalised Sobol indices
            self.CI_gen_sobol_i = self.bootstrapping(
                self.compute_first_order_generalised_sobol_indices,
                estimator_inputs,
                computed_indices["gen_sobol_i"],
                num_bootstrap_samples,
                confidence_level,
            )

            self.logger.info(
                "UQpy: Confidence intervals for First order Generalised Sobol indices computed successfully.\n"
            )

            # Total order generalised Sobol indices
            self.CI_gen_sobol_total_i = self.bootstrapping(
                self.compute_total_order_generalised_sobol_indices,
                estimator_inputs,
                computed_indices["gen_sobol_total_i"],
                num_bootstrap_samples,
                confidence_level,
            )

            self.logger.info(
                "UQpy: Confidence intervals for Total order Sobol Generalised indices computed successfully.\n"
            )

            # Store the indices in the dictionary
            computed_indices["CI_gen_sobol_i"] = self.CI_gen_sobol_i
            computed_indices["CI_gen_sobol_total_i"] = self.CI_gen_sobol_total_i

        return computed_indices

    @staticmethod
    def compute_first_order_generalised_sobol_indices(
        A_model_evals, B_model_evals, C_i_model_evals
    ):

        """
        Compute the generalised Sobol indices for models with multiple outputs.

        :param A_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_samples, n_outputs)`.
        :param B_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_samples, n_outputs)`.
        :param C_i_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_outputs, n_samples, num_vars)`.

        :return: First order generalised Sobol indices, :class:`numpy.ndarray` of shape :code:`(n_outputs, num_vars)`.

        """

        num_vars = C_i_model_evals.shape[2]
        n_outputs = A_model_evals.shape[1]

        # store generalised Sobol indices
        gen_sobol_i = np.zeros((num_vars, 1))

        for i in range(num_vars):

            all_Y_i = A_model_evals.T  # shape: (n_outputs, n_samples)
            all_Y_i_tilde = B_model_evals.T  # shape: (n_outputs, n_samples)
            all_Y_i_u = C_i_model_evals[:, :, i]  # shape: (n_outputs, n_samples)

            # compute the mean using all model evaluations
            # shape: (n_outputs, 1)
            mean = (
                np.mean(all_Y_i, axis=1, keepdims=1)
                + np.mean(all_Y_i_u, axis=1, keepdims=1)
                + np.mean(all_Y_i_tilde, axis=1, keepdims=1)
            ) / 3

            # center the evaluations since mean is available
            all_Y_i = all_Y_i - mean
            all_Y_i_tilde = all_Y_i_tilde - mean
            all_Y_i_u = all_Y_i_u - mean

            # compute the variance matrix using all available model evaluations
            # shape: (n_outputs, n_outputs)
            C = (np.cov(all_Y_i) + np.cov(all_Y_i_u) + np.cov(all_Y_i_tilde)) / 3

            # compute covariance btw. RVs 'X' and 'Y'
            # shape: (2*n_outputs, 2*n_outputs)
            # It contains the following 4 block matrices:
            #   (1, 1) variance of 'X'
            #  *(1, 2) covariance between 'X' and 'Y' (a.k.a. cross-covariance)
            #   (2, 1) covariance between 'Y' and 'X' (a.k.a. cross-covariance)
            #   (2, 2) variance of 'Y'
            _cov_1 = np.cov(all_Y_i_u, all_Y_i)  # for first order indices

            # We need the cross-covariance between 'X' and 'Y'
            # Extract *(1, 2) (upper right block)
            # shape: (n_outputs, n_outputs)
            C_u = _cov_1[0:n_outputs, n_outputs : 2 * n_outputs]

            denominator = np.trace(C)

            # Generalised Sobol indices
            gen_sobol_i[i] = np.trace(C_u) / denominator

        return gen_sobol_i

    @staticmethod
    def compute_total_order_generalised_sobol_indices(
        A_model_evals, B_model_evals, C_i_model_evals
    ):

        """
        Compute the generalised Sobol indices for models with multiple outputs.

        :param A_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_samples, n_outputs)`.
        :param B_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_samples, n_outputs)`.
        :param C_i_model_evals: Model evaluations, :class:`numpy.ndarray` of shape :code:`(n_outputs, n_samples, num_vars)`.

        :return: Total order generalised Sobol indices, :class:`numpy.ndarray` of shape :code:`(n_outputs, num_vars)`.

        """

        num_vars = C_i_model_evals.shape[2]
        n_outputs = A_model_evals.shape[1]

        # store generalised Sobol indices
        gen_sobol_total_i = np.zeros((num_vars, 1))

        for i in range(num_vars):

            all_Y_i = A_model_evals.T  # shape: (n_outputs, n_samples)
            all_Y_i_tilde = B_model_evals.T  # shape: (n_outputs, n_samples)
            all_Y_i_u = C_i_model_evals[:, :, i]  # shape: (n_outputs, n_samples)

            # compute the mean using all model evaluations
            # shape: (n_outputs, 1)
            mean = (
                np.mean(all_Y_i, axis=1, keepdims=1)
                + np.mean(all_Y_i_u, axis=1, keepdims=1)
                + np.mean(all_Y_i_tilde, axis=1, keepdims=1)
            ) / 3

            # center the evaluations since mean is available
            all_Y_i = all_Y_i - mean
            all_Y_i_tilde = all_Y_i_tilde - mean
            all_Y_i_u = all_Y_i_u - mean

            # compute the variance matrix using all available model evaluations
            # shape: (n_outputs, n_outputs)
            C = (np.cov(all_Y_i) + np.cov(all_Y_i_u) + np.cov(all_Y_i_tilde)) / 3

            # compute covariance btw. RVs 'X' and 'Y'
            # shape: (2*n_outputs, 2*n_outputs)
            # It contains the following 4 block matrices:
            #   (1, 1) variance of 'X'
            #  *(1, 2) covariance between 'X' and 'Y' (a.k.a. cross-covariance)
            #   (2, 1) covariance between 'Y' and 'X' (a.k.a. cross-covariance)
            #   (2, 2) variance of 'Y'
            _cov_2 = np.cov(all_Y_i_u, all_Y_i_tilde)  # for total order indices

            # We need the cross-covariance between 'X' and 'Y'
            # Extract *(1, 2) (upper right block)
            # shape: (n_outputs, n_outputs)
            C_u_tilde = _cov_2[0:n_outputs, n_outputs : 2 * n_outputs]
            denominator = np.trace(C)

            # Generalised Sobol indices
            gen_sobol_total_i[i] = 1 - np.trace(C_u_tilde) / denominator

        return gen_sobol_total_i


# src/UQpy/sensitivity/sobol.py

"""

The Sobol class computes the Sobol indices for single output and multi-output
models. The Sobol indices can be computed using various pick-and-freeze 
schemes.

The schemes implemented are listed below:

# First order indices: 
- Sobol1993 [1]: Requires n_samples*(num_vars + 1) model evaluations
- Saltelli2002 [3]: Requires n_samples*(2*num_vars + 1) model evaluations
- Janon2014 [4]: Requires n_samples*(num_vars + 1) model evaluations

# Second order indices:
- Saltelli2002 [3]: Requires n_samples*(2*num_vars + 1) model evaluations

# Total order indices:
- Homma1996: Requires n_samples*(num_vars + 1) model evaluations
- Saltelli2002 [3]: Requires n_samples*(2*num_vars + 1) model evaluations

For more details on "Saltelli2002" refer to [3].
    
Note: Apart from second order indices, the Saltelli2002 scheme provides 
      more accurate estimates of all indices, as opposed to Homma1996 or Sobol1993. 
      Because this method efficiently utilizes the higher number of model evaluations.

Additionally, we can compute the confidence intervals for the Sobol indices 
using bootstrapping [2].


References
----------

.. [1] Sobol, I.M. (1993) Sensitivity Estimates for Nonlinear Mathematical Models.  
       Mathematical Modelling and Computational Experiments, 4, 407-414.

.. [2] Jeremy Orloff and Jonathan Bloom (2014), Bootstrap confidence intervals, 
       Introduction to Probability and Statistics, MIT OCW. 

.. [3] Saltelli, A. (2002). Making best use of model evaluations to
       compute sensitivity indices.

.. [4] Janon, Alexander; Klein, Thierry; Lagnoux, Agnes; Nodet, Maëlle; 
       Prior, Clementine. Asymptotic normality and efficiency of two Sobol index 
       estimators. ESAIM: Probability and Statistics, Volume 18 (2014), pp. 342-364. 
       doi:10.1051/ps/2013040. http://www.numdam.org/articles/10.1051/ps/2013040/

"""

import math
import logging
import itertools

import numpy as np

from UQpy.sensitivity.baseclass.sensitivity import Sensitivity
from UQpy.sensitivity.baseclass.pickfreeze import generate_pick_freeze_samples
from UQpy.utilities.UQpyLoggingFormatter import UQpyLoggingFormatter

# TODO: Sampling strategies


class Sobol(Sensitivity):
    """
    Compute Sobol sensitivity indices using the pick
    and freeze algorithm. For models with multiple outputs
    (vector-valued response), the sensitivity indices are computed for each
    output separately.
    For time-series models, the sensitivity indices are computed for each
    time instant separately. (Pointwise-in-time Sobol indices)

    :param runmodel_object: The computational model. It should be of type :class:`.RunModel`. \
        The output QoI can be a scalar or vector of length :code:`ny`, then the sensitivity \
        indices of all :code:`ny` outputs are computed independently.

    :param distributions: List of :class:`.Distribution` objects corresponding to each \
        random variable, or :class:`.JointIndependent` object \
        (multivariate RV with independent marginals).

    :param random_state: Random seed used to initialize the pseudo-random number \
        generator. Default is :any:`None`.

    **Methods:**
    """

    def __init__(
        self, runmodel_object, dist_object, random_state=None, **kwargs
    ) -> None:

        super().__init__(runmodel_object, dist_object, random_state, **kwargs)

        # Create logger with the same name as the class
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.ERROR)
        frmt = UQpyLoggingFormatter()

        # create console handler with a higher log level
        ch = logging.StreamHandler()
        ch.setFormatter(frmt)

        # add the handler to the logger
        self.logger.addHandler(ch)

        self.sobol_i = None
        "First order Sobol indices, :class:`numpy.ndarray` of shape `(num_vars, n_outputs)`"

        self.sobol_total_i = None
        "Total order Sobol indices, :class:`numpy.ndarray` of shape `(num_vars, n_outputs)`"

        self.sobol_ij = None
        "Second order Sobol indices, :class:`numpy.ndarray` of shape `(num_second_order_terms, n_outputs)`"

        self.CI_sobol_i = None
        "Confidence intervals for the first order Sobol indices, :class:`numpy.ndarray` of shape `(num_vars, 2)`"

        self.CI_sobol_total_i = None
        "Confidence intervals for the total order Sobol indices, :class:`numpy.ndarray` of shape `(num_vars, 2)`"

        self.CI_sobol_ij = None
        "Confidence intervals for the second order Sobol indices, :class:`numpy.ndarray` of shape `(num_second_order_terms, 2)`"

        self.n_samples = None
        "Number of samples used to compute the sensitivity indices, :class:`int`"

        self.num_vars = None
        "Number of model input variables, :class:`int`"

        self.multioutput = None
        "True if the model has multiple outputs, :class:`bool`"

    def run(
        self,
        n_samples=1_000,
        num_bootstrap_samples=None,
        confidence_level=0.95,
        estimate_second_order=False,
        first_order_scheme="Janon2014",
        total_order_scheme="Homma1996",
        second_order_scheme="Saltelli2002",
    ):

        """
        Compute the sensitivity indices and confidence intervals.

        :param n_samples: Number of samples used to compute the sensitivity indices. \
            Default is 1,000.

        :param num_bootstrap_samples: Number of bootstrap samples used to compute the \
            confidence intervals. Default is :any:`None`.

        :param confidence_interval: Confidence level used to compute the confidence \
            intervals. Default is 0.95.

        :param estimate_second_order: If True, the second order Sobol indices are \
            estimated. Default is False.

        :param first_order_scheme: Scheme used to compute the first order Sobol \
            indices. Default is "Janon2014".

        :param total_order_scheme: Scheme used to compute the total order Sobol \
            indices. Default is "Homma1996".

        :param second_order_scheme: Scheme used to compute the second order \
            Sobol indices. Default is "Saltelli2002".

        :return: A :class:`dict` with the following keys: \
            :code:`sobol_i` of shape :code:`(num_vars, 1)`, \
            :code:`sobol_total_i` of shape :code:`(num_vars, 1)`, \
            :code:`sobol_ij` of shape :code:`(num_second_order_terms, 1)`, \
            :code:`CI_sobol_i` of shape :code:`(num_vars, 2)`, \
            if multioutput: Shape: `(n_outputs, num_vars, 2)`, \
            :code:`CI_sobol_total_i` of shape :code:`(num_vars, 2)`, \
            if multioutput: Shape: `(n_outputs, num_vars, 2)`, \
            :code:`CI_sobol_ij` of shape :code:`(num_second_order_terms, 2)`
            if multioutput: Shape: `(n_outputs, num_second_order_terms, 2)`, \

        """
        # Check n_samples data type
        self.n_samples = n_samples
        if not isinstance(self.n_samples, int):
            raise TypeError("UQpy: n_samples should be an integer.")

        # Check num_bootstrap_samples data type
        if num_bootstrap_samples is not None:
            if not isinstance(num_bootstrap_samples, int):
                raise TypeError("UQpy: num_bootstrap_samples should be an integer.")
        elif num_bootstrap_samples is None:
            self.logger.info(
                "UQpy: num_bootstrap_samples is set to None, confidence intervals will not be computed."
            )

        ################## GENERATE SAMPLES ##################

        (
            A_samples,
            B_samples,
            C_i_generator,
            D_i_generator,
        ) = generate_pick_freeze_samples(
            self.dist_object, self.n_samples, self.random_state
        )

        self.logger.info("UQpy: Generated samples using the pick-freeze scheme.")

        self.num_vars = A_samples.shape[1]  # Number of variables

        ################# MODEL EVALUATIONS ####################

        A_model_evals = self._run_model(A_samples)  # shape: (n_samples, n_outputs)

        self.logger.info("UQpy: Model evaluations A completed.")

        B_model_evals = self._run_model(B_samples)  # shape: (n_samples, n_outputs)

        self.logger.info("UQpy: Model evaluations B completed.")

        # Check the number of outputs of the model
        try:
            self.n_outputs = A_model_evals.shape[1]
        except:
            self.n_outputs = 1

        # multioutput flag
        self.multioutput = True if self.n_outputs > 1 else False

        if not self.multioutput:
            A_model_evals = A_model_evals.reshape(-1, 1)
            B_model_evals = B_model_evals.reshape(-1, 1)

        C_i_model_evals = np.zeros((self.n_outputs, self.n_samples, self.num_vars))

        for i, C_i in enumerate(C_i_generator):
            C_i_model_evals[:, :, i] = self._run_model(C_i).T

        self.logger.info("UQpy: Model evaluations C completed.")

        # Compute D_i_model_evals only if needed
        if estimate_second_order or total_order_scheme == "Saltelli2002":

            D_i_model_evals = np.zeros((self.n_outputs, self.n_samples, self.num_vars))

            for i, D_i in enumerate(D_i_generator):
                D_i_model_evals[:, :, i] = self._run_model(D_i).T

            self.logger.info("UQpy: Model evaluations D completed.")

        else:
            D_i_model_evals = None

        self.logger.info("UQpy: All model evaluations computed successfully.")

        ######################### STORAGE ########################

        # Create dictionary to store the sensitivity indices
        computed_indices = {}

        ################## COMPUTE SOBOL INDICES ##################

        # First order Sobol indices
        self.sobol_i = compute_first_order(
            A_model_evals,
            B_model_evals,
            C_i_model_evals,
            D_i_model_evals,
            scheme=first_order_scheme,
        )

        self.logger.info("UQpy: First order Sobol indices computed successfully.")

        computed_indices["sobol_i"] = self.sobol_i

        # Total order Sobol indices
        self.sobol_total_i = compute_total_order(
            A_model_evals,
            B_model_evals,
            C_i_model_evals,
            D_i_model_evals,
            scheme=total_order_scheme,
        )

        self.logger.info("UQpy: Total order Sobol indices computed successfully.")

        computed_indices["sobol_total_i"] = self.sobol_total_i

        if estimate_second_order:

            # Second order Sobol indices
            self.sobol_ij = compute_second_order(
                A_model_evals,
                B_model_evals,
                C_i_model_evals,
                D_i_model_evals,
                computed_indices["sobol_i"],
                scheme=second_order_scheme,
            )

            self.logger.info("UQpy: Second order Sobol indices computed successfully.")

            computed_indices["sobol_ij"] = self.sobol_ij

        ################## CONFIDENCE INTERVALS ####################

        if num_bootstrap_samples is not None:

            self.logger.info("UQpy: Computing confidence intervals ...")

            estimator_inputs = [
                A_model_evals,
                B_model_evals,
                C_i_model_evals,
                D_i_model_evals,
            ]

            # First order Sobol indices
            self.CI_sobol_i = self.bootstrapping(
                compute_first_order,
                estimator_inputs,
                computed_indices["sobol_i"],
                num_bootstrap_samples,
                confidence_level,
                scheme=first_order_scheme,
            )

            self.logger.info(
                "UQpy: Confidence intervals for First order Sobol indices computed successfully."
            )

            computed_indices["CI_sobol_i"] = self.CI_sobol_i

            # Total order Sobol indices
            self.CI_sobol_total_i = self.bootstrapping(
                compute_total_order,
                estimator_inputs,
                computed_indices["sobol_total_i"],
                num_bootstrap_samples,
                confidence_level,
                scheme=total_order_scheme,
            )

            self.logger.info(
                "UQpy: Confidence intervals for Total order Sobol indices computed successfully."
            )

            computed_indices["CI_sobol_total_i"] = self.CI_sobol_total_i

            # Second order Sobol indices
            if estimate_second_order:
                self.CI_sobol_ij = self.bootstrapping(
                    compute_second_order,
                    estimator_inputs,
                    computed_indices["sobol_ij"],
                    num_bootstrap_samples,
                    confidence_level,
                    first_order_sobol=computed_indices["sobol_i"],
                    scheme=second_order_scheme,
                )

                self.logger.info(
                    "UQpy: Confidence intervals for Second order Sobol indices computed successfully."
                )

                computed_indices["CI_sobol_ij"] = self.CI_sobol_ij

        return computed_indices


###################### Pick and Freeze Methods #####################

"""

These methods are also called by other sensitivity methods (such as Chatterjee, 
Cramer-von Mises) to estimate the Sobol indices and therefore are implemented as 
functions and not static methods in the Sobol class.


#! Saltelli2002
--------------------------------------------------------------------------------

Sobol indices estimated as per Theorem 2 in [3]_. Refer page 7 in
[3]_ for details.

Since there are several sets of function evaluations available,
there are several ways to estimate E[Y]^2 and V[Y].
Below we summarise the evaluations to be used as given in Theorem 2.

# First-order indices:
    - E[Y]^2 : f_A, f_B
    - V[Y] : f_A
    - S_i = ( <f_A, f_C_i>/N - E[Y]^2 ) / V[Y]


# Second-order indices:
    - Estimate 1:
        - E[Y]^2 : f_C_l, f_D_l -> l = max(i,j)
        - V[Y] : f_C_j or f_D_i
        - V^c_ij = f_D_i, f_C_j

    - Estimate 2:
        - E[Y]^2: f_C_l, f_D_l -> l = min(i,j)
        - V[Y] : f_C_i or f_D_j
        - V^c_ij = f_D_j, f_C_i

    where:
    S_ij = S^c_ij - S_i - S_j
    S^c_ij = ( <f_A, f_C_ij>/N - E[Y]^2 ) / V[Y] # Esimate 1
            = ( <f_B, f_D_ij>/N - E[Y]^2 ) / V[Y] # Esimate 2

# Total-order indices:
        - E[Y]^2 : f_B
        - V[Y] : f_B
        - S_T_i = 1 - ( <f_B, f_C_i>/N - E[Y]^2 ) / V[Y]

For m=5, the Sobol indices are estimated as follows:
First order indices: 2 estimates
Second order indices: 2 estimates
Total order indices: 2 estimates
S_{-ij}: 2 estimates
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
|       |   f_B  |  f_C_1  |  f_C_2  |  f_C_3  |  f_C_4  |  f_C_5 |  f_D_1  |  f_D_2  |  f_D_3  |  f_D_4  | f_D_5 |  f_A |
+=======+========+=========+=========+=========+=========+========+=========+=========+=========+=========+=======+======+
|  f_B  |  V[Y]  |         |         |         |         |        |         |         |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_C_1 |  S_T_1 |   V[Y]  |         |         |         |        |         |         |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_C_2 |  S_T_2 | V^c_-12 |   V[Y]  |         |         |        |         |         |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_C_3 |  S_T_3 | V^c_-13 | V^c_-23 |   V[Y]  |         |        |         |         |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_C_4 |  S_T_4 | V^c_-14 | V^c_-24 | V^c_-34 |   V[Y]  |        |         |         |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_C_5 |  S_T_5 | V^c_-15 | V^c_-25 | V^c_-35 | V^c_-45 |  V[Y]  |         |         |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_D_1 |   S_1  |  E^2[Y] |  V^c_12 |  V^c_13 |  V^c_14 | V^c_15 |   V[Y]  |         |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_D_2 |   S_2  |  V^c_12 |  E^2[Y] |  V^c_23 |  V^c_24 | V^c_25 | V^c_-12 |   V[Y]  |         |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_D_3 |   S_3  |  V^c_13 |  V^c_23 |  E^2[Y] |  V^c_34 | V^c_35 | V^c_-13 | V^c_-23 |   V[Y]  |         |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_D_4 |   S_4  |  V^c_14 |  V^c_24 |  V^c_34 |  E^2[Y] | V^c_45 | V^c_-14 | V^c_-24 | V^c_-34 |   V[Y]  |       |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
| f_D_5 |   S_5  |  V^c_15 |  V^c_25 |  V^c_35 |  V^c_45 | E^2[Y] | V^c_-15 | V^c_-25 | V^c_-35 | V^c_-45 |  V[Y] |      |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+
|  f_A  | E^2[Y] |   S_1   |   S_2   |   S_3   |   S_4   |   S_5  |  S_T_1  |  S_T_2  |  S_T_3  |  S_T_4  | S_T_5 | V[Y] |
+-------+--------+---------+---------+---------+---------+--------+---------+---------+---------+---------+-------+------+

For m>5, we can follow the same procedure as above.

For m = 4, the Sobol indices are estimated as follows:
First order indices: 2 estimates
Second order indices: 4 estimates
Total order indices: 2 estimates
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
|       |   f_B  |  f_C_1 |  f_C_2 |  f_C_3 |  f_C_4 |  f_D_1 |  f_D_2 |  f_D_3 | f_D_4 |  f_A |
+=======+========+========+========+========+========+========+========+========+=======+======+
|  f_B  |  V[Y]  |        |        |        |        |        |        |        |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_C_1 |  S_T_1 |  V[Y]  |        |        |        |        |        |        |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_C_2 |  S_T_2 | V^c_34 |  V[Y]  |        |        |        |        |        |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_C_3 |  S_T_3 | V^c_24 | V^c_14 |  V[Y]  |        |        |        |        |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_C_4 |  S_T_4 | V^c_23 | V^c_13 | V^c_12 |  V[Y]  |        |        |        |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_D_1 |   S_1  | E^2[Y] | V^c_12 | V^c_13 | V^c_14 |  V[Y]  |        |        |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_D_2 |   S_2  | V^c_12 | E^2[Y] | V^c_23 | V^c_24 | V^c_34 |  V[Y]  |        |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_D_3 |   S_3  | V^c_13 | V^c_23 | E^2[Y] | V^c_34 | V^c_25 | V^c_14 |  V[Y]  |       |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
| f_D_4 |   S_4  | V^c_14 | V^c_24 | V^c_34 | E^2[Y] | V^c_23 | V^c_13 | V^c_12 |  V[Y] |      |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+
|  f_A  | E^2[Y] |   S_1  |   S_2  |   S_3  |   S_4  |  S_T_1 |  S_T_2 |  S_T_3 | S_T_4 | V[Y] |
+-------+--------+--------+--------+--------+--------+--------+--------+--------+-------+------+

For m = 3, the Sobol indices are estimated as follows:
First order indices: 4 estimates
Second order indices: 2 estimates
Total order indices: 2 estimates
+-------+--------+--------+--------+--------+-------+-------+-------+------+
|       |   f_B  |  f_C_1 |  f_C_2 |  f_C_3 | f_D_1 | f_D_2 | f_D_3 |  f_A |
+=======+========+========+========+========+=======+=======+=======+======+
|  f_B  |  V[Y]  |        |        |        |       |       |       |      |
+-------+--------+--------+--------+--------+-------+-------+-------+------+
| f_C_1 |  S_T_1 |  V[Y]  |        |        |       |       |       |      |
+-------+--------+--------+--------+--------+-------+-------+-------+------+
| f_C_2 |  S_T_2 |   S_3  |  V[Y]  |        |       |       |       |      |
+-------+--------+--------+--------+--------+-------+-------+-------+------+
| f_C_3 |  S_T_3 |   S_2  |   S_1  |  V[Y]  |       |       |       |      |
+-------+--------+--------+--------+--------+-------+-------+-------+------+
| f_D_1 |   S_1  | E^2[Y] | V^c_12 | V^c_13 |  V[Y] |       |       |      |
+-------+--------+--------+--------+--------+-------+-------+-------+------+
| f_D_2 |   S_2  | V^c_12 | E^2[Y] | V^c_23 |  S_3  |  V[Y] |       |      |
+-------+--------+--------+--------+--------+-------+-------+-------+------+
| f_D_3 |   S_3  | V^c_13 | V^c_23 | E^2[Y] |  S_2  |  S_1  |  V[Y] |      |
+-------+--------+--------+--------+--------+-------+-------+-------+------+
|  f_A  | E^2[Y] |   S_1  |   S_2  |   S_3  | S_T_1 | S_T_2 | S_T_3 | V[Y] |
+-------+--------+--------+--------+--------+-------+-------+-------+------+

"""


def compute_first_order(
    A_model_evals,
    B_model_evals,
    C_i_model_evals,
    D_i_model_evals=None,
    scheme="Janon2014",
):

    """
    Compute first order Sobol' indices using the Pick-and-Freeze scheme.

    For the Sobol1996 scheme:
    For computing the first order Sobol' indices, only f_A_model_evals and
    f_C_i_model_evals are required. The other inputs are optional.
        f_B_model_evals is set to None if f_B_model_evals is not provided.

    **Inputs:**

    * **A_model_evals** (`ndarray`):
        Shape: `(n_samples, n_outputs)`.

    * **B_model_evals** (`ndarray`):
        If not available, pass `None`.
        Shape: `(n_samples, n_outputs)`.

    * **C_i_model_evals** (`ndarray`):
        Shape: `(n_outputs, n_samples, num_vars)`.

    * **D_i_model_evals** (`ndarray`, optional):
        Shape: `(n_outputs, n_samples, num_vars)`.

    * **scheme** (`str`, optional):
        Scheme to use for computing the first order Sobol' indices.
        Default: 'Sobol1993'.

    **Outputs:**

    * **first_order_sobol** (`ndarray`):
        First order Sobol' indices.
        Shape: `(num_vars, n_outputs)`.

    """

    n_samples = A_model_evals.shape[0]
    n_outputs = A_model_evals.shape[1]
    num_vars = C_i_model_evals.shape[2]

    # Store first order Sobol' indices
    first_order_sobol = np.zeros((num_vars, n_outputs))

    if scheme == "Sobol1993":

        for output_j in range(n_outputs):

            f_A = A_model_evals[:, output_j]
            f_B = B_model_evals[:, output_j] if B_model_evals is not None else None

            # combine all model evaluations
            # to improve accuracy of the estimator
            _all_model_evals = np.append(f_A, f_B) if f_B is not None else f_A
            f_0 = np.mean(_all_model_evals)  # scalar

            f_0_square = f_0**2
            total_variance = np.var(_all_model_evals, ddof=1)

            for var_i in range(num_vars):

                f_C_i = C_i_model_evals[output_j, :, var_i]

                S_i = (np.dot(f_A, f_C_i) / n_samples - f_0_square) / total_variance

                first_order_sobol[var_i, output_j] = S_i

    elif scheme == "Janon2014":

        for output_j in range(n_outputs):

            f_A = A_model_evals[:, output_j]

            for var_i in range(num_vars):

                f_C_i = C_i_model_evals[output_j, :, var_i]

                # combine all model evaluations
                # to improve accuracy of the estimator
                _all_model_evals = np.append(f_A, f_C_i)
                f_0 = np.mean(_all_model_evals)

                f_0_square = f_0**2
                total_variance = np.mean(_all_model_evals**2) - f_0_square

                S_i = (np.dot(f_A, f_C_i) / n_samples - f_0_square) / total_variance

                first_order_sobol[var_i, output_j] = S_i

    elif scheme == "Saltelli2002":

        """
        Number of estimates for first order indices is 4 if
        num_vars is 3, else 2.

        """

        for output_j in range(n_outputs):

            f_A = A_model_evals[:, output_j]
            f_B = B_model_evals[:, output_j]
            f_0_square = np.dot(f_A, f_B) / n_samples
            total_variance = np.var(f_A, ddof=1)

            for var_i in range(num_vars):

                f_C_i = C_i_model_evals[output_j, :, var_i]
                f_D_i = D_i_model_evals[output_j, :, var_i]

                # (Estimate 1)
                est_1 = (np.dot(f_A, f_C_i) / n_samples - f_0_square) / total_variance

                # (Estimate 2)
                est_2 = (np.dot(f_B, f_D_i) / n_samples - f_0_square) / total_variance

                if num_vars == 3:

                    # list of variable indices
                    list_vars = list(range(num_vars))
                    list_vars.remove(var_i)
                    # combination of all remaining variables indices
                    rem_vars_perm = list(itertools.permutations(list_vars, 2))

                    # (Estimate 3)
                    var_a, var_b = rem_vars_perm[0]
                    f_C_a = C_i_model_evals[output_j, :, var_a]
                    f_C_b = C_i_model_evals[output_j, :, var_b]
                    est_3 = (
                        np.dot(f_C_a, f_C_b) / n_samples - f_0_square
                    ) / total_variance

                    # (Estimate 4)
                    var_a, var_b = rem_vars_perm[1]
                    f_D_a = D_i_model_evals[output_j, :, var_a]
                    f_D_b = D_i_model_evals[output_j, :, var_b]
                    est_4 = (
                        np.dot(f_D_a, f_D_b) / n_samples - f_0_square
                    ) / total_variance

                    first_order_sobol[var_i, output_j] = (
                        est_1 + est_2 + est_3 + est_4
                    ) / 4

                else:
                    first_order_sobol[var_i, output_j] = (est_1 + est_2) / 2

    return first_order_sobol


def compute_total_order(
    A_model_evals,
    B_model_evals,
    C_i_model_evals,
    D_i_model_evals=None,
    scheme="Homma1996",
):

    """
    Compute total order Sobol' indices using the Pick-and-Freeze scheme.

    For the Homma1996 scheme:
    For computing the first order Sobol' indices, only f_B_model_evals and
    f_C_i_model_evals are required.
        f_A_model_evals is set to None if f_A_model_evals is not provided.

    **Inputs:**

    * **A_model_evals** (`ndarray`):
        If not available, pass `None`.
        Shape: `(n_samples, n_outputs)`.

    * **B_model_evals** (`ndarray`):
        Shape: `(n_samples, n_outputs)`.

    * **C_i_model_evals** (`ndarray`):
        Shape: `(n_outputs, n_samples, num_vars)`.

    * **D_i_model_evals** (`ndarray`, optional):
        Shape: `(n_outputs, n_samples, num_vars)`.

    * **scheme** (`str`, optional):
        Scheme to use for computing the total order Sobol' indices.
        Default: 'Homma1996'.

    **Outputs:**

    * **total_order_sobol** (`ndarray`):
        Total order Sobol' indices.
        Shape: `(num_vars, n_outputs)`.

    """

    n_samples = A_model_evals.shape[0]
    n_outputs = A_model_evals.shape[1]
    num_vars = C_i_model_evals.shape[2]

    # Store total order Sobol' indices
    total_order_sobol = np.zeros((num_vars, n_outputs))

    if scheme == "Homma1996":

        for output_j in range(n_outputs):

            f_A = A_model_evals[:, output_j] if A_model_evals is not None else None
            f_B = B_model_evals[:, output_j]

            # combine all model evaluations
            # to improve accuracy of the estimator
            _all_model_evals = np.append(f_A, f_B) if f_A is not None else f_B
            f_0 = np.mean(_all_model_evals)  # scalar

            f_0_square = f_0**2
            total_variance = np.var(_all_model_evals, ddof=1)

            for var_i in range(num_vars):

                f_C_i = C_i_model_evals[output_j, :, var_i]

                S_T_i = (
                    1 - (np.dot(f_B, f_C_i) / n_samples - f_0_square) / total_variance
                )

                total_order_sobol[var_i, output_j] = S_T_i

    elif scheme == "Saltelli2002":

        for output_j in range(n_outputs):

            f_A = A_model_evals[:, output_j]
            f_B = B_model_evals[:, output_j]
            f_0_square = np.mean(f_B) ** 2
            total_variance = np.var(f_B, ddof=1)

            for var_i in range(num_vars):

                f_C_i = C_i_model_evals[output_j, :, var_i]
                f_D_i = D_i_model_evals[output_j, :, var_i]

                # (Estimate 1)
                est_1 = (
                    1 - (np.dot(f_B, f_C_i) / n_samples - f_0_square) / total_variance
                )

                # (Estimate 2)
                est_2 = (
                    1 - (np.dot(f_A, f_D_i) / n_samples - f_0_square) / total_variance
                )

                total_order_sobol[var_i, output_j] = (est_1 + est_2) / 2

    return total_order_sobol


def compute_second_order(
    A_model_evals,
    B_model_evals,
    C_i_model_evals,
    D_i_model_evals,
    first_order_sobol=None,  # None to make it a make keyword argument
    scheme="Saltelli2002",
):
    """
    Compute the second order Sobol indices using the Pick-and-Freeze scheme.

    NOTE:
    - Number of estimates for second order indices is 4 if
        num_vars is 4, else 2.

    - Although the B_model_evals are not being used currently, they are
        included for use in estimate 3 and 4 for case num_vars = 4.

    **Inputs:**

    * **A_model_evals** (`ndarray`):
        Shape: `(n_samples, n_outputs)`.

    * **B_model_evals** (`ndarray`):
        Shape: `(n_samples, n_outputs)`.

    * **C_i_model_evals** (`ndarray`):
        Shape: `(n_outputs, n_samples, num_vars)`.

    * **D_i_model_evals** (`ndarray`, optional):
        Shape: `(n_outputs, n_samples, num_vars)`.

    * **first_order_sobol** (`ndarray`):
        First order Sobol' indices.
        Shape: `(num_vars, n_outputs)`.

    * **scheme** (`str`, optional):
        Scheme to use for computing the first order Sobol' indices.
        Default: 'Sobol1993'.

    **Outputs:**

    * **second_order_sobol** (`ndarray`):
        Second order Sobol indices.
        Shape: `(num_second_order_terms, n_outputs)`.
    """

    n_samples = A_model_evals.shape[0]
    n_outputs = A_model_evals.shape[1]
    num_vars = C_i_model_evals.shape[2]

    second_order_terms = itertools.combinations(range(num_vars), 2)
    second_order_terms = list(second_order_terms)
    num_second_order_terms = math.comb(num_vars, 2)

    # Store second order Sobol' indices
    second_order_sobol = np.zeros((num_second_order_terms, n_outputs))

    if scheme == "Saltelli2002":

        for output_j in range(n_outputs):

            for k in range(num_second_order_terms):

                var_a, var_b = second_order_terms[k]
                S_a = first_order_sobol[var_a, output_j]
                S_b = first_order_sobol[var_b, output_j]

                # (Estimate 1)
                var_c = np.max([var_a, var_b])
                f_C_c = C_i_model_evals[output_j, :, var_c]
                f_D_c = D_i_model_evals[output_j, :, var_c]
                f_0_square = np.dot(f_D_c, f_C_c) / n_samples
                total_variance = np.var(f_D_c, ddof=1)

                f_C_a = C_i_model_evals[output_j, :, var_a]
                f_D_b = D_i_model_evals[output_j, :, var_b]
                S_c_ab_1 = (
                    np.dot(f_C_a, f_D_b) / n_samples - f_0_square
                ) / total_variance

                est_1 = S_c_ab_1 - S_a - S_b

                # (Estimate 2)
                var_c = np.min([var_a, var_b])
                f_C_c = C_i_model_evals[output_j, :, var_c]
                f_D_c = D_i_model_evals[output_j, :, var_c]
                f_0_square = np.dot(f_D_c, f_C_c) / n_samples
                total_variance = np.var(f_D_c, ddof=1)

                f_D_a = D_i_model_evals[output_j, :, var_a]
                f_C_b = C_i_model_evals[output_j, :, var_b]
                S_c_ab_2 = (
                    np.dot(f_D_a, f_C_b) / n_samples - f_0_square
                ) / total_variance

                est_2 = S_c_ab_2 - S_a - S_b

                if num_vars == 4:

                    # (Estimate 3)
                    # TODO: How to compute this?

                    # (Estimate 4)
                    # TODO: How to compute this?

                    # second_order_sobol[k, output_j] = (
                    #     est_1 + est_2 + est_3 + est_4
                    # ) / 4

                    pass

                else:
                    second_order_sobol[k, output_j] = (est_1 + est_2) / 2

    return second_order_sobol


# src/UQpy/stochastic_process/BispectralRepresentation.py

import itertools
from UQpy.utilities import *


class BispectralRepresentation:
    def __init__(
            self,
            n_samples: int,
            power_spectrum: Union[list, np.ndarray],
            bispectrum: Union[list, np.ndarray],
            time_interval: Union[list, np.ndarray],
            frequency_interval: Union[list, np.ndarray],
            n_time_intervals: Union[list, np.ndarray],
            n_frequency_intervals: Union[list, np.ndarray],
            case="uni",
            random_state: RandomStateType = None,
    ):
        """
        A class to simulate non-Gaussian stochastic processes from a given power spectrum and bispectrum based on the
        3-rd order Spectral Representation Method. This class can simulate uni-variate, one-dimensional and
        multi-dimensional stochastic processes.

        :param n_samples: Number of samples of the stochastic process to be simulated.
         The :meth:`run` method is automatically called if `n_samples` is provided. If `n_samples` is not
         provided, then the :class:`.BispectralRepresentation` object is created but samples are not generated.
        :param power_spectrum: The discretized power spectrum.
         - For uni-variate, one-dimensional processes `power_spectrum` will be :class:`list` or :class:`numpy.ndarray`
         of length `n_frequency_intervals`.

         - For uni-variate, multi-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(n_frequency_intervals[0], ..., n_frequency_intervals[n_dimensions-1])`

        :param bispectrum: The prescribed bispectrum.
         - For uni-variate, one-dimensional processes, `bispectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(n_frequency_intervals, n_frequency_intervals)`

         - For uni-variate, multi-dimensional processes, `bispectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(n_frequency_intervals[0], ..., n_frequency_intervals[n_dimensions-1], n_frequency_intervals[0], ..., n_frequency_intervals[n_dimensions-1])`

        :param time_interval: Length of time discretizations (:math:`\Delta t`) for each dimension of size
         `n_dimensions`.
        :param frequency_interval: Length of frequency discretizations (:math:`\Delta \omega`) for each dimension of
         size `n_dimensions`.
        :param n_time_intervals: Number of time discretizations for each dimensions of size `n_dimensions`.
        :param n_frequency_intervals: Number of frequency discretizations for each dimension of size
         `n_dimensions`.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
         If an :any:`int` is provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise,
         the object itself can be passed directly.
        """
        self.n_samples = n_samples
        self.n_frequency_intervals = np.array(n_frequency_intervals)
        self.n_time_intervals = np.array(n_time_intervals)
        self.frequency_interval = np.array(frequency_interval)
        self.time_interval = np.array(time_interval)
        self.n_dimensions: int = len(power_spectrum.shape)
        """The dimensionality of the stochastic process."""
        self.power_spectrum = power_spectrum
        self.bispectrum = bispectrum

        # Error checks
        t_u = (2 * np.pi / (2 * self.n_frequency_intervals * self.frequency_interval))
        if (self.time_interval > t_u).any():
            raise RuntimeError("UQpy: Aliasing might occur during execution")

        self.random_state = random_state
        if isinstance(self.random_state, int):
            np.random.seed(self.random_state)
        elif not isinstance(self.random_state, (type(None), np.random.RandomState)):
            raise TypeError("UQpy: random_state must be None, an int or an np.random.RandomState object.")

        self.logger = logging.getLogger(__name__)

        self.bispectrum_amplitude: float = np.absolute(bispectrum)
        """The amplitude of the bispectrum."""
        self.bispectrum_real: float = np.real(bispectrum)
        """The real part of the bispectrum."""
        self.bispectrum_imaginary: float = np.imag(bispectrum)
        """The imaginary part of the bispectrum."""
        self.biphase: NumpyFloatArray = np.arctan2(self.bispectrum_imaginary, self.bispectrum_real)
        """The biphase values of the bispectrum."""
        self.biphase[np.isnan(self.biphase)] = 0

        self.phi: NumpyFloatArray = None
        """The random phase angles used in the simulation of the stochastic process.
        The shape of the phase angles :code:`(n_samples, n_variables, n_frequency_intervals[0], ...,
        n_frequency_intervals[n_dimensions-1])`"""
        self.samples: NumpyFloatArray = None
        """Generated samples.
        The shape of the samples is :code:`(n_samples, n_variables, n_time_intervals[0], ...,
        n_time_intervals[n_dimensions-1])`"""

        self.case = case

        if self.n_dimensions == len(self.power_spectrum.shape):
            self.case = "uni"
            self._compute_bicoherence_uni()
        else:
            self.n_variables: int = self.power_spectrum.shape[0]
            """Number of variables in the stochastic process."""
            self.case = "multi"

        if self.n_samples is not None:
            self.run(n_samples=self.n_samples)

    def _compute_bicoherence_uni(self):
        self.logger.info("UQpy: Stochastic Process: Computing the partial bicoherence values.")
        self.bc2 = np.zeros_like(self.bispectrum_real)
        """The bicoherence values of the power spectrum and bispectrum."""
        self.pure_power_sepctrum = np.zeros_like(self.power_spectrum)
        """The pure part of the power spectrum."""
        self.sum_bc2 = np.zeros_like(self.power_spectrum)
        """The sum of the bicoherence values for single frequencies."""

        if self.n_dimensions == 1:
            self.pure_power_sepctrum[0] = self.power_spectrum[0]
            self.pure_power_sepctrum[1] = self.power_spectrum[1]

        if self.n_dimensions == 2:
            self.pure_power_sepctrum[0, :] = self.power_spectrum[0, :]
            self.pure_power_sepctrum[1, :] = self.power_spectrum[1, :]
            self.pure_power_sepctrum[:, 0] = self.power_spectrum[:, 0]
            self.pure_power_sepctrum[:, 1] = self.power_spectrum[:, 1]

        if self.n_dimensions == 3:
            self.pure_power_sepctrum[0, :, :] = self.power_spectrum[0, :, :]
            self.pure_power_sepctrum[1, :, :] = self.power_spectrum[1, :, :]
            self.pure_power_sepctrum[:, 0, :] = self.power_spectrum[:, 0, :]
            self.pure_power_sepctrum[:, 1, :] = self.power_spectrum[:, 1, :]
            self.pure_power_sepctrum[:, :, 0] = self.power_spectrum[:, :, 0]
            self.pure_power_sepctrum[:, 0, 1] = self.power_spectrum[:, :, 1]

        self.ranges = [range(self.n_frequency_intervals[i]) for i in range(self.n_dimensions)]

        for i in itertools.product(*self.ranges):
            wk = np.array(i)
            for j in itertools.product(*[range(np.int32(k)) for k in np.ceil((wk + 1) / 2)]):
                wj = np.array(j)
                wi = wk - wj
                if (self.bispectrum_amplitude[(*wi, *wj)] > 0 and self.pure_power_sepctrum[(*wi, *[])]
                        * self.pure_power_sepctrum[(*wj, *[])] != 0):
                    self.bc2[(*wi, *wj)] = (self.bispectrum_amplitude[(*wi, *wj)] ** 2
                                            / (self.pure_power_sepctrum[(*wi, *[])]
                                               * self.pure_power_sepctrum[(*wj, *[])]
                                               * self.power_spectrum[(*wk, *[])]) * np.prod(self.frequency_interval))
                    self.sum_bc2[(*wk, *[])] = (self.sum_bc2[(*wk, *[])] + self.bc2[(*wi, *wj)])
                else:
                    self.bc2[(*wi, *wj)] = 0
            if self.sum_bc2[(*wk, *[])] > 1:
                self.logger.info("UQpy: Stochastic Process: Results may not be as expected as sum of partial "
                                 "bicoherences is greater than 1")
                for j in itertools.product(*[range(k) for k in np.ceil((wk + 1) / 2, dtype=np.int32)]):
                    wj = np.array(j)
                    wi = wk - wj
                    self.bc2[(*wi, *wj)] = (self.bc2[(*wi, *wj)] / self.sum_bc2[(*wk, *[])])
                self.sum_bc2[(*wk, *[])] = 1
            self.pure_power_sepctrum[(*wk, *[])] = self.power_spectrum[(*wk, *[])] * (1 - self.sum_bc2[(*wk, *[])])

    def _simulate_bsrm_uni(self, phi):
        coeff = np.sqrt((2 ** (self.n_dimensions + 1)) * self.power_spectrum * np.prod(self.frequency_interval))
        phi_e = np.exp(phi * 1.0j)
        biphase_e = np.exp(self.biphase * 1.0j)
        b = np.sqrt(1 - self.sum_bc2) * phi_e
        bc = np.sqrt(self.bc2)

        phi_e = np.einsum("i...->...i", phi_e)
        b = np.einsum("i...->...i", b)

        for i in itertools.product(*self.ranges):
            wk = np.array(i)
            for j in itertools.product(*[range(np.int32(k)) for k in np.ceil((wk + 1) / 2)]):
                wj = np.array(j)
                wi = wk - wj
                b[(*wk, *[])] = (
                        b[(*wk, *[])]
                        + bc[(*wi, *wj)]
                        * biphase_e[(*wi, *wj)]
                        * phi_e[(*wi, *[])]
                        * phi_e[(*wj, *[])])

        b = np.einsum("...i->i...", b)
        b = b * coeff
        b[np.isnan(b)] = 0
        samples = np.fft.fftn(b, self.n_time_intervals)
        samples = samples[:, np.newaxis]
        return np.real(samples)

    def run(self, n_samples: int):
        """
        Execute the random sampling in the :class:`.BispectralRepresentation` class.

        The :meth:`run` method is the function that performs random sampling in the :class:`.BispectralRepresentation`
        class. If `n_samples` is provided, the :meth:`run` method is automatically called when the
        :class:`.BispectralRepresentation` object is defined. The user may also call the :meth:`run` method directly to
        generate samples. The :meth:`run` method of the :class:`.BispectralRepresentation` class can be invoked many
        times and each time the generated samples are appended to the existing samples.

        :param n_samples: Number of samples of the stochastic process to be simulated.
         If the :meth:`run` method is invoked multiple times, the newly generated samples will be appended to the
         existing samples.

        The :meth:`run` method has no returns, although it creates and/or appends the :py:attr:`samples` attribute of
        the :class:`.BispectralRepresentation` class.
        """
        if n_samples is None:
            raise ValueError("UQpy: Stochastic Process: Number of samples must be defined.")
        if not isinstance(n_samples, int):
            raise ValueError("UQpy: Stochastic Process: n_samples should be an integer.")

        self.logger.info("UQpy: Stochastic Process: Running 3rd-order Spectral Representation Method.")

        samples = None
        phi = None

        if self.case == "uni":
            self.logger.info("UQpy: Stochastic Process: Starting simulation of uni-variate Stochastic Processes.")
            self.logger.info("UQpy: The number of dimensions is %i:", self.n_dimensions)
            phi = (np.random.uniform(
                size=np.append(self.n_samples,
                               np.ones(self.n_dimensions, dtype=np.int32)
                               * self.n_frequency_intervals, )) * 2 * np.pi)
            samples = self._simulate_bsrm_uni(phi)

        if self.samples is None:
            self.samples = samples
            self.phi = phi
        else:
            self.samples = np.concatenate((self.samples, samples), axis=0)
            self.phi = np.concatenate((self.phi, phi), axis=0)

        self.logger.info("UQpy: Stochastic Process: 3rd-order Spectral Representation Method Complete.")


# src/UQpy/stochastic_process/InverseTranslation.py

from UQpy.distributions import *
from UQpy.utilities import *
from UQpy.stochastic_process.supportive import (
    inverse_wiener_khinchin_transform,
    wiener_khinchin_transform,
)


class InverseTranslation:
    def __init__(
        self,
        distributions: Distribution,
        time_interval: Union[list, np.ndarray],
        frequency_interval: Union[list, np.ndarray],
        n_time_intervals: Union[list, np.ndarray],
        n_frequency_intervals: Union[list, np.ndarray],
        correlation_function_non_gaussian: Union[list, np.ndarray] = None,
        power_spectrum_non_gaussian: Union[list, np.ndarray] = None,
        samples_non_gaussian: Union[list, np.ndarray] = None,
        percentage_error: float = 5.0
    ):
        """
        A class to perform Iterative Translation Approximation Method to find the underlying  Gaussian Stochastic
        Processes which upon translation would yield the necessary non-Gaussian Stochastic Processes.

        :param distributions: An instance of the :py:mod:`UQpy` :class:`.Distributions` class defining the marginal
         distribution of the non-Gaussian stochastic process.
        :param time_interval: The value of time discretization.
        :param frequency_interval: The value of frequency discretization.
        :param n_time_intervals: The number of time discretizations.
        :param n_frequency_intervals: The number of frequency discretizations.
        :param correlation_function_non_gaussian: The auto correlation function of the non-Gaussian stochastic
         processes. Either the power spectrum or the auto correlation function of the Gaussian stochastic process needs
         to be defined.
        :param power_spectrum_non_gaussian: The power spectrum of the non-Gaussian stochastic processes.
        :param samples_non_gaussian: Samples of non-Gaussian stochastic processes.
         `samples_non_gaussian` is optional. If no samples are passed, the :class:`.InverseTranslation` class will
         compute the underlying Gaussian correlation using the ITAM.
        :param percentage_error:
        """
        self.samples_gaussian = None
        """The inverse translated Gaussian samples from the non-Gaussian samples."""
        self.samples_non_gaussian = None
        self.samples_shape = None
        self.distributions = distributions
        self.frequency = np.arange(0, n_frequency_intervals) * frequency_interval
        self.time = np.arange(0, n_time_intervals) * time_interval
        self.error = percentage_error
        self.logger = logging.getLogger(__name__)
        if correlation_function_non_gaussian is None and power_spectrum_non_gaussian is None:
            self.logger.info("Either the Power Spectrum or the Autocorrelation function should be specified")
        if correlation_function_non_gaussian is None:
            self.power_spectrum_non_gaussian = power_spectrum_non_gaussian
            self.correlation_function_non_gaussian = wiener_khinchin_transform(power_spectrum_non_gaussian,
                                                                               self.frequency, self.time)
        elif power_spectrum_non_gaussian is None:
            self.correlation_function_non_gaussian = correlation_function_non_gaussian
            self.power_spectrum_non_gaussian = inverse_wiener_khinchin_transform(correlation_function_non_gaussian,
                                                                                 self.frequency, self.time)
        self.num = self.correlation_function_non_gaussian.shape[0]
        self.dim = len(self.correlation_function_non_gaussian.shape)

        if samples_non_gaussian is not None:
            self.run(samples_non_gaussian)

        self.power_spectrum_gaussian: NumpyFloatArray = self._itam_power_spectrum()
        """The power spectrum of the inverse translated Gaussian stochastic processes"""
        self.auto_correlation_function_gaussian = wiener_khinchin_transform(
            self.power_spectrum_gaussian, self.frequency, self.time)
        self.correlation_function_gaussian: NumpyFloatArray = (
            self.auto_correlation_function_gaussian / self.auto_correlation_function_gaussian[0])
        """The correlation function of the inverse translated Gaussian stochastic processes."""

    def run(self, samples_non_gaussian):
        """

        :param samples_non_gaussian: Samples of non-Gaussian stochastic processes. If samples are provided at the
         initialization, then the run method is executed automatically. If no samples are passed at the
         initialization, the :class:`.InverseTranslation` class will compute the underlying Gaussian correlation using
         the ITAM and the run method needs to be manually executed by the user.
        """

        self.samples_shape = samples_non_gaussian.shape
        self.samples_non_gaussian = samples_non_gaussian.flatten()[:, np.newaxis]
        self.samples_gaussian: NumpyFloatArray = self._inverse_translate_non_gaussian_samples().reshape(
            self.samples_shape)


    def _inverse_translate_non_gaussian_samples(self):
        if not hasattr(self.distributions, "cdf"):
            raise AttributeError("UQpy: The marginal dist_object needs to have an inverse cdf defined.")
        non_gaussian_cdf = getattr(self.distributions, "cdf")
        samples_cdf = non_gaussian_cdf(self.samples_non_gaussian)
        return Normal(loc=0.0, scale=1.0).icdf(samples_cdf)

    def _itam_power_spectrum(self):
        target_S = self.power_spectrum_non_gaussian
        i_converge = 0
        max_iter = 500
        target_R = wiener_khinchin_transform(target_S, self.frequency, self.time)
        R_g_iterate = target_R
        S_g_iterate = target_S
        R_ng_iterate = np.zeros_like(R_g_iterate)
        r_ng_iterate = np.zeros_like(R_g_iterate)
        S_ng_iterate = np.zeros_like(S_g_iterate)
        non_gaussian_moments = getattr(self.distributions, 'moments')()

        for _ in range(max_iter):
            R_g_iterate = wiener_khinchin_transform(S_g_iterate, self.frequency, self.time)
            for i in range(len(target_R)):
                r_ng_iterate[i] = correlation_distortion(dist_object=self.distributions,
                                                         rho=R_g_iterate[i] / R_g_iterate[0])
            R_ng_iterate = r_ng_iterate * non_gaussian_moments[1] + non_gaussian_moments[0] ** 2
            S_ng_iterate = inverse_wiener_khinchin_transform(R_ng_iterate, self.frequency, self.time)

            err1 = np.sum((target_S - S_ng_iterate) ** 2)
            err2 = np.sum(target_S ** 2)

            if 100 * np.sqrt(err1 / err2) < self.error:
                i_converge = 1

            ratio = target_S / S_ng_iterate
            ratio = np.nan_to_num(ratio, nan=0.0, posinf=0.0, neginf=0.0)

            S_g_next_iterate = (ratio ** 1.3) * S_g_iterate

            # Eliminate Numerical error of Upgrading Scheme
            S_g_next_iterate[S_g_next_iterate < 0] = 0
            S_g_iterate = S_g_next_iterate

            if i_converge:
                break

        return S_g_iterate


# src/UQpy/stochastic_process/KarhunenLoeveExpansion.py

from scipy.linalg import sqrtm
import numpy as np
from UQpy.utilities import *


class KarhunenLoeveExpansion:
    # TODO: Test this for non-stationary processes.

    def __init__(
        self,
        n_samples: int,
        correlation_function: np.ndarray,
        time_interval: Union[np.ndarray, float],
        threshold: int = None,
        random_state: RandomStateType = None,
    ):
        """
        A class to simulate stochastic processes from a given auto-correlation function based on the Karhunen-Loeve
        Expansion

        :param n_samples: Number of samples of the stochastic process to be simulated.
         The :meth:`run` method is automatically called if `n_samples` is provided. If `n_samples` is not
         provided, then the :class:`.KarhunenLoeveExpansion` object is created but samples are not generated.
        :param correlation_function: The correlation function of the stochastic process of size
         :code:`(n_time_intervals, n_time_intervals)`
        :param time_interval: The length of time discretization.
        :param threshold: The threshold number of eigenvalues to be used in the expansion.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.

        """
        self.correlation_function = correlation_function
        self.time_interval = time_interval
        self.n_eigenvalues = threshold or len(self.correlation_function[0])
        self.random_state = random_state
        if isinstance(self.random_state, int):
            np.random.seed(self.random_state)
        elif not isinstance(self.random_state, (type(None), np.random.RandomState)):
            raise TypeError("UQpy: random_state must be None, an int or an np.random.RandomState object.")

        self.logger = logging.getLogger(__name__)
        self.n_samples = n_samples

        self.samples: NumpyFloatArray = None
        """Array of generated samples."""
        self.xi: NumpyFloatArray = None
        """The independent gaussian random variables used in the expansion."""

        if self.n_samples is not None:
            self.run(n_samples=self.n_samples)

    def _simulate(self, xi):
        lam, phi = np.linalg.eig(self.correlation_function)
        lam = lam[: self.n_eigenvalues]
        lam = np.diag(lam)
        self.phi = np.real(phi[:, : self.n_eigenvalues])
        self.lam = lam.astype(np.float64)
        samples = np.dot(self.phi, np.dot(sqrtm(self.lam), xi))
        samples = np.real(samples)
        samples = samples.T
        samples = samples[:, np.newaxis]
        return samples

    def run(self, n_samples: int):
        """
        Execute the random sampling in the :class:`.KarhunenLoeveExpansion` class.

        The :meth:`run` method is the function that performs random sampling in the :class:`.KarhunenLoeveExpansion`
        class. If `n_samples` is provided when the :class:`.KarhunenLoeveExpansion` object is defined, the
        :meth:`run` method is automatically called. The user may also call the :meth:`run` method directly to generate
        samples. The :meth:`run`` method of the :class:`.KarhunenLoeveExpansion` class can be invoked many times and
        each time the generated samples are appended to the existing samples.

        :param n_samples: Number of samples of the stochastic process to be simulated.
         If the :meth:`run` method is invoked multiple times, the newly generated samples will be appended to the
         existing samples.

        The :meth:`run` method has no returns, although it creates and/or appends the :py:attr:`samples` attribute of
        the :class:`KarhunenLoeveExpansion` class.
        """
        if n_samples is None:
            raise ValueError("UQpy: Stochastic Process: Number of samples must be defined.")
        if not isinstance(n_samples, int):
            raise ValueError("UQpy: Stochastic Process: n_samples should be an integer.")

        self.logger.info("UQpy: Stochastic Process: Running Karhunen Loeve Expansion.")

        self.logger.info("UQpy: Stochastic Process: Starting simulation of Stochastic Processes.")
        xi = np.random.normal(size=(self.n_eigenvalues, self.n_samples))
        samples = self._simulate(xi)

        if self.samples is None:
            self.samples = samples
            self.xi = xi
        else:
            self.samples = np.concatenate((self.samples, samples), axis=0)
            self.xi = np.concatenate((self.xi, xi), axis=0)

        self.logger.info("UQpy: Stochastic Process: Karhunen-Loeve Expansion Complete.")


# src/UQpy/stochastic_process/SpectralRepresentation.py

from UQpy.utilities import *
import numpy as np


class SpectralRepresentation:
    def __init__(
            self,
            n_samples: int,
            power_spectrum: Union[list, np.ndarray, float],
            time_interval: Union[list, np.ndarray, float],
            frequency_interval: Union[list, np.ndarray, float],
            n_time_intervals: Union[list, np.ndarray, float],
            n_frequency_intervals: Union[list, np.ndarray, float],
            random_state: RandomStateType = None,
    ):
        """
        A class to simulate stochastic processes from a given power spectrum density using the Spectral Representation
        Method. This class can simulate uni-variate, multi-variate, and multi-dimensional stochastic processes. The
        class uses Singular Value Decomposition, as opposed to Cholesky Decomposition, to ensure robust, near-positive
        definite multi-dimensional power spectra.

        :param n_samples: Number of samples of the stochastic process to be simulated. The :py:meth:`run` method
         is automatically called if `n_samples` is provided. If `n_samples` is not provided, then the
         :class:`.SpectralRepresentation` object is created but samples are not generated.
        :param power_spectrum: The discretized power spectrum.

         * For uni-variate, one-dimensional processes `power_spectrum` will be :class:`list` or :class:`numpy.ndarray` of length `n_frequency_intervals`.

         * For multi-variate, one-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(n_of_variables, n_variables, n_frequency_intervals)`.

         * For uni-variate, multi-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(n_frequency_intervals[0], ..., n_frequency_intervals[n_dimensions-1])`

         * For multi-variate, multi-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(n_variables, n_variables, n_frequency_intervals[0],...,n_frequency_intervals[n_dimensions-1])`.

        :param time_interval: Length of time discretizations (:math:`\Delta t`) for each
         dimension of size `n_dimensions`.
        :param frequency_interval: Length of frequency discretizations
         (:math:`\Delta \omega`) for each dimension of size `n_dimensions`.
        :param n_time_intervals: Number of time discretizations for each dimensions of
         size `n_dimensions`.
        :param n_frequency_intervals: Number of frequency discretizations for each
         dimension of size `n_dimensions`.
        :param random_state: Random seed used to initialize the pseudo-random number generator. Default is :any:`None`.
         If an :any:`int` is provided, this sets the seed for an object of :class:`numpy.random.RandomState`. Otherwise,
         the object itself can be passed directly.
        """
        self.power_spectrum = power_spectrum
        if (
                isinstance(time_interval, float)
                and isinstance(frequency_interval, float)
                and isinstance(n_time_intervals, int)
                and isinstance(n_frequency_intervals, int)
        ):
            time_interval = [time_interval]
            frequency_interval = [frequency_interval]
            n_time_intervals = [n_time_intervals]
            n_frequency_intervals = [n_frequency_intervals]
        self.time_interval = np.array(time_interval)
        self.frequency_interval = np.array(frequency_interval)
        self.number_time_intervals = np.array(n_time_intervals)
        self.number_frequency_intervals = np.array(n_frequency_intervals)
        self.n_samples = n_samples

        # Error checks
        t_u = 2 * np.pi / (2 * self.number_frequency_intervals * self.frequency_interval)

        if (self.time_interval > t_u).any():
            raise RuntimeError("UQpy: Aliasing might occur during execution")

        self.logger = logging.getLogger(__name__)

        self.random_state = random_state
        if isinstance(self.random_state, int):
            np.random.seed(self.random_state)
        elif not isinstance(self.random_state, (type(None), np.random.RandomState)):
            raise TypeError("UQpy: random_state must be None, an int or an np.random.RandomState object.")

        self.samples: NumpyFloatArray = None
        """Generated samples.
        The shape of the samples is :code:`(n_samples, n_variables, n_time_intervals[0], ...,
        n_time_intervals[n_dimensions-1])`"""
        self.n_variables: int = None
        """Number of variables in the stochastic process."""
        self.n_dimensions: int = len(self.number_frequency_intervals)
        """The dimensionality of the stochastic process."""
        self.phi: NumpyFloatArray = None
        """The random phase angles used in the simulation of the stochastic process.
        The shape of the phase angles :code:`(n_samples, n_variables, n_frequency_intervals[0], ...,
        n_frequency_intervals[n_dimensions-1])`"""

        if self.n_dimensions == len(self.power_spectrum.shape):
            self.case = "uni"
        else:
            self.n_variables = self.power_spectrum.shape[0]
            self.case = "multi"

        # Run Spectral Representation Method
        if self.n_samples is not None:
            self.run(n_samples=self.n_samples)

    def run(self, n_samples):
        """
        Execute the random sampling in the :class:`.SpectralRepresentation` class.

        The :meth:`run` method is the function that performs random sampling in the :class:`.SpectralRepresentation`
        class. If `n_samples` is provided when the :class:`.SpectralRepresentation` object is defined, the
        :meth:`run` method is automatically called. The user may also call the :meth:`run` method directly to generate
        samples. The :meth:`run` method of the :class:`.SpectralRepresentation` class can be invoked many times and each
        time the generated samples are appended to the existing samples.

        :param int n_samples: Number of samples of the stochastic process to be simulated.
         If the :meth:`run` method is invoked multiple times, the newly generated samples will be appended to the
         existing samples.

        The :meth:`run` method has no returns, although it creates and/or appends the :py:attr:`samples` attribute of
        the :class:`.SpectralRepresentation` class.
        """
        if n_samples is None:
            raise ValueError("UQpy: Stochastic Process: Number of samples must be defined.")
        if not isinstance(n_samples, int):
            raise ValueError("UQpy: Stochastic Process: nsamples should be an integer.")

        self.logger.info("UQpy: Stochastic Process: Running Spectral Representation Method.")

        samples = None
        phi = None

        if self.case == "uni":
            self.logger.info("UQpy: Stochastic Process: Starting simulation of uni-variate Stochastic Processes.")
            self.logger.info("UQpy: The number of dimensions is %i:", self.n_dimensions)
            phi = (np.random.uniform(size=np.append(self.n_samples, np.ones(self.n_dimensions, dtype=np.int32)
                                                    * self.number_frequency_intervals, )) * 2 * np.pi)
            samples = self._simulate_uni(phi)

        elif self.case == "multi":
            self.logger.info("UQpy: Stochastic Process: Starting simulation of multi-variate Stochastic Processes.")
            self.logger.info("UQpy: Stochastic Process: The number of variables is %i:", self.n_variables)
            self.logger.info("UQpy: Stochastic Process: The number of dimensions is  %i:", self.n_dimensions)
            phi = (np.random.uniform(size=
                                     np.append(self.n_samples, np.append(np.ones(self.n_dimensions,
                                                                                 dtype=np.int32)
                                                                         * self.number_frequency_intervals,
                                                                         self.n_variables, ), )) * 2 * np.pi)
            samples = self._simulate_multi(phi)

        if self.samples is None:
            self.samples = samples
            self.phi = phi
        else:
            self.samples = np.concatenate((self.samples, samples), axis=0)
            self.phi = np.concatenate((self.phi, phi), axis=0)

        self.logger.info("UQpy: Stochastic Process: Spectral Representation Method Complete.")

    def _simulate_uni(self, phi):
        fourier_coefficient = np.exp(phi * 1.0j) * np.sqrt(
            2 ** (self.n_dimensions + 1) * self.power_spectrum * np.prod(self.frequency_interval))
        samples = np.fft.fftn(fourier_coefficient, s=self.number_time_intervals)
        samples = np.real(samples)
        samples = samples[:, np.newaxis]
        return samples

    def _simulate_multi(self, phi):
        power_spectrum = np.einsum("ij...->...ij", self.power_spectrum)
        coefficient = np.sqrt(2 ** (self.n_dimensions + 1)) * np.sqrt(np.prod(self.frequency_interval))
        u, s, v = np.linalg.svd(power_spectrum)
        power_spectrum_decomposed = np.einsum("...ij,...j->...ij", u, np.sqrt(s))
        fourier_coefficient = coefficient * np.einsum(
            "...ij,n...j -> n...i", power_spectrum_decomposed, np.exp(phi * 1.0j))
        fourier_coefficient[np.isnan(fourier_coefficient)] = 0
        samples = np.real(np.fft.fftn(fourier_coefficient, s=self.number_time_intervals,
                                      axes=tuple(np.arange(1, 1 + self.n_dimensions))))
        samples = np.einsum("n...m->nm...", samples)
        return samples


# src/UQpy/stochastic_process/Translation.py

import itertools

import numpy as np
from scipy.stats import norm

from UQpy.distributions.baseclass import Distribution
from UQpy.utilities import *
from UQpy.stochastic_process.supportive import (
    inverse_wiener_khinchin_transform,
    wiener_khinchin_transform,
    scaling_correlation_function,
)


class Translation:
    def __init__(
            self,
            distributions: Distribution,
            time_interval: Union[list, np.ndarray, float],
            frequency_interval: Union[list, np.ndarray, float],
            n_time_intervals: Union[list, np.ndarray, float],
            n_frequency_intervals: Union[list, np.ndarray, float],
            power_spectrum_gaussian: np.ndarray = None,
            correlation_function_gaussian: np.ndarray = None,
            samples_gaussian: np.ndarray = None,
    ):
        """
        A class to translate Gaussian Stochastic Processes to non-Gaussian Stochastic Processes

        :param distributions: An instance of the UQpy :class:`.Distribution` class defining the marginal distribution to
         which the Gaussian stochastic process should be translated to.
        :param time_interval: The value of time discretization.
        :param frequency_interval: The value of frequency discretization.
        :param n_time_intervals: The number of time discretizations.
        :param n_frequency_intervals: The number of frequency discretizations.
        :param power_spectrum_gaussian: The power spectrum of the gaussian stochastic process to be translated.
         `power_spectrum_gaussian` must be of size :code:`(n_frequency_intervals)`.
        :param correlation_function_gaussian: The auto correlation function of the Gaussian stochastic process to be
         translated. Either the power spectrum or the auto correlation function of the gaussian stochastic process needs
         to be defined. `correlation_function_gaussian` must be of size :code:`(n_time_intervals)`.
        :param samples_gaussian: Samples of Gaussian stochastic process to be translated.
         `samples_gaussian` is optional. If no samples are passed, the :class:`.Translation` class will compute the
         correlation distortion.
        """
        self.samples_non_gaussian = None
        self.samples_gaussian = None
        """Translated non-Gaussian stochastic process from Gaussian samples."""
        self.samples_shape = None
        self.distributions = distributions
        self.time_interval = time_interval
        self.frequency_interval = frequency_interval
        self.n_time_intervals = n_time_intervals
        self.n_frequency_intervals = n_frequency_intervals
        self.correlation_function_non_gaussian: callable = None
        """The correlation function of the translated non-Gaussian stochastic processes obtained by distorting the 
        Gaussian correlation function."""
        self.scaled_correlation_function_non_gaussian: callable = None
        """This obtained by scaling the correlation function of the non-Gaussian stochastic processes to make the
        correlation at '0' lag to be 1"""
        if correlation_function_gaussian is None and power_spectrum_gaussian is None:
            print("Either the Power Spectrum or the Autocorrelation function should be specified")
        if correlation_function_gaussian is None:
            self.power_spectrum_gaussian = power_spectrum_gaussian
            self.correlation_function_gaussian = wiener_khinchin_transform(
                power_spectrum_gaussian,
                np.arange(0, self.n_frequency_intervals) * self.frequency_interval,
                np.arange(0, self.n_time_intervals) * self.time_interval,)
        elif power_spectrum_gaussian is None:
            self.correlation_function_gaussian = correlation_function_gaussian
            self.power_spectrum_gaussian = inverse_wiener_khinchin_transform(
                correlation_function_gaussian,
                np.arange(0, self.n_frequency_intervals) * self.frequency_interval,
                np.arange(0, self.n_time_intervals) * self.time_interval,)
        self.shape = self.correlation_function_gaussian.shape
        self.dim = len(self.correlation_function_gaussian.shape)
        if samples_gaussian is not None:
            self.run(samples_gaussian)

        (self.correlation_function_non_gaussian, self.scaled_correlation_function_non_gaussian,) \
            = self._autocorrelation_distortion()
        self.power_spectrum_non_gaussian: NumpyFloatArray = inverse_wiener_khinchin_transform(
            self.correlation_function_non_gaussian,
            np.arange(0, self.n_frequency_intervals) * self.frequency_interval,
            np.arange(0, self.n_time_intervals) * self.time_interval,)
        """The power spectrum of the translated non-Gaussian stochastic processes."""

    def run(self, samples_gaussian):
        """

        :param samples_gaussian: Samples of Gaussian stochastic process to be translated.
         `samples_gaussian` is optional. If samples are provided at the object initialization then the run method is
         executed automatically. If no samples are passed, the :class:`.Translation` class will compute the
         correlation distortion and the use needs to execute the run method manually in order to compute non-gaussian
         samples.
        """
        self.samples_shape = samples_gaussian.shape
        self.samples_gaussian: NumpyFloatArray = samples_gaussian.flatten()[:, np.newaxis]
        self.samples_non_gaussian = self._translate_gaussian_samples().reshape(self.samples_shape)

    def _translate_gaussian_samples(self):
        standard_deviation = np.sqrt(self.correlation_function_gaussian[0])
        samples_cdf = norm.cdf(self.samples_gaussian, scale=standard_deviation)
        if not hasattr(self.distributions, "icdf"):
            raise AttributeError("UQpy: The marginal dist_object needs to have an inverse cdf defined.")
        non_gaussian_icdf = getattr(self.distributions, "icdf")
        return non_gaussian_icdf(samples_cdf)

    def _autocorrelation_distortion(self):
        correlation_function_gaussian = scaling_correlation_function(self.correlation_function_gaussian)
        correlation_function_gaussian = np.clip(correlation_function_gaussian, -0.999, 0.999)
        correlation_function_non_gaussian = np.zeros_like(correlation_function_gaussian)
        for i in itertools.product(*[range(s) for s in self.shape]):
            correlation_function_non_gaussian[i] = correlation_distortion(
                self.distributions, correlation_function_gaussian[i])
        if hasattr(self.distributions, "moments"):
            non_gaussian_moments = getattr(self.distributions, "moments")()
        else:
            raise AttributeError("UQpy: The marginal dist_object needs to have defined moments.")
        scaled_correlation_function_non_gaussian = correlation_function_non_gaussian * non_gaussian_moments[1] + \
                                                   non_gaussian_moments[0] ** 2
        return correlation_function_non_gaussian, scaled_correlation_function_non_gaussian


# src/UQpy/stochastic_process/__init__.py

from UQpy.stochastic_process.BispectralRepresentation import BispectralRepresentation
from UQpy.stochastic_process.InverseTranslation import InverseTranslation
from UQpy.stochastic_process.KarhunenLoeveExpansion import KarhunenLoeveExpansion
from UQpy.stochastic_process.SpectralRepresentation import SpectralRepresentation
from UQpy.stochastic_process.Translation import Translation

from .supportive import *

from . import supportive


# src/UQpy/surrogates/__init__.py

from UQpy.surrogates.polynomial_chaos import *
from UQpy.surrogates.stochastic_reduced_order_models import *
from UQpy.surrogates.kriging import *
from UQpy.surrogates.gaussian_process import *
from UQpy.surrogates.baseclass import *

from . import polynomial_chaos, stochastic_reduced_order_models, gaussian_process, baseclass


# src/UQpy/transformations/Correlate.py

import numpy as np
from beartype import beartype
from scipy.linalg import cholesky

from UQpy.utilities.ValidationTypes import NumpyFloatArray


class Correlate:

    @beartype
    def __init__(self, samples_u: np.ndarray, corr_z: np.ndarray):
        """
        A class to induce correlation to standard normal random variables.

        :param samples_u: Uncorrelated  standard normal vector of shape ``(nsamples, n_dimensions)``.
        :param corr_z: The correlation  matrix (:math:`\mathbf{C_Z}`) of the standard normal random vector **Z** .
        """
        self.samples_y = samples_u
        self.corr_z = corr_z
        self.H: NumpyFloatArray = cholesky(self.corr_z, lower=True)
        """The lower diagonal matrix resulting from the Cholesky decomposition of the correlation  matrix
        (:math:`\mathbf{C_Z}`)."""
        self.samples_z: NumpyFloatArray = (self.H @ samples_u.T).T
        """Correlated standard normal vector of shape ``(nsamples, n_dimensions)``."""


# src/UQpy/transformations/Decorrelate.py

import numpy as np
from beartype import beartype
from scipy.linalg import cholesky

from UQpy.utilities.ValidationTypes import NumpyFloatArray


class Decorrelate:
    @beartype
    def __init__(self, samples_z: np.ndarray, corr_z: np.ndarray):
        """
        A class to remove correlation from correlated standard normal random variables.

        :param samples_z: Correlated standard normal vector of shape ``(n_samples, n_dimensions)``.
        :param corr_z: The correlation  matrix (:math:`\mathbf{C_Z}`) of the standard normal random vector **Z** .
        """
        self.samples_z = samples_z
        self.corr_z = corr_z
        self.H: NumpyFloatArray = cholesky(self.corr_z, lower=True)
        """The lower diagonal matrix resulting from the Cholesky decomposition of the correlation  matrix
        (:math:`\mathbf{C_Z}`)."""
        self.samples_u: NumpyFloatArray = np.linalg.solve(self.H, samples_z.T.squeeze()).T
        """Uncorrelated standard normal vector of shape ``(n_samples, n_dimensions)``."""


# src/UQpy/transformations/Nataf.py

import logging
from typing import Union, Annotated
from beartype import beartype
from beartype.vale import Is

from UQpy.distributions import *
import numpy as np
import scipy.stats as stats
from UQpy.utilities.Utilities import nearest_psd, calculate_gauss_quadrature_2d
from UQpy.utilities.Utilities import bi_variate_normal_pdf
from scipy.linalg import cholesky
from UQpy.utilities.ValidationTypes import PositiveInteger, NumpyFloatArray

DistributionList = Annotated[
    list[Distribution],
    Is[lambda lst: all(isinstance(item, Distribution) for item in lst)],
]


class Nataf:

    @beartype
    def __init__(
        self,
        distributions: Union[Distribution, DistributionList],
        samples_x: Union[None, np.ndarray] = None,
        samples_z: Union[None, np.ndarray] = None,
        jacobian: bool = False,
        corr_z: Union[None, np.ndarray] = None,
        corr_x: Union[None, np.ndarray] = None,
        itam_beta: Union[float, int] = 1.0,
        itam_threshold1: Union[float, int] = 0.001,
        itam_threshold2: Union[float, int] = 0.1,
        itam_max_iter: int = 100,
    ):
        """
        Transform random variables using the Nataf or Inverse Nataf transformation

        :param distributions: Probability distribution of each random variable. Must be an object of type
         :class:`.DistributionContinuous1D` or :class:`.JointIndependent`.
        :param samples_x: Random vector of shape ``(nsamples, n_dimensions)`` with prescribed probability
         distributions.
         If `samples_x` is provided, the :class:`.Nataf` class transforms them to `samples_z`.
        :param samples_z: Standard normal random vector of shape ``(nsamples_, n_dimensions)``
         If `samples_z` is provided, the :class:`.Nataf` class transforms them to `samples_x`.
        :param jacobian: A boolean whether to return the jacobian of the transformation. Default: :any:`False`
        :param corr_z: The correlation  matrix (:math:`\mathbf{C_Z}`) of the standard normal random vector **Z** .
         Default: The identity matrix.
         If ``corr_z`` is specified, the :class:`.Nataf` class computes the correlation distortion integral above to
         solve for ``corr_x``.
        :param corr_x: The correlation  matrix (:math:`\mathbf{C_X}`) of the random vector **X** .
         Default: The identity matrix.
         If ``corr_x`` is specified, the :class:`.Nataf` class invokes the ITAM to compute ``corr_z``.
        :param itam_beta: A parameter selected to optimize convergence speed and desired accuracy of the ITAM method.
         Default: :math:`1.0`
        :param itam_threshold1: If ``corr_x`` is specified, this specifies the threshold value for the error in the
         `ITAM` method (see :py:mod:`.utilities` module) given by
         :math:`\epsilon_1 = ||\mathbf{C_X}^{target} - \mathbf{C_X}^{computed}||`
         Default: :math:`0.001`
        :param itam_threshold2: If ``corr_x`` is specified, this specifies the threshold value for the error difference
         between iterations in the `ITAM` method (see :py:mod:`.utilities` module) given by
         :math:`\epsilon_1^{i} - \epsilon_1^{i-1}`
         for iteration :math:`i`.
         Default: :math:`0.01`
        :param itam_max_iter: Maximum number of iterations for the ITAM method. Default: :math:`100`
        """
        self.n_dimensions = 0
        if isinstance(distributions, list):
            for i in range(len(distributions)):
                self.update_dimensions(distributions[i])
        else:
            self.update_dimensions(distributions)
        self.dist_object = distributions
        self.samples_x: NumpyFloatArray = samples_x
        """Random vector of shape ``(nsamples, n_dimensions)`` with prescribed probability distributions."""
        self.samples_z:NumpyFloatArray = samples_z
        """Standard normal random vector of shape ``(nsamples, n_dimensions)``"""
        self.jacobian = jacobian
        self.jzx: NumpyFloatArray = None
        """The Jacobian of the transformation of shape ``(n_dimensions, n_dimensions)``."""
        self.jxz: NumpyFloatArray = None
        """The Jacobian of the transformation of shape ``(n_dimensions, n_dimensions)``."""
        self.itam_max_iter = itam_max_iter
        self.itam_beta = float(itam_beta)
        self.itam_threshold1 = float(itam_threshold1)
        self.itam_threshold2 = float(itam_threshold2)
        self.logger = logging.getLogger(__name__)

        if corr_x is None and corr_z is None:
            self.corr_x: NumpyFloatArray = np.eye(self.n_dimensions)
            """Distorted correlation matrix (:math:`\mathbf{C_X}`) of the random vector **X**."""
            self.corr_z: NumpyFloatArray = np.eye(self.n_dimensions)
            """Distorted correlation matrix (:math:`\mathbf{C_Z}`) of the standard normal vector **Z**."""
        elif corr_x is not None:
            self.corr_x = corr_x
            if np.all(np.equal(self.corr_x, np.eye(self.n_dimensions))):
                self.corr_z = self.corr_x
            elif all(isinstance(x, Normal) for x in distributions):
                self.corr_z = self.corr_x
            else:
                self.corr_z, self.itam_error1, self.itam_error2 =\
                    self.itam(self.dist_object, self.corr_x, self.itam_max_iter, self.itam_beta,
                              self.itam_threshold1, self.itam_threshold2,)
        elif corr_z is not None:
            self.corr_z = corr_z
            if np.all(np.equal(self.corr_z, np.eye(self.n_dimensions))):
                self.corr_x = self.corr_z
            elif all(isinstance(x, Normal) for x in distributions):
                self.corr_x = self.corr_z
            else:
                self.corr_x = self.distortion_z2x(self.dist_object, self.corr_z)

        self.H: NumpyFloatArray = cholesky(self.corr_z, lower=True)
        """The lower triangular matrix resulting from the Cholesky decomposition of the correlation matrix
        :math:`\mathbf{C_Z}`."""

        if self.samples_x is not None or self.samples_z is not None:
            self.run(self.samples_x, self.samples_z, self.jacobian)

    @beartype
    def run(
        self,
        samples_x: Union[None, np.ndarray] = None,
        samples_z: Union[None, np.ndarray] = None,
        jacobian: bool = False,
    ):
        """
        Execute the Nataf transformation or its inverse.
        If `samples_x` is provided, the :meth:`run` method performs the Nataf transformation. If `samples_z` is
        provided, the :meth:`run` method performs the inverse Nataf transformation.

        :param samples_x: Random vector **X**  with prescribed probability distributions or standard normal random
         vector **Z** of shape :code:`(nsamples, n_dimensions)`.
        :param samples_z: Standard normal random vector of shape ``(nsamples, n_dimensions)``
         If `samples_z` is provided, the :class:`.Nataf` class transforms them to `samples_x`.
        :param jacobian: The jacobian of the transformation of shape ``(n_dimensions, n_dimensions)``. Default:
         :any:`False`
        """
        self.jacobian = jacobian

        if samples_x is not None:
            if len(samples_x.shape) != 2:
                self.samples_x = np.atleast_2d(samples_x).T
            else:
                self.samples_x = samples_x
            if not self.jacobian:
                self.samples_z = self._transform_x2z(self.samples_x)
            elif self.jacobian:
                self.samples_z, self.jxz = self._transform_x2z(self.samples_x, jacobian=self.jacobian)

        if samples_z is not None:
            if len(samples_z.shape) != 2:
                self.samples_z = np.atleast_2d(samples_z).T
            else:
                self.samples_z = samples_z
            if not self.jacobian:
                self.samples_x = self._transform_z2x(self.samples_z)
            elif self.jacobian:
                self.samples_x, self.jzx = self._transform_z2x(self.samples_z, jacobian=self.jacobian)

    @staticmethod
    def itam(
        distributions: Union[
            DistributionContinuous1D,
            JointIndependent,
            list[Union[DistributionContinuous1D, JointIndependent]]],
        corr_x,
        itam_max_iter: int = 100,
        itam_beta: Union[float, int] = 1.0,
        itam_threshold1: Union[float, int] = 0.001,
        itam_threshold2: Union[float, int] = 0.01,
    ):
        """
        Calculate the correlation matrix :math:`\mathbf{C_Z}` of the standard normal random vector
        :math:`\mathbf{Z}` given the correlation matrix :math:`\mathbf{C_X}` of the random vector :math:`\mathbf{X}`
        using the `ITAM` method :cite:`StochasticProcess13`.

        :param distributions:  Probability distribution of each random variable. Must be an object of type
         :class:`.DistributionContinuous1D` or :class:`.JointIndependent`.
         `distributions` must have a ``cdf`` method.
        :param corr_x: The correlation  matrix (:math:`\mathbf{C_X}`) of the random vector **X** .
         Default: The ``identity`` matrix.
        :param itam_max_iter: Maximum number of iterations for the ITAM method. Default: :math:`100`
        :param itam_beta: A parameters selected to optimize convergence speed and desired accuracy of the ITAM method
         (see :cite:`Nataf2`). Default: :math:`1.0`
        :param itam_threshold1: A threshold value for the relative difference between the non-Gaussian correlation
         function and the underlying Gaussian. Default: :math:`0.001`
        :param itam_threshold2: If ``corr_x`` is specified, this specifies the threshold value for the error difference
         between iterations in the `ITAM` method (see :py:mod:`.utilities` module) given by:
         :math:`\epsilon_1^{i} - \epsilon_1^{i-1}`
         for iteration :math:`i`. Default: :math:`0.01`
        :return: Distorted correlation matrix (:math:`\mathbf{C_Z}`) of the standard normal vector **Z**,
         List of ITAM errors for each iteration, List of ITAM difference errors for each iteration
        """
        # Initial Guess
        corr_z0 = corr_x
        corr_z = np.zeros_like(corr_z0)
        # Iteration Condition
        itam_error1 = []
        itam_error2 = []
        itam_error1.append(100.0)
        itam_error2.append(abs(itam_error1[0] - 0.1) / 0.1)

        logger = logging.getLogger(__name__)

        logger.info("UQpy: Initializing Iterative Translation Approximation Method (ITAM)")

        for k in range(itam_max_iter):
            error0 = itam_error1[k]

            corr0 = Nataf.distortion_z2x(distributions, corr_z0)

            max_ratio = np.amax(np.ones((len(corr_x), len(corr_x))) / abs(corr_z0))

            corr_z = np.nan_to_num((corr_x / corr0) ** itam_beta * corr_z0)

            # Do not allow off-diagonal correlations to equal or exceed one
            corr_z[corr_z < -1.0] = (max_ratio + 1) / 2 * corr_z0[corr_z < -1.0]
            corr_z[corr_z > 1.0] = (max_ratio + 1) / 2 * corr_z0[corr_z > 1.0]

            corr_z = np.array(nearest_psd(corr_z))

            corr_z0 = corr_z.copy()

            itam_error1.append(np.linalg.norm(corr_x - corr0))
            itam_error2.append(abs(itam_error1[-1] - error0) / error0)

            logger.info("UQpy: ITAM iteration number %i", k)
            logger.info("UQpy: Current error, %e", (itam_error1[-1], itam_error2[-1]))

            if itam_error1[k] <= itam_threshold1 and itam_error2[k] <= itam_threshold2:
                break

        logger.info("UQpy: ITAM Done.")

        return corr_z, itam_error1, itam_error2

    @staticmethod
    def distortion_z2x(distributions:  Union[Distribution, list[Distribution]], corr_z: np.ndarray):
        """
        This is a method to calculate the correlation matrix :math:`\mathbf{C_x}` of the random vector
        :math:`\mathbf{x}`  given the correlation matrix :math:`\mathbf{C_z}` of the standard normal random vector
        :math:`\mathbf{z}`.

        :param distributions: This is a method to calculate the correlation matrix :math:`\mathbf{C_x}` of the random vector
         :math:`\mathbf{x}`  given the correlation matrix :math:`\mathbf{C_z}` of the standard normal random vector
         :math:`\mathbf{z}`.
         This method is part of the :class:`.Nataf` class.
        :param corr_z: The correlation  matrix (:math:`\mathbf{C_z}`) of the standard normal vector **Z** .
         Default: The ``identity`` matrix.
        :return: Distorted correlation matrix (:math:`\mathbf{C_X}`) of the random vector **X**.
        """
        logger = logging.getLogger(__name__)
        z_max = 8
        z_min = -z_max
        ng = 128

        eta, w2d, xi = calculate_gauss_quadrature_2d(ng, z_max, z_min)

        corr_x = np.ones_like(corr_z)
        logger.info("UQpy: Computing Nataf correlation distortion...")

        is_joint = isinstance(distributions, JointIndependent)
        marginals = distributions.marginals if is_joint else distributions
        corr_x = Nataf.calculate_corr_x(corr_x, corr_z, marginals, eta, w2d, xi, is_joint)
        return corr_x

    @staticmethod
    def calculate_corr_x(corr_x, corr_z, marginals, eta, w2d, xi, is_joint):
        if all(hasattr(m, "moments") for m in marginals) and all(
            hasattr(m, "icdf") for m in marginals
        ):
            for i in range(len(marginals)):
                i_cdf_i = marginals[i].icdf
                mi = marginals[i].moments()
                if not (np.isfinite(mi[0]) and np.isfinite(mi[1])):
                    raise RuntimeError("UQpy: The marginal distributions need to have finite mean and variance.")
                for j in range(i + 1, len(marginals)):
                    i_cdf_j = marginals[j].icdf
                    mj = marginals[j].moments()
                    if not (np.isfinite(mj[0]) and np.isfinite(mj[1])):
                        raise RuntimeError("UQpy: The marginal distributions need to have finite mean and variance.")
                    term1 = mj[0] ** 2 if is_joint else mj[0]
                    term2 = mi[0] ** 2 if is_joint else mi[0]
                    tmp_f_xi = i_cdf_j(np.atleast_2d(stats.norm.cdf(xi)).T) - term1
                    tmp_f_eta = i_cdf_i(np.atleast_2d(stats.norm.cdf(eta)).T) - term2

                    phi2 = bi_variate_normal_pdf(xi, eta, corr_z[i, j])

                    corr_x[i, j] = (1 / (np.sqrt(mj[1]) * np.sqrt(mi[1])) * np.sum(tmp_f_xi * tmp_f_eta * w2d * phi2))
                    corr_x[j, i] = corr_x[i, j]
        return corr_x

    def _transform_x2z(self, samples_x, jacobian=False):
        """
        This is a method to transform a vector :math:`\mathbf{x}` of  samples with marginal distributions
        :math:`f_i(x_i)` and cumulative distributions :math:`F_i(x_i)` to a vector :math:`\mathbf{z}` of standard normal
        samples  according to: :math:`Z_{i}=\Phi^{-1}(F_i(X_{i}))`, where :math:`\Phi` is the cumulative
        distribution function of a standard  normal variable.

        :param samples_x: Random vector of shape ``(nsamples, n_dimensions)`` with prescribed probability distributions.
        :param jacobian: A boolean whether to return the jacobian of the transformation.
         Default: False
        :return: Standard normal random vector of shape ``(nsamples, n_dimensions)``, the jacobian of the transformation
         of shape ``(n_dimensions, n_dimensions)``.
        """

        m, n = np.shape(samples_x)
        samples_z = None

        if isinstance(self.dist_object, JointIndependent):
            if all(hasattr(m, "cdf") for m in self.dist_object.marginals):
                samples_z = np.zeros_like(samples_x)
                for j in range(len(self.dist_object.marginals)):
                    samples_z[:, j] = stats.norm.ppf(self.dist_object.marginals[j].cdf(samples_x[:, j]))
        elif isinstance(self.dist_object, DistributionContinuous1D):
            samples_z = stats.norm.ppf(self.dist_object.cdf(samples_x))
        else:
            samples_z = np.zeros_like(samples_x)
            for j in range(n):
                samples_z[:, j] = stats.norm.ppf(
                    self.dist_object[j].cdf(samples_x[:, j])
                )

        if not jacobian:
            return samples_z
        jac = np.zeros(shape=(n, n))
        jxz = [None] * m
        for i in range(m):
            for j in range(n):
                xi = np.array([samples_x[i, j]])
                zi = np.array([samples_z[i, j]])
                jac[j, j] = stats.norm.pdf(zi) / self.dist_object[j].pdf(xi)
            jxz[i] = np.linalg.solve(jac, self.H)

        return samples_z, jxz

    def _transform_z2x(self, samples_z, jacobian=False):
        """
        This is a method to transform a standard normal vector :math:`\mathbf{z}` to a vector
        :math:`\mathbf{x}` of samples with marginal distributions :math:`f_i(x_i)` and cumulative distributions
        :math:`F_i(x_i)` to samples  according to: :math:`Z_{i}=\Phi^{-1}(F_i(X_{i}))`, where :math:`\Phi` is the
        cumulative distribution function of a standard  normal variable.

        :param samples_z: Standard normal random vector of shape ``(nsamples, n_dimensions)``
        :param jacobian: A boolean whether to return the jacobian of the transformation. Default: False
        :return: Random vector of shape ``(nsamples, n_dimensions)`` with prescribed probability distributions,
         The jacobian of the transformation of shape ``(n_dimensions, n_dimensions)``.
        """
        m, n = np.shape(samples_z)
        h = cholesky(self.corr_z, lower=True)
        # samples_z = (h @ samples_y.T).T
        samples_x = np.zeros_like(samples_z)
        if isinstance(self.dist_object, JointIndependent):
            if all(hasattr(m, "icdf") for m in self.dist_object.marginals):
                for j in range(len(self.dist_object.marginals)):
                    samples_x[:, j] = self.dist_object.marginals[j].icdf(stats.norm.cdf(samples_z[:, j]))

        elif isinstance(self.dist_object, DistributionContinuous1D):
            samples_x = self.dist_object.icdf(stats.norm.cdf(samples_z))
        elif isinstance(self.dist_object, list):
            for j in range(samples_x.shape[1]):
                samples_x[:, j] = self.dist_object[j].icdf(stats.norm.cdf(samples_z[:, j]))

        if not jacobian:
            return samples_x
        jac = np.zeros(shape=(n, n))
        jzx = [None] * m
        for i in range(m):
            for j in range(n):
                xi = np.array([samples_x[i, j]])
                zi = np.array([samples_z[i, j]])
                jac[j, j] = self.dist_object[j].pdf(xi) / stats.norm.pdf(zi)
            jzx[i] = np.linalg.solve(h, jac)

        return samples_x, jzx

    @beartype
    def rvs(self, nsamples: PositiveInteger):
        """
        Generate realizations from the joint pdf of the random vector **X**.

        :param nsamples: Number of samples to generate.
        :return: Random vector in the parameter space of shape ``(nsamples, n_dimensions)``.
        """
        h = cholesky(self.corr_z, lower=True)
        n = int(nsamples)
        m = np.size(self.dist_object)
        y = np.random.randn(nsamples, m)
        z = np.dot(h, y.T).T
        samples_x = np.zeros([n, m])
        for i in range(m):
            samples_x[:, i] = self.dist_object[i].icdf(stats.norm.cdf(z[:, i]))
        return samples_x

    def update_dimensions(self, dist_object):
        if isinstance(dist_object, DistributionContinuous1D):
            self.n_dimensions += 1
        elif isinstance(dist_object, JointIndependent):
            self.n_dimensions += len(dist_object.marginals)
        else:
            raise TypeError("UQpy: A  ``DistributionContinuous1D``  or ``JointIndependent`` object must be provided.")


# src/UQpy/transformations/__init__.py

from .Nataf import Nataf
from .Correlate import Correlate
from .Decorrelate import Decorrelate


# src/UQpy/utilities/Constants.py

SKLEARN_STRING = "<class 'sklearn.gaussian_process._gpr.GaussianProcessRegressor'>"


# src/UQpy/utilities/DistanceMetric.py

from enum import Enum


class DistanceMetric(Enum):
    BRAYCURTIS = 1
    CANBERRA = 2
    CHEBYSHEV = 3
    CITYBLOCK = 4
    CORRELATION = 5
    COSINE = 6
    DICE = 7
    EUCLIDEAN = 8
    HAMMING = 9
    JACCARD = 10
    KULSINKSI = 11
    MAHALANOBIS = 12
    MATCHING = 13
    MINKOWSKI = 14
    ROGERSTANIMOTO = 15
    RUSSELLRAO = 16
    SEUCLIDEAN = 17
    SOKALMICHENER = 18
    SOKALSNEATH = 19
    SQEUCLIDEAN = 20


# src/UQpy/utilities/FminCobyla.py

from scipy.optimize import fmin_cobyla
import logging


class FminCobyla:

    def __init__(self):
        # super().__init__(None)
        self.logger = logging.getLogger(__name__)
        self.optimization = fmin_cobyla
        self.method = 'cobyla'

        self.constraints = {}
        self.arguments = {}

    def optimize(self, function, initial_guess, args=(), jac=False):
        if self.constraints:
            return fmin_cobyla(function, initial_guess, cons=self.constraints,
                               args=args, consargs=self.arguments, rhobeg=1.0, rhoend=0.0001, maxfun=1000,
                               disp=None, catol=0.0002)
        else:
            return fmin_cobyla(function, initial_guess, args=args,
                               rhobeg=1.0, rhoend=0.0001, maxfun=1000,
                               disp=None, catol=0.0002)

    def apply_constraints(self, constraints):
        self.constraints = constraints

    def apply_constraints_argument(self, arguments):
        self.arguments = arguments

    def supports_jacobian(self):
        return self.method.lower() in ['cg', 'bfgs', 'newton-cg', 'l-bfgs-b', 'tnc', 'slsqp', 'dogleg', 'trust-ncg',
                                       'trust-krylov', 'trust-exact', 'trust-constr']

    def update_bounds(self, bounds):
        pass


# src/UQpy/utilities/GrassmannPoint.py

from typing import Annotated

from beartype import beartype
from beartype.vale import Is

from UQpy.utilities.ValidationTypes import Numpy2DFloatArrayOrthonormal, Numpy2DFloatArray


class GrassmannPoint:
    @beartype
    def __init__(self, data: Numpy2DFloatArrayOrthonormal):
        """
        :param data: Matrix representing the point on the Grassmann manifold.
        """
        self._data = data

    @property
    def data(self) -> Numpy2DFloatArray:
        """
        The matrix containing the Grassmann point
        """
        return self._data


# src/UQpy/utilities/MinimizeOptimizer.py

from scipy.optimize import minimize
import logging


class MinimizeOptimizer:

    def __init__(self, method: str = 'bfgs', bounds=None):
        # super().__init__(bounds)
        self._bounds = None
        self.logger = logging.getLogger(__name__)
        self.optimization = minimize
        self.method = method
        self.save_bounds(bounds)
        self.constraints = {}

    def save_bounds(self, bounds):
        if self.method.lower() in ['nelder-mead', 'l-bfgs-b', 'tnc', 'slsqp', 'powell', 'trust-constr']:
            self._bounds = bounds
        else:
            self.logger.warning("The selected optimizer method does not support bounds and thus will be ignored.")

    def optimize(self, function, initial_guess, args=(), jac=False):
        if self.constraints:
            return minimize(function, initial_guess, args=args,
                            method=self.method, bounds=self._bounds,
                            constraints=self.constraints, jac=jac,
                            options={'disp': True, 'maxiter': 10000, 'catol': 0.002})
        else:
            return minimize(function, initial_guess, args=args,
                            method=self.method, bounds=self._bounds, jac=jac,
                            options={'disp': True, 'maxiter': 10000, 'catol': 0.002})

    def apply_constraints(self, constraints):
        if self.method.lower() in ['cobyla', 'slsqp', 'trust-constr']:
            self.constraints = constraints
        else:
            self.logger.warning("The selected optimizer method does not support constraints and thus will be ignored.")

    def update_bounds(self, bounds):
        self.save_bounds(bounds)

    def supports_jacobian(self):
        return self.method.lower() in ['cg', 'bfgs', 'newton-cg', 'l-bfgs-b', 'tnc', 'slsqp', 'dogleg', 'trust-ncg',
                                       'trust-krylov', 'trust-exact', 'trust-constr']


# src/UQpy/utilities/NoPublicConstructor.py

"Code retrieved from: https://stackoverflow.com/a/64682734/5647511"
from typing import Type, Any, TypeVar


T = TypeVar("T")


class NoPublicConstructor(type):
    """Metaclass that ensures a private constructor

    If a class uses this metaclass like this:

        class SomeClass(metaclass=NoPublicConstructor):
            pass

    If you try to instantiate your class (`SomeClass()`),
    a `TypeError` will be thrown.
    """

    def __call__(cls, *args, **kwargs):
        raise TypeError(
            f"{cls.__module__}.{cls.__qualname__} has no public constructor. "
            f"Use one of the create methods instead."
        )

    def _create(cls: Type[T], *args: Any, **kwargs: Any) -> T:
        return super().__call__(*args, **kwargs)  # type: ignore


# src/UQpy/utilities/UQpyLoggingFormatter.py

import logging

"""Based on https://stackoverflow.com/questions/14844970/
modifying-logging-message-format-based-on-message-logging-level-in-python3"""


class UQpyLoggingFormatter(logging.Formatter):
    def __init__(self):
        super().__init__(fmt="%(levelno)d: %(msg)s", datefmt=None, style="%")
        self.format_dictionary = {
            logging.NOTSET: "%(message)s",
            logging.DEBUG: "%(message)s",
            logging.INFO: "[%(levelname)s] - %(asctime)s - %(message)s",
            logging.WARNING: "[%(levelname)s] - %(asctime)s - File: %(filename)s - %(message)s",
            logging.WARN: "[%(levelname)s] - %(asctime)s - File: %(filename)s - %(message)s",
            logging.ERROR: "[%(levelname)s] - %(asctime)s - File: %(filename)s - Method: %(funcName)s -"
            " Line: %(lineno)s - %(message)s",
            logging.CRITICAL: "[%(levelname)s] - %(asctime)s - File: %(filename)s - Method: %(funcName)s -"
            " Line: %(lineno)s - %(message)s",
        }

    def format(self, record):

        # Save the original format configured by the user
        # when the logger formatter was instantiated
        format_orig = self._style._fmt

        # Replace the original format with one customized by logging level
        self._style._fmt = self.format_dictionary[record.levelno]

        # Call the original formatter class to do the grunt work
        result = logging.Formatter.format(self, record)

        # Restore the original format configured by the user
        self._style._fmt = format_orig

        return result


# src/UQpy/utilities/Utilities.py

import numpy as np
import scipy.stats as stats
from UQpy.utilities.ValidationTypes import RandomStateType

from UQpy.run_model.RunModel import RunModel


def svd(matrix, rank=None, tol=None):
    """
    Compute the singular value decomposition (SVD) of a matrix.

    **Inputs:**

    * **matrix** (`ndarray`):
        Matrix of ``shape=(m, n)`` to perform the factorization using thin SVD

    * **tol** (`float`):
        Tolerance to estimate the rank of the matrix.

        Default: Machine precision

    * **iterations** (`rank`):
        Number of eigenvalues to keep.

        Default: None

    **Output/Returns:**

    * **u** (`ndarray`):
        Matrix of left eigenvectors of ``shape=(m, rank)``.

    * **v** (`ndarray`):
        Matrix of right eigenvectors of ``shape=(rank, n)``.

    * **s** (`ndarray`):
        Matrix of eigenvalues ``shape=(rank, rank)``.

    """
    ui, si, vi = np.linalg.svd(matrix, full_matrices=True, hermitian=False)
    si = np.diag(si)
    vi = vi.T
    if rank is None:
        if tol is not None:
            rank = np.linalg.matrix_rank(si, tol=tol)
        else:
            rank = np.linalg.matrix_rank(si)
        u = ui[:, :rank]
        s = si[:rank, :rank]
        v = vi[:, :rank]
    else:
        u = ui[:, :rank]
        s = si[:rank, :rank]
        v = vi[:, :rank]

    return u, s, v


def nearest_psd(input_matrix, iterations=10):
    """
    A function to compute the nearest positive semi-definite matrix of a given matrix :cite:`Utilities3`.

    :param numpy.ndarray input_matrix: Matrix to find the nearest PD.
    :param iterations: Number of iterations to perform. Default: 10
    :return: Nearest PSD matrix to input_matrix.
    """
    n = input_matrix.shape[0]
    w = np.identity(n)
    # w is the matrix used for the norm (assumed to be Identity matrix here)
    # the algorithm should work for any diagonal W
    delta_s = 0
    psd_matrix = input_matrix.copy()
    for k in range(iterations):

        r_k = psd_matrix - delta_s
        x_k = _get_ps(r_k, w=w)
        delta_s = x_k - r_k
        psd_matrix = _get_pu(x_k, w=w)

    return psd_matrix


def nearest_pd(input_matrix):
    """
    This is a method to find the nearest positive-definite matrix to input (:cite:`Utilities1`, :cite:`nearest_psd`).

    :param numpy.ndarray input_matrix: Matrix to find the nearest PD.
    :return: Nearest PD matrix to input_matrix.
    """
    b = (input_matrix + input_matrix.T) / 2
    _, s, v = np.linalg.svd(b)

    h = np.dot(v.T, np.dot(np.diag(s), v))

    a2 = (b + h) / 2

    pd_matrix = (a2 + a2.T) / 2

    if _is_pd(pd_matrix):
        return pd_matrix

    spacing = np.spacing(np.linalg.norm(pd_matrix))
    k = 1
    while not _is_pd(pd_matrix):
        min_eig = np.min(np.real(np.linalg.eigvals(pd_matrix)))
        pd_matrix += np.eye(input_matrix.shape[0]) * (-min_eig * k ** 2 + spacing)
        k += 1

    return pd_matrix


def _is_pd(input_matrix):
    try:
        _ = np.linalg.cholesky(input_matrix)
        return True
    except np.linalg.LinAlgError:
        return False


def run_parallel_python(model_script, model_object_name, sample, dict_kwargs=None):
    """
    Method needed by :class:`.RunModel` to execute a python model in parallel
    """
    _ = sample
    exec("from " + model_script[:-3] + " import " + model_object_name)

    if dict_kwargs is None:
        par_res = eval(model_object_name + "(sample)")
    else:
        par_res = eval(model_object_name + "(sample, **dict_kwargs)")

    return par_res


def gradient(runmodel_object=None, point=None, order="first", df_step=None):
    """
    This method estimates the gradients (1st, 2nd, mixed) of a function using a finite difference scheme in the
    standard normal space. First order gradients are calculated using central finite differences.

    :param runmodel_object: The numerical model. It should be of type :class:`.RunModel` or a
     `callable`.
    :param point: The point to evaluate the gradient with shape ``point``.shape=(1, dimension)
    :param order: Order of the gradient. Available options: 'first', 'second', 'mixed'. Default: 'First'.
    :param df_step: Finite difference step. Default: 0.001.
    :return: Vector of first-order gradients (if order = 'first'). Vector of second-order gradients
     (if order = 'second'). Vector of mixed gradients (if order = 'mixed').
    """
    point = np.atleast_2d(point)

    dimension = point.shape[1]

    if df_step is None:
        df_step = [0.001] * dimension
    elif isinstance(df_step, float):
        df_step = [df_step] * dimension
    elif isinstance(df_step, list):
        if len(df_step) == 1:
            df_step = [df_step[0]] * dimension

    if not callable(runmodel_object) and not isinstance(runmodel_object, RunModel):
        raise RuntimeError("A RunModel object or callable function must be provided as model.")

    def func(m):
        def func_eval(x):
            if isinstance(m, RunModel):
                m.run(samples=x, append_samples=False)
                return np.array(m.qoi_list).flatten()
            else:
                return m(x).flatten()

        return func_eval

    f_eval = func(m=runmodel_object)

    if order.lower() == "first":
        du_dj = np.zeros([point.shape[0], dimension])

        for ii in range(dimension):
            eps_i = df_step[ii]
            u_i1_j = point.copy()
            u_i1_j[:, ii] = u_i1_j[:, ii] + eps_i
            u_1i_j = point.copy()
            u_1i_j[:, ii] = u_1i_j[:, ii] - eps_i

            qoi_plus = f_eval(u_i1_j)
            qoi_minus = f_eval(u_1i_j)

            du_dj[:, ii] = (qoi_plus - qoi_minus) / (2 * eps_i)

        return du_dj

    elif order.lower() == "second":
        # print('Calculating second order derivatives..')
        d2u_dj = np.zeros([point.shape[0], dimension])
        for ii in range(dimension):
            u_i1_j = point.copy()
            u_i1_j[:, ii] = u_i1_j[:, ii] + df_step[ii]
            u_1i_j = point.copy()
            u_1i_j[:, ii] = u_1i_j[:, ii] - df_step[ii]

            qoi_plus = f_eval(u_i1_j)
            qoi = f_eval(point)
            qoi_minus = f_eval(u_1i_j)

            d2u_dj[:, ii] = (qoi_plus - 2 * qoi + qoi_minus) / (
                df_step[ii] * df_step[ii]
            )

        return d2u_dj

    elif order.lower() == "mixed":

        import itertools

        range_ = list(range(dimension))
        d2u_dij = np.zeros([point.shape[0], int(dimension * (dimension - 1) / 2)])
        count = 0
        for i in itertools.combinations(range_, 2):
            u_i1_j1 = point.copy()
            u_i1_1j = point.copy()
            u_1i_j1 = point.copy()
            u_1i_1j = point.copy()

            eps_i1_0 = df_step[i[0]]
            eps_i1_1 = df_step[i[1]]

            u_i1_j1[:, i[0]] += eps_i1_0
            u_i1_j1[:, i[1]] += eps_i1_1

            u_i1_1j[:, i[0]] += eps_i1_0
            u_i1_1j[:, i[1]] -= eps_i1_1

            u_1i_j1[:, i[0]] -= eps_i1_0
            u_1i_j1[:, i[1]] += eps_i1_1

            u_1i_1j[:, i[0]] -= eps_i1_0
            u_1i_1j[:, i[1]] -= eps_i1_1

            print("hi")
            qoi_0 = f_eval(u_i1_j1)
            qoi_1 = f_eval(u_i1_1j)
            qoi_2 = f_eval(u_1i_j1)
            qoi_3 = f_eval(u_1i_1j)

            d2u_dij[:, count] = (qoi_0 + qoi_3 - qoi_1 - qoi_2) / (
                4 * eps_i1_0 * eps_i1_1
            )

            count += 1
        return d2u_dij


def bi_variate_normal_pdf(x1, x2, rho):
    return (
        1
        / (2 * np.pi * np.sqrt(1 - rho ** 2))
        * np.exp(-1 / (2 * (1 - rho ** 2)) * (x1 ** 2 - 2 * rho * x1 * x2 + x2 ** 2))
    )

def _get_a_plus(a):
    eig_val, eig_vec = np.linalg.eig(a)
    q = np.matrix(eig_vec)
    x_diagonal = np.matrix(np.diag(np.maximum(eig_val, 0)))

    return q * x_diagonal * q.T


def _get_ps(a, w=None):
    w05 = np.matrix(w ** 0.5)

    return w05.I * _get_a_plus(w05 * a * w05) * w05.I


def _get_pu(a, w=None):
    a_ret = np.array(a.copy())
    a_ret[w > 0] = np.array(w)[w > 0]
    return np.matrix(a_ret)


def _nn_coord(x, k):
    """
    Select k elements close to x.
    Select k elements close to x to be used to construct a sparse kernel
    matrix to be used in the diffusion maps.

    :param Union[list, numpy.ndarray] x: Matrices to be tested.
    :param int k: Number of points close to x.
    :return: Indices of the closer points.
    :rtype: int
    """
    if isinstance(x, list):
        x = np.array(x)

    dim = np.shape(x)

    if len(dim) != 1:
        raise ValueError("k MUST be a vector.")

    if not isinstance(k, int):
        raise TypeError("k MUST be integer.")

    if k < 1:
        raise ValueError("k MUST be larger than or equal to 1.")

    idx = x.argsort()[: len(x) - k]
    return idx


def correlation_distortion(dist_object, rho):
    """
    This method computes the correlation distortion from Gaussian distribution to any other distribution defined in
    UQpy.distributions

    :param Distribution dist_object: The object of the Distribution the correlation needs to be calculated.
    :param float rho: The Gaussian  correlation value.
    :return: The distorted correlation value.
    :rtype: float
    """

    if rho == 1.0:
        rho = 0.999
    n = 512
    zmax = 8
    zmin = -zmax
    eta, w2d, xi = calculate_gauss_quadrature_2d(n, zmax, zmin)
    tmp_f_xi = dist_object.icdf(stats.norm.cdf(xi[:, np.newaxis]))
    tmp_f_eta = dist_object.icdf(stats.norm.cdf(eta[:, np.newaxis]))
    coef = tmp_f_xi * tmp_f_eta * w2d
    phi2 = bi_variate_normal_pdf(xi, eta, rho)
    rho_non = np.sum(coef * phi2)
    rho_non = (
        rho_non - dist_object.moments(moments2return="m") ** 2
    ) / dist_object.moments(moments2return="v")
    return rho_non


def calculate_gauss_quadrature_2d(n, zmax, zmin):
    points, weights = np.polynomial.legendre.leggauss(n)
    points = -(0.5 * (points + 1) * (zmax - zmin) + zmin)
    weights = weights * (0.5 * (zmax - zmin))
    xi = np.tile(points, [n, 1])
    xi = xi.flatten(order="F")
    eta = np.tile(points, n)
    first = np.tile(weights, n)
    first = np.reshape(first, [n, n])
    second = np.transpose(first)
    weights2d = first * second
    w2d = weights2d.flatten()
    return eta, w2d, xi


def process_random_state(random_state: RandomStateType):
    if isinstance(random_state, (int, type(None))):
        return np.random.RandomState(random_state)
    else:
        return random_state


def calculate_gradient(krig_object, step_size, x, y, xt):
    """
    Estimating gradients with a Kriging metamodel (surrogate).

    :param krig_object:
    :param step_size:
    :param numpy.ndarray x: Samples in the training data.
    :param numpy.ndarray y: Function values evaluated at the samples in the training data.
    :param numpy.ndarray xt: Samples where gradients need to be evaluated.
    :return: First-order gradient evaluated at the points 'xt' using central difference.
    :rtype: numpy.ndarray
    """

    if krig_object is not None:
        krig_object.fit(x, y)
        krig_object.nopt = 1
        tck = krig_object.predict
    else:
        from scipy.interpolate import LinearNDInterpolator

        tck = LinearNDInterpolator(x, y, fill_value=0).__call__

    gr = gradient(point=xt, runmodel_object=tck, order="first", df_step=step_size)
    return gr


# src/UQpy/utilities/ValidationTypes.py

from typing import Union, Annotated

import numpy as np
from beartype.vale import Is

RandomStateType = Union[None, int, np.random.RandomState]
PositiveInteger = Annotated[int, Is[lambda number: number > 0]]
PositiveFloat = Annotated[float, Is[lambda number: number > 0]]
Numpy2DFloatArray = Annotated[
    np.ndarray,
    Is[lambda array: array.ndim == 2 and np.issubdtype(array.dtype, float)],
]
NumpyFloatArray = Annotated[
    np.ndarray,
    Is[lambda array: np.issubdtype(array.dtype, float)],
]
Numpy2DFloatArrayOrthonormal = Annotated[
    np.ndarray,
    Is[lambda array: array.ndim == 2 and np.issubdtype(array.dtype, float) and
                     np.allclose(array.T @ array, np.eye(array.shape[1]))],
]



# src/UQpy/utilities/__init__.py

from UQpy.utilities.Constants import *
from UQpy.utilities.NoPublicConstructor import *
from UQpy.utilities.UQpyLoggingFormatter import *
from UQpy.utilities.Utilities import *
from UQpy.utilities.ValidationTypes import *
from UQpy.utilities.DistanceMetric import *
from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.utilities.distances import *
from UQpy.utilities.kernels import *


# docs/source/conf.py

# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import os
import sys

sys.path.insert(0, os.path.abspath("./src/UQpy/docs"))

# -- Project information -----------------------------------------------------

project = "UQpy"
copyright = "2021, Michael D. Shields"
author = (
    "Michael D. Shields, Dimitrs G. Giovanis, Audrey Olivier, B.S. Aakash, Mohit Singh Chauhan, Lohit Vandanapu, "
    "Ketson RM dos Santos"
)

# The full version, including alpha/beta/rc tags
release = "v4.0.0"

# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    "sphinx.ext.autodoc",
    # "sphinx.ext.napoleon",
    # "sphinx.ext.doctest",
    "sphinx.ext.intersphinx",
    # "sphinx.ext.coverage",
    "sphinx.ext.mathjax",
    "sphinx.ext.viewcode",
    # "sphinx.ext.githubpages",
    # "sphinx.ext.autosummary",
    # "sphinx.ext.autosectionlabel",
    "sphinx_autodoc_typehints",
    "sphinxcontrib.bibtex",
    "sphinx_gallery.gen_gallery",
]

# nbsphinx_custom_formats={
#     ".md": ["jupytext.reads", {"fmt": "mystnb"}]
# }
autoclass_content = "init"
add_module_names = False
autodoc_member_order = "bysource"
napoleon_use_param = True
# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]
bibtex_bibfiles = ["bibliography.bib"]
bibtex_default_style = "unsrt"

# Try to remove duplicate labels
autosectionlabel_prefix_document = True

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "Model_Runs**"]

sphinx_gallery_conf = {
    "examples_dirs": [
        "../code/dimension_reduction/diffusion_maps",
        "../code/dimension_reduction/pod",
        "../code/dimension_reduction/grassmann",
        "../code/distributions/continuous_1d",
        "../code/distributions/discrete_1d",
        "../code/distributions/multivariate",
        "../code/distributions/user_defined",
        "../code/sampling/adaptive_kriging",
        "../code/sampling/importance_sampling",
        "../code/sampling/monte_carlo",
        "../code/sampling/latin_hypercube",
        "../code/sampling/mcmc",
        "../code/sampling/simplex",
        "../code/sampling/true_stratified_sampling",
        "../code/sampling/refined_stratified_sampling",
        "../code/inference/mle",
        "../code/inference/info_model_selection",
        "../code/inference/bayes_parameter_estimation",
        "../code/inference/bayes_model_selection",
        "../code/transformations/nataf",
        "../code/sensitivity/morris",
        "../code/sensitivity/sobol",
        "../code/sensitivity/cramer_von_mises",
        "../code/sensitivity/chatterjee",
        "../code/sensitivity/generalised_sobol",
        "../code/stochastic_processes/bispectral",
        "../code/stochastic_processes/karhunen_loeve",
        "../code/stochastic_processes/spectral",
        "../code/stochastic_processes/translation",
        "../code/reliability/form",
        "../code/reliability/sorm",
        "../code/reliability/subset_simulation",
        "../code/surrogates/srom",
        "../code/surrogates/gpr",
        "../code/surrogates/pce",
        "../code/RunModel",
    ],  # path to your example scripts,
    "gallery_dirs": [
        "auto_examples/dimension_reduction/diffusion_maps",
        "auto_examples/dimension_reduction/pod",
        "auto_examples/dimension_reduction/grassmann",
        "auto_examples/distributions/continuous_1d",
        "auto_examples/distributions/discrete_1d",
        "auto_examples/distributions/multivariate",
        "auto_examples/distributions/user_defined",
        "auto_examples/sampling/adaptive_kriging",
        "auto_examples/sampling/importance_sampling",
        "auto_examples/sampling/monte_carlo",
        "auto_examples/sampling/latin_hypercube",
        "auto_examples/sampling/mcmc",
        "auto_examples/sampling/simplex",
        "auto_examples/sampling/true_stratified_sampling",
        "auto_examples/sampling/refined_stratified_sampling",
        "auto_examples/inference/mle",
        "auto_examples/inference/info_model_selection",
        "auto_examples/inference/bayes_parameter_estimation",
        "auto_examples/inference/bayes_model_selection",
        "auto_examples/transformations/nataf",
        "auto_examples/sensitivity/morris",
        "auto_examples/sensitivity/sobol",
        "auto_examples/sensitivity/cramer_von_mises",
        "auto_examples/sensitivity/chatterjee",
        "auto_examples/sensitivity/generalised_sobol",
        "auto_examples/stochastic_processes/bispectral",
        "auto_examples/stochastic_processes/karhunen_loeve",
        "auto_examples/stochastic_processes/spectral",
        "auto_examples/stochastic_processes/translation",
        "auto_examples/reliability/form",
        "auto_examples/reliability/sorm",
        "auto_examples/reliability/subset_simulation",
        "auto_examples/surrogates/srom",
        "auto_examples/surrogates/gpr",
        "auto_examples/surrogates/pce",
        "auto_examples/RunModel",
    ],  # path to where to save gallery generated output
    "binder": {
        # Required keys
        "org": "SURGroup",
        "repo": "UQpy",
        "branch": "master",  # Can be any branch, tag, or commit hash. Use a branch that hosts your docs.
        "binderhub_url": "https://mybinder.org",
        # Any URL of a binderhub deployment. Must be full URL (e.g. https://mybinder.org).
        "dependencies": "./binder/requirements.txt",
        "notebooks_dir": "notebooks",
        "use_jupyter_lab": True
        # Jupyter notebooks for Binder will be copied to this directory (relative to built documentation root).
    },
    "ignore_pattern": "/local_",
}

# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
language = None
pygments_style = None

html_theme = "sphinx_rtd_theme"

html_theme_options = {
    "logo_only": True,
    "style_nav_header_background": "#F0F0F0",
    "vcs_pageview_mode": "view",
}

github_url = "https://github.com/SURGroup/UQpy"
html_logo = "_static/logo.png"

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]

html_sidebars = {
    "**": [
        "about.html",
        "navigation.html",
        "relations.html",
        "searchbox.html",
    ]
}

# The suffix(es) of source filenames.
# You can specify multiple suffix as a list of string:
#
source_suffix = [".rst", ".md"]
# source_suffix = ".rst"

# The master toctree document.
master_doc = "index"

# -- Options for HTMLHelp output ---------------------------------------------

# Output file base name for HTML help builder.
htmlhelp_basename = "uqpydoc"

# autodoc_default_flags = ['members',  'private-members']

# -- Options for LaTeX output ------------------------------------------------

latex_elements = {
    # The paper size ('letterpaper' or 'a4paper').
    #
    # 'papersize': 'letterpaper',
    # The font size ('10pt', '11pt' or '12pt').
    #
    # 'pointsize': '10pt',
    # Additional stuff for the LaTeX preamble.
    #
    # 'preamble': '',
    # Latex figure (float) alignment
    #
    # 'figure_align': 'htbp',
}

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title,
#  author, documentclass [howto, manual, or own class]).
# latex_documents = [
#     (master_doc, "uqpy.tex", "UQpy Documentation", "Michael D. Shields", "manual"),
# ]

# -- Options for manual page output ------------------------------------------

# One entry per manual page. List of tuples
# (source start file, name, description, authors, manual section).
# man_pages = [(master_doc, "uqpy", "UQpy Documentation", [author], 1)]

# -- Options for Texinfo output ----------------------------------------------

# Grouping the document tree into Texinfo files. List of tuples
# (source start file, target name, title, author,
#  dir menu entry, description, category)
# texinfo_documents = [
#     (
#         master_doc,
#         "uqpy",
#         "UQpy Documentation",
#         author,
#         "uqpy",
#         "One line description of project.",
#         "Miscellaneous",
#     ),
# ]

# -- Options for Epub output -------------------------------------------------

# Bibliographic Dublin Core info.
# epub_title = project

# The unique identifier of the text. This can be a ISBN number
# or the project homepage.
#
# epub_identifier = ''

# A unique identification for the text.
#
# epub_uid = ''

# A list of files that should not be packed into the epub file.
epub_exclude_files = ["search.html"]

# Example configuration for intersphinx: refer to the Python standard library.
intersphinx_mapping = {
    "python": ("https://docs.python.org/3", None),
    "numpy": ("https://docs.scipy.org/doc/numpy/", None),
    "scipy": ("https://docs.scipy.org/doc/scipy/", None),
}


# src/UQpy/__init__.py

"""
Uncertainty Quantification with Python
========================================
"""

import pkg_resources

import UQpy.distributions
import UQpy.sampling
import UQpy.dimension_reduction
import UQpy.inference
import UQpy.reliability
import UQpy.run_model.RunModel
import UQpy.sampling
import UQpy.stochastic_process
import UQpy.surrogates
import UQpy.sensitivity
import UQpy.transformations
import UQpy.utilities.Utilities


from UQpy.dimension_reduction import *
from UQpy.distributions import *
from UQpy.inference import *
from UQpy.reliability import *
from UQpy.run_model import *
from UQpy.sampling import *
from UQpy.stochastic_process import *
from UQpy.surrogates import *
from UQpy.transformations import *
from UQpy.utilities.Utilities import *
from UQpy.sensitivity import *
from UQpy.utilities.UQpyLoggingFormatter import *
import logging
from beartype.roar import BeartypeDecorHintPep585DeprecationWarning
from warnings import filterwarnings
filterwarnings("ignore", category=BeartypeDecorHintPep585DeprecationWarning)

logger = logging.getLogger(__name__)
logger.setLevel(logging.ERROR)

formatter = UQpyLoggingFormatter()

console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)

logger.addHandler(console_handler)

logging.logThreads = 0
logging.logProcesses = 0

try:
    __version__ = pkg_resources.get_distribution("UQpy").version
except pkg_resources.DistributionNotFound:
    __version__ = None

try:
    __version__ = pkg_resources.get_distribution("UQpy").version
except pkg_resources.DistributionNotFound:
    __version__ = None


# tests/unit_tests/sampling/MCMC/test_mcmc_algorithms.py

from UQpy.sampling.mcmc import *
import UQpy.distributions as Distributions


# Tests for parent MCMC and MH algorithm
def test_mh_1d_target_pdf():
    target = Distributions.Normal().pdf
    x = MetropolisHastings(dimension=1, pdf_target=target, n_chains=1, random_state=123, nsamples=10)
    assert round(float(x.samples[-1]), 3) == -1.291


def test_mh_1d_samples_per_chain():
    target = Distributions.Normal().pdf
    x = MetropolisHastings(dimension=1, pdf_target=target, n_chains=2, random_state=123,
                           nsamples_per_chain=5)
    assert round(float(x.samples[-1]), 3) == 0.474


def test_mh_1d_acceptance_rate():
    target = Distributions.Normal().pdf
    x = MetropolisHastings(dimension=1, pdf_target=target, n_chains=1, random_state=123, nsamples=100)
    assert round(float(x.acceptance_rate[0]), 3) == 0.707


def test_mh_1d_save_log_pdf():
    target = Distributions.Normal().pdf
    x = MetropolisHastings(dimension=1, pdf_target=target, n_chains=1, random_state=123, save_log_pdf=True,
                           nsamples=10)
    assert round(float(x.log_pdf_values[-1]), 3) == -1.752


def test_mh_1d_target_log_pdf():
    target = Distributions.Normal().log_pdf
    x = MetropolisHastings(dimension=1, log_pdf_target=target, n_chains=1, random_state=123, nsamples=10)
    assert round(float(x.samples[-1]), 3) == -1.291


def test_mh_2d():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    x = MetropolisHastings(dimension=2, pdf_target=target, n_chains=1, random_state=123, nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.406, -1.217]


def test_mh_2d_burn_jump():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    x = MetropolisHastings(dimension=2, log_pdf_target=target, burn_length=10, jump=2, n_chains=1,
                           random_state=123, nsamples=10)
    assert x.iterations_number == 30


def test_mh_2d_nsamples_check():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    x = MetropolisHastings(dimension=2, pdf_target=target, n_chains=2, random_state=123, nsamples=60)
    assert x.nsamples_per_chain + x.samples_counter == 90


def test_mh_2d_2chains():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    x = MetropolisHastings(dimension=2, pdf_target=target, n_chains=2, random_state=123, nsamples=60)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.064, -0.533]


def test_mh_2d_2chains_non_concatenated():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    x = MetropolisHastings(dimension=2, pdf_target=target, concatenate_chains=False, n_chains=2, random_state=123,
                           nsamples=60)
    assert [[round(float(x.samples[-1][0][0]), 3), round(float(x.samples[-1][0][1]), 3)],
            [round(float(x.samples[-1][1][0]), 3), round(float(x.samples[-1][1][1]), 3)]] == [[1.767, 1.465],
                                                                                              [-0.064, -0.533]]


def test_mh_2d_seed():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    x = MetropolisHastings(pdf_target=target, seed=[0., 0.], n_chains=1, random_state=123, nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.406, -1.217]


def test_mh_1d_symmetric_proposal_pdf():
    target = Distributions.Normal().pdf
    proposal = Distributions.Normal()
    x = MetropolisHastings(dimension=1, pdf_target=target, proposal=proposal, proposal_is_symmetric=True,
                           n_chains=1, random_state=123, nsamples=10)
    assert round(float(x.samples[-1]), 3) == -1.291


def test_mh_1d_asymmetric_proposal_pdf():
    target = Distributions.Normal().pdf
    proposal = Distributions.Normal()
    x = MetropolisHastings(dimension=1, pdf_target=target, proposal=proposal, proposal_is_symmetric=False,
                           n_chains=1, random_state=123, nsamples=10)
    assert round(float(x.samples[-1]), 3) == -1.291


def test_mmh_1d_burn_jump():
    target = Distributions.Normal().pdf
    x = ModifiedMetropolisHastings(dimension=1, pdf_target=target, burn_length=10,
                                   jump=2, n_chains=1, random_state=123, nsamples=10)
    assert round(float(x.samples[-1]), 3) == 0.497


def test_mmh_2d_list_target_pdf():
    target = [Distributions.Normal().pdf, Distributions.Normal().pdf]
    x = ModifiedMetropolisHastings(dimension=2, pdf_target=target, n_chains=1, random_state=123, nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.810, 0.173]


def test_mmh_2d_list_target_log_pdf():
    target = [Distributions.Normal().log_pdf, Distributions.Normal().log_pdf]
    x = ModifiedMetropolisHastings(dimension=2, log_pdf_target=target, n_chains=1, random_state=123,
                                   nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.810, 0.173]


def test_mmh_2d_joint_proposal():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    proposal = Distributions.JointIndependent(marginals=[Distributions.Normal(scale=0.2),
                                                         Distributions.Normal(scale=0.2)])
    x = ModifiedMetropolisHastings(dimension=2, pdf_target=target, n_chains=1, proposal=proposal, random_state=123,
                                   nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.783, -0.195]


def test_mmh_2d_list_proposal():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    proposal = [Distributions.Normal(scale=0.2), Distributions.Normal(scale=0.2)]
    x = ModifiedMetropolisHastings(dimension=2, pdf_target=target, n_chains=1, proposal=proposal, random_state=123,
                                   nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.783, -0.195]


def test_mmh_2d_single1d_proposal():
    target = Distributions.MultivariateNormal([0., 0.]).pdf
    proposal = Distributions.Normal(scale=0.2)
    x = ModifiedMetropolisHastings(dimension=2, pdf_target=target, n_chains=1, proposal=proposal, random_state=123,
                                   nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.783, -0.195]


def test_mmh_2d_list_proposal_log_target():
    target = [Distributions.Normal().log_pdf, Distributions.Normal().log_pdf]
    proposal = [Distributions.Normal(scale=0.2), Distributions.Normal(scale=0.2)]
    x = ModifiedMetropolisHastings(dimension=2, log_pdf_target=target, n_chains=1, proposal=proposal,
                                   random_state=123, nsamples=10)
    assert [round(float(x.samples[-1][0]), 3), round(float(x.samples[-1][1]), 3)] == [-0.783, -0.195]


def test_dram_1d_burn_jump():
    target = Distributions.Normal().pdf
    x = DRAM(dimension=1, pdf_target=target, burn_length=10, jump=2, n_chains=1, random_state=123,
             nsamples=10)
    assert round(float(x.samples[-1]), 3) == 0.935


def test_dream_1d_burn_jump():
    target = Distributions.Normal().pdf
    x = DREAM(pdf_target=target, burn_length=10, jump=2, dimension=1, n_chains=10, random_state=123,
              nsamples=20)
    assert round(float(x.samples[-1]), 3) == 0.0


def test_dream_1d_check_chains():
    target = Distributions.Normal().pdf
    x = DREAM(pdf_target=target, burn_length=0, jump=2, save_log_pdf=True, dimension=1, check_chains=(1000, 1),
              n_chains=20, random_state=123, nsamples=2000)
    assert (round(float(x.samples[-1]), 3) == 0.593)


def test_dream_1d_adapt_chains():
    target = Distributions.Normal().pdf
    x = DREAM(pdf_target=target, burn_length=1000, jump=2, save_log_pdf=True, dimension=1,
              crossover_adaptation=(1000, 1), n_chains=20, random_state=123, nsamples=2000)
    assert (round(float(x.samples[-1]), 3) == -0.446)


def test_stretch_1d_burn_jump():
    target = Distributions.Normal().pdf
    x = Stretch(pdf_target=target, burn_length=10, jump=2, dimension=1, n_chains=2, random_state=123,
                nsamples=10)
    assert round(float(x.samples[-1]), 3) == -0.961


def test_unconcatenate_chains_mcmc():
    target = Distributions.Normal().pdf
    x = ModifiedMetropolisHastings(dimension=1, pdf_target=target, burn_length=10, jump=2, n_chains=2,
                                   save_log_pdf=True, random_state=123)
    x.run(nsamples=5)
    x.run(nsamples=5)
    assert (round(float(x.samples[-1]), 3) == -0.744)


# setup.py

#!/usr/bin/env python
import sys
version = sys.argv[1]
del sys.argv[1]
from setuptools import setup, find_packages

setup(
    name='UQpy',
    version=version,
    url='https://github.com/SURGroup/UQpy',
    description="UQpy is a general purpose toolbox for Uncertainty Quantification",
    author="Michael D. Shields, Dimitris G. Giovanis, Audrey Olivier, Aakash Bangalore-Satish, Mohit Chauhan, "
           "Lohit Vandanapu, Ketson R.M. dos Santos",
    author_email="UQpy.info@gmail.com",
    license='MIT',
    platforms=["OSX", "Windows", "Linux"],
    packages=find_packages("src"),
    package_dir={"": "src"},
    package_data={"": ["*.pdf"]},
    install_requires=[
        "numpy", "scipy", "matplotlib", "scikit-learn", 'fire',
        "beartype==0.9.1",
    ],
    classifiers=[
        'Programming Language :: Python :: 3',
        'Intended Audience :: Science/Research',
        'Topic :: Scientific/Engineering :: Mathematics',
        'License :: OSI Approved :: MIT License',
        'Natural Language :: English',
    ],
)


# tests/unit_tests/dimension_reduction/test_POD.py

import numpy as np

# Define a dataset to test DirectPOD, SnapshotPOD and HOSVD methods
import pytest
from beartype.roar import BeartypeCallHintPepParamException

from UQpy.dimension_reduction.pod.DirectPOD import DirectPOD
from UQpy.dimension_reduction.pod.SnapshotPOD import SnapshotPOD
from UQpy.dimension_reduction.hosvd.HigherOrderSVD import HigherOrderSVD

Data = np.zeros((2, 2, 3))

Data[:, :, 0] = [
    [0.9073, 1.7842],
    [2.1488, 4.2495]]

Data[:, :, 1] = [
    [6.7121, 0.5334],
    [0.3054, 0.3207]]

Data[:, :, 2] = [
    [-0.3698, 0.0151],
    [2.3753, 4.7146]]


def test_DirectPOD_listData():
    list_data = list(Data)
    pod_dir = DirectPOD(solution_snapshots=list_data, reconstruction_percentage=100)
    actual_result = pod_dir.reconstructed_solution[0][1][1]
    expected_result = 0.3054
    assert expected_result == round(actual_result, 6)

def test_DirectPOD():
    pod_dir = DirectPOD(solution_snapshots=Data, n_modes=1)
    reconstructed_solutions = pod_dir.reconstructed_solution
    assert round(reconstructed_solutions[0][1][1], 6) == 0.761704


def test_SnapshotPOD():
    pod_snap = SnapshotPOD(solution_snapshots=Data, n_modes=1)
    reconstructed_solutions = pod_snap.reconstructed_solution
    assert round(reconstructed_solutions[0][1][1], 6) == -0.181528


def test_SnapshotPOD_listData():
    list_data = list(Data)
    pod_dir = SnapshotPOD(solution_snapshots=list_data, reconstruction_percentage=100)
    actual_result = pod_dir.reconstructed_solution[0][1][1]
    expected_result = 0.3054
    assert expected_result == round(actual_result, 6)


def test_POD_unfold():
    list_data = list(Data)
    pod_output = HigherOrderSVD.unfold3d(list_data)
    actual_result = pod_output[0][0][1]
    expected_result = 6.7121
    assert expected_result == round(actual_result, 6)


def test_HOSVD():
    hosvd = HigherOrderSVD(solution_snapshots=Data, reconstruction_percentage=90)
    reconstructed_solutions = HigherOrderSVD.reconstruct(hosvd.u1, hosvd.u2, hosvd.u3hat, hosvd.s3hat)
    assert round(reconstructed_solutions[0][1][1], 6) == 0.714928


def test_DirectPOD_modes_less_than_zero():
    with pytest.raises(BeartypeCallHintPepParamException):
        pod_dir = DirectPOD(solution_snapshots=Data, n_modes=-1)


def test_DirectPOD_mode_non_integer():
    with pytest.raises(BeartypeCallHintPepParamException):
        pod_dir = DirectPOD(solution_snapshots=Data, n_modes=1.5)


def test_DirectPOD_reconstr_perc_less_than_zero():
    with pytest.raises(BeartypeCallHintPepParamException):
        pod_dir = DirectPOD(solution_snapshots=Data, reconstruction_percentage=-1)


def test_DirectPOD_both_modes_and_reconstr_error():
    with pytest.raises(ValueError):
        pod_dir = DirectPOD(solution_snapshots=Data, n_modes=1, reconstruction_percentage=50)


def test_HOSVD_modes_less_than_zero():
    with pytest.raises(BeartypeCallHintPepParamException):
        hosvd = HigherOrderSVD(solution_snapshots=Data, modes=-1)


def test_HOSVD_mode_non_integer():
    with pytest.raises(BeartypeCallHintPepParamException):
        hosvd = HigherOrderSVD(solution_snapshots=Data, modes=1.5)


def test_HOSVD_reconstr_perc_less_than_zero():
    with pytest.raises(BeartypeCallHintPepParamException):
        hosvd = HigherOrderSVD(solution_snapshots=Data, reconstruction_percentage=-1)


def test_HOSVD_both_modes_and_reconstr_error():
    with pytest.raises(ValueError):
        hosvd = HigherOrderSVD(solution_snapshots=Data, modes=1, reconstruction_percentage=50)


def test_SnapshotPOD_modes_less_than_zero():
    with pytest.raises(BeartypeCallHintPepParamException):
        snap = SnapshotPOD(solution_snapshots=Data, n_modes=-1)


def test_SnapshotPOD_mode_non_integer():
    with pytest.raises(BeartypeCallHintPepParamException):
        snap = SnapshotPOD(solution_snapshots=Data, n_modes=1.5)


def test_SnapshotPOD_reconstr_perc_less_than_zero():
    with pytest.raises(BeartypeCallHintPepParamException):
        snap = SnapshotPOD(solution_snapshots=Data, reconstruction_percentage=-1)


def test_SnapshotPOD_both_modes_and_reconstr_error():
    with pytest.raises(ValueError):
        snap = SnapshotPOD(solution_snapshots=Data, n_modes=1, reconstruction_percentage=50)


# tests/unit_tests/dimension_reduction/test_distances.py

from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.dimension_reduction.grassmann_manifold.projections.SVDProjection import SVDProjection
from UQpy.utilities.distances.euclidean_distances import L2Distance
from UQpy.utilities.distances.grassmannian_distances import AsimovDistance, BinetCauchyDistance, FubiniStudyDistance, \
    GeodesicDistance, ProcrustesDistance, ProjectionDistance, SpectralDistance
from UQpy.utilities.distances import MartinDistance
from UQpy.utilities.DistanceMetric import DistanceMetric
import numpy as np
import sys


def test_euclidean_distance_2points():
    x = np.array([[2.0, 3.1], [4.0, 1.25]])
    d = L2Distance()
    distance = np.round(d.compute_distance(xi=np.array([2.0, 3.1]), xj=np.array([4.0, 1.25])), 3)

    assert distance == 2.724


def test_grassmann_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(GeodesicDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 1.491253


def test_fs_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(FubiniStudyDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 1.491253


def test_procrustes_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(ProcrustesDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 1.356865


def test_projection_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(ProjectionDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 0.996838


def test_binet_cauchy_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(BinetCauchyDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 0.996838


def test_asimov_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(AsimovDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 1.491253


def test_martin_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(MartinDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 2.25056


def test_spectral_distance():
    xi = np.array([[-np.sqrt(2)/2, -np.sqrt(2)/4], [np.sqrt(2)/2, -np.sqrt(2)/4], [0, -np.sqrt(3)/2]])
    xj = np.array([[0, np.sqrt(2)/2], [1, 0], [0, -np.sqrt(2)/2]])
    distance = np.round(SpectralDistance().compute_distance(GrassmannPoint(xi), GrassmannPoint(xj)), 6)
    assert distance == 1.356865



# tests/unit_tests/dimension_reduction/test_dmaps.py

from UQpy.utilities.kernels.GaussianKernel import GaussianKernel
from UQpy.dimension_reduction.diffusion_maps.DiffusionMaps import DiffusionMaps


def test_dmaps_swiss_roll():
    import numpy as np
    from UQpy.dimension_reduction.diffusion_maps.DiffusionMaps import DiffusionMaps

    # set parameters
    length_phi = 15  # length of swiss roll in angular direction
    length_Z = 15  # length of swiss roll in z direction
    sigma = 0.1  # noise strength
    m = 4000  # number of samples

    np.random.seed(1111)
    # create dataset
    phi = length_phi * np.random.rand(m)
    xi = np.random.rand(m)
    Z0 = length_Z * np.random.rand(m)
    X0 = 1. / 6 * (phi + sigma * xi) * np.sin(phi)
    Y0 = 1. / 6 * (phi + sigma * xi) * np.cos(phi)

    swiss_roll = np.array([X0, Y0, Z0]).transpose()
    dmaps = DiffusionMaps(data=swiss_roll, kernel=GaussianKernel(epsilon=0.5),
                          alpha=0.5, n_eigenvectors=3, is_sparse=True, n_neighbors=100)

    evals = dmaps.eigenvalues
    assert round(evals[0], 9) == 1.0
    assert round(evals[1], 9) == 0.997460971
    assert round(evals[2], 9) == 0.987499519


def test_dmaps_circular():
    import numpy as np
    np.random.seed(1111)
    a = 6
    b = 1
    k = 10
    u = np.linspace(0, 2 * np.pi, 1000)

    v = k * u

    x0 = (a + b * np.cos(0.8 * v)) * (np.cos(u))
    y0 = (a + b * np.cos(0.8 * v)) * (np.sin(u))
    z0 = b * np.sin(0.8 * v)

    rox = 0.2
    roy = 0.2
    roz = 0.2
    x = x0 + rox * np.random.normal(0, 1, len(x0))
    y = y0 + roy * np.random.normal(0, 1, len(y0))
    z = z0 + roz * np.random.normal(0, 1, len(z0))

    X = np.array([x, y, z]).transpose()

    dmaps = DiffusionMaps(data=X, alpha=1, n_eigenvectors=3,
                          kernel=GaussianKernel(epsilon=0.3))

    evals = dmaps.eigenvalues
    assert np.round(evals[0], 5) == 1.0
    assert np.round(evals[1], 5) == 0.99962
    assert np.round(evals[2], 5) == 0.99961


# tests/unit_tests/dimension_reduction/test_grassman.py

import sys

import numpy as np
import scipy

from UQpy.utilities.distances.grassmannian_distances import GeodesicDistance
from UQpy.dimension_reduction.grassmann_manifold.projections.SVDProjection import SVDProjection
from UQpy.dimension_reduction.grassmann_manifold.GrassmannInterpolation import GrassmannInterpolation

sol0 = np.array([[0.61415, 1.03029, 1.02001, 0.57327, 0.79874, 0.73274],
                 [0.56924, 0.91700, 0.88841, 0.53737, 0.68676, 0.67751],
                 [0.51514, 0.87898, 0.87779, 0.47850, 0.69085, 0.61525],
                 [0.63038, 1.10822, 1.12313, 0.58038, 0.89142, 0.75429],
                 [0.69666, 1.03114, 0.95037, 0.67211, 0.71184, 0.82522],
                 [0.66595, 1.03789, 0.98690, 0.63420, 0.75416, 0.79110]])

sol1 = np.array([[1.05134, 1.37652, 0.95634, 0.85630, 0.47570, 1.22488],
                 [0.16370, 0.63105, 0.14533, 0.81030, 0.44559, 0.43358],
                 [1.23478, 2.10342, 1.04698, 1.68755, 0.92792, 1.73277],
                 [0.90538, 1.64067, 0.62027, 1.17577, 0.63644, 1.34925],
                 [0.58210, 0.75795, 0.65519, 0.65712, 0.37251, 0.65740],
                 [0.99174, 1.59375, 0.63724, 0.89107, 0.47631, 1.36581]])

sol2 = np.array([[1.04142, 0.91670, 1.47962, 1.23350, 0.94111, 0.61858],
                 [1.00464, 0.65684, 1.35136, 1.11288, 0.96093, 0.42340],
                 [1.05567, 1.33192, 1.56286, 1.43412, 0.77044, 0.97182],
                 [0.89812, 0.86136, 1.20204, 1.17892, 0.83788, 0.61160],
                 [0.46935, 0.39371, 0.63534, 0.57856, 0.47615, 0.26407],
                 [1.14102, 0.80869, 1.39123, 1.33076, 0.47719, 0.68170]])

sol3 = np.array([[0.60547, 0.11492, 0.78956, 0.13796, 0.76685, 0.41661],
                 [0.32771, 0.11606, 0.67630, 0.15208, 0.44845, 0.34840],
                 [0.58959, 0.10156, 0.72623, 0.11859, 0.73671, 0.38714],
                 [0.36283, 0.07979, 0.52824, 0.09760, 0.46313, 0.27906],
                 [0.87487, 0.22452, 1.30208, 0.30189, 1.22015, 0.62918],
                 [0.56006, 0.16879, 1.09635, 0.20431, 0.69439, 0.60317]])


def test_solution_reconstruction():
    nodes = np.array([[0, 0], [1, 0], [1, 1], [0, 1]])
    point = np.array([0.1, 0.1])  # Point to interpolate.

    D1 = 6
    r0 = 2  # rank sample 0
    r1 = 3  # rank sample 1
    r2 = 4  # rank sample 2
    r3 = 3  # rank sample 2

    np.random.seed(1111)  # For reproducibility.

    # Creating a list of solutions.
    Solutions = [sol0, sol1, sol2, sol3]

    manifold_projection = SVDProjection(Solutions, p="max")

    interpolation = GrassmannInterpolation(interpolation_method=None,
                                           manifold_data=manifold_projection.v,
                                           coordinates=nodes,
                                           distance=GeodesicDistance())

    interpolated_solution = interpolation.interpolate_manifold(point=point)
    # assert round(interpolated_solution.data[0, 0], 9) == -0.353239531
    assert round(interpolated_solution.data[0, 0], 9) == -0.316807309

def test_parsimonious():
    from UQpy.utilities.kernels.GaussianKernel import GaussianKernel
    from UQpy.dimension_reduction.diffusion_maps.DiffusionMaps import DiffusionMaps
    from sklearn.datasets import make_s_curve

    n = 4000
    X, X_color = make_s_curve(n, random_state=3, noise=0)
    kernel = GaussianKernel()

    dmaps_object = DiffusionMaps(data=X, alpha=1.0, n_eigenvectors=9,
                                 is_sparse=True, n_neighbors=100,
                                 kernel=kernel)

    dmaps_object.parsimonious(dim=2)

    assert dmaps_object.parsimonious_indices[0] == 1
    assert dmaps_object.parsimonious_indices[1] == 5


# tests/unit_tests/dimension_reduction/test_karcher.py

import copy

import numpy as np
import numpy.random
from beartype import beartype

from UQpy.utilities.distances.grassmannian_distances import GeodesicDistance
from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.dimension_reduction.grassmann_manifold.projections.SVDProjection import SVDProjection
from UQpy.dimension_reduction.grassmann_manifold.GrassmannOperations import GrassmannOperations
import sys
from numpy.random import RandomState

sol0 = np.array([[0.61415, 1.03029, 1.02001, 0.57327, 0.79874, 0.73274],
                 [0.56924, 0.91700, 0.88841, 0.53737, 0.68676, 0.67751],
                 [0.51514, 0.87898, 0.87779, 0.47850, 0.69085, 0.61525],
                 [0.63038, 1.10822, 1.12313, 0.58038, 0.89142, 0.75429],
                 [0.69666, 1.03114, 0.95037, 0.67211, 0.71184, 0.82522],
                 [0.66595, 1.03789, 0.98690, 0.63420, 0.75416, 0.79110]])

sol1 = np.array([[1.05134, 1.37652, 0.95634, 0.85630, 0.47570, 1.22488],
                 [0.16370, 0.63105, 0.14533, 0.81030, 0.44559, 0.43358],
                 [1.23478, 2.10342, 1.04698, 1.68755, 0.92792, 1.73277],
                 [0.90538, 1.64067, 0.62027, 1.17577, 0.63644, 1.34925],
                 [0.58210, 0.75795, 0.65519, 0.65712, 0.37251, 0.65740],
                 [0.99174, 1.59375, 0.63724, 0.89107, 0.47631, 1.36581]])

sol2 = np.array([[1.04142, 0.91670, 1.47962, 1.23350, 0.94111, 0.61858],
                 [1.00464, 0.65684, 1.35136, 1.11288, 0.96093, 0.42340],
                 [1.05567, 1.33192, 1.56286, 1.43412, 0.77044, 0.97182],
                 [0.89812, 0.86136, 1.20204, 1.17892, 0.83788, 0.61160],
                 [0.46935, 0.39371, 0.63534, 0.57856, 0.47615, 0.26407],
                 [1.14102, 0.80869, 1.39123, 1.33076, 0.47719, 0.68170]])

sol3 = np.array([[0.60547, 0.11492, 0.78956, 0.13796, 0.76685, 0.41661],
                 [0.32771, 0.11606, 0.67630, 0.15208, 0.44845, 0.34840],
                 [0.58959, 0.10156, 0.72623, 0.11859, 0.73671, 0.38714],
                 [0.36283, 0.07979, 0.52824, 0.09760, 0.46313, 0.27906],
                 [0.87487, 0.22452, 1.30208, 0.30189, 1.22015, 0.62918],
                 [0.56006, 0.16879, 1.09635, 0.20431, 0.69439, 0.60317]])


def test_karcher():
    # Creating a list of matrices.
    matrices = [sol0, sol1, sol2, sol3]
    manifold_projection = SVDProjection(matrices, p="max")

    # optimization_method = GradientDescent(acceleration=True, error_tolerance=1e-4, max_iterations=1000)
    psi_mean = GrassmannOperations.karcher_mean(grassmann_points=manifold_projection.u,
                                                optimization_method="GradientDescent",
                                                distance=GeodesicDistance())

    phi_mean = GrassmannOperations.karcher_mean(grassmann_points=manifold_projection.v,
                                                optimization_method="GradientDescent",
                                                distance=GeodesicDistance())

    assert round(psi_mean.data[0, 0], 9) == -0.387197957
    # assert round(psi_mean.data[0, 0], 9) == -0.418075902
    # assert round(phi_mean.data[0, 0], 9) == -0.353239531


def test_karcher_stochastic():
    matrices = [sol0, sol1, sol2, sol3]
    manifold_projection = SVDProjection(matrices, p="max")

    # optimization_method = GradientDescent(acceleration=True, error_tolerance=1e-4, max_iterations=1000)
    psi_mean = GrassmannOperations.karcher_mean(grassmann_points=manifold_projection.u,
                                                optimization_method="StochasticGradientDescent",
                                                distance=GeodesicDistance())

    phi_mean = GrassmannOperations.karcher_mean(grassmann_points=manifold_projection.v,
                                                optimization_method="StochasticGradientDescent",
                                                distance=GeodesicDistance())


# tests/unit_tests/dimension_reduction/test_kernel.py

import sys

from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.dimension_reduction.grassmann_manifold.projections.SVDProjection import SVDProjection
from UQpy.utilities.kernels.ProjectionKernel import ProjectionKernel
from UQpy.utilities.kernels.BinetCauchyKernel import BinetCauchyKernel
from UQpy.utilities.kernels.GaussianKernel import GaussianKernel
import numpy as np


def test_kernel_projection():
    xi = GrassmannPoint(np.array([[-np.sqrt(2) / 2, -np.sqrt(2) / 4], [np.sqrt(2) / 2,
                                                                       -np.sqrt(2) / 4], [0, -np.sqrt(3) / 2]]))
    xj = GrassmannPoint(np.array([[0, np.sqrt(2) / 2], [1, 0], [0, -np.sqrt(2) / 2]]))
    xk = GrassmannPoint(np.array([[-0.69535592, -0.0546034], [-0.34016974, -0.85332868],
                                  [-0.63305978, 0.51850616]]))
    points = [xi, xj, xk]
    k = ProjectionKernel()
    k.calculate_kernel_matrix(points)
    kernel = np.matrix.round(k.kernel_matrix, 4)

    assert np.allclose(kernel, np.array([[2, 1.0063, 1.2345], [1.0063, 2, 1.0101], [1.2345, 1.0101, 2]]))


def test_kernel_binet_cauchy():
    xi = GrassmannPoint(np.array([[-np.sqrt(2) / 2, -np.sqrt(2) / 4], [np.sqrt(2) / 2, -np.sqrt(2) / 4],
                                  [0, -np.sqrt(3) / 2]]))
    xj = GrassmannPoint(np.array([[0, np.sqrt(2) / 2], [1, 0], [0, -np.sqrt(2) / 2]]))
    xk = GrassmannPoint(np.array([[-0.69535592, -0.0546034], [-0.34016974, -0.85332868], [-0.63305978, 0.51850616]]))
    points = [xi, xj, xk]

    kernel = BinetCauchyKernel()
    kernel.calculate_kernel_matrix(points)
    kernel = np.matrix.round(kernel.kernel_matrix, 4)

    assert np.allclose(kernel, np.array([[1, 0.0063, 0.2345], [0.0063, 1, 0.0101], [0.2345, 0.0101, 1]]))


def test_kernel_gaussian_1d():
    xi = np.array([1, 2, 3, 4])
    xj = np.array([0.2, -1, 3, 5])
    xk = np.array([1, 2, 3, 4])
    points = [xi, xj, xk]
    gaussian = GaussianKernel(epsilon=2.0)
    gaussian.calculate_kernel_matrix(points)

    assert np.allclose(np.matrix.round(gaussian.kernel_matrix, 4),
                       np.array([[1., 0.26447726, 1.], [0.26447726, 1., 0.26447726], [1, 0.26447726, 1]]),
                       atol=1e-04)
    assert np.round(gaussian.epsilon, 4) == 2


def test_kernel_gaussian_2d():
    xi = np.array([[-np.sqrt(2) / 2, -np.sqrt(2) / 4], [np.sqrt(2) / 2, -np.sqrt(2) / 4], [0, -np.sqrt(3) / 2]])
    xj = np.array([[0, np.sqrt(2) / 2], [1, 0], [0, -np.sqrt(2) / 2]])
    xk = np.array([[-0.69535592, -0.0546034], [-0.34016974, -0.85332868], [-0.63305978, 0.51850616]])
    points = [xi, xj, xk]
    gaussian = GaussianKernel()
    gaussian.calculate_kernel_matrix(points)

    assert np.allclose(np.matrix.round(gaussian.kernel_matrix, 4), np.array([[1., 0.39434829, 0.15306655],
                                                                             [0.39434829, 1., 0.06422136],
                                                                             [0.15306655, 0.06422136, 1.]]), atol=1e-4)
    assert np.round(gaussian.epsilon, 4) == 1.0


sol0 = np.array([[0.61415, 1.03029, 1.02001, 0.57327, 0.79874, 0.73274],
                 [0.56924, 0.91700, 0.88841, 0.53737, 0.68676, 0.67751],
                 [0.51514, 0.87898, 0.87779, 0.47850, 0.69085, 0.61525],
                 [0.63038, 1.10822, 1.12313, 0.58038, 0.89142, 0.75429],
                 [0.69666, 1.03114, 0.95037, 0.67211, 0.71184, 0.82522],
                 [0.66595, 1.03789, 0.98690, 0.63420, 0.75416, 0.79110]])

sol1 = np.array([[1.05134, 1.37652, 0.95634, 0.85630, 0.47570, 1.22488],
                 [0.16370, 0.63105, 0.14533, 0.81030, 0.44559, 0.43358],
                 [1.23478, 2.10342, 1.04698, 1.68755, 0.92792, 1.73277],
                 [0.90538, 1.64067, 0.62027, 1.17577, 0.63644, 1.34925],
                 [0.58210, 0.75795, 0.65519, 0.65712, 0.37251, 0.65740],
                 [0.99174, 1.59375, 0.63724, 0.89107, 0.47631, 1.36581]])

sol2 = np.array([[1.04142, 0.91670, 1.47962, 1.23350, 0.94111, 0.61858],
                 [1.00464, 0.65684, 1.35136, 1.11288, 0.96093, 0.42340],
                 [1.05567, 1.33192, 1.56286, 1.43412, 0.77044, 0.97182],
                 [0.89812, 0.86136, 1.20204, 1.17892, 0.83788, 0.61160],
                 [0.46935, 0.39371, 0.63534, 0.57856, 0.47615, 0.26407],
                 [1.14102, 0.80869, 1.39123, 1.33076, 0.47719, 0.68170]])

sol3 = np.array([[0.60547, 0.11492, 0.78956, 0.13796, 0.76685, 0.41661],
                 [0.32771, 0.11606, 0.67630, 0.15208, 0.44845, 0.34840],
                 [0.58959, 0.10156, 0.72623, 0.11859, 0.73671, 0.38714],
                 [0.36283, 0.07979, 0.52824, 0.09760, 0.46313, 0.27906],
                 [0.87487, 0.22452, 1.30208, 0.30189, 1.22015, 0.62918],
                 [0.56006, 0.16879, 1.09635, 0.20431, 0.69439, 0.60317]])


def test_kernel():

    np.random.seed(1111)  # For reproducibility.
    from numpy.random import RandomState
    rnd = RandomState(0)

    # Creating a list of solutions.
    Solutions = [sol0, sol1, sol2, sol3]
    manifold_projection = SVDProjection(Solutions, p="max")
    kernel = ProjectionKernel()

    kernel.calculate_kernel_matrix(manifold_projection.u)

    assert np.round(kernel.kernel_matrix[0, 1], 8) == 6.0


# tests/unit_tests/distributions/test__independent_distributions.py

# Test all distributions available in UQpy, using the cdf method or pdf method for multivariate distributions

from UQpy.distributions import *
import numpy as np


def test_beta():
    result = Beta(a=1., b=2.).cdf(x=0.8)
    assert result == 0.96


def test_cauchy():
    assert np.round(Cauchy().cdf(x=0.8), 3) == 0.715


def test_chi_square():
    assert np.round(ChiSquare(df=5.).cdf(x=0.8), 3) == 0.023


def test_exponential():
    assert np.round(Exponential().cdf(x=0.8), 3) == 0.551


def test_gamma():
    assert np.round(Gamma(a=2.).cdf(x=0.8), 3) == 0.191


def test_gen_extreme():
    assert GeneralizedExtreme(c=2.).cdf(x=0.8) == 1.


def test_inverse_gauss():
    assert np.round(InverseGauss(mu=2.).cdf(x=0.8), 3) == 0.411


def test_laplace():
    assert np.round(Laplace().cdf(x=0.8), 3) == 0.775


def test_levy():
    assert np.round(Levy().cdf(x=0.8), 3) == 0.264


def test_logistic():
    assert np.round(Logistic().cdf(x=0.8), 3) == 0.690


def test_lognormal():
    assert np.round(Lognormal(s=2.).cdf(x=0.8), 3) == 0.456


def test_maxwell():
    assert np.round(Maxwell().cdf(x=0.8), 3) == 0.113


def test_normal():
    assert np.round(Normal().cdf(x=0.8), 3) == 0.788


def test_pareto():
    assert np.round(Pareto(b=2.).cdf(x=1.1), 3) == 0.174


def test_poisson():
    assert np.round(Poisson(mu=2.).cdf(x=1.), 3) == 0.406


def test_rayleigh():
    assert np.round(Rayleigh().cdf(x=0.8), 3) == 0.274


def test_truncated_normal():
    assert np.round(TruncatedNormal(a=-1., b=1.).cdf(x=0.8), 3) == 0.922


# For multinomial, mvnormal, more tests are needed


def test_multinomial_1():
    assert Multinomial(n=5, p=[0.2, 0.3, 0.5]).pmf(x=[1, 1, 3]) == 0.15


def test_multinomial_2():
    assert np.round(Multinomial(n=5, p=[0.2, 0.3, 0.5]).log_pmf(x=[1, 1, 3]), 3) == -1.897


def test_multinomial_3():
    samples = Multinomial(n=5, p=[0.2, 0.3, 0.5]).rvs(nsamples=2, random_state=123)
    assert np.all(samples == np.array([[1, 1, 3], [0, 2, 3]]))


def test_multinomial_4():
    multinomial = Multinomial(n=5, p=[0.2, 0.3, 0.5])
    moments = multinomial.moments(moments2return='m')
    assert np.all(moments == [1., 1.5, 2.5])


def test_multinomial_5():
    cov = Multinomial(n=5, p=[0.2, 0.3, 0.5]).moments(moments2return='v')
    assert np.all(np.round(cov, 2) == np.array([[0.80, -0.30, -0.50], [-0.30, 1.05, -0.75], [-0.50, -0.75, 1.25]]))


def test_multinomial_6():
    moments = Multinomial(n=5, p=[0.2, 0.3, 0.5]).moments(moments2return='mv')
    true_values = (np.array([1., 1.5, 2.5]),
                   np.array([[0.80, -0.30, -0.50], [-0.30, 1.05, -0.75], [-0.50, -0.75, 1.25]]))
    assert np.all(moments[0] == true_values[0]) and np.all(np.round(moments[1], 2) == true_values[1])


def test_mvnormal_1():
    assert np.round(MultivariateNormal(mean=[1., 2.], cov=3.).cdf(x=[0.8, 0.8]), 3) == 0.111


def test_mvnormal_2():
    assert np.round(MultivariateNormal(mean=[1., 2.], cov=3.).pdf(x=[0.8, 0.8]), 3) == 0.041


def test_mvnormal_3():
    assert np.round(MultivariateNormal(mean=[1., 2.], cov=3.).log_pdf(x=[0.8, 0.8]), 3) == -3.183


def test_mvnormal_4():
    data = np.array([[0., 0.9], [0.1, 1.], [-0.1, 1.1]])
    true_mean = np.array([0., 1.])
    true_cov = np.array([[0.010, -0.005], [-0.005, 0.010]])
    dict_fit = MultivariateNormal(mean=None, cov=None).fit(data=data)
    assert np.all(dict_fit['mean'] == true_mean) and np.all(np.round(dict_fit['cov'], 3) == true_cov)


def test_mvnormal_5():
    samples = MultivariateNormal(mean=[1., 2.], cov=1.).rvs(nsamples=3, random_state=123)
    assert np.all(np.round(samples, 3) == np.array([[-0.086, 2.997], [1.283, 0.494], [0.421, 3.651]]))


def test_mvnormal_6():
    assert np.all(MultivariateNormal(mean=[1., 2.], cov=3.).moments(moments2return='m') == [1., 2.])


def test_mvnormal_7():
    assert np.all(MultivariateNormal(mean=[1., 2.], cov=3.).moments(moments2return='v') == 3.)


def test_mvnormal_8():
    moments = MultivariateNormal(mean=[1., 2.], cov=3.).moments(moments2return='mv')
    assert np.all(moments[0] == [1., 2.]) and moments[1] == 3.


# Check copulas
unif = np.array([0.4, 0.9]).reshape((1, 2))


def test_clayton():
    assert np.round(Clayton(theta=2.).evaluate_cdf(unit_uniform_samples=unif), 3) == 0.393


def test_frank():
    assert np.round(Frank(theta=2.).evaluate_cdf(unit_uniform_samples=unif), 3) == 0.379


def test_gumbel_1():
    assert np.round(Gumbel(theta=2.).evaluate_cdf(unit_uniform_samples=unif), 3) == 0.398


def test_gumbel_2():
    assert np.round(Gumbel(theta=2.).evaluate_pdf(unit_uniform_samples=unif), 3) == 0.261


# Check JointInd and JointCopula

marginals = [Normal(loc=2., scale=2.), Lognormal(s=1., loc=0., scale=np.exp(1))]
dist_joint = JointIndependent(marginals=marginals)
dist_joint_copula = JointCopula(marginals=marginals, copula=Gumbel(theta=2.))


def test_joint_ind_1():
    marginals_ = [Normal(loc=2., scale=2.), Lognormal(s=1., loc=0., scale=np.exp(1))]
    dist_joint_ = JointIndependent(marginals=marginals_)
    dist_joint_.update_parameters(loc_0=3.)
    assert dist_joint_.get_parameters()['loc_0'] == 3.


def test_joint_ind_2():
    samples = dist_joint.rvs(nsamples=1, random_state=123)
    assert np.all(np.round(samples, 3) == [[-0.171, 0.918]])


def test_joint_ind_3():
    x = np.array([0.5, 0.5]).reshape((1, 2))
    assert np.round(dist_joint.pdf(x=x), 3) == 0.029


def test_joint_ind_4():
    x = np.array([0.5, 0.5]).reshape((1, 2))
    assert np.round(dist_joint.log_pdf(x=x), 3) == -3.553


def test_joint_ind_5():
    x = np.array([0.5, 0.5]).reshape((1, 2))
    assert np.round(dist_joint.cdf(x=x), 3) == 0.010


def test_joint_ind_6():
    assert np.all(np.round(dist_joint.moments(moments2return='m'), 3) == [2., 4.482])


def test_joint_ind_7():
    marginals_ = [Normal(loc=None, scale=2.), Lognormal(s=1., loc=0., scale=np.exp(1))]
    dist_joint_ = JointIndependent(marginals=marginals_)
    data = np.array([[-0.17126121, 0.91793325], [3.99469089, 7.36946747], [2.565957, 3.60736828]])
    mle_fit = dist_joint_.fit(data=data)
    assert np.round(mle_fit['loc_0'], 3) == 2.130


def test_joint_copula_1():
    marginals_ = [Normal(loc=2., scale=2.), Lognormal(s=1., loc=0., scale=np.exp(1))]
    dist_joint_ = JointCopula(marginals=marginals_, copula=Gumbel(theta=3.))
    dist_joint_.update_parameters(theta_c=2.)
    assert dist_joint_.get_parameters()['theta_c'] == 2.


def test_joint_copula_3():
    x = np.array([0.5, 0.5]).reshape((1, 2))
    assert np.round(dist_joint_copula.pdf(x=x), 3) == 0.045


def test_joint_copula_4():
    x = np.array([0.5, 0.5]).reshape((1, 2))
    assert np.round(dist_joint_copula.log_pdf(x=x), 3) == -3.092


def test_joint_copula_5():
    x = np.array([0.5, 0.5]).reshape((1, 2))
    assert np.round(dist_joint_copula.cdf(x=x), 3) == 0.032


# tests/unit_tests/distributions/test_distribution_methods.py

from UQpy.distributions import *
import numpy as np

# Test all functions for one type of continuous distribution: uniform
dist_continuous = Uniform(loc=1., scale=2.)


def test_get_params():
    assert dist_continuous.get_parameters()['loc'] == 1.


def test_update_params():
    dist = Uniform(loc=1., scale=2.)
    dist.update_parameters(loc=2.)
    assert dist.get_parameters()['loc'] == 2.


def test_continuous_pdf():
    assert dist_continuous.pdf(x=1.5) == 0.5


def test_continuous_cdf():
    assert dist_continuous.cdf(x=1.5) == 0.25


def test_continuous_log_pdf():
    assert np.round(dist_continuous.log_pdf(x=1.5), 3) == -0.693


def test_continuous_icdf():
    assert dist_continuous.icdf(x=0.9) == 2.8


def test_continuous_rvs():
    samples = dist_continuous.rvs(nsamples=2, random_state=123)
    assert np.all(np.round(samples, 3) == np.array([2.393, 1.572]).reshape((2, 1)))


def test_continuous_fit():
    dict_fit = Uniform(loc=None, scale=None).fit(data=[1.5, 2.5, 3.5])
    assert dict_fit == {'loc': 1.5, 'scale': 2.0}


def test_continuous_moments():
    assert dist_continuous.moments(moments2return='m') == 2.


# Test all functions for one type of discrete distribution: binomial
dist_discrete = Binomial(n=5, p=0.2)


def test_discrete_pmf():
    assert np.round(dist_discrete.pmf(x=2.), 3) == 0.205


def test_discrete_cdf():
    assert np.round(dist_discrete.cdf(x=2.), 3) == 0.942


def test_discrete_log_pmf():
    assert np.round(dist_discrete.log_pmf(x=2.), 3) == -1.586


def test_discrete_icdf():
    assert dist_discrete.icdf(0.9) == 2.


def test_discrete_rvs():
    samples = dist_discrete.rvs(nsamples=2, random_state=123)
    assert np.all(np.round(samples, 3) == np.array([1., 0.]).reshape((2, 1)))


def test_discrete_moments():
    assert dist_discrete.moments(moments2return='m') == 1.


# Test functions for Copula


def test_update_params_copula():
    copula = Gumbel(theta=2.)
    copula.update_parameters(theta=1.)
    assert copula.get_parameters()['theta'] == 1.


# tests/unit_tests/inference/pfn_cubic.py

import numpy as np

def model_cubic(theta):
    domain = np.linspace(0, 10, 50)
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain + inpt[1] * domain ** 2 + inpt[2] * domain ** 3

# tests/unit_tests/inference/pfn_linear.py

import numpy as np

def model_linear(theta):
    domain = np.linspace(0, 10, 50)
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain


# tests/unit_tests/inference/pfn_models.py

import numpy as np

domain = np.linspace(0, 10, 4)

def model_quadratic1(theta):
    # this one takes one parameter vector theta and return one qoi
    domain = np.linspace(0, 10, 50)
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain + inpt[1] * domain ** 2

def model_quadratic(theta):
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain + inpt[1] * domain ** 2


def model_quadratic_vectorized(theta):
    # this one takes several parameter vector theta - vectorized computations to return appropriate qois
    inpts = np.tile(theta[..., np.newaxis], [1, 1, 50])
    return inpts[:, 0, :] * domain + inpts[:, 1, :] * domain ** 2


def model_linear(theta):
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain


def model_cubic(theta):
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain + inpt[1] * domain ** 2 + inpt[2] * domain ** 3

# tests/unit_tests/inference/pfn_quadratic.py

import numpy as np

def model_quadratic(theta):
    domain = np.linspace(0, 10, 50)
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain + inpt[1] * domain ** 2

# tests/unit_tests/inference/test_bayes_model_selection.py

from UQpy.run_model import PythonModel, RunModel
from UQpy.inference import *
from UQpy.inference.inference_models.ComputationalModel import ComputationalModel
from UQpy.sampling.mcmc import MetropolisHastings
from UQpy.distributions.collection.Normal import Normal
from UQpy.distributions.collection.JointIndependent import JointIndependent
import shutil

import os

dir_path = os.path.dirname(os.path.realpath(__file__))
# print(dir_path)
# print(os.getcwd())
# os.chdir(dir_path)
# print(os.getcwd())


# os.chdir("~/test/unit_tests/inference")


def test_models():
    a = os.getcwd()
    if "inference" not in a:
        os.chdir("../inference")
    data_ex1 = np.loadtxt('data_ex1a.txt')

    model = PythonModel(model_script='pfn_linear.py', model_object_name='model_linear', var_names=['theta_0'])
    runmodel4 = RunModel(model=model)

    model1 = PythonModel(model_script='pfn_quadratic.py', model_object_name='model_quadratic', var_names=['theta_0', 'theta_1'])
    runmodel5 = RunModel(model=model1)

    model2 = PythonModel(model_script='pfn_cubic.py', model_object_name='model_cubic',
                         var_names=['theta_0', 'theta_1', 'theta_2'])
    runmodel6 = RunModel(model=model2)

    prior1 = Normal()
    prior2 = JointIndependent(marginals=[Normal(), Normal()])
    prior3 = JointIndependent(marginals=[Normal(), Normal(), Normal()])

    model_n_params = [1, 2, 3]
    model1 = ComputationalModel(n_parameters=1, runmodel_object=runmodel4, prior=prior1,
                                error_covariance=np.ones(50), name='model_linear')
    model2 = ComputationalModel(n_parameters=2, runmodel_object=runmodel5, prior=prior2,
                                error_covariance=np.ones(50), name='model_quadratic')
    model3 = ComputationalModel(n_parameters=3, runmodel_object=runmodel6, prior=prior3,
                                error_covariance=np.ones(50), name='model_cubic')

    proposals = [Normal(0, 10),
                 JointIndependent([Normal(0, 1), Normal(0, 1)]),
                 JointIndependent([Normal(0, 1), Normal(0, 2), Normal(0.025)])]

    # sampling =
    mh1 = MetropolisHastings(args_target=(data_ex1,),
                             log_pdf_target=model1.evaluate_log_posterior,
                             jump=1, burn_length=500,
                             proposal=proposals[0], random_state=0, seed=[0.])
    mh2 = MetropolisHastings(args_target=(data_ex1,),
                             log_pdf_target=model2.evaluate_log_posterior,
                             jump=1, burn_length=500,
                             proposal=proposals[1], random_state=0, seed=[0., 0.])
    mh3 = MetropolisHastings(args_target=(data_ex1,),
                             log_pdf_target=model3.evaluate_log_posterior,
                             jump=1, burn_length=500,
                             proposal=proposals[2], random_state=0, seed=[0., 0., 0.])

    e1 = BayesParameterEstimation(inference_model=model1, data=data_ex1, sampling_class=mh1)
    e2 = BayesParameterEstimation(inference_model=model2, data=data_ex1, sampling_class=mh2)
    e3 = BayesParameterEstimation(inference_model=model3, data=data_ex1, sampling_class=mh3)

    selection = BayesModelSelection(parameter_estimators=[e1, e2, e3],
                                    prior_probabilities=[1. / 3., 1. / 3., 1. / 3.],
                                    nsamples=[2000, 2000, 2000])

    selection.sort_models()
    assert selection.probabilities[0] == 1.0
    assert selection.probabilities[1] == 0.0
    assert selection.probabilities[2] == 0.0

    assert selection.candidate_models[0].name == 'model_quadratic'
    assert selection.candidate_models[1].name == 'model_cubic'
    assert selection.candidate_models[2].name == 'model_linear'



# tests/unit_tests/inference/test_bayes_parameter_estimation.py

import numpy as np
from sklearn.neighbors import KernelDensity  # for the plots

from UQpy.distributions.collection import JointIndependent, Uniform, Lognormal, Normal
from UQpy.inference.inference_models.DistributionModel import DistributionModel
from UQpy.inference.BayesParameterEstimation import BayesParameterEstimation
from UQpy.sampling.ImportanceSampling import ImportanceSampling
from UQpy.sampling.mcmc.MetropolisHastings import MetropolisHastings


def pdf_from_kde(domain, samples1d):
    bandwidth = 1.06 * np.std(samples1d) * samples1d.size ** (-1 / 5)
    kde = KernelDensity(bandwidth=bandwidth).fit(samples1d.reshape((-1, 1)))
    log_dens = kde.score_samples(domain)
    return np.exp(log_dens)


def test_probability_model_importance_sampling():
    # Generate data from a probability model, here a Gaussian pdf, then learn its parameters,
    # mean and covariance, from this data

    np.random.seed(100)
    mu, sigma = 10, 1  # true mean and standard deviation
    np.random.seed(1)
    data = np.random.normal(mu, sigma, 100).reshape((-1, 1))

    p0 = Uniform(loc=0., scale=15)
    p1 = Lognormal(s=1., loc=0., scale=1.)
    prior = JointIndependent(marginals=[p0, p1])

    # create an instance of class Model
    candidate_model = DistributionModel(distributions=Normal(loc=None, scale=None),
                                        n_parameters=2, prior=prior)

    sampling = ImportanceSampling(random_state=1)

    bayes_estimator = BayesParameterEstimation(sampling_class=sampling,
                                               inference_model=candidate_model,
                                               data=data,
                                               nsamples=10000)
    bayes_estimator.sampler.resample()
    s_posterior = bayes_estimator.sampler.unweighted_samples

    assert s_posterior[0, 1] == 0.8616126410951304
    assert s_posterior[9999, 0] == 10.02449120238032


def test_probability_model_mcmc():
    np.random.seed(100)
    mu, sigma = 10, 1  # true mean and standard deviation
    np.random.seed(1)
    data = np.random.normal(mu, sigma, 100).reshape((-1, 1))

    p0 = Uniform(loc=0., scale=15)
    p1 = Lognormal(s=1., loc=0., scale=1.)
    prior = JointIndependent(marginals=[p0, p1])

    # create an instance of class Model
    candidate_model = DistributionModel(distributions=Normal(loc=None, scale=None),
                                        n_parameters=2, prior=prior)

    sampling = MetropolisHastings(jump=10, burn_length=10, seed=[1.0, 0.2], random_state=1)
    bayes_estimator = BayesParameterEstimation(sampling_class=sampling,
                                               inference_model=candidate_model,
                                               data=data,
                                               nsamples=5)
    s = bayes_estimator.sampler.samples

    assert s[0, 1] == 3.5196936384257835
    assert s[1, 0] == 11.143811671048994
    assert s[2, 0] == 10.162512455643435
    assert s[3, 1] == 0.8541521389437781
    assert s[4, 1] == 1.0095454025762525


# tests/unit_tests/inference/test_inference_distribution.py

import numpy as np
from UQpy.inference import *
from UQpy.distributions import *
from UQpy.inference.inference_models.DistributionModel import DistributionModel

# data used throughout
from UQpy.inference.information_criteria import BIC, AICc
from UQpy.sampling.mcmc import MetropolisHastings

data = [0., 1., -1.5, -0.2]
# first candidate model, 1-dimensional
prior = Lognormal(s=1., loc=0., scale=1.)
dist = Normal(loc=0., scale=None)
candidate_model = DistributionModel(n_parameters=1, distributions=dist, prior=prior)
candidate_model_no_prior = DistributionModel(n_parameters=1, distributions=dist)
# second candidate model, 2-dimensional
prior2 = JointIndependent([Uniform(loc=0., scale=0.5), Lognormal(s=1., loc=0., scale=1.)])
dist2 = Uniform(loc=None, scale=None)
candidate_model2 = DistributionModel(n_parameters=2, distributions=dist2, prior=prior2)


def test_mle():
    ml_estimator = MLE(inference_model=candidate_model, data=data, n_optimizations=3)
    assert round(ml_estimator.mle[0], 3) == 0.907


def test_info_model_selection_bic():
    mle1 = MLE(inference_model=candidate_model, data=data)
    mle2 = MLE(inference_model=candidate_model2, data=data)
    selector = InformationModelSelection(parameter_estimators=[mle1, mle2], criterion=BIC(), n_optimizations=[5]*2)
    assert round(selector.probabilities[0], 3) == 0.284


def test_info_model_selection_aic():
    mle1 = MLE(inference_model=candidate_model, data=data)
    mle2 = MLE(inference_model=candidate_model2, data=data)
    selector = InformationModelSelection(parameter_estimators=[mle1, mle2], criterion=AIC(), n_optimizations=[5]*2)
    selector.sort_models()
    assert round(selector.probabilities[0], 3) == 0.650


def test_info_model_selection_aicc():
    mle1 = MLE(inference_model=candidate_model, data=data)
    mle2 = MLE(inference_model=candidate_model2, data=data)
    selector = InformationModelSelection(parameter_estimators=[mle1, mle2], criterion=AICc(), n_optimizations=[5]*2)
    assert round(selector.probabilities[0], 3) == 0.988


def test_bayes_mcmc():
    mh1 = MetropolisHastings(args_target=(data, ),
                             log_pdf_target=candidate_model_no_prior.evaluate_log_posterior,
                             jump=2, burn_length=5, seed=[1., ], random_state=123)
    bayes_estimator = BayesParameterEstimation(sampling_class=mh1, inference_model=candidate_model_no_prior,
                                               data=data, nsamples=50)
    assert np.round(np.mean(bayes_estimator.sampler.samples), 3) == 1.275


def test_bayes_is():
    is1 = ImportanceSampling(args_target=(data, ),
                             log_pdf_target=candidate_model.evaluate_log_posterior,
                             proposal=candidate_model.prior,
                             random_state=123)
    bayes_estimator = BayesParameterEstimation(sampling_class=is1, inference_model=candidate_model,
                                               data=data, nsamples=100)
    assert np.round(np.mean(bayes_estimator.sampler.samples), 3) == 1.873


def test_bayes_selection():
    mh1 = MetropolisHastings(args_target=(data, ),
                             log_pdf_target=candidate_model.evaluate_log_posterior,
                             random_state=123, n_chains=2, dimension=1)
    mh2 = MetropolisHastings(args_target=(data, ),
                             log_pdf_target=candidate_model2.evaluate_log_posterior,
                             random_state=123, n_chains=2, dimension=1)
    parameter_estimator = BayesParameterEstimation(inference_model=candidate_model, data=data, sampling_class=mh1)
    parameter_estimator1 = BayesParameterEstimation(inference_model=candidate_model2, data=data, sampling_class=mh2)
    selection = BayesModelSelection(parameter_estimators=[parameter_estimator, parameter_estimator1], nsamples=[50, 50])
    assert round(selection.probabilities[0], 3) == 1.000


def test_bayes_selection2():
    mh1 = MetropolisHastings(args_target=(data, ),
                             log_pdf_target=candidate_model.evaluate_log_posterior,
                             random_state=123, n_chains=2, dimension=1)
    mh2 = MetropolisHastings(args_target=(data, ),
                             log_pdf_target=candidate_model2.evaluate_log_posterior,
                             random_state=123, n_chains=2, dimension=1)
    parameter_estimator = BayesParameterEstimation(inference_model=candidate_model, data=data, sampling_class=mh1)
    parameter_estimator1 = BayesParameterEstimation(inference_model=candidate_model2, data=data, sampling_class=mh2)
    selection = BayesModelSelection(parameter_estimators=[parameter_estimator, parameter_estimator1], nsamples=[50, 50])
    selection.sort_models()
    assert round(selection.probabilities[0], 3) == 1.000


# tests/unit_tests/inference/test_inference_runmodel.py

from UQpy.inference.inference_models.ComputationalModel import *
from UQpy.inference.inference_models.LogLikelihoodModel import *
from UQpy.inference.MLE import *
from UQpy.distributions import *
import pytest
from UQpy.run_model.model_execution.PythonModel import PythonModel
import numpy as np

data = [0., 1., -1.5, -0.2]


@pytest.fixture
def setup():
    model = PythonModel(model_script='pfn_models.py', model_object_name='model_quadratic',
                      var_names=['theta_0', 'theta_1'], delete_files=True)
    h_func = RunModel(model=model)
    yield h_func


def user_log_likelihood(data, model_outputs, params=None):
    return np.sum(Normal().log_pdf(np.array(data)-model_outputs[0])).reshape((-1,))


def user_log_likelihood_uniform(data, params):
    return [np.sum(Uniform(loc=p[0], scale=p[1]).log_pdf(data)) for p in params]


def test_mle(setup):
    h_func = setup
    candidate_model = ComputationalModel(n_parameters=2, runmodel_object=h_func, error_covariance=1.)
    optimizer = MinimizeOptimizer(method='nelder-mead', bounds=((-2, 2), (-2, 2)))
    ml_estimator = MLE(inference_model=candidate_model, data=data, n_optimizations=2,
                       optimizer=optimizer, random_state=123)
    assert round(ml_estimator.mle[0], 3) == -0.039


def test_mle_optimizer(setup):
    h_func = setup
    candidate_model = ComputationalModel(n_parameters=2, runmodel_object=h_func, error_covariance=np.ones(4))
    ml_estimator = MLE(inference_model=candidate_model, data=data, initial_parameters=[0., 0.], random_state=123)
    assert round(ml_estimator.mle[0], 3) == -0.039


def test_user_loglike(setup):
    h_func = setup
    candidate_model = ComputationalModel(
        n_parameters=2, runmodel_object=h_func, log_likelihood=user_log_likelihood)
    optimizer = MinimizeOptimizer(method='nelder-mead', bounds=((-2, 2), (-2, 2)))
    ml_estimator = MLE(inference_model=candidate_model, optimizer=optimizer,
                       data=data, n_optimizations=2, random_state=123)
    assert round(ml_estimator.mle[0], 3) == -0.039


def test_user_loglike_uniform():
    candidate_model = LogLikelihoodModel(n_parameters=2, log_likelihood=user_log_likelihood_uniform)
    optimizer = MinimizeOptimizer(method='nelder-mead', bounds=((-2, 2), (-2, 2)))
    ml_estimator = MLE(inference_model=candidate_model, data=data, n_optimizations=2,
                       optimizer=optimizer, random_state=123)
    assert round(ml_estimator.mle[0], 3) == 0.786


# tests/unit_tests/inference/test_info_model_selection.py

from UQpy.inference.MLE import MLE
from UQpy.inference.information_criteria import *
from UQpy.distributions.collection import Gamma, Exponential, ChiSquare
from UQpy.inference.inference_models.DistributionModel import DistributionModel
from UQpy.inference.InformationModelSelection import InformationModelSelection


def test_aic():
    data = Gamma(a=2, loc=0, scale=2).rvs(nsamples=500, random_state=12)

    m0 = DistributionModel(distributions=Gamma(a=None, loc=None, scale=None), n_parameters=3, name='gamma')
    m1 = DistributionModel(distributions=Exponential(loc=None, scale=None), n_parameters=2, name='exponential')
    m2 = DistributionModel(distributions=ChiSquare(df=None, loc=None, scale=None),
                           n_parameters=3, name='chi-square')

    candidate_models = [m0, m1, m2]
    mle1 = MLE(inference_model=m0, random_state=0, data=data)
    mle2 = MLE(inference_model=m1, random_state=0, data=data)
    mle3 = MLE(inference_model=m2, random_state=0, data=data)
    selector = InformationModelSelection(parameter_estimators=[mle1, mle2, mle3], criterion=AIC(),
                                         n_optimizations=[5]*3)
    selector.sort_models()
    assert 2285.9685816790425 == selector.criterion_values[0]
    assert 2285.9685821390594 == selector.criterion_values[1]
    assert 2368.9477307542193 == selector.criterion_values[2]


def test_bic():
    data = Gamma(a=2, loc=0, scale=2).rvs(nsamples=500, random_state=12)
    m0 = DistributionModel(distributions=Gamma(a=None, loc=None, scale=None), n_parameters=3, name='gamma')
    m1 = DistributionModel(distributions=Exponential(loc=None, scale=None), n_parameters=2, name='exponential')
    m2 = DistributionModel(distributions=ChiSquare(df=None, loc=None, scale=None),
                           n_parameters=3, name='chi-square')

    candidate_models = [m0, m1, m2]
    mle1 = MLE(inference_model=m0, random_state=0, data=data)
    mle2 = MLE(inference_model=m1, random_state=0, data=data)
    mle3 = MLE(inference_model=m2, random_state=0, data=data)
    selector = InformationModelSelection(parameter_estimators=[mle1, mle2, mle3], criterion=BIC(),
                                         n_optimizations=[5]*3)
    selector.sort_models()
    assert 0.5000000575021204 == selector.probabilities[0]
    assert 0.4999999424978796 == selector.probabilities[1]
    assert 3.939737591540338e-18 == selector.probabilities[2]


def test_aicc():
    data = Gamma(a=2, loc=0, scale=2).rvs(nsamples=500, random_state=12)
    m0 = DistributionModel(distributions=Gamma(a=None, loc=None, scale=None), n_parameters=3, name='gamma')
    m1 = DistributionModel(distributions=Exponential(loc=None, scale=None), n_parameters=2, name='exponential')
    m2 = DistributionModel(distributions=ChiSquare(df=None, loc=None, scale=None),
                           n_parameters=3, name='chi-square')

    candidate_models = [m0, m1, m2]
    mle1 = MLE(inference_model=m0, random_state=0, data=data)
    mle2 = MLE(inference_model=m1, random_state=0, data=data)
    mle3 = MLE(inference_model=m2, random_state=0, data=data)
    selector = InformationModelSelection(parameter_estimators=[mle1, mle2, mle3], criterion=AICc(),
                                         n_optimizations=[5]*3)

    selector.sort_models()
    assert 2286.0169687758166 == selector.criterion_values[0]
    assert 2286.0169692358336 == selector.criterion_values[1]
    assert 2368.9718756234347 == selector.criterion_values[2]


# tests/unit_tests/inference/test_mle.py

import numpy as np

from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from UQpy.distributions.collection import Normal
from UQpy.inference.inference_models.DistributionModel import DistributionModel
from UQpy.inference.inference_models.ComputationalModel import ComputationalModel
from UQpy.inference import MLE
from UQpy.run_model.RunModel import RunModel
import shutil


def test_simple_probability_model():
    np.random.seed(1)
    mu, sigma = 0, 0.1  # true mean and standard deviation
    data_1 = np.random.normal(mu, sigma, 1000).reshape((-1, 1))
    # set parameters to be learnt as None
    dist = Normal(loc=None, scale=None)
    candidate_model = DistributionModel(distributions=dist, n_parameters=2)

    ml_estimator = MLE(inference_model=candidate_model, data=data_1, n_optimizations=3, random_state=1)

    assert ml_estimator.mle[0] == 0.003881247615960185
    assert ml_estimator.mle[1] == 0.09810041339322118


def test_regression_model():
    param_true = np.array([1.0, 2.0]).reshape((1, -1))
    from UQpy.run_model.model_execution.PythonModel import PythonModel
    model = PythonModel(model_script='pfn_models.py', model_object_name='model_quadratic',
                      var_names=['theta_0', 'theta_1'])
    h_func = RunModel(model=model)
    h_func.run(samples=param_true)

    # Add noise
    error_covariance = 1.
    data_clean = np.array(h_func.qoi_list[0])
    noise = Normal(loc=0., scale=np.sqrt(error_covariance)).rvs(nsamples=4, random_state=1).reshape((4,))
    data_3 = data_clean + noise

    candidate_model = ComputationalModel(n_parameters=2, runmodel_object=h_func, error_covariance=error_covariance)

    optimizer = MinimizeOptimizer(method='nelder-mead')
    ml_estimator = MLE(inference_model=candidate_model, data=data_3, n_optimizations=1, random_state=1,
                       optimizer=optimizer)

    assert ml_estimator.mle[0] == 0.8689097631871134
    assert ml_estimator.mle[1] == 2.0030767805841143



# tests/unit_tests/reliability/Resonance_pfn.py

import numpy as np


class RunPythonModel:

    def __init__(self, samples=None):
        self.samples = samples
        self.qoi = [0] * self.samples.shape[0]

        self.omega = 6.
        self.epsilon = 0.0001

        for i in range(self.samples.shape[0]):
            add = self.samples[i][1] - self.samples[i][0] * (self.omega + self.epsilon) ** 2
            diff = self.samples[i][0] * (self.omega - self.epsilon) ** 2 - self.samples[i][1]
            self.qoi[i] = np.maximum(add, diff)

# tests/unit_tests/reliability/Rosenbrock.py

from UQpy.distributions import DistributionND
import numpy as np

class Rosenbrock(DistributionND):
    def __init__(self, p=20.):
        super().__init__(p=p)

    def pdf(self, x):
        return np.exp(-(100*(x[:, 1]-x[:, 0]**2)**2+(1-x[:, 0])**2) / self.parameters['p'])

    def log_pdf(self, x):
        return -(100*(x[:, 1]-x[:, 0]**2)**2+(1-x[:, 0])**2)/self.parameters['p']

# tests/unit_tests/reliability/Rosenbrock_pfn.py

class RunPythonModel:

    def __init__(self, samples=None):

        self.samples = samples
        self.qoi = [0]*self.samples.shape[0]

        for i in range(self.samples.shape[0]):
            self.qoi[i] = 120 - self.samples[i][1] - 3*self.samples[i][0]

# tests/unit_tests/reliability/pfn.py

import numpy as np













def example2(samples=None):
    import numpy as np
    d = 2
    beta = 3.0902
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        g[i] = -1 / np.sqrt(d) * (samples[i, 0] + samples[i, 1]) + beta
    return g


def example3(samples=None):
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        g[i] = 6.2 * samples[i, 0] - samples[i, 1] * samples[i, 2] ** 2
    return g


def example4(samples=None):
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        g[i] = samples[i, 0] * samples[i, 1] - 80
    return g


def RunPythonModel(samples, b_eff, d):

    qoi = list()
    for i in range(samples.shape[0]):
        qoi.append(b_eff * np.sqrt(d) - np.sum(samples[i, :]))
    return qoi



# tests/unit_tests/reliability/pfn1.py

def model_i(samples):
    resistance = samples[0, 0]
    stress = samples[0, 1]
    return resistance - stress

# tests/unit_tests/reliability/pfn2.py

import numpy as np


def model_j(samples):
    d0 = 3
    e = 30000000
    l = 100
    w = 2
    t = 4
    return d0 - 4 * l ** 3 / (e * w * t) * np.sqrt((samples[0, 1] / t ** 2) ** 2 + (samples[0, 0] / w ** 2) ** 2)

# tests/unit_tests/reliability/pfn3.py

import numpy as np


def example1(samples=None):
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        R = samples[i, 0]
        S = samples[i, 1]
        g[i] = R - S
    return g

# tests/unit_tests/reliability/pfn4.py


def model_k(samples):
    return samples[0, 0] * samples[0, 1] - 80

# tests/unit_tests/reliability/pfn5.py

import numpy as np


def Example1(samples=None):

    x = np.zeros(samples.shape[0])

    omega = 6.
    epsilon = 0.0001

    for i in range(samples.shape[0]):
        add = samples[i][1] - samples[i][0]*(omega+epsilon)**2
        diff = samples[i][0]*(omega-epsilon)**2 - samples[i][1]
        x[i] = np.maximum(add, diff)

    return x

# tests/unit_tests/reliability/test_form.py

from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.run_model.RunModel import RunModel
from UQpy.distributions import *
from UQpy.reliability import FORM
import glob
import shutil
import numpy as np
import pytest
import os


@pytest.fixture
def setup():
    model = PythonModel(model_script='pfn1.py', model_object_name='model_i', delete_files=True)
    h_func = RunModel(model=model)
    yield h_func
    # shutil.rmtree(h_func.model_dir)


def test_seeds_xu_is_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.u_record[0][0][0], [0., 0.], rtol=1e-02)


def test_seeds_x_is_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1])
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_tol1_is_not_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1], tol1=1.0e-3)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_tol2_is_not_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1], tol2=1.0e-3)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_tol3_is_not_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1], tol3=1.0e-3)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_tol12_is_not_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1], tol1=1.0e-3, tol2=1.0e-3)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_tol13_is_not_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1], tol1=1.0e-3, tol3=1.0e-3)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_tol23_is_not_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1], tol3=1.0e-3, tol2=1.0e-3)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_tol123_is_not_none(setup):
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup, seed_u=[1, 1], tol1=1.0e-3, tol3=1.0e-3, tol2=1.0e-3)
    form_obj.run()
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(form_obj.failure_probability, 0.0126, rtol=1e-02)


def test_form_example():
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    model = PythonModel(model_script='pfn3.py', model_object_name='example1', delete_files=True)
    RunModelObject = RunModel(model=model)
    dist1 = Normal(loc=200., scale=20.)
    dist2 = Normal(loc=150, scale=10.)
    Q = FORM(distributions=[dist1, dist2], runmodel_object=RunModelObject,
             tol1=1e-5, tol2=1e-5)
    Q.run()

    # print results
    np.allclose(Q.DesignPoint_U, np.array([-2., 1.]))
    np.allclose(Q.DesignPoint_X, np.array([160., 160.]))
    assert Q.beta[0] == 2.236067977499917
    assert Q.failure_probability[0] == 0.012673659338729965
    np.allclose(Q.dg_u_record, np.array([0., 0.]))



# tests/unit_tests/reliability/test_sorm.py

from UQpy.distributions import Normal
from UQpy.reliability import FORM, SORM
from UQpy.run_model.RunModel import RunModel
from UQpy.run_model.model_execution.PythonModel import PythonModel
import glob
import shutil
import numpy as np
import pytest
import os


@pytest.fixture
def setup():
    path = os.path.abspath(os.path.dirname(__file__))
    os.chdir(path)
    model = PythonModel(model_script='pfn4.py', model_object_name='model_k', delete_files=True)
    h_func = RunModel(model=model)
    yield h_func


def test_sorm(setup):
    dist1 = Normal(loc=500, scale=100)
    dist2 = Normal(loc=1000, scale=100)
    dist = [dist1, dist2]
    form_obj = FORM(distributions=dist, runmodel_object=setup)
    form_obj.run()
    sorm_obj = SORM(form_object=form_obj)
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(sorm_obj.failure_probability, 2.8803e-7, rtol=1e-02)

def test_form_obj():
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    with pytest.raises(Exception):
        assert SORM(form_object='form')






# tests/unit_tests/reliability/test_subset.py

import shutil

import numpy as np

from UQpy.run_model.RunModel import RunModel
from UQpy.distributions.collection.Normal import Normal
from UQpy.distributions.collection.JointIndependent import JointIndependent
from UQpy.sampling.MonteCarloSampling import MonteCarloSampling
from UQpy.reliability.SubsetSimulation import SubsetSimulation


def test_subset():  # Define the structural problem
    n_variables = 2
    model = 'pfn5.py'
    Example = 'Example1'

    omega = 6
    epsilon = 0.01
    mu_m = 5
    sigma_m = 1
    mu_k = 125
    sigma_k = 25
    m = np.linspace(mu_m - 3 * sigma_m, mu_m + 3 * sigma_m, 101)
    d_m = Normal(loc=mu_m, scale=sigma_m)
    d_k = Normal(loc=mu_k, scale=sigma_k)
    dist_nominal = JointIndependent(marginals=[d_m, d_k])

    from UQpy.sampling import Stretch

    n_samples_set = 1000
    p_cond = 0.1
    n_chains = int(n_samples_set * p_cond)

    mc = MonteCarloSampling(distributions=dist_nominal,
                            nsamples=n_samples_set,
                            random_state=1)

    init_sus_samples = mc.samples
    from UQpy.run_model.RunModel import RunModel
    from UQpy.run_model.model_execution.PythonModel import PythonModel

    model = PythonModel(model_script=model, model_object_name=Example)
    RunModelObject_SuS = RunModel(model=model)


    sampling = Stretch(pdf_target=dist_nominal.pdf, dimension=2, n_chains=1000, random_state=0)

    SuS_object = SubsetSimulation(sampling=sampling, runmodel_object=RunModelObject_SuS, conditional_probability=p_cond,
                                nsamples_per_subset=n_samples_set, samples_init=init_sus_samples)

    print(SuS_object.failure_probability)
    assert SuS_object.failure_probability == 3.1200000000000006e-05








# tests/unit_tests/reliability/test_taylor_series.py

from UQpy.run_model import PythonModel, RunModel
from UQpy.distributions import Normal
from UQpy.reliability import TaylorSeries
from UQpy.transformations import Nataf
import glob
import shutil
import numpy as np
import pytest
import os

path = os.path.abspath(os.path.dirname(__file__))
os.chdir(path)

def model_i(samples):
    qoi_list = [0] * samples.shape[0]
    for i in range(samples.shape[0]):
        resistance = samples[i, 0]
        stress = samples[i, 1]
        qoi_list[i] = resistance - stress
    return qoi_list


########################################################################################################################
# Tests for the derivatives function of the TaylorSeries class.


def test_derivatives_1_no_samples():
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    with pytest.raises(Exception):
        assert TaylorSeries._derivatives(nataf_object=ntf_obj, runmodel_object=model_i)


def test_derivatives_3_no_nataf():
    point_u = np.array([-2, 1])
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    with pytest.raises(Exception):
        assert TaylorSeries._derivatives(point_u=point_u, runmodel_object=model_i)

@pytest.fixture
def setup():
    model = PythonModel(model_script='pfn1.py', model_object_name='model_i', delete_files=True)
    h_func = RunModel(model=model)
    yield h_func
    # shutil.rmtree(h_func.model_dir)

def test_derivatives_4_callable_model(setup):
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    point_u = np.array([-2, 1])
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    gradient, qoi, array_of_samples = TaylorSeries._derivatives(point_u=point_u,
                                                                runmodel_object=setup,
                                                                nataf_object=ntf_obj)
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(array_of_samples[0], [160, 160], rtol=1e-09)
    np.testing.assert_allclose(gradient, [20, -10], rtol=1e-09)


def test_derivatives_5_run_model(setup):
    dist1 = Normal(loc=200, scale=20)
    dist2 = Normal(loc=150, scale=10)
    point_u = np.array([-2, 1])
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    gradient, qoi, array_of_samples = TaylorSeries._derivatives(point_u=point_u,
                                                                runmodel_object=setup,
                                                                nataf_object=ntf_obj)
    for file_name in glob.glob("Model_Runs_*"):
        shutil.rmtree(file_name)
    np.testing.assert_allclose(array_of_samples[0], [160, 160], rtol=1e-09)
    np.testing.assert_allclose(gradient, [20, -10], rtol=1e-09)


def test_derivatives_6_second():
    model = PythonModel(model_script='pfn2.py', model_object_name='model_j', delete_files=True)
    h_func = RunModel(model=model)
    dist1 = Normal(loc=500, scale=100)
    dist2 = Normal(loc=1000, scale=100)
    point_u = np.array([1.73673009, 0.16383283])
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    hessian = TaylorSeries._derivatives(point_u=point_u, runmodel_object=h_func,
                                        nataf_object=ntf_obj, order='second')
    np.testing.assert_allclose(hessian, [[-0.00720754, 0.00477726], [0.00477726, -0.00316643]], rtol=1e-04)




# tests/unit_tests/run_model/process_third_party_output.py

import numpy as np


def read_output(index):
    x = np.load("./OutputFiles/oupt_%d.npy" % index)
    return x


# tests/unit_tests/run_model/process_third_party_output_blank.py

import numpy as np



# tests/unit_tests/run_model/process_third_party_output_class.py

import numpy as np


class ReadOutput:
    def __init__(self, index):
        self.qoi = np.load("./OutputFiles/oupt_%d.npy" % index)


# tests/unit_tests/run_model/python_model.py

import numpy as np


def sum_rvs(samples=None):
    x = np.sum(samples, axis=1)
    return x


def sum_rvs_vec(samples=None):
    x = np.sum(samples, axis=2)
    return x


class SumRVs:
    def __init__(self, samples=None):

        self.qoi = np.sum(samples, axis=1)

class SumRVsVec:
    def __init__(self, samples=None):

        self.qoi = np.sum(samples, axis=2)


def det_rvs(samples=None):

    x = samples[:][0] * np.linalg.det(samples[:][1])
    return x


def det_rvs_par(samples=None):
    x = samples[0][0] * np.linalg.det(samples[0][1])
    return x


class DetRVs:
    def __init__(self, samples=None):

        self.qoi = samples[0][0] * np.linalg.det(samples[0][1])


def det_rvs_fixed(samples=None, coeff=None):

    x = coeff * np.linalg.det(samples[:])
    return x


# tests/unit_tests/run_model/python_model_blank.py

import numpy as np


# tests/unit_tests/run_model/python_model_class.py

import numpy as np


class SumRVs:
    def __init__(self, samples=None):

        self.qoi = np.sum(samples, axis=1)


# tests/unit_tests/run_model/python_model_function.py

import numpy as np


def sum_rvs(samples=None):
    x = np.sum(samples, axis=1)
    return x


# tests/unit_tests/run_model/python_model_sum_scalar.py

import fire
import os


def python(index):
    command1 = "cp ./InputFiles/sum_scalar_" + str(index) + ".py ."
    # The user will need to modify command2 to point to the Matlab application on their system.
    command2 = "python3 sum_scalar_" + str(index) + ".py"
    command3 = "mv ./OutputFiles/oupt.npy ./OutputFiles/oupt_" + str(index) + ".npy"
    command4 = "rm sum_scalar_" + str(index) + ".py"
    os.system(command1)
    os.system(command2)
    os.system(command3)
    os.system(command4)


if __name__ == '__main__':
    fire.Fire(python)


# tests/unit_tests/run_model/python_model_sum_scalar_default.py

import fire
import os


def python(index):
    command1 = "cp ./InputFiles/sum_scalar_default_" + str(index) + ".py ."
    # The user will need to modify command2 to point to the Matlab application on their system.
    command2 = "python3 sum_scalar_default_" + str(index) + ".py"
    command3 = "mv ./OutputFiles/oupt.npy ./OutputFiles/oupt_" + str(index) + ".npy"
    command4 = "rm sum_scalar_" + str(index) + ".py"
    os.system(command1)
    os.system(command2)
    os.system(command3)
    os.system(command4)


if __name__ == '__main__':
    fire.Fire(python)


# tests/unit_tests/run_model/sum_scalar.py

import numpy as np
import os

x = np.zeros(3)
x[0] = <var1>
x[1] = <var11>
x[2] = <var111>
output = sum(x)

if not os.path.isdir('OutputFiles'):
    os.mkdir('OutputFiles')

np.save('OutputFiles/oupt.npy', output)


# tests/unit_tests/run_model/sum_scalar_default.py

import numpy as np
import os

x = np.zeros(3)
x[0] = <x0>
x[1] = <x1>
x[2] = <x2>
output = sum(x)

if not os.path.isdir('OutputFiles'):
    os.mkdir('OutputFiles')

np.save('OutputFiles/oupt.npy', output)


# tests/unit_tests/run_model/test_RunModel.py

import shutil

from beartype.roar import BeartypeCallHintPepParamException

from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.run_model import ThirdPartyModel, RunModel
from UQpy.sampling import MonteCarloSampling
from UQpy.run_model.RunModel import RunModel
from UQpy.distributions import Normal
import pytest
import os
import numpy as np

d = Normal(loc=0, scale=1)
x_mcs = MonteCarloSampling(distributions=[d, d, d], nsamples=5, random_state=1234)
x_mcs_new = MonteCarloSampling(distributions=[d, d, d], nsamples=5, random_state=2345)
verbose_parameter = True


# def test_div_zero():
#     with pytest.raises(TypeError):
#         model = PythonModel(model_script='python_model.py', model_object_name='SumRVs', fmt=20,
#                             delete_files=True)
#         runmodel_object = RunModel_New(model=model)
#
#
# def test_fmt_1():
#     with pytest.raises(TypeError):
#         model = PythonModel(model_script='python_model.py', model_object_name='SumRVs', fmt=20,
#                             delete_files=True)
#
#
# def test_fmt_2():
#     with pytest.raises(ValueError):
#         model = PythonModel(model_script='python_model.py', model_object_name='SumRVs', fmt="random_string",
#                             delete_files=True)
#         runmodel_object = RunModel_New(model=model)


def test_var_names():
    with pytest.raises(BeartypeCallHintPepParamException):
        model = PythonModel(model_script='python_model.py', model_object_name='SumRVs', var_names=[20],
                            delete_files=True)
        runmodel_object = RunModel(model=model)



def test_model_script():
    with pytest.raises(ValueError):
        model = PythonModel(model_script='random_file_name', model_object_name='SumRVs', delete_files=True)
        runmodel_object = RunModel(model=model)


def test_samples():
    with pytest.raises(BeartypeCallHintPepParamException):
        model = PythonModel(model_script='python_model.py', model_object_name='SumRVs',
                         delete_files=True)
        runmodel_object = RunModel(model=model, samples="samples_string")



def test_python_serial_workflow_class_vectorized():
    model = PythonModel(model_script='python_model.py', model_object_name='SumRVs')
    model_python_serial_class = RunModel(model=model)
    model_python_serial_class.run(samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_serial_class.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))



def test_python_serial_workflow_class():
    model = PythonModel(model_script='python_model.py', model_object_name='SumRVs')
    model_python_serial_class = RunModel(model=model)
    model_python_serial_class.run(samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_serial_class.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))


def test_direct_samples():
    model = PythonModel(model_script='python_model.py', model_object_name='SumRVs')
    model_python_serial_class = RunModel(model=model, samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_serial_class.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))


def test_append_samples_true():
    model = PythonModel(model_script='python_model.py', model_object_name='SumRVs')
    model_python_serial_class = RunModel(model=model, samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_serial_class.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))
    model_python_serial_class.run(x_mcs_new.samples, append_samples=True)
    assert np.allclose(np.array(model_python_serial_class.qoi_list).flatten(),
                       np.sum(np.vstack((x_mcs.samples, x_mcs_new.samples)), axis=1))


def test_append_samples_false():
    model = PythonModel(model_script='python_model.py', model_object_name='SumRVs')
    model_python_serial_class = RunModel(model=model, samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_serial_class.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))
    model_python_serial_class.run(x_mcs_new.samples, append_samples=False)
    assert np.allclose(np.array(model_python_serial_class.qoi_list).flatten(), np.sum(x_mcs_new.samples, axis=1))


def test_python_serial_workflow_function_vectorized():
    model = PythonModel(model_script='python_model.py', model_object_name='sum_rvs')
    model_python_serial_function = RunModel(model=model)
    model_python_serial_function.run(samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_serial_function.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))


def test_python_serial_workflow_function():
    model = PythonModel(model_script='python_model.py', model_object_name='sum_rvs')
    model_python_serial_function = RunModel(model=model)
    model_python_serial_function.run(samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_serial_function.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))


# def test_python_serial_workflow_function_no_object_name():
#     model_python_serial_function = RunModel(ntasks=1, model_script='python_model_function.py', vec=False)
#     model_python_serial_function.run(samples=x_mcs.samples)
#     assert np.allclose(np.array(model_python_serial_function.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))
#     shutil.rmtree(model_python_serial_function.model_dir)


# def test_python_serial_workflow_class_no_object_name():
#     model_python_serial_function = RunModel(ntasks=1, model_script='python_model_class.py', vec=False)
#     model_python_serial_function.run(samples=x_mcs.samples)
#     assert np.allclose(np.array(model_python_serial_function.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))

@pytest.mark.skip()
def test_python_parallel_workflow_class():
    model = PythonModel(model_script='python_model.py', model_object_name='SumRVs')
    model_python_parallel_class = RunModel(model=model, samples=x_mcs.samples, ntasks=3)
    model_python_parallel_class.run(samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_parallel_class.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))
    shutil.rmtree(model_python_parallel_class.model_dir)

@pytest.mark.skip()
def test_python_parallel_workflow_function():
    model = PythonModel(model_script='python_model.py', model_object_name='sum_rvs')
    model_python_parallel_function = RunModel(model=model, ntasks=3)
    model_python_parallel_function.run(samples=x_mcs.samples)
    assert np.allclose(np.array(model_python_parallel_function.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1))
    shutil.rmtree(model_python_parallel_function.model_dir)


# def test_third_party_serial():
#     names = ['var1', 'var11', 'var111']
#     model = ThirdPartyModel(model_script='python_model_sum_scalar.py',
#                             input_template='sum_scalar.py', var_names=names, model_object_name="matlab",
#                             output_script='process_third_party_output.py', output_object_name='read_output',
#                             fmt="{:>10.4f}", delete_files=True)
#     m = RunModel_New(model=model)
#     m.run(x_mcs.samples)
#     assert np.allclose(np.array(m.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1), atol=1e-4)
#     shutil.rmtree(m.model.model_dir)


# def test_third_party_serial_output_class():
#     names = ['var1', 'var11', 'var111']
#     m = RunModel(ntasks=1, model_script='python_model_sum_scalar.py',
#                  input_template='sum_scalar.py', var_names=names, model_object_name="matlab",
#                  output_script='process_third_party_output_class.py', output_object_name='ReadOutput',
#                  resume=False, fmt="{:>10.4f}", verbose=verbose_parameter, delete_files=True)
#     m.run(x_mcs.samples)
#     assert np.allclose(np.array(m.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1), atol=1e-4)
#     shutil.rmtree(m.model_dir)


# def test_third_party_serial_no_output_class():
#     names = ['var1', 'var11', 'var111']
#     m = RunModel(ntasks=1, model_script='python_model_sum_scalar.py', input_template='sum_scalar.py', var_names=names,
#                  model_object_name="matlab", output_script='process_third_party_output_class.py', resume=False,
#                  fmt="{:>10.4f}", verbose=verbose_parameter, delete_files=True)
#     m.run(x_mcs.samples)
#     assert np.allclose(np.array(m.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1), atol=1e-4)
#     shutil.rmtree(m.model_dir)


# def test_third_party_serial_no_output_function():
#     names = ['var1', 'var11', 'var111']
#     model = ThirdPartyModel(model_script='python_model_sum_scalar.py',
#                             input_template='sum_scalar.py', var_names=names, model_object_name="matlab",
#                             output_script='process_third_party_output.py', output_object_name='read_output',
#                             fmt="{:>10.4f}", delete_files=True)
#     m = RunModel_New(model=model)
#     m.run(x_mcs.samples)
#     assert np.allclose(np.array(m.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1), atol=1e-4)
#     shutil.rmtree(m.model_dir)


@pytest.mark.skip()
def test_third_party_parallel():
    names = ['var1', 'var11', 'var111']
    model = ThirdPartyModel(model_script='python_model_sum_scalar.py', fmt="{:>10.4f}", delete_files=True,
                 input_template='sum_scalar.py', var_names=names, model_object_name="matlab",
                 output_script='process_third_party_output.py', output_object_name='read_output')
    m = RunModel(model=model, ntasks=3)
    m.run(x_mcs.samples)
    assert np.allclose(np.array(m.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1), atol=1e-4)
    shutil.rmtree(m.model.model_dir)


@pytest.mark.skip()
def test_third_party_default_var_names():
    model = ThirdPartyModel(model_script='python_model_sum_scalar.py', fmt="{:>10.4f}", delete_files=True,
                            input_template='sum_scalar.py', model_object_name="matlab",
                            output_script='process_third_party_output.py', output_object_name='read_output')
    model_third_party_default_names = RunModel(model=model, ntasks=3, samples=x_mcs.samples)
    assert np.allclose(np.array(model_third_party_default_names.qoi_list).flatten(), np.sum(x_mcs.samples, axis=1),
                       atol=1e-4)
    shutil.rmtree(model_third_party_default_names.model.model_dir)


def test_third_party_var_names():
    names = ['var1', 'var11', 'var111', 'var1111']
    with pytest.raises(TypeError):
        model = ThirdPartyModel(model_script='python_model_sum_scalar.py', fmt="{:>10.4f}", delete_files=True,
                                input_template='sum_scalar.py', model_object_name="matlab",
                                output_script='process_third_party_output.py', output_object_name='read_output')
        model_third_party_default_names = RunModel(model=model, ntasks=3, samples=x_mcs.samples)


def test_python_serial_workflow_function_object_name_error():
    with pytest.raises(TypeError):
        model = PythonModel(model_script='python_model.py')
        model = RunModel(model=model)
        model.run(x_mcs.samples)


def test_python_serial_workflow_function_wrong_object_name():
    with pytest.raises(AttributeError):
        model = PythonModel(model_script='python_model.py', model_object_name="random_model_name")
        model = RunModel(model=model)
        model.run(x_mcs.samples)


def test_python_serial_workflow_function_no_objects():
    with pytest.raises(TypeError):
        model = PythonModel(model_script='python_model_blank.py')
        model = RunModel(model=model)
        model.run(x_mcs.samples)


# tests/unit_tests/sampling/BraninHoo.py

import numpy as np


def function(z, a=1, b=5.1/(4*np.pi**2), c=5/np.pi, r=6, s=10, t=1/(8*np.pi)):
    f = a*(z[:, 1] - b*z[:, 0]**2 + c*z[:, 0] - r)**2 + s*(1 - t)*np.cos(z[:, 0]) + s + 5*z[:, 0]
    return f

# tests/unit_tests/sampling/python_model_1Dfunction.py

def y_func(z):
    import numpy as np
    return np.sin(z)


# tests/unit_tests/sampling/python_model_function.py

def y_func(z):
    return 1/(6.2727*(abs(0.3-z[:, 0]**2-z[:, 1]**2)+0.01))

# tests/unit_tests/sampling/series.py

import numpy as np


def series(z, k=7):
    t1 = 3 + 0.1 * (z[:, 1] - z[:, 0]) ** 2 - (z[:, 1] + z[:, 0]) / np.sqrt(2)
    t2 = 3 + 0.1 * (z[:, 1] - z[:, 0]) ** 2 + (z[:, 1] + z[:, 0]) / np.sqrt(2)
    t3 = z[:, 1] - z[:, 0] + k / np.sqrt(2)
    t4 = z[:, 0] - z[:, 1] + k / np.sqrt(2)
    return min([t1, t2, t3, t4])

# tests/unit_tests/sampling/test_adaptive_kriging.py

import pytest

from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

from UQpy.surrogates.kriging.Kriging import Kriging
from UQpy.sampling import MonteCarloSampling, AdaptiveKriging
from UQpy.run_model.RunModel import RunModel
from UQpy.distributions.collection import Normal
from UQpy.sampling.adaptive_kriging_functions import *
import shutil


def test_akmcs_weighted_u():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=0)
    model = PythonModel(model_script='series.py', model_object_name="series")
    rmodel = RunModel(model=model)
    regression_model = LinearRegression()
    correlation_model = ExponentialCorrelation()
    K = Kriging(regression_model=regression_model, correlation_model=correlation_model,
                optimizer=MinimizeOptimizer('l-bfgs-b'),
                optimizations_number=10, correlation_model_parameters=[1, 1], random_state=1)
    # OPTIONS: 'U', 'EFF', 'Weighted-U'
    learning_function = WeightedUFunction(weighted_u_stop=2)
    a = AdaptiveKriging(distributions=marginals, runmodel_object=rmodel, surrogate=K,
                        learning_nsamples=10**3, n_add=1, learning_function=learning_function,
                        random_state=2)
    a.run(nsamples=25, samples=x.samples)

    assert a.samples[23, 0] == 1.083176685073489
    assert a.samples[20, 1] == 0.20293978126855253



def test_akmcs_u():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    model = PythonModel(model_script='series.py', model_object_name="series")
    rmodel = RunModel(model=model)
    regression_model = LinearRegression()
    correlation_model = ExponentialCorrelation()
    K = Kriging(regression_model=regression_model, correlation_model=correlation_model,
                optimizer=MinimizeOptimizer('l-bfgs-b'),
                optimizations_number=10, correlation_model_parameters=[1, 1], random_state=0)
    # OPTIONS: 'U', 'EFF', 'Weighted-U'
    learning_function = UFunction(u_stop=2)
    a = AdaptiveKriging(distributions=marginals, runmodel_object=rmodel, surrogate=K,
                        learning_nsamples=10**3, n_add=1, learning_function=learning_function,
                        random_state=2)
    a.run(nsamples=25, samples=x.samples)

    assert a.samples[23, 0] == -4.141979058326188
    assert a.samples[20, 1] == -1.6476534435429009



def test_akmcs_expected_feasibility():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    model = PythonModel(model_script='series.py', model_object_name="series")
    rmodel = RunModel(model=model)
    regression_model = LinearRegression()
    correlation_model = ExponentialCorrelation()
    K = Kriging(regression_model=regression_model, correlation_model=correlation_model,
                optimizations_number=10, correlation_model_parameters=[1, 1],
                optimizer=MinimizeOptimizer('l-bfgs-b'),)
    # OPTIONS: 'U', 'EFF', 'Weighted-U'
    learning_function = ExpectedFeasibility(eff_a=0, eff_epsilon=2, eff_stop=0.001)
    a = AdaptiveKriging(distributions=marginals, runmodel_object=rmodel, surrogate=K,
                        learning_nsamples=10**3, n_add=1, learning_function=learning_function,
                        random_state=2)
    a.run(nsamples=25, samples=x.samples)

    assert a.samples[23, 0] == 1.366058523912817
    assert a.samples[20, 1] == -12.914668932772358



def test_akmcs_expected_improvement():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    model = PythonModel(model_script='series.py', model_object_name="series")
    rmodel = RunModel(model=model)
    regression_model = LinearRegression()
    correlation_model = ExponentialCorrelation()
    K = Kriging(regression_model=regression_model, correlation_model=correlation_model,
                optimizations_number=10, correlation_model_parameters=[1, 1],
                optimizer=MinimizeOptimizer('l-bfgs-b'),)
    # OPTIONS: 'U', 'EFF', 'Weighted-U'
    learning_function = ExpectedImprovement()
    a = AdaptiveKriging(distributions=marginals, runmodel_object=rmodel, surrogate=K,
                        learning_nsamples=10**3, n_add=1, learning_function=learning_function,
                        random_state=2)
    a.run(nsamples=25, samples=x.samples)

    assert a.samples[23, 0] == 4.553078100499578
    assert a.samples[20, 1] == -3.508949564718469



def test_akmcs_expected_improvement_global_fit():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    model = PythonModel(model_script='series.py', model_object_name="series")
    rmodel = RunModel(model=model)
    regression_model = LinearRegression()
    correlation_model = ExponentialCorrelation()
    K = Kriging(regression_model=regression_model, correlation_model=correlation_model,
                optimizations_number=10, correlation_model_parameters=[1, 1],
                optimizer=MinimizeOptimizer('l-bfgs-b'),)
    # OPTIONS: 'U', 'EFF', 'Weighted-U'
    learning_function = ExpectedImprovementGlobalFit()
    a = AdaptiveKriging(distributions=marginals, runmodel_object=rmodel, surrogate=K,
                        learning_nsamples=10**3, n_add=1, learning_function=learning_function,
                        random_state=2)
    a.run(nsamples=25, samples=x.samples)

    assert a.samples[23, 0] == 11.939859785098493
    assert a.samples[20, 1] == -8.429899469300118


def test_akmcs_samples_error():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=0)
    model = PythonModel(model_script='series.py', model_object_name="series")
    rmodel = RunModel(model=model)
    regression_model = LinearRegression()
    correlation_model = ExponentialCorrelation()
    K = Kriging(regression_model=regression_model, correlation_model=correlation_model,
                optimizer=MinimizeOptimizer('l-bfgs-b'),
                optimizations_number=10, correlation_model_parameters=[1, 1], random_state=1)
    # OPTIONS: 'U', 'EFF', 'Weighted-U'
    learning_function = WeightedUFunction(weighted_u_stop=2)
    with pytest.raises(NotImplementedError):
        a = AdaptiveKriging(distributions=[Normal(loc=0., scale=4.)]*3, runmodel_object=rmodel, surrogate=K,
                            learning_nsamples=10**3, n_add=1, learning_function=learning_function,
                            random_state=2, samples=x.samples)


def test_akmcs_u_run_from_init():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    model = PythonModel(model_script='series.py', model_object_name="series")
    rmodel = RunModel(model=model)
    regression_model = LinearRegression()
    correlation_model = ExponentialCorrelation()
    K = Kriging(regression_model=regression_model, correlation_model=correlation_model,
                optimizer=MinimizeOptimizer('l-bfgs-b'),
                optimizations_number=10, correlation_model_parameters=[1, 1], random_state=0)
    # OPTIONS: 'U', 'EFF', 'Weighted-U'
    learning_function = UFunction(u_stop=2)
    a = AdaptiveKriging(distributions=marginals, runmodel_object=rmodel, surrogate=K,
                        learning_nsamples=10**3, n_add=1, learning_function=learning_function,
                        random_state=2, nsamples=25, samples=x.samples)

    assert a.samples[23, 0] == -4.141979058326188
    assert a.samples[20, 1] == -1.6476534435429009


# tests/unit_tests/sampling/test_importance_sampling.py

import numpy as np
from UQpy.sampling import ImportanceSampling
from UQpy.distributions import JointIndependent, Uniform


def log_rosenbrock(x, param):
    return -(100 * (x[:, 1] - x[:, 0] ** 2) ** 2 + (1 - x[:, 0]) ** 2) / param


def rosenbrock(x):
    return np.exp(-(100 * (x[:, 1] - x[:, 0] ** 2) ** 2 + (1 - x[:, 0]) ** 2) / 20)


proposal = JointIndependent([Uniform(loc=-8, scale=16), Uniform(loc=-10, scale=60)])
proposal2 = JointIndependent([Uniform(loc=-8, scale=16), Uniform(loc=-10, scale=60)])
del proposal2.log_pdf


def test_pdf_target():
    w = ImportanceSampling(pdf_target=rosenbrock,
                           proposal=proposal,
                           random_state=123,
                           nsamples=2000)
    assert (w.weights.shape == (2000, ) and np.all(np.round(w.samples[-1], 3) == [-6.434, 27.373]))


def test_log_pdf_target():
    w = ImportanceSampling(log_pdf_target=log_rosenbrock,
                           args_target=(20, ),
                           proposal=proposal,
                           random_state=123,
                           nsamples=2000)
    assert (w.weights.shape == (2000, ) and np.all(np.round(w.samples[-1], 3) == [-6.434, 27.373]))


def test_resampling():
    w = ImportanceSampling(log_pdf_target=log_rosenbrock, args_target=(20,),
                           proposal=proposal, random_state=123,
                           nsamples=2000)
    w.resample(nsamples=1000)
    result = w.unweighted_samples[-1]
    assert np.all(np.round(result, 3) == [-4.912, 23.106])


def test_resampling2():
    w = ImportanceSampling(log_pdf_target=log_rosenbrock, args_target=(20,),
                           proposal=proposal, random_state=123, nsamples=2000)
    w.resample()
    assert w.unweighted_samples.shape == (2000, 2)


def test_rerun():
    w = ImportanceSampling(log_pdf_target=log_rosenbrock, args_target=(20,),
                           proposal=proposal, random_state=123)
    w.run(nsamples=1000)
    w.resample(nsamples=1000)
    w.run(nsamples=1000)
    assert w.samples.shape == (2000, 2)


def test_proposal():
    w = ImportanceSampling(log_pdf_target=log_rosenbrock, args_target=(20,),
                           proposal=proposal2, random_state=123, nsamples=2000)
    assert (w.weights.shape == (2000,) and np.all(np.round(w.samples[-1], 3) == [-6.434, 27.373]))



# tests/unit_tests/sampling/test_latin_hypercube.py

from UQpy import Centered, MaxiMin
from UQpy.sampling.stratified_sampling.LatinHypercubeSampling import *
from UQpy.distributions.collection.Uniform import *
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria.MinCorrelation import *
import numpy as np

distribution = Uniform(1, 1)
distribution1 = Uniform(1, 2)
distribution2 = Uniform(3, 4)

def test_lhs_random_criterion():
    random_criterion = Random()
    latin_hypercube_sampling = \
        LatinHypercubeSampling(distributions=distribution, nsamples=2,
                               criterion=random_criterion, random_state=1)
    actual_samples = latin_hypercube_sampling._samples.flatten()
    expected_samples = np.array([[1.208511], [1.86016225]]).flatten()
    np.testing.assert_allclose(expected_samples, actual_samples, rtol=1e-6)


def test_lhs_centered_criterion():
    centered_criterion = Centered()
    latin_hypercube_sampling = \
        LatinHypercubeSampling(distributions=distribution, nsamples=2,
                               criterion=centered_criterion, random_state=1)
    actual_samples = latin_hypercube_sampling._samples.flatten()
    expected_samples = np.array([[1.25], [1.75]]).flatten()
    np.testing.assert_allclose(expected_samples, actual_samples, rtol=1e-6)


def test_lhs_maximin_criterion():
    maximin_criterion = MaxiMin()
    latin_hypercube_sampling = \
        LatinHypercubeSampling(distributions=distribution, nsamples=2,
                               criterion=maximin_criterion, random_state=1)
    actual_samples = latin_hypercube_sampling._samples.flatten()
    expected_samples = np.array([[1.208511], [1.86016225]]).flatten()
    np.testing.assert_allclose(expected_samples, actual_samples, rtol=1e-6)

def test_lhs_maximin_criterion_joint_independent():
    maximin_criterion = MaxiMin()
    latin_hypercube_sampling = \
        LatinHypercubeSampling(distributions=JointIndependent(marginals=[distribution, distribution1]),
                               nsamples=2, criterion=maximin_criterion, random_state=1)
    actual_samples = latin_hypercube_sampling._samples.flatten()
    expected_samples = np.array([1.86016225, 1.00011437, 1.208511,   2.30233257]).flatten()
    np.testing.assert_allclose(expected_samples, actual_samples, rtol=1e-6)


from scipy.spatial.distance import pdist

dist1 = Uniform(loc=0., scale=1.)
dist2 = Uniform(loc=0., scale=1.)
joint_dist = JointIndependent(marginals=[dist1, dist2])
x1b = LatinHypercubeSampling(distributions=joint_dist, criterion=MinCorrelation(), nsamples=5, random_state=123)

x1d = LatinHypercubeSampling(distributions=dist1, criterion=Centered(), random_state=123, nsamples=5)

x1e = LatinHypercubeSampling(distributions=dist1, criterion=Centered(), nsamples=5)
x1g = LatinHypercubeSampling(distributions=[dist1, dist2], nsamples=5)


def d_func(x): return pdist(x, metric='euclidean')


x1h = LatinHypercubeSampling(distributions=[dist1, dist2], criterion=MaxiMin(), nsamples=5, random_state=789)


def test_samples4():
    """ Check the samples attribute, when dist_object is a list of distribution class object and criterion is
    'centered'."""
    expected_samples = np.array([0.1, 0.9, 0.5, 0.3, 0.7])
    assert (x1d._samples.round(2) == expected_samples).all()


def test_samples5():
    """ Check the samples attribute, when dist_object is a list of distribution class object, criterion is
    'maximin' and metric is callable."""
    expected_samples = np.array([[0.56, 0.2 ], [0.25, 0.62], [1.  , 0.4 ], [0.72, 0.91], [0.06, 0.15]])
    assert (x1h._samples.round(2) == expected_samples).all()


# tests/unit_tests/sampling/test_monte_carlo.py

import numpy as np
import pytest
from beartype.roar import BeartypeCallHintPepParamException

from UQpy.distributions import Normal, MultivariateNormal
from UQpy.sampling import MonteCarloSampling

dist1 = Normal(loc=0., scale=1.)
dist2 = Normal(loc=0., scale=1.)

x = MonteCarloSampling(distributions=dist1, nsamples=5, random_state=np.random.RandomState(123))
x.transform_u01()
y = MonteCarloSampling(distributions=[dist1, dist2])
y.run(nsamples=5, random_state=123)
y.transform_u01()

# Call run method multiple time, to cover lines where samples are append to existing ones
z1 = MonteCarloSampling(distributions=dist1, nsamples=2, random_state=123)
z1.run(nsamples=2)
z2 = MonteCarloSampling(distributions=[dist1, dist2], nsamples=2, random_state=np.random.RandomState(123))
z2.run(nsamples=2)
# Same object as z2, just to cover lines where, random_state is an integer
z3 = MonteCarloSampling(distributions=[dist1, dist2], nsamples=2, random_state=123)

z4 = MonteCarloSampling(distributions=[MultivariateNormal([0, 0])], nsamples=2,
                        random_state=np.random.RandomState(123))
z4.run(nsamples=2)
z4.transform_u01()
dist3 = Normal(loc=0., scale=1.)
del dist3.rvs
z5 = MonteCarloSampling(distributions=[dist3], random_state=np.random.RandomState(123))


def test_dist_object():
    """Validate dist_object, when dist_object is a distribution object."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions='abc', random_state=np.random.RandomState(123))


def test_dist_object2():
    """Validate dist_object, when dist_object is not a list of distribution object."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions=['abc'], random_state=np.random.RandomState(123))


def test_dist_object3():
    """Create a MCS object using DistributionND class object. Validate dist_object, when dist_object is not a
    list of distribution object."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions=['abc'], random_state=np.random.RandomState(123))


def test_distribution_nd():
    """Validate error check, when distribution object doesn't have 'rvs' method."""
    with pytest.raises(ValueError):
        z5.run(nsamples=2)


def test_samples1():
    """ Check the samples attribute, when dist_object is a distribution class object."""
    expected_samples = np.array([[-1.09], [1.], [0.28], [-1.51], [-0.58]])
    assert (x.samples.round(2) == expected_samples).all()


def test_samples2():
    """ Check the samples attribute, when dist_object is list of multiple distribution class object."""
    expected_samples = np.array([[-1.09,  1.65], [1., -2.43], [0.28, -0.43], [-1.51,  1.27], [-0.58, -0.87]])
    assert (y.samples.round(2) == expected_samples).all()


def test_samples3():
    """ Check the samples attribute, when 'run' method is called twice and samples are appended, and dist_object is a
    distribution class object.."""
    expected_samples = np.array([[-1.09], [1.], [0.28], [-1.51]])
    assert (z1.samples.round(2) == expected_samples).all()


def test_samples4():
    """ Check the samples attribute, when 'run' method is called twice and samples are appended, and dist_object is a
    list of multiple distribution class object."""
    expected_samples = np.array([[-1.09,  0.28], [1., -1.51], [-0.58, -2.43], [1.65, -0.43]])
    assert (z2.samples.round(2) == expected_samples).all()


def test_samples5():
    """ Check the samples attribute, when dist_object is a list of distributionND class object."""
    expected_samples = np.array([[[-1.09,  1.]], [[0.28, -1.51]], [[-0.58,  1.65]], [[-2.43, -0.43]]])
    assert (z4.samples.round(2) == expected_samples).all()


def test_samples_u01():
    """ Check the samplesU01 attribute, when dist_object is a distribution class object."""
    expected_samples_u01 = np.array([[0.14], [0.84], [0.61], [0.07], [0.28]])
    assert (x.samplesU01.round(2) == expected_samples_u01).all()


def test_samples2_u01():
    """ Check the samplesU01 attribute, when dist_object is list of multiple distribution class object."""
    expected_samples_u01 = np.array([[0.14, 0.95], [0.84, 0.01], [0.61, 0.33], [0.07, 0.9], [0.28, 0.19]])
    assert (y.samplesU01.round(2) == expected_samples_u01).all()


def test_samples2_u03():
    """ Check the samplesU01 attribute, when dist_object is list of distributionND class objects."""
    assert [x_.round(2)[0, 0] for x_ in z4.samplesU01] == [0.12, 0.04, 0.27, 0.0]


def test_nsamples_none():
    """Validate error check, when nsamples is None, while calling 'run' method."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions=dist1, random_state=np.random.RandomState(123)).run(nsamples=None)


def test_nsamples_not_integer():
    """Validate error check, when nsamples is not an integer, while calling 'run' method."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions=dist1, random_state=np.random.RandomState(123)).run(nsamples='abc')


def test_random_state1():
    """Check if random_state attribute is an integer, np.random.RandomState object or None, when dist_object is a
    distribution class object."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions=dist1, random_state='abc')



def test_random_state2():
    """Check if random_state attribute is not an integer, np.random.RandomState object or None, when dist_object is a
    list of multiple distribution class object."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions=[dist1, dist2], random_state='abc')


def test_run_random_state():
    """Check if random_state attribute is not an integer, np.random.RandomState object or None, when when 'run' method
    is called."""
    with pytest.raises(BeartypeCallHintPepParamException):
        MonteCarloSampling(distributions=dist1).run(nsamples=5, random_state='abc')


def test_cdf_method1():
    """Check, if dist_object attribute has 'cdf' method, when dist_object is a distribution class object."""
    del x.dist_object.cdf
    with pytest.raises(ValueError):
        x.transform_u01()


def test_cdf_method2():
    """Check, if dist_object attribute has 'cdf' method,  when dist_object is list of multiple distribution class
    object."""
    y.dist_object, y.array = [1, 1], True
    with pytest.raises(ValueError):
        y.transform_u01()


def test_cdf_method3():
    """Check, if dist_object attribute has 'cdf' method,  when dist_object is a distributionND class object."""

    z4.dist_object, z4.list = [1, 2], True
    with pytest.raises(ValueError):
        z4.transform_u01()


# tests/unit_tests/sampling/test_refined_stratified.py

import pytest
from beartype.roar import BeartypeCallHintPepParamException

from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from UQpy.sampling.stratified_sampling.refinement.GradientEnhancedRefinement import GradientEnhancedRefinement
from UQpy.distributions.collection.Uniform import Uniform
from UQpy.sampling.stratified_sampling.RefinedStratifiedSampling import *
from UQpy.sampling.stratified_sampling.refinement.RandomRefinement import *
from UQpy.sampling.stratified_sampling.strata.VoronoiStrata import *
from UQpy.run_model.RunModel import *
from UQpy.surrogates.kriging.Kriging import Kriging


def test_rss_simple_rectangular():
    marginals = [Uniform(loc=0., scale=1.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[4, 4])
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata,
                               nsamples_per_stratum=1, random_state=1)
    algorithm = RandomRefinement(strata)
    y = RefinedStratifiedSampling(stratified_sampling=x,
                                  nsamples=18,
                                  samples_per_iteration=2,
                                  refinement_algorithm=algorithm,
                                  random_state=2)
    assert y.samples[16, 0] == 0.06614276178462988
    assert y.samples[16, 1] == 0.7836449863362334
    assert y.samples[17, 0] == 0.1891972651582183
    assert y.samples[17, 1] == 0.2961099664117288


def test_rss_simple_voronoi():
    marginals = [Uniform(loc=0., scale=1.), Uniform(loc=0., scale=1.)]
    strata = VoronoiStrata(seeds_number=16, dimension=2, random_state=1)
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata,
                               nsamples_per_stratum=1, random_state=1)
    algorithm = RandomRefinement(strata)
    y = RefinedStratifiedSampling(stratified_sampling=x,
                                  nsamples=18,
                                  samples_per_iteration=2,
                                  refinement_algorithm=algorithm,
                                  random_state=2)
    assert np.round(y.samples[16, 0], 6) == 0.363793
    assert np.round(y.samples[16, 1], 6) == 0.467625
    assert np.round(y.samples[17, 0], 6) == 0.424586
    assert np.round(y.samples[17, 1], 6) == 0.217301


def test_rect_rss():
    """
    Test the 6 samples generated by RSS using rectangular stratification
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[2, 2], random_state=1)
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1, )
    y = RefinedStratifiedSampling(stratified_sampling=x, nsamples=6, samples_per_iteration=2, random_state=2,
                                  refinement_algorithm=RandomRefinement(strata=strata))
    assert np.allclose(y.samples, np.array([[0.417022, 0.36016225], [1.00011437, 0.15116629],
                                            [0.14675589, 0.5461693], [1.18626021, 0.67278036],
                                            [0.77483124, 0.7176612], [1.7101839, 0.66516741]]))
    assert np.allclose(np.array(y.samplesU01), np.array([[0.208511, 0.36016225], [0.50005719, 0.15116629],
                                                         [0.07337795, 0.5461693], [0.59313011, 0.67278036],
                                                         [0.38741562, 0.7176612], [0.85509195, 0.66516741]]))


def test_rect_gerss():
    """
    Test the 6 samples generated by GE-RSS using rectangular stratification
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[2, 2], random_state=1)
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1)
    model = PythonModel(model_script='python_model_function.py', model_object_name="y_func")
    rmodel = RunModel(model=model)
    from UQpy.surrogates.kriging.regression_models import LinearRegression
    from UQpy.surrogates.kriging.correlation_models import ExponentialCorrelation

    K = Kriging(regression_model=LinearRegression(), correlation_model=ExponentialCorrelation(), optimizations_number=20, random_state=0,
                correlation_model_parameters=[1, 1], optimizer=MinimizeOptimizer('l-bfgs-b'), )
    K.fit(samples=x.samples, values=rmodel.qoi_list)
    refinement = GradientEnhancedRefinement(strata=x.strata_object, runmodel_object=rmodel,
                                            surrogate=K, nearest_points_number=4)
    z = RefinedStratifiedSampling(stratified_sampling=x, random_state=2, refinement_algorithm=refinement)
    z.run(nsamples=6)
    assert np.allclose(z.samples, np.array([[0.417022, 0.36016225], [1.00011437, 0.15116629],
                                            [0.14675589, 0.5461693], [1.18626021, 0.67278036],
                                            [1.51296312, 0.77483124], [0.74237455, 0.66026822]]))
    # assert np.allclose(z.samples, np.array([[0.417022, 0.36016225], [1.00011437, 0.15116629],
    #                                         [0.14675589, 0.5461693], [1.18626021, 0.67278036],
    #                                         [1.59254104, 0.96577043], [1.97386531, 0.24237455]]))
    # assert np.allclose(z.samples, np.array([[0.417022, 0.36016225], [1.00011437, 0.15116629],
    #                                         [0.14675589, 0.5461693], [1.18626021, 0.67278036],
    #                                         [1.59254104, 0.96577043], [1.7176612, 0.2101839]]))
    # assert np.allclose(z.samplesU01, np.array([[0.208511, 0.36016225], [0.50005719, 0.15116629],
    #                                            [0.07337795, 0.5461693], [0.59313011, 0.67278036],
    #                                            [0.79627052, 0.96577043], [0.98693265, 0.24237455]]))
    # assert np.allclose(z.samplesU01, np.array([[0.208511, 0.36016225], [0.50005719, 0.15116629],
    #                                            [0.07337795, 0.5461693], [0.59313011, 0.67278036],
    #                                            [0.79627052, 0.96577043], [0.8588306 , 0.2101839]]))


def test_vor_rss():
    """
    Test the 6 samples generated by RSS using voronoi stratification
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata_vor = VoronoiStrata(seeds_number=4, dimension=2, random_state=10)
    x_vor = TrueStratifiedSampling(distributions=marginals, strata_object=strata_vor, nsamples_per_stratum=1, )
    y_vor = RefinedStratifiedSampling(stratified_sampling=x_vor, nsamples=6, samples_per_iteration=2,
                                      refinement_algorithm=RandomRefinement(strata=x_vor.strata_object))
    assert np.allclose(y_vor.samples, np.array([[1.78345908, 0.01640854], [1.46201137, 0.70862104],
                                                [0.4021338, 0.05290083], [0.1062376, 0.88958226],
                                                [0.61246269, 0.47160095], [0.85778034, 0.72123075]]))

    assert np.allclose(y_vor.samplesU01, np.array([[0.89172954, 0.01640854], [0.73100569, 0.70862104],
                                                   [0.2010669, 0.05290083], [0.0531188, 0.88958226],
                                                   [0.30623134, 0.47160095], [0.42889017, 0.72123075]]))


def test_vor_gerss():
    """
    Test the 6 samples generated by GE-RSS using voronoi stratification
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata_vor = VoronoiStrata(seeds_number=4, dimension=2, random_state=10)
    x_vor = TrueStratifiedSampling(distributions=marginals, strata_object=strata_vor, nsamples_per_stratum=1, )
    from UQpy.surrogates.kriging.regression_models.LinearRegression import LinearRegression
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import ExponentialCorrelation
    model = PythonModel(model_script='python_model_function.py', model_object_name="y_func")
    rmodel = RunModel(model=model)
    K_ = Kriging(regression_model=LinearRegression(), correlation_model=ExponentialCorrelation(), optimizations_number=20,
                 optimizer=MinimizeOptimizer('l-bfgs-b'), random_state=0,
                 correlation_model_parameters=[1, 1])

    K_.fit(samples=x_vor.samples, values=rmodel.qoi_list)
    z_vor = RefinedStratifiedSampling(stratified_sampling=x_vor, nsamples=6, random_state=x_vor.random_state,
                                      refinement_algorithm=GradientEnhancedRefinement(strata=x_vor.strata_object,
                                                                                      runmodel_object=rmodel,
                                                                                      surrogate=K_,
                                                                                      nearest_points_number=4))
    assert np.allclose(z_vor.samples, np.array([[1.78345908, 0.01640854], [1.46201137, 0.70862104],
                                                [0.4021338, 0.05290083], [0.1062376, 0.88958226],
                                                [0.61246269, 0.47160095], [1.16609055, 0.30832536]]))
    assert np.allclose(z_vor.samplesU01, np.array([[0.89172954, 0.01640854], [0.73100569, 0.70862104],
                                                   [0.2010669, 0.05290083], [0.0531188, 0.88958226],
                                                   [0.30623134, 0.47160095], [0.58304527, 0.30832536]]))


def test_rss_random_state():
    """
        Check 'random_state' is an integer or RandomState object.
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[2, 2])
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1, random_state=1)
    with pytest.raises(BeartypeCallHintPepParamException):
        RefinedStratifiedSampling(stratified_sampling=x, samples_number=6, samples_per_iteration=2, random_state='abc',
                                  refinement_algorithm=RandomRefinement(x.strata_object))


def test_rss_runmodel_object():
    """
        Check 'runmodel_object' should be a UQpy.RunModel class object.
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[2, 2])
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1, random_state=1)
    from UQpy.surrogates.kriging.regression_models import LinearRegression
    from UQpy.surrogates.kriging.correlation_models import ExponentialCorrelation

    K = Kriging(regression_model=LinearRegression(), correlation_model=ExponentialCorrelation(), optimizations_number=20,
                correlation_model_parameters=[1, 1], optimizer=MinimizeOptimizer('l-bfgs-b'), )
    model = PythonModel(model_script='python_model_function.py', model_object_name="y_func")
    rmodel = RunModel(model=model)
    K.fit(samples=x.samples, values=rmodel.qoi_list)
    with pytest.raises(BeartypeCallHintPepParamException):
        refinement = GradientEnhancedRefinement(strata=x.strata_object, runmodel_object='abc',
                                                surrogate=K)
        RefinedStratifiedSampling(stratified_sampling=x, samples_number=6, samples_per_iteration=2,
                                  refinement_algorithm=refinement)


def test_rss_kriging_object():
    """
        Check 'kriging_object', it should have 'fit' and 'predict' methods.
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[2, 2])
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1, random_state=1)
    model = PythonModel(model_script='python_model_function.py', model_object_name="y_func")
    rmodel = RunModel(model=model)
    with pytest.raises(NotImplementedError):
        refinement = GradientEnhancedRefinement(strata=x.strata_object, runmodel_object=rmodel,
                                                surrogate="abc")
        RefinedStratifiedSampling(stratified_sampling=x, nsamples=6, samples_per_iteration=2,
                                  refinement_algorithm=refinement)


def test_nsamples():
    """
        Check 'nsamples' attributes, it should be an integer.
    """
    marginals = [Uniform(loc=0., scale=2.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[2, 2])
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1, random_state=1)
    with pytest.raises(BeartypeCallHintPepParamException):
        RefinedStratifiedSampling(stratified_sampling=x, nsamples='a', samples_per_iteration=2,
                                  refinement_algorithm=RandomRefinement(x.strata_object))


# tests/unit_tests/sampling/test_true_stratified.py

import pytest
from beartype.roar import BeartypeCallHintPepParamException

from UQpy.distributions.collection import *
from UQpy.sampling.stratified_sampling.TrueStratifiedSampling import *
from UQpy.sampling.stratified_sampling.strata import SamplingCriterion
from UQpy.sampling.stratified_sampling.strata.VoronoiStrata import *


def test_rectangular_sts():
    marginals = [Uniform(loc=0., scale=1.), Uniform(loc=0., scale=1.)]
    strata = RectangularStrata(strata_number=[4, 4], random_state=1)
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata,
                               nsamples_per_stratum=1, )
    assert x.samples[6, 0] == 0.5511130624328794
    assert x.samples[12, 1] == 0.9736516658759619
    assert x.samples[2, 0] == 0.5366889727042783
    assert x.samples[9, 1] == 0.5495253722712197


def test_delaunay_sts():
    marginals = [Exponential(loc=1., scale=1.), Exponential(loc=1., scale=1.)]
    seeds = np.array([[0, 0], [0.4, 0.8], [1, 0], [1, 1]])
    strata_obj = DelaunayStrata(seeds=seeds, )
    sts_obj = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj, random_state=1,
                                     nsamples_per_stratum=1, )
    assert np.round(sts_obj.samples[2, 0], 8) == 1.90258174


def test_voronoi_sts():
    marginals = [Exponential(loc=1., scale=1.), Exponential(loc=1., scale=1.)]
    strata = VoronoiStrata(seeds_number=8, dimension=2, random_state=3)
    x = TrueStratifiedSampling(distributions=marginals, strata_object=strata,
                               nsamples_per_stratum=3)
    assert np.round(x.samples[7, 0], 8) == 3.69284409
    assert np.round(x.samples[20, 1], 8) == 1.15559632
    assert np.round(x.samples[1, 0], 8) == 1.8393015
    assert np.round(x.samples[15, 1], 8) == 2.11772762


# Rectangular
marginals = [Exponential(loc=1., scale=1.), Exponential(loc=1., scale=1.)]
strata = RectangularStrata(strata_number=[3, 3])

nsamples_per_stratum = [1] * 9
nsamples_per_stratum[4] = 0
x_sts = TrueStratifiedSampling(distributions=marginals, strata_object=strata,
                               nsamples_per_stratum=nsamples_per_stratum, random_state=1)
strata1 = RectangularStrata(strata_number=[3, 3], sampling_criterion=SamplingCriterion.CENTERED,
                            random_state=1)
x_sts1 = TrueStratifiedSampling(distributions=marginals, strata_object=strata1, nsamples_per_stratum=1, )

# Voronoi
strata_vor = VoronoiStrata(seeds_number=8, dimension=2, random_state=3)
sts_vor = TrueStratifiedSampling(distributions=marginals, strata_object=strata_vor)
sts_vor.run(nsamples_per_stratum=1)

strata_vor1 = VoronoiStrata(seeds_number=8, dimension=2, random_state=3)
sts_vor2 = TrueStratifiedSampling(distributions=marginals, strata_object=strata_vor1, )
sts_vor2.run(nsamples=8, nsamples_per_stratum=1)
sts_vor2.run()

# Delaunay
seeds = np.array([[0, 0], [0.4, 0.8], [1, 0], [1, 1]])
strata_del = DelaunayStrata(seeds=seeds, random_state=2)
sts_del = TrueStratifiedSampling(distributions=marginals, strata_object=strata_del,
                                 nsamples_per_stratum=2)


# Unit tests
def test_rect_random():
    """
    Test the samples generated by 3x3 rectangular strata, with random points
    """
    assert np.allclose(x_sts.samples, np.array([[1.14966929, 1.27457918], [1.4055223, 1.1062248],
                                                [2.25732188, 1.03126317], [1.0640978, 1.59515015],
                                                [2.60406483, 1.71936573], [1.15051073, 3.254492],
                                                [1.51330216, 4.20330958], [2.12638191, 3.20869262]]))
    assert np.allclose(x_sts.samplesU01, np.array([[0.13900733, 0.24010816], [0.33337146, 0.10077752],
                                                   [0.7155853, 0.03077953], [0.06208674, 0.44852024],
                                                   [0.79892249, 0.51293891], [0.1397315, 0.89507317],
                                                   [0.40148408, 0.95937248], [0.67579586, 0.89015584]]))


def test_rect_centered():
    """
    Test the samples generated by 3x3 rectangular strata, with centered points
    """
    tmp1 = (np.round(x_sts1.samples, 3) == np.array([[1.182, 1.182], [1.693, 1.182], [2.792, 1.182], [1.182, 1.693],
                                                     [1.693, 1.693], [2.792, 1.693], [1.182, 2.792], [1.693, 2.792],
                                                     [2.792, 2.792]])).all()
    tmp2 = (np.round(x_sts1.samplesU01, 3) == np.array([[0.167, 0.167], [0.5, 0.167], [0.833, 0.167], [0.167, 0.5],
                                                        [0.5, 0.5], [0.833, 0.5], [0.167, 0.833], [0.5, 0.833],
                                                        [0.833, 0.833]])).all()
    assert tmp1 and tmp2


def test_rect_sts_criterion():
    """
        Test the 'sts_criterion' attribute for RectangularSTS class.
    """
    with pytest.raises(AttributeError):
        RectangularStrata(sampling_criterion=SamplingCriterion.aaa)


def test_rect_strata_object():
    """
        Test type of strata_object. It should be a RectangularStrata object.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        TrueStratifiedSampling(distributions=marginals, strata_object=None, nsamples_per_stratum=1)


def test_rect_nsamples_check1():
    """
        In case of centered sampling, nsamples should be equal to number of strata in strata_object.
    """
    with pytest.raises(ValueError):
        strata.sampling_criterion = SamplingCriterion.CENTERED
        TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1,
                               nsamples=10)


def test_vor_nsamples_per_strarum_not_0():
    """
        Test samples and weights of the class object, when all stratum has atleast 1 sample.
    """
    assert np.allclose(sts_vor.samples, np.array([[1.5185188,  2.50692166], [1.19961165, 2.53070353],
                                                  [4.64208735, 3.15996289], [1.1938551,  1.30464598],
                                                  [1.22346669, 1.46101267], [1.00251352, 1.68941482],
                                                  [1.58543262, 1.07725168], [2.12618711, 1.66531153]]))
    assert np.allclose(sts_vor.weights, np.array([0.16067610144427008, 0.15776137825840353, 0.09704553679275912,
                                                  0.12806452403454796, 0.0316694475356838, 0.053596335374910736,
                                                  0.2494276251010699, 0.12175905145835474]))



def test_vor_strata_object():
    """
        Test type of strata_object. It should be a VoronoiStrata object.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        TrueStratifiedSampling(distributions=marginals, strata_object=None)


def test_vor_random_state():
    """
        Check 'random_state' is an integer or RandomState object.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        TrueStratifiedSampling(distributions=marginals, strata_object=strata_vor, random_state='abc')


def test_vor_dist_object():
    """
        Check 'dist_object' is a Distribution object.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        TrueStratifiedSampling(distributions=[2, 1], strata_object=strata_vor)


def test_vor_dist_object2():
    """
        Check 'dist_object' is a Distribution object.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        TrueStratifiedSampling(distributions=2, strata_object=strata_vor)


def test_voronoi_nsamples_check():
    """
        In case of centered sampling, nsamples should be equal to number of strata in strata_object.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        sts_vor.run(nsamples='abc')


def test_voronoi_nsamples_per_stratum_check():
    """
        Check length of nsamples_per_stratum should be equal to number of strata in strata_object.
    """
    with pytest.raises(ValueError):
        sts_vor.run(nsamples_per_stratum=[2, 1, 0, 1])


def test_voronoi_nsamples_per_stratum_check2():
    """
        Check nsamples_per_stratum should an integer or list.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        sts_vor.run(nsamples_per_stratum='abc')


def test_del_nsamples_per_strarum_not_0():
    """
        Test samples and weights of the class object, when all stratum has atleast 1 sample.
    """
    assert np.allclose(sts_del.samples, np.array([[1.14604405, 3.46440724], [1.10919352, 1.98225373],
                                                  [1.74510444, 1.4267083], [1.95844896, 1.14825374],
                                                  [1.87318612, 3.40220043], [1.60852035, 4.16117521],
                                                  [3.23430303, 2.49568182], [3.89471882, 2.03363426]]))
    assert np.allclose(sts_del.weights, np.array([0.1,  0.1,  0.2,  0.2,  0.05, 0.05, 0.15, 0.15]))


def test_del_strata_object():
    """
        Test type of strata_object. It should be a DelaunayStrata object.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        TrueStratifiedSampling(distributions=marginals, strata_object=None)



# tests/unit_tests/stochastic_process/test_bispectral_1d.py

from UQpy.stochastic_process import BispectralRepresentation
import numpy as np

n_sim = 100  # Num of samples
n = 1  # Num of dimensions
# Input parameters
T = 60  # Time(1 / T = dw)
nt = 256  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nf = 128  # Num of Discretized Freq.
# # Generation of Input Data(Stationary)
dt = T / nt
t = np.linspace(0, T - dt, nt)
df = F / nf
f = np.linspace(0, F - df, nf)

S = 32 * 1 / np.sqrt(2 * np.pi) * np.exp(-1 / 2 * f ** 2)
# Generating the 2 dimensional mesh grid
fx = f
fy = f
Fx, Fy = np.meshgrid(f, f)

b = 95 * 2 * 1 / (2 * np.pi) * np.exp(2 * (-1 / 2 * (Fx ** 2 + Fy ** 2)))
B_Real = b
B_Imag = b

B_Real[0, :] = 0
B_Real[:, 0] = 0
B_Imag[0, :] = 0
B_Imag[:, 0] = 0

B_Complex = B_Real + 1j * B_Imag
B_Ampl = np.absolute(B_Complex)

BSRM_object = BispectralRepresentation(n_sim, S, B_Complex, [dt], [df], [nt], [nf], random_state=128)
samples_1d = BSRM_object.samples


def test_samples_1d_shape():
    assert samples_1d.shape == (n_sim, 1, nt)


def test_samples_1d_values():
    assert (samples_1d[65, 0, 42], -4.203128646610985)

# tests/unit_tests/stochastic_process/test_inverse_translation.py

import numpy as np
from UQpy.distributions import Uniform
from UQpy.stochastic_process import SpectralRepresentation, Translation, InverseTranslation

n_sim = 100  # Num of samples
T = 100  # Time(1 / T = dw)
nt = 256  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nw = 128  # Num of Discretized Freq.
dt = T / nt
t = np.linspace(0, T - dt, nt)
dw = F / nw
w = np.linspace(0, F - dw, nw)
S = 125 / 4 * w ** 2 * np.exp(-5 * w)
SRM_object = SpectralRepresentation(n_sim, S, dt, dw, nt, nw, random_state=128)
samples = SRM_object.samples


def S_to_R(S, w, t):
    dw = w[1] - w[0]
    fac = np.ones(len(w))
    fac[1: len(w) - 1: 2] = 4
    fac[2: len(w) - 2: 2] = 2
    fac = fac * dw / 3
    R = np.zeros(len(t))
    for i in range(len(t)):
        R[i] = 2 * np.dot(fac, S * np.cos(w * t[i]))
    return R


R = S_to_R(S, w, t)
distribution = Uniform(0, 1)

Translate_object = Translation(distributions=distribution, time_interval=dt, frequency_interval=dw,
                               n_time_intervals=nt, n_frequency_intervals=nw, correlation_function_gaussian=R,
                               samples_gaussian=samples)

samples_ng = Translate_object.samples_non_gaussian
R_ng = Translate_object.scaled_correlation_function_non_gaussian

InverseTranslate_object = InverseTranslation(distributions=distribution, time_interval=dt, frequency_interval=dw,
                                             n_time_intervals=nt, n_frequency_intervals=nw,
                                             correlation_function_non_gaussian=R_ng, samples_non_gaussian=samples_ng,
                                             percentage_error=5.0)
samples_g = InverseTranslate_object.samples_gaussian
S_g = InverseTranslate_object.power_spectrum_gaussian
R_g = InverseTranslate_object.auto_correlation_function_gaussian
r_g = InverseTranslate_object.correlation_function_gaussian


def test_samples_shape():
    assert samples_g.shape == samples_ng.shape


def test_samples_g_value():
    assert np.isclose(samples_g[25, 0, 43], 0.2544126816395569)


def test_R_g_value():
    assert np.isclose(R_g[42], 0.06893298630483506)



# tests/unit_tests/stochastic_process/test_karhunen_loeve.py

from UQpy.stochastic_process import KarhunenLoeveExpansion
import numpy as np

n_sim = 100  # Num of samples
m = 400 + 1
T = 1000
dt = T / (m - 1)
t = np.linspace(0, T, m)

# Target Covariance(ACF)
R = np.zeros([m, m])
for i in range(m):
    for j in range(m):
        R[i, j] = 2 * np.exp(-((t[j] - t[i]) / 281) ** 2)

KLE_Object = KarhunenLoeveExpansion(n_samples=n_sim, correlation_function=R,
                                    time_interval=dt, random_state=128)
samples = KLE_Object.samples


def test_samples_shape():
    assert samples.shape == (n_sim, 1, len(t))


def test_samples_values():
    assert np.isclose(samples[27, 0, 246], 0.22392952712490516, rtol=0.01)


# tests/unit_tests/stochastic_process/test_spectral_1d_1v.py

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np


n_sim = 100  # Num of samples
n = 1  # Num of dimensions
m = 1  # Num of variables
T = 100  # Time(1 / T = dw)
nt = 256  # Num of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nw = 128  # Num of Discretized Freq.

# # Generation of Input Data(Stationary)
dt = T / nt
t = np.linspace(0, T - dt, nt)
dw = F / nw
w = np.linspace(0, F - dw, nw)
t_u = 2 * np.pi / 2 / F

S_1d_1v = 125 / 4 * w ** 2 * np.exp(-5 * w)
SRM_object = SpectralRepresentation(n_sim, S_1d_1v, dt, dw, nt, nw, random_state=128)
samples_1d_1v = SRM_object.samples


def test_samples_1d_1v_shape():
    assert samples_1d_1v.shape == (n_sim, 1, nt)


def test_samples_1d_1v_value():
    assert np.isclose(samples_1d_1v[53, 0, 134], -0.9143690244714813)


# tests/unit_tests/stochastic_process/test_spectral_1d_mv.py

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np

np.random.RandomState(1024)

n_sim = 100  # Num of samples
n = 1  # Num of dimensions
m = 3  # Num of variables

T = 10  # Time(1 / T = dw)
nt = 256  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nf = 128  # Num of Discretized Freq.

# # Generation of Input Data(Stationary)
dt = T / nt
t = np.linspace(0, T - dt, nt)
df = F / nf
f = np.linspace(0, F - df, nf)

t_u = 2 * np.pi / 2 / F

S_11 = 38.3 / (1 + 6.19 * f) ** (5 / 3)
S_22 = 43.4 / (1 + 6.98 * f) ** (5 / 3)
S_33 = 135 / (1 + 21.8 * f) ** (5 / 3)

g_12 = np.exp(-0.1757 * f)
g_13 = np.exp(-3.478 * f)
g_23 = np.exp(-3.392 * f)

S_list = np.array([S_11, S_22, S_33])
g_list = np.array([g_12, g_13, g_23])

# Assembly of S_jk
S_sqrt = np.sqrt(S_list)
S_jk = np.einsum('i...,j...->ij...', S_sqrt, S_sqrt)
# Assembly of g_jk
g_jk = np.zeros_like(S_jk)
counter = 0
for i in range(m):
    for j in range(i + 1, m):
        g_jk[i, j] = g_list[counter]
        counter = counter + 1
g_jk = np.einsum('ij...->ji...', g_jk) + g_jk

for i in range(m):
    g_jk[i, i] = np.ones_like(S_jk[0, 0])
S_1d_mv = S_jk * g_jk

SRM_object = SpectralRepresentation(n_sim, S_1d_mv, dt, df, nt, nf, random_state=128)
samples_1d_mv = SRM_object.samples


def test_samples_1d_mv_shape():
    assert samples_1d_mv.shape == (n_sim, m, nt)


def test_samples_1d_mv_values():
    assert np.isclose(-6.292191903354104, samples_1d_mv[43, 2, 67])


# tests/unit_tests/stochastic_process/test_spectral_nd_1v.py

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np

n_sim = 10  # Num of samples
n = 2  # Num of dimensions
m = 1  # Num of variables
T = 10
nt = 200
dt = T / nt
t = np.linspace(0, T - dt, nt)
# Frequency
W = np.array([1.0, 1.5])
nw = 100
dw = W / nw
x_list = [np.linspace(0, W[i] - dw[i], nw) for i in range(n)]
xy_list = np.array(np.meshgrid(*x_list, indexing='ij'))

S_nd_1v = 125 / 4 * np.linalg.norm(xy_list, axis=0) ** 2 * np.exp(-5 * np.linalg.norm(xy_list, axis=0))
SRM_object = SpectralRepresentation(n_sim, S_nd_1v, [dt, dt], dw, [nt, nt], [nw, nw], random_state=128)
samples_nd_1v = SRM_object.samples


def test_samples_nd_1v_shape():
    assert samples_nd_1v.shape == (n_sim, 1, nt, nt)


def test_samples_nd_1v_values():
    assert np.isclose(1.0430071116540038, samples_nd_1v[4, 0, 107, 59])

# tests/unit_tests/stochastic_process/test_spectral_nd_mv.py

from UQpy.stochastic_process import SpectralRepresentation
import numpy as np

n_sim = 10  # Number of samples
n = 2  # Number of Dimensions
m = 3  # Number of Variables

# Time
T = 10  # Simulation Time
dt = 0.1
nt = int(T / dt) + 1
t = np.linspace(0, T, nt)

# Frequency
nw = 100
W = np.array([1.5, 2.5])
dw = W / (nw - 1)

x_list = [np.linspace(dw[i], W[i], nw) for i in range(n)]
xy_list = np.array(np.meshgrid(*x_list, indexing='ij'))

S_11 = 125 / 4 * np.linalg.norm(xy_list, axis=0) ** 2 * np.exp(-5 * np.linalg.norm(xy_list, axis=0))
S_22 = 125 / 4 * np.linalg.norm(xy_list, axis=0) ** 2 * np.exp(-3 * np.linalg.norm(xy_list, axis=0))
S_33 = 125 / 4 * np.linalg.norm(xy_list, axis=0) ** 2 * np.exp(-7 * np.linalg.norm(xy_list, axis=0))

g_12 = np.exp(-0.1757 * np.linalg.norm(xy_list, axis=0))
g_13 = np.exp(-3.478 * np.linalg.norm(xy_list, axis=0))
g_23 = np.exp(-3.392 * np.linalg.norm(xy_list, axis=0))

S_list = np.array([S_11, S_22, S_33])
g_list = np.array([g_12, g_13, g_23])

# Assembly of S_jk
S_sqrt = np.sqrt(S_list)
S_jk = np.einsum('i...,j...->ij...', S_sqrt, S_sqrt)
# Assembly of g_jk
g_jk = np.zeros_like(S_jk)
counter = 0
for i in range(m):
    for j in range(i + 1, m):
        g_jk[i, j] = g_list[counter]
        counter = counter + 1
g_jk = np.einsum('ij...->ji...', g_jk) + g_jk

for i in range(m):
    g_jk[i, i] = np.ones_like(S_jk[0, 0])
S_nd_mv = S_jk * g_jk

SRM_object = SpectralRepresentation(n_sim, S_nd_mv, [dt, dt], dw, [nt, nt], [nw, nw], random_state=128)
samples_nd_mv = SRM_object.samples


def test_samples_nd_mv_shape():
    assert samples_nd_mv.shape == (n_sim, m, nt, nt)


def test_samples_nd_mv_values():
    assert np.isclose(samples_nd_mv[3, 1, 31, 79], 0.7922504882569233)


# tests/unit_tests/stochastic_process/test_translation.py

import numpy as np
from UQpy import Translation, SpectralRepresentation
from UQpy.distributions import Uniform

n_sim = 100  # Num of samples
T = 100  # Time(1 / T = dw)
nt = 256  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nw = 128  # Num of Discretized Freq.
dt = T / nt
t = np.linspace(0, T - dt, nt)
dw = F / nw
w = np.linspace(0, F - dw, nw)
S = 125 * w ** 2 * np.exp(-5 * w)
SRM_object = SpectralRepresentation(n_sim, S, dt, dw, nt, nw, random_state=128)
samples = SRM_object.samples


def S_to_R(S, w, t):
    dw = w[1] - w[0]
    fac = np.ones(len(w))
    fac[1: len(w) - 1: 2] = 4
    fac[2: len(w) - 2: 2] = 2
    fac = fac * dw / 3
    R = np.zeros(len(t))
    for i in range(len(t)):
        R[i] = 2 * np.dot(fac, S * np.cos(w * t[i]))
    return R


R_g = S_to_R(S, w, t)

distribution = Uniform(0, 1)

Translate_object = Translation(distributions=distribution, time_interval=dt, frequency_interval=dw,
                               n_time_intervals=nt, n_frequency_intervals=nw,
                               correlation_function_gaussian=R_g, samples_gaussian=samples)
samples_ng = Translate_object.samples_non_gaussian


def test_samples_ng_shape():
    assert samples_ng.shape == samples.shape


def test_samples_ng_values():
    assert np.isclose(samples_ng[18, 0, 153], 0.12719431113548574)


def test_correlation_values():
    assert np.isclose(Translate_object.correlation_function_non_gaussian[52], -0.00604468129285296)



# tests/unit_tests/surrogates/python_model_1Dfunction.py

def y_func(z):
    import numpy as np
    return np.sin(z)

# tests/unit_tests/surrogates/python_model_function.py

def y_func(z):
    return 1/(6.2727*(abs(0.3-z[:, 0]**2-z[:, 1]**2)+0.01))

# tests/unit_tests/surrogates/test_gpr.py

import pytest
from beartype.roar import BeartypeCallHintPepParamException

from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from UQpy.utilities.FminCobyla import FminCobyla
from UQpy.surrogates.gaussian_process.GaussianProcessRegression import GaussianProcessRegression
# # from UQpy.utilities.strata.Rectangular import Rectangular
# # from UQpy.sampling.StratifiedSampling import StratifiedSampling
# # from UQpy.RunModel import RunModel
# # from UQpy.distributions.collection.Uniform import Uniform
import numpy as np
import shutil
from UQpy.surrogates.gaussian_process.constraints import NonNegative
from UQpy.surrogates.gaussian_process.regression_models import LinearRegression, ConstantRegression, QuadraticRegression
from UQpy.surrogates.gaussian_process.kernels import RBF, Matern


samples = np.linspace(0, 5, 20).reshape(-1, 1)
values = np.cos(samples)
optimizer = MinimizeOptimizer(method="L-BFGS-B", bounds=[[0.1, 5], [0.1, 3]])
gpr = GaussianProcessRegression(kernel=RBF(), hyperparameters=[2.852, 2.959], random_state=1)
gpr.fit(samples=samples, values=values)

optimizer1 = MinimizeOptimizer(method="L-BFGS-B")
gpr2 = GaussianProcessRegression(kernel=Matern(nu=0.5), hyperparameters=[10, 4], optimizer=optimizer1,
                                 bounds=[[0.01, 100], [0.1, 10]], optimizations_number=10, random_state=2)
gpr2.fit(samples=samples, values=values, optimizations_number=10)

points11 = np.array([[0], [1], [2]])
points12 = np.array([[-1], [0], [-1]])
hyperparameters1 = np.array([1.5, 10])

points21 = np.array([[0, 0], [1, 1], [2, 2]])
points22 = np.array([[-1, 0], [1, 2], [0, -2]])
hyperparameters2 = np.array([3, 1, 2])
rbf_kernel = RBF()
matern_kernel_1_2 = Matern(nu=0.5)
matern_kernel_3_2 = Matern(nu=1.5)
matern_kernel_5_2 = Matern(nu=2.5)
matern_kernel_inf = Matern(nu=np.inf)
matern_kernel_2_1 = Matern(nu=2)

constant_reg = ConstantRegression()
linear_reg = LinearRegression()
quadratic_reg = QuadraticRegression()

non_negative = NonNegative(constraint_points=np.array([[1], [2], [3]]), observed_error=0.1, z_value=2)
optimizer3 = FminCobyla()
gpr3 = GaussianProcessRegression(regression_model=constant_reg, kernel=RBF(), hyperparameters=[2.852, 2.959],
                                 random_state=1, optimizer=optimizer3, optimize_constraints=non_negative,
                                 bounds=[[0.1, 5], [0.1, 3]])
gpr3.fit(samples=samples, values=values+1.05)

gpr4 = GaussianProcessRegression(kernel=RBF(), hyperparameters=[2.852, 2.959, 0.001], random_state=1,
                                 optimizer=optimizer3, bounds=[[0.1, 5], [0.1, 3], [1e-6, 1e-1]], noise=True)
gpr4.fit(samples=samples, values=values)



def test_predict1():
    prediction = np.round(gpr.predict([[1], [np.pi/2], [np.pi]], True), 3)
    expected_prediction = np.array([[0.54, 0., -1.], [0.,  0.,  0.]])
    assert (expected_prediction == prediction).all()


def test_predict2():
    prediction = np.round(gpr2.predict([[1], [2*np.pi], [np.pi]], True), 3)
    expected_prediction = np.array([[0.537,  0.238, -0.998], [0.077,  0.39,  0.046]])
    assert (expected_prediction == prediction).all()


def test_predict3():
    prediction = np.round(gpr3.predict([[1], [2*np.pi], [np.pi]], True), 3)
    expected_prediction = np.array([[1.59, 2.047, 0.05], [0., 0.006, 0.]])
    assert (expected_prediction == prediction).all()


def test_predict4():
    prediction = np.round(gpr4.predict([[1], [2 * np.pi], [np.pi]], True), 3)
    expected_prediction = np.array([[0.54,  0.983, -1.], [0.,  0.04,  0.]])
    assert np.isclose(prediction, expected_prediction, atol=0.05).all()


def test_rbf():
    """
    Test RBF kernel for 1-d input model
    """
    cov = rbf_kernel.c(points11, points12, hyperparameters1)
    covariance_matrix = np.round(cov, 3)
    expected_covariance_matrix = np.array([[80.074, 100.,  80.074], [41.111,  80.074,  41.111],
                                           [13.534,  41.111,  13.534]])
    assert (expected_covariance_matrix == covariance_matrix).all()


def test_rbf2():
    """
    Test RBF kernel for 2-d input model
    """
    cov = rbf_kernel.c(points21, points22, hyperparameters2)
    covariance_matrix2 = np.round(cov, 3)
    expected_covariance_matrix2 = np.array([[3.784e+00, 5.120e-01, 5.410e-01], [1.943e+00, 2.426e+00, 4.200e-02],
                                            [3.280e-01, 3.784e+00, 1.000e-03]])
    assert (expected_covariance_matrix2 == covariance_matrix2).all()


def test_matern_inf():
    """
    Test Matern kernel (nu=inf) for 1-d input model, should concide with RBF kernel
    """
    cov = matern_kernel_inf.c(points11, points12, hyperparameters1)
    covariance_matrix = np.round(cov, 3)
    expected_covariance_matrix = np.array([[80.074, 100.,  80.074], [41.111,  80.074,  41.111],
                                           [13.534,  41.111,  13.534]])
    assert (expected_covariance_matrix == covariance_matrix).all()


def tets_matern_inf2():
    """
    Test Matern kernel (nu=inf) for 2-d input model
    """
    cov = matern_kernel_inf.c(points21, points22, hyperparameters2)
    covariance_matrix2 = np.round(cov, 3)
    expected_covariance_matrix2 = np.array([[3.784e+00, 5.120e-01, 5.410e-01], [1.943e+00, 2.426e+00, 4.200e-02],
                                            [3.280e-01, 3.784e+00, 1.000e-03]])
    assert (expected_covariance_matrix2 == covariance_matrix2).all()


def test_matern_1_2():
    """
    Test Matern kernel (nu=0.5) for 1-d input model, should concide with absolute exponential kernel
    """
    cov = matern_kernel_1_2.c(points11, points12, hyperparameters1)
    covariance_matrix = np.round(cov, 3)
    expected_covariance_matrix = np.array([[51.342, 100.,  51.342], [26.36,  51.342,  26.36],
                                           [13.534,  26.36,  13.534]])
    assert (expected_covariance_matrix == covariance_matrix).all()


def test_matern_1_2_2():
    """
    Test Matern kernel (nu=0.5) for 2-d input model
    """
    cov = matern_kernel_1_2.c(points21, points22, hyperparameters2)
    covariance_matrix2 = np.round(cov, 3)
    expected_covariance_matrix2 = np.array([[2.866, 0.527, 0.541], [1.203, 1.472, 0.196], [0.428, 2.866, 0.069]])
    assert (expected_covariance_matrix2 == covariance_matrix2).all()


def test_matern_3_2():
    """
    Test Matern kernel (nu=1.5) for 1-d input model
    """
    cov = matern_kernel_3_2.c(points11, points12, hyperparameters1)
    covariance_matrix = np.round(cov, 3)
    expected_covariance_matrix = np.array([[67.906, 100.,  67.906], [32.869,  67.906,  32.869],
                                           [13.973,  32.869,  13.973]])
    assert (expected_covariance_matrix == covariance_matrix).all()


def test_matern_5_2():
    """
    Test Matern kernel (nu=2.5) for 1-d input model
    """
    cov = matern_kernel_5_2.c(points11, points12, hyperparameters1)
    covariance_matrix = np.round(cov, 3)
    expected_covariance_matrix = np.array([[72.776, 100.,  72.776], [35.222,  72.776,  35.222],
                                           [13.866,  35.222,  13.866]])
    assert (expected_covariance_matrix == covariance_matrix).all()


def test_matern_2_1():
    """
    Test Matern kernel (nu=2) for 1-d input model
    """
    cov = matern_kernel_2_1.c(points11, points12, hyperparameters1)
    covariance_matrix = np.round(cov, 3)
    expected_covariance_matrix = np.array([[70.894, np.nan,  70.894], [34.25, 70.894, 34.25],
                                           [13.921, 34.25, 13.921]])
    assert np.array_equal(expected_covariance_matrix, covariance_matrix, equal_nan=True)


def test_constant_regress():
    """
    Verify Constant Regression Model
    """
    tmp = constant_reg.r(np.array([[1, 0, -10]]))
    expected_tmp = np.array([[1.]])
    tmp1 = constant_reg.r(np.array([[1], [0], [-10]]))
    expected_tmp1 = np.array([[1.], [1.], [1.]])
    assert (tmp == expected_tmp).all() and (tmp1 == expected_tmp1).all()


def test_linear_regress():
    """
    Verify Linear Regression Model
    """
    tmp = linear_reg.r(np.array([[1, 0, -10]]))
    expected_tmp = np.array([[1., 1., 0., -10.]])
    tmp1 = linear_reg.r(np.array([[1], [0], [-10]]))
    expected_tmp1 = np.array([[1., 1.], [1., 0.], [1., -10]])
    assert (tmp == expected_tmp).all() and (tmp1 == expected_tmp1).all()


def test_quadratic_regress():
    """
    Verify Quadratic Regression Model
    """
    tmp = quadratic_reg.r(np.array([[1, 2, 3]]))
    expected_tmp = np.array([[1., 1., 2., 3., 1., 2., 3., 4., 6., 9.]])
    tmp1 = quadratic_reg.r(np.array([[1], [0], [-10]]))
    expected_tmp1 = np.array([[1., 1., 1], [1., 0., 0.], [1., -10, 100]])
    assert (tmp == expected_tmp).all() and (tmp1 == expected_tmp1).all()


def test_hyperparameters_length():
    """
        Raises an error if shape/length of hyperparameters is not correct.
    """
    with pytest.raises(RuntimeError):
        gpr5 = GaussianProcessRegression(regression_model=constant_reg, kernel=RBF(), hyperparameters=[2.852, 2.959],
                                         random_state=1, optimizer=optimizer3, bounds=[[0.1, 5], [0.1, 3]], noise=True)
        gpr5.fit(samples=samples, values=values)
#
#
# def test_jacobian():
#     jacobian = np.round(gpr.jacobian([[np.pi], [np.pi/2]]), 3)
#     expected_jacobian = np.array([-0., -1.])
#     assert (expected_jacobian == jacobian).all()
#
#
# def test_jacobian1():
#     jacobian = np.round(krig3.jacobian([[np.pi], [np.pi/2]]), 3)
#     expected_jacobian = np.array([0., -1.])
#     assert (expected_jacobian == jacobian).all()
#
#
# def test_regress():
#     from UQpy.surrogates.kriging.regression_models import Constant, Linear, Quadratic
#     krig.regression_model = Constant()
#     tmp = krig.regression_model.r([[0], [1]])
#     tmp_test1 = (tmp[0] == np.array([[1.], [1.]])).all() and (tmp[1] == np.array([[[0.]], [[0.]]])).all()
#
#     krig.regression_model = Linear()
#     tmp = krig.regression_model.r([[0], [1]])
#     tmp_test2 = (tmp[0] == (np.array([[1., 0.], [1., 1.]]))).all() and \
#                 (tmp[1] == np.array([[[0., 1.]], [[0., 1.]]])).all()
#
#     krig.regression_model = Quadratic()
#     tmp = krig.regression_model.r([[-1, 1], [2, -0.5]])
#     tmp_test3 = (tmp[0] == np.array([[1., -1.,  1.,  1., -1.,  1.], [1.,  2., -0.5,  4., -1.,  0.25]])).all() and \
#                 (tmp[1] == np.array([[[0., 1., 0., -2., 1., 0.],
#                                       [0., 0., 1., 0., -1., 2.]],
#                                      [[0., 1., 0., 4., -0.5, 0.],
#                                       [0., 0., 1., 0., 2., -1.]]])).all()
#
#     assert tmp_test1 and tmp_test2 and tmp_test3
#
#
# def test_corr():
#     from UQpy.surrogates.kriging.correlation_models import Exponential, Linear, Spherical, Cubic, Spline
#     krig.correlation_model = Exponential()
#     rx_exponential = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1])), 3) ==
#                       np.array([[0.135], [0.368], [1.]])).all()
#     drdt_exponential = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1]), dt=True)[1], 3) ==
#                         np.array([[[-0.271]], [[-0.368]], [[0.]]])).all()
#     drdx_exponential = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1]), dx=True)[1], 3) ==
#                         np.array([[[0.135]], [[0.368]], [[0.]]])).all()
#     expon = rx_exponential and drdt_exponential and drdx_exponential
#
#     krig.correlation_model = Linear()
#     rx_linear = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1])), 3) ==
#                  np.array([[0.], [0.], [1.]])).all()
#     drdt_linear = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3) ==
#                    np.array([[[-0.1]], [[-0.]], [[-0.1]]])).all()
#     drdx_linear = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3) ==
#                    np.array([[[1.]], [[-0.]], [[-1.]]])).all()
#     linear = rx_linear and drdt_linear and drdx_linear
#
#     krig.correlation_model = Spherical()
#     rx_spherical = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1])), 3) ==
#                     np.array([[0.], [0.], [1.]])).all()
#     drdt_spherical = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3)
#                       == np.array([[[-0.148]], [[-0.]], [[-0.148]]])).all()
#     drdx_spherical = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3)
#                       == np.array([[[1.485]], [[-0.]], [[-1.485]]])).all()
#     spherical = rx_spherical and drdt_spherical and drdx_spherical
#
#     krig.correlation_model = Cubic()
#     rx_cubic = (np.round(krig.correlation_model.c([[0.2], [0.5], [1]], [[0.5]], np.array([1])), 3) ==
#                 np.array([[0.784], [1.], [0.5]])).all()
#     drdt_cubic = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3) ==
#                   np.array([[[-0.054]], [[0.]], [[-0.054]]])).all()
#     drdx_cubic = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3) ==
#                   np.array([[[0.54]], [[0.]], [[-0.54]]])).all()
#     cubic = rx_cubic and drdt_cubic and drdx_cubic
#
#     krig.correlation_model = Spline()
#     rx_spline = (np.round(krig.correlation_model.c([[0.2], [0.5], [1]], [[0.5]], np.array([1])), 3) ==
#                  np.array([[0.429], [1.], [0.156]])).all()
#     drdt_spline = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3) ==
#                    np.array([[[-0.21]], [[0.]], [[-0.21]]])).all()
#     drdx_spline = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3) ==
#                    np.array([[[2.1]], [[0.]], [[-2.1]]])).all()
#     spline = rx_spline and drdt_spline and drdx_spline
#
#     assert expon and linear and spherical and cubic and spline
#
#
# def test_reg_model():
#     """
#         Raises an error if reg_model is not callable or a string of an in-built model.
#     """
#     with pytest.raises(BeartypeCallHintPepParamException):
#         Kriging(regression_model='A', correlation_model=Gaussian(), correlation_model_parameters=[1])
#
#
# def test_corr_model():
#     """
#         Raises an error if corr_model is not callable or a string of an in-built model.
#     """
#     with pytest.raises(BeartypeCallHintPepParamException):
#         Kriging(regression_model=Linear(), correlation_model='A', correlation_model_parameters=[1])
#
#
# def test_corr_model_params():
#     """
#         Raises an error if corr_model_params is not defined.
#     """
#     with pytest.raises(TypeError):
#         Kriging(regression_model=Linear(), correlation_model=Gaussian(), bounds=[[0.01, 5]],
#                 optimizations_number=100, random_state=1)
#
#
# def test_optimizer():
#     """
#         Raises an error if corr_model_params is not defined.
#     """
#     with pytest.raises(BeartypeCallHintPepParamException):
#         Kriging(regression_model=Linear(), correlation_model=Gaussian(),
#                 correlation_model_parameters=[1], optimizer='A')
#
#
# def test_random_state():
#     """
#         Raises an error if type of random_state is not correct.
#     """
#     with pytest.raises(BeartypeCallHintPepParamException):
#         Kriging(regression_model=Linear(), correlation_model=Gaussian(),
#                 correlation_model_parameters=[1], random_state='A')
#
#
# def test_loglikelihood():
#     prediction = np.round(krig3.log_likelihood(np.array([1.5]),
#                                                krig3.correlation_model, np.array([[1], [2]]),
#                                                np.array([[1], [1]]), np.array([[1], [2]]), return_grad=False), 3)
#     expected_prediction = 1.679
#     assert (expected_prediction == prediction).all()
#
#
# def test_loglikelihood_derivative():
#     prediction = np.round(krig3.log_likelihood(np.array([1.5]), krig3.correlation_model, np.array([[1], [2]]),
#                                                np.array([[1], [1]]), np.array([[1], [2]]), return_grad=True)[1], 3)
#     expected_prediction = np.array([-0.235])
#     assert (expected_prediction == prediction).all()
#
# def test_example():
#     from UQpy.surrogates import Kriging
#     from UQpy.utilities.strata import Rectangular
#     from UQpy.sampling import StratifiedSampling, MonteCarloSampling
#     from UQpy.RunModel import RunModel
#     from UQpy.distributions import Uniform
#     import numpy as np
#     import matplotlib.pyplot as plt
#
#     marginals = [Uniform()]
#
#     strata = Rectangular(strata_number=[10])
#
#     x = StratifiedSampling(distributions=marginals, strata_object=strata,
#                            samples_per_stratum_number=1, random_state=2)
#
#     def func(x):
#         return 1 / (1 + (10 * x) ** 4) + 0.5 * np.exp(-100 * (x - 0.5) ** 2)+0.05
#
#     samples = x.samples.copy()
#     output = np.zeros([x.samples.shape[0], 1])
#     for j in range(x.samples.shape[0]):
#         output[j, 0] = func(x.samples[j, 0])
#
#     from UQpy.surrogates.kriging.regression_models import Linear
#     from UQpy.surrogates.kriging.correlation_models import Gaussian
#
#     from UQpy.utilities.optimization.MinimizeOptimizer import MinimizeOptimizer
#     from UQpy.surrogates.kriging.constraints.Nonnegative import Nonnegative
#     optimizer = MinimizeOptimizer(method="cobyla", bounds=[[1, 1000]])
#
#     K = Kriging(regression_model=Linear(), correlation_model=Gaussian(), optimizer=optimizer,
#                 optimizations_number=20, correlation_model_parameters=[1], random_state=2,
#                 optimize_constraints=Nonnegative(np.linspace(min(x.samples), max(x.samples), 30)))
#
#     K.fit(samples=samples, values=output)
#     print(K.correlation_model_parameters)
#
#     num = 1000
#     x1 = np.linspace(min(x.samples), max(x.samples), num)
#
#     y, y_sd = K.predict(x1.reshape([num, 1]), return_std=True)
#     y_grad = K.jacobian(x1.reshape([num, 1]))
#
#     y_act = np.zeros([num, 1])
#     for i in range(num):
#         y_act[i, 0] = func(x1[i, 0])
#
#     fig = plt.figure(figsize=(10, 8))
#     ax = plt.subplot(111)
#     plt.plot(x1, y_act, label='Actual fucntion')
#     plt.plot(x1, y, label='Surrogate')
#     # plt.plot(x1, y_grad, label='Gradient')
#     plt.scatter(K.samples, K.values, label='Data')
#     plt.fill(np.concatenate([x1, x1[::-1]]), np.concatenate([y - 1.9600 * y_sd,
#                                                              (y + 1.9600 * y_sd)[::-1]]),
#              alpha=.5, fc='y', ec='None', label='95% CI')
#     box = ax.get_position()
#     ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
#
#     # Put a legend to the right of the current axis
#     ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
#     plt.grid()
#     plt.show()

# tests/unit_tests/surrogates/test_kriging.py

import pytest
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
from beartype.roar import BeartypeCallHintPepParamException

from UQpy.surrogates.kriging.Kriging import Kriging
import numpy as np
from UQpy.surrogates.kriging.regression_models import LinearRegression, ConstantRegression
from UQpy.surrogates.kriging.correlation_models import GaussianCorrelation


samples = np.linspace(0, 5, 20).reshape(-1, 1)
values = np.cos(samples)
optimizer = MinimizeOptimizer(method="L-BFGS-B")
krig = Kriging(regression_model=LinearRegression(), correlation_model=GaussianCorrelation(), optimizer=optimizer,
               correlation_model_parameters=[0.14], optimize=False, random_state=1)
krig.fit(samples=samples, values=values, correlation_model_parameters=[0.3])

optimizer = MinimizeOptimizer(method="L-BFGS-B")
krig2 = Kriging(regression_model=ConstantRegression(), correlation_model=GaussianCorrelation(), optimizer=optimizer,
                correlation_model_parameters=[0.3], bounds=[[0.01, 5]],
                optimize=False, optimizations_number=100, normalize=False,
                random_state=2)
krig2.fit(samples=samples, values=values)


# Using the in-built linear regression model as a function
linear_regression_model = Kriging(regression_model=LinearRegression(), correlation_model=GaussianCorrelation(), optimizer=optimizer,
                                  correlation_model_parameters=[1]).regression_model
optimizer = MinimizeOptimizer(method="L-BFGS-B")
gaussian_corrleation_model = Kriging(regression_model=LinearRegression(), correlation_model=GaussianCorrelation(), optimizer=optimizer,
                                     correlation_model_parameters=[1]).correlation_model

optimizer = MinimizeOptimizer(method="L-BFGS-B")
krig3 = Kriging(regression_model=linear_regression_model, correlation_model=gaussian_corrleation_model,
                optimizer=optimizer,
                correlation_model_parameters=[1], optimize=False, normalize=False, random_state=0)
krig3.fit(samples=samples, values=values)


def test_predict():
    prediction = np.round(krig.predict([[1], [np.pi/2], [np.pi]], True), 3)
    expected_prediction = np.array([[0.54,  0., -1.], [0.,  0.,  0.]])
    assert (expected_prediction == prediction).all()


def test_predict1():
    prediction = np.round(krig2.predict([[1], [2*np.pi], [np.pi]], True), 3)
    expected_prediction = np.array([[0.54,  1.009, -1.], [0.,  0.031,  0.]])
    assert (expected_prediction == prediction).all()


def test_predict2():
    prediction = np.round(krig3.predict([[1], [np.pi/2], [np.pi]]), 3)
    expected_prediction = np.array([[0.54, -0., -1.]])
    assert (expected_prediction == prediction).all()


def test_jacobian():
    jacobian = np.round(krig.jacobian([[np.pi], [np.pi/2]]), 3)
    expected_jacobian = np.array([-0., -1.])
    assert (expected_jacobian == jacobian).all()


def test_jacobian1():
    jacobian = np.round(krig3.jacobian([[np.pi], [np.pi/2]]), 3)
    expected_jacobian = np.array([0., -1.])
    assert (expected_jacobian == jacobian).all()


def test_regression_models():
    from UQpy.surrogates.kriging.regression_models import ConstantRegression, LinearRegression, QuadraticRegression
    krig.regression_model = ConstantRegression()
    tmp = krig.regression_model.r([[0], [1]])
    tmp_test1 = (tmp[0] == np.array([[1.], [1.]])).all() and (tmp[1] == np.array([[[0.]], [[0.]]])).all()

    krig.regression_model = LinearRegression()
    tmp = krig.regression_model.r([[0], [1]])
    tmp_test2 = (tmp[0] == (np.array([[1., 0.], [1., 1.]]))).all() and \
                (tmp[1] == np.array([[[0., 1.]], [[0., 1.]]])).all()

    krig.regression_model = QuadraticRegression()
    tmp = krig.regression_model.r([[-1, 1], [2, -0.5]])
    tmp_test3 = (tmp[0] == np.array([[1., -1.,  1.,  1., -1.,  1.], [1.,  2., -0.5,  4., -1.,  0.25]])).all() and \
                (tmp[1] == np.array([[[0., 1., 0., -2., 1., 0.],
                                      [0., 0., 1., 0., -1., 2.]],
                                     [[0., 1., 0., 4., -0.5, 0.],
                                      [0., 0., 1., 0., 2., -1.]]])).all()

    assert tmp_test1 and tmp_test2 and tmp_test3


def test_correlation_models():
    from UQpy.surrogates.kriging.correlation_models import ExponentialCorrelation, LinearCorrelation, SphericalCorrelation, CubicCorrelation, SplineCorrelation
    krig.correlation_model = ExponentialCorrelation()
    rx_exponential = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1])), 3) ==
                      np.array([[0.135], [0.368], [1.]])).all()
    drdt_exponential = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1]), dt=True)[1], 3) ==
                        np.array([[[-0.271]], [[-0.368]], [[0.]]])).all()
    drdx_exponential = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1]), dx=True)[1], 3) ==
                        np.array([[[0.135]], [[0.368]], [[0.]]])).all()
    expon = rx_exponential and drdt_exponential and drdx_exponential

    krig.correlation_model = LinearCorrelation()
    rx_linear = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1])), 3) ==
                 np.array([[0.], [0.], [1.]])).all()
    drdt_linear = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3) ==
                   np.array([[[-0.1]], [[-0.]], [[-0.1]]])).all()
    drdx_linear = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3) ==
                   np.array([[[1.]], [[-0.]], [[-1.]]])).all()
    linear = rx_linear and drdt_linear and drdx_linear

    krig.correlation_model = SphericalCorrelation()
    rx_spherical = (np.round(krig.correlation_model.c([[0], [1], [2]], [[2]], np.array([1])), 3) ==
                    np.array([[0.], [0.], [1.]])).all()
    drdt_spherical = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3)
                      == np.array([[[-0.148]], [[-0.]], [[-0.148]]])).all()
    drdx_spherical = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3)
                      == np.array([[[1.485]], [[-0.]], [[-1.485]]])).all()
    spherical = rx_spherical and drdt_spherical and drdx_spherical

    krig.correlation_model = CubicCorrelation()
    rx_cubic = (np.round(krig.correlation_model.c([[0.2], [0.5], [1]], [[0.5]], np.array([1])), 3) ==
                np.array([[0.784], [1.], [0.5]])).all()
    drdt_cubic = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3) ==
                  np.array([[[-0.054]], [[0.]], [[-0.054]]])).all()
    drdx_cubic = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3) ==
                  np.array([[[0.54]], [[0.]], [[-0.54]]])).all()
    cubic = rx_cubic and drdt_cubic and drdx_cubic

    krig.correlation_model = SplineCorrelation()
    rx_spline = (np.round(krig.correlation_model.c([[0.2], [0.5], [1]], [[0.5]], np.array([1])), 3) ==
                 np.array([[0.429], [1.], [0.156]])).all()
    drdt_spline = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dt=True)[1], 3) ==
                   np.array([[[-0.21]], [[0.]], [[-0.21]]])).all()
    drdx_spline = (np.round(krig.correlation_model.c([[0.4], [0.5], [0.6]], [[0.5]], np.array([1]), dx=True)[1], 3) ==
                   np.array([[[2.1]], [[0.]], [[-2.1]]])).all()
    spline = rx_spline and drdt_spline and drdx_spline

    assert expon and linear and spherical and cubic and spline


def test_wrong_regression_model():
    """
        Raises an error if reg_model is not callable or a string of an in-built model.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        Kriging(regression_model='A', correlation_model=GaussianCorrelation(), correlation_model_parameters=[1])


def test_wrong_correlation_model():
    """
        Raises an error if corr_model is not callable or a string of an in-built model.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        Kriging(regression_model=LinearRegression(), correlation_model='A', correlation_model_parameters=[1])


def test_missing_correlation_model_parameters():
    """
        Raises an error if corr_model_params is not defined.
    """
    with pytest.raises(TypeError):
        Kriging(regression_model=LinearRegression(), correlation_model=GaussianCorrelation(), bounds=[[0.01, 5]],
                optimizations_number=100, random_state=1)


def test_optimizer():
    """
        Raises an error if corr_model_params is not defined.
    """
    with pytest.raises(ValueError):
        Kriging(regression_model=LinearRegression(), correlation_model=GaussianCorrelation(),
                correlation_model_parameters=[1], optimizer='A')


def test_random_state():
    """
        Raises an error if type of random_state is not correct.
    """
    with pytest.raises(BeartypeCallHintPepParamException):
        Kriging(regression_model=LinearRegression(), correlation_model=GaussianCorrelation(),
                correlation_model_parameters=[1], random_state='A')


def test_log_likelihood():
    prediction = np.round(krig3.log_likelihood(np.array([1.5]),
                                               krig3.correlation_model, np.array([[1], [2]]),
                                               np.array([[1], [1]]), np.array([[1], [2]]), return_grad=False), 3)
    expected_prediction = 1.679
    assert (expected_prediction == prediction).all()


def test_log_likelihood_derivative():
    prediction = np.round(krig3.log_likelihood(np.array([1.5]), krig3.correlation_model, np.array([[1], [2]]),
                                               np.array([[1], [1]]), np.array([[1], [2]]), return_grad=True)[1], 3)
    expected_prediction = np.array([-0.235])
    assert (expected_prediction == prediction).all()



# tests/unit_tests/surrogates/test_pce.py

import pytest

from UQpy.distributions import JointIndependent, Normal
from UQpy.sampling import MonteCarloSampling
from UQpy.distributions import Uniform
from UQpy.sensitivity.PceSensitivity import PceSensitivity
from UQpy.surrogates import *
import numpy as np

from UQpy.surrogates.polynomial_chaos.polynomials.TotalDegreeBasis import TotalDegreeBasis
from UQpy.surrogates.polynomial_chaos.polynomials.TensorProductBasis import TensorProductBasis

np.random.seed(1)
max_degree, n_samples = 2, 10
dist = Uniform(loc=0, scale=10)


def func(x):
    return x * np.sin(x) / 10


x = dist.rvs(n_samples)
x_test = dist.rvs(n_samples)
y = func(x)


# Unit tests
def test_1():
    """
    Test td basis
    """
    polynomials = TotalDegreeBasis(dist, max_degree).polynomials
    value = polynomials[1].evaluate(x)[0]
    assert round(value, 4) == -0.2874


def test_2():
    """
    Test tp basis
    """
    polynomial_basis = TensorProductBasis(distributions=dist,
                                                                   max_degree=max_degree).polynomials
    value = polynomial_basis[1].evaluate(x)[0]
    assert round(value, 4) == -0.2874


def test_3():
    """
    Test PCE coefficients w/ lasso
    """
    polynomial_basis = TensorProductBasis(dist, max_degree)
    lasso = LassoRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=lasso)
    pce.fit(x, y)
    assert round(pce.coefficients[0][0], 4) == 0.0004


#
def test_4():
    """
    Test PCE coefficients w/ ridge
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    ridge = RidgeRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=ridge)
    pce.fit(x, y)
    assert round(pce.coefficients[0][0], 4) == 0.0276


#
def test_5():
    """
    Test PCE coefficients w/ lstsq
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    assert round(pce.coefficients[0][0], 4) == 0.2175


#
def test_6():
    """
    Test PCE prediction
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    y_test = pce.predict(x_test)
    assert round(y_test[0][0], 4) == -0.1607


#
def test_7():
    """
    Test Sobol indices
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    pce_sensitivity = PceSensitivity(pce)
    first_order_sobol = pce_sensitivity.calculate_first_order_indices()
    assert round(first_order_sobol[0][0], 3) == 1.0


#
def test_8():
    """
    Test Sobol indices
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    pce_sensitivity = PceSensitivity(pce)

    total_order_sobol = pce_sensitivity.calculate_total_order_indices()
    assert round(total_order_sobol[0][0], 3) == 1.0


#
def test_9():
    """
    Test Sobol indices
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    pce_sensitivity = PceSensitivity(pce)


#
def test_10():
    """
    Test Sobol indices
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    pce_sensitivity = PceSensitivity(pce)
    with pytest.raises(ValueError):
        generalized_total_order_sobol = pce_sensitivity.calculate_generalized_total_order_indices()


#
def test_11():
    """
    PCE mean
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    mean, _ = pce.get_moments()
    assert round(mean, 3) == 0.218


#
def test_12():
    """
    PCE variance
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, y)
    _, variance = pce.get_moments()
    assert round(variance, 3) == 0.185


def function(x):
    # without square root
    u1 = x[:, 4] * np.cos(x[:, 0])
    u2 = x[:, 4] * np.cos(x[:, 0]) + x[:, 5] * np.cos(np.sum(x[:, :2], axis=1))
    u3 = x[:, 4] * np.cos(x[:, 0]) + x[:, 5] * np.cos(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.cos(
        np.sum(x[:, :3], axis=1))
    u4 = x[:, 4] * np.cos(x[:, 0]) + x[:, 5] * np.cos(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.cos(
        np.sum(x[:, :3], axis=1)) + x[:, 7] * np.cos(np.sum(x[:, :4], axis=1))

    v1 = x[:, 4] * np.sin(x[:, 0])
    v2 = x[:, 4] * np.sin(x[:, 0]) + x[:, 5] * np.sin(np.sum(x[:, :2], axis=1))
    v3 = x[:, 4] * np.sin(x[:, 0]) + x[:, 5] * np.sin(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.sin(
        np.sum(x[:, :3], axis=1))
    v4 = x[:, 4] * np.sin(x[:, 0]) + x[:, 5] * np.sin(np.sum(x[:, :2], axis=1)) + x[:, 6] * np.sin(
        np.sum(x[:, :3], axis=1)) + x[:, 7] * np.sin(np.sum(x[:, :4], axis=1))

    return (u1 + u2 + u3 + u4) ** 2 + (v1 + v2 + v3 + v4) ** 2


dist_1 = Uniform(loc=0, scale=2 * np.pi)
dist_2 = Uniform(loc=0, scale=1)

marg = [dist_1] * 4
marg_1 = [dist_2] * 4
marg.extend(marg_1)

joint = JointIndependent(marginals=marg)

n_samples_2 = 10
mcs_2 = MonteCarloSampling(distributions=joint, nsamples=n_samples_2, random_state=0)
x_2 = mcs_2.samples
y_2 = function(x_2)

polynomial_basis = TotalDegreeBasis(joint, 2)
least_squares = LeastSquareRegression()
pce_2 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
pce_2.fit(x_2, y_2)


def test_17():
    """
    Test Sobol indices for vector-valued quantity of interest on the random inputs
    """
    pce_sensitivity = PceSensitivity(pce_2)
    pce_sensitivity.run()
    generalized_first_sobol = pce_sensitivity.generalized_first_order_indices
    assert round(generalized_first_sobol[0], 4) == 0.0137


def test_18():
    """
    Test Sobol indices for vector-valued quantity of interest on the random inputs
    """
    pce_sensitivity = PceSensitivity(pce_2)
    pce_sensitivity.run()
    generalized_total_sobol = pce_sensitivity.generalized_total_order_indices
    assert round(generalized_total_sobol[0], 4) == 0.4281


def test_19():
    """
    Test Higher statistical moments on Uniform distribution (skewness=0, kurtosis=1.8 from definition)
    """
    polynomial_basis = TotalDegreeBasis(dist, max_degree)
    least_squares = LeastSquareRegression()
    pce = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pce.fit(x, x)
    mean, var, skew, kurtosis = pce.get_moments(True)

    assert round(kurtosis, 3) == 1.8 and round(skew, 3) == 0


def test_20():
    """
    Test Higher statistical moments on Gaussian distribution (skewness=0, kurtosis=3 from definition)
    """
    dist_Gauss = Normal(loc=0, scale=1)

    mcs = MonteCarloSampling(dist, nsamples=n_samples, random_state=1)

    polynomial_basis = TotalDegreeBasis(dist_Gauss, max_degree)
    least_squares = LeastSquareRegression()
    pceGauss = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
    pceGauss.fit(x, x)
    mean, var, skew, kurtosis = pceGauss.get_moments(True)

    assert round(kurtosis, 3) == 3 and round(skew, 3) == 0


def functionLAR(x):
    u1 = x[:, 0] ** 2 + x[:, 2] ** 2 + x[:, 3] ** 2

    return u1


n_samples_2 = 100
x_2 = joint.rvs(n_samples_2)
y_2 = functionLAR(x_2)

polynomial_basis = TotalDegreeBasis(joint, 6)
least_squares = LeastSquareRegression()
pce_2 = PolynomialChaosExpansion(polynomial_basis=polynomial_basis, regression_method=least_squares)
pce_2.fit(x_2, y_2)


def test_21():
    """
    Test Model Selection Algorithm based on Least Angle Regression (select only important basis functions from large set)
    """

    pce_lar = polynomial_chaos.regressions.LeastAngleRegression.model_selection(pce_2)
    pce2_lar_sens = PceSensitivity(pce_lar)

    assert all((np.argwhere(np.round(pce2_lar_sens.calculate_generalized_total_order_indices(), 3) > 0)
                == [[0], [2], [3]]))


# tests/unit_tests/transformations/test_correlate.py

# Test the Correlation module

from UQpy.transformations import Correlate
import numpy as np
import pytest


def test_samples():
    samples_z = np.array([[0.3, 0.2], [0.2, 2.4]])
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Correlate(samples_u=samples_z, corr_z=rz)
    np.testing.assert_allclose(ntf_obj.samples_z, [[0.3, 0.36], [0.2, 1.6]], rtol=1e-09)


def test_samples_u():
    samples_z = np.array([[0.3, 0.2], [0.2, 2.4]])
    with pytest.raises(Exception):
        assert Correlate(samples_u=samples_z)


def test_corr_z():
    rz = np.array([[1.0, 0.0], [0.0, 1.0]])
    with pytest.raises(Exception):
        assert Correlate(corr_z=rz)



# tests/unit_tests/transformations/test_decorrelate.py

# Test the Correlation module

from UQpy.transformations import Decorrelate
import numpy as np
import pytest


def test_samples():
    samples_z = np.array([[0.3, 0.36], [0.2, 1.6]])
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Decorrelate(samples_z=samples_z, corr_z=rz)
    np.testing.assert_allclose(ntf_obj.samples_u, [[0.3, 0.19999999999999998], [0.2, 2.4000000000000004]], rtol=1e-09)


def test_samples_u():
    samples_z = np.array([[0.3, 0.36], [0.2, 1.6]])
    with pytest.raises(Exception):
        assert Decorrelate(samples_z=samples_z)


def test_corr_z():
    rz = np.array([[1.0, 0.0], [0.0, 1.0]])
    with pytest.raises(Exception):
        assert Decorrelate(corr_z=rz)


def test_h():
    samples_z = np.array([[0.3, 0.36], [0.2, 1.6]])
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Decorrelate(samples_z=samples_z, corr_z=rz)
    np.testing.assert_allclose(ntf_obj.H, [[1.0, 0.], [0.8, 0.6]], rtol=1e-09)


def test_samples1():
    samples_z = np.array([[0.3, 0.36], [0.2, 1.6]])
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Decorrelate(samples_z=samples_z, corr_z=rz)
    assert (ntf_obj.samples_z == np.array([[0.3, 0.36], [0.2, 1.6]])).all()


def test_corr_z_1():
    samples_z = np.array([[0.3, 0.36], [0.2, 1.6]])
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Decorrelate(samples_z=samples_z, corr_z=rz)
    assert (ntf_obj.corr_z == np.array([[1.0, 0.8], [0.8, 1.0]])).all()


# tests/unit_tests/transformations/test_nataf.py

# Test the Nataf transformation

from UQpy.distributions import *
from UQpy.transformations import Nataf
import numpy as np
import pytest


def test_correct_distribution_dimensions():
    dist1 = Normal(loc=0.0, scale=1.0)
    dist2 = Normal(loc=0.0, scale=1.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    assert ntf_obj.n_dimensions == 2


def test_wrong_distribution_in_list():
    dist1 = Normal(loc=0.0, scale=1.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    with pytest.raises(Exception):
        assert Nataf(distributions=[dist1, 'Beta'], corr_x=rx)


def test_wrong_distribution():
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    with pytest.raises(Exception):
        assert Nataf(distributions='Normal', corr_x=rx)


def test_identity_correlation_x_normal():
    dist1 = Normal(loc=0.0, scale=1.0)
    dist2 = Normal(loc=0.0, scale=1.0)
    ntf_obj = Nataf(distributions=[dist1, dist2])
    assert np.all(np.equal(ntf_obj.corr_x, np.eye(2)))


def test_identity_correlation_z_normal():
    dist1 = Normal(loc=0.0, scale=1.0)
    dist2 = Normal(loc=0.0, scale=1.0)
    ntf_obj = Nataf(distributions=[dist1, dist2])
    assert np.all(np.equal(ntf_obj.corr_z, np.eye(2)))


def test_identity_correlation_uniform_z():
    dist1 = Uniform(loc=0.0, scale=1.0)
    dist2 = Uniform(loc=0.0, scale=1.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    assert np.all(np.equal(ntf_obj.corr_z, rx))


def test_identity_correlation_normal_z():
    dist1 = Normal(loc=0.0, scale=1.0)
    dist2 = Normal(loc=0.0, scale=1.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    assert np.all(np.equal(ntf_obj.corr_z, rx))


def test_non_identity_correlation_normal():
    dist1 = Normal(loc=0.0, scale=1.0)
    dist2 = Normal(loc=0.0, scale=1.0)
    rx = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    assert np.all(np.equal(ntf_obj.corr_z, rx))


def test_non_identity_correlation_uniform_z():
    dist1 = Uniform(loc=0.0, scale=1.0)
    dist2 = Uniform(loc=0.0, scale=1.0)
    rx = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_x=rx)
    np.testing.assert_allclose(ntf_obj.corr_z, [[1., 0.8134732861515996], [0.8134732861515996, 1.]], rtol=1e-09)


def test_non_identity_correlation_uniform_x():
    dist1 = Uniform(loc=0.0, scale=1.0)
    dist2 = Uniform(loc=0.0, scale=1.0)
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_z=rz)
    assert (ntf_obj.corr_x == [[1., 0.7859392826067285], [0.7859392826067285, 1.]]).all()


def test_attribute_h():
    dist1 = Normal(loc=0.0, scale=1.0)
    dist2 = Normal(loc=0.0, scale=1.0)
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    ntf_obj = Nataf(distributions=[dist1, dist2], corr_z=rz)
    np.testing.assert_allclose(ntf_obj.H, [[1., 0.], [0.8, 0.6]], rtol=1e-09)


def test_samples_x():
    dist1 = Normal(loc=0.0, scale=1.0)
    ntf_obj = Nataf(distributions=[dist1])
    assert ntf_obj.samples_x is None


def test_samples_z():
    dist1 = Normal(loc=0.0, scale=1.0)
    ntf_obj = Nataf(distributions=[dist1])
    assert ntf_obj.samples_z is None


def test_samples_x_shape():
    dist1 = Normal(loc=0.0, scale=5.0)
    samples_x = np.array([0.3, 1.2, 3.5])
    ntf_obj = Nataf(distributions=[dist1])
    ntf_obj.run(samples_x)
    assert ntf_obj.samples_x.shape == (3, 1)


def test_samples_z_shape():
    dist1 = Normal(loc=0.0, scale=5.0)
    ntf_obj = Nataf(distributions=[dist1])
    samples_z = np.array([0.3, 1.2, 3.5])
    ntf_obj.run(samples_z)
    assert ntf_obj.samples_z.shape == (3, 1)


def test_samples_x_jxz1():
    dist1 = Normal(loc=0.0, scale=5.0)
    ntf_obj = Nataf(distributions=[dist1])
    samples_x = np.array([0.3, 1.2, 3.5])
    ntf_obj.run(samples_x=samples_x, jacobian=False)
    assert ntf_obj.jxz is None


def test_samples_x_jxz2():
    dist1 = Uniform(loc=0.0, scale=5.0)
    dist2 = Uniform(loc=0.0, scale=3.0)
    ntf_obj = Nataf(distributions=[dist1, dist2])
    samples_x = np.array([[0.3, 1.2, 3.5], [0.2, 2.4, 0.9]]).T
    ntf_obj.run(samples_x=samples_x, jacobian=True)
    g = []
    for i in range(3):
        if i == 0:
            g.append((ntf_obj.jxz[i] == np.array([[1.6789373877365803, 0.0], [0.0, 2.577850090371836]])).all())
        elif i == 1:
            g.append((ntf_obj.jxz[i] == np.array([[0.6433491348614259, 0.0], [0.0, 1.1906381155257868]])).all())
        else:
            g.append((ntf_obj.jxz[i] == np.array([[0.5752207318528584, 0.0], [0.0, 0.958701219754764]])).all())
    assert np.all(g)


def test_samples_x1():
    dist1 = Uniform(loc=0.0, scale=5.0)
    dist2 = Uniform(loc=0.0, scale=3.0)
    ntf_obj = Nataf(distributions=[dist1, dist2])
    samples_x = np.array([[0.3, 1.2, 3.5], [0.2, 2.4, 0.9]]).T
    ntf_obj.run(samples_x=samples_x, jacobian=True)
    g = []
    for i in range(3):
        if i == 0:
            g.append((ntf_obj.samples_z[i] == np.array([-1.5547735945968535, -1.501085946044025])).all())
        elif i == 1:
            g.append((ntf_obj.samples_z[i] == np.array([-0.7063025628400874, 0.841621233572914])).all())
        else:
            g.append((ntf_obj.samples_z[i] == np.array([0.5244005127080407, -0.5244005127080409])).all())
    assert np.all(g)


def test_samples_z_jzx1():
    dist1 = Normal(loc=0.0, scale=5.0)
    ntf_obj = Nataf(distributions=[dist1])
    samples_z = np.array([0.3, 1.2, 3.5])
    ntf_obj.run(samples_z=samples_z, jacobian=False)
    assert ntf_obj.jzx is None


def test_samples_z_jzx2():
    dist1 = Uniform(loc=0.0, scale=5.0)
    dist2 = Uniform(loc=0.0, scale=3.0)
    ntf_obj = Nataf(distributions=[dist1, dist2])
    samples_z = np.array([[0.3, 1.2], [0.2, 2.4]]).T
    ntf_obj.run(samples_z=samples_z, jacobian=True)
    g = []
    for i in range(2):
        if i == 0:
            g.append((ntf_obj.jzx[i] == np.array([[0.524400601939789, 0.0], [0.0, 0.8524218415758338]])).all())
        else:
            g.append((ntf_obj.jzx[i] == np.array([[1.0299400748281828, 0.0], [0.0, 14.884586948005541]])).all())
    assert np.all(g)


def test_samples_z2():
    dist1 = Uniform(loc=0.0, scale=5.0)
    dist2 = Uniform(loc=0.0, scale=3.0)
    ntf_obj = Nataf(distributions=[dist1, dist2])
    samples_z = np.array([[0.3, 1.2], [0.2, 2.4]]).T
    ntf_obj.run(samples_z=samples_z, jacobian=True)
    g = []
    for i in range(2):
        if i == 0:
            g.append((ntf_obj.samples_x[i] == np.array([3.089557110944763, 1.737779128317309])).all())
        elif i == 1:
            g.append((ntf_obj.samples_x[i] == np.array([4.424651648891459, 2.9754073922262116])).all())
    assert np.all(g)


def test_itam_beta():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx)
    assert ntf_obj.itam_beta == 1.0


def test_itam_max_iteration():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx)
    assert ntf_obj.itam_max_iter == 100


def test_itam_threshold1():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx)
    assert ntf_obj.itam_threshold1 == 0.001


def test_itam_threshold2():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx)
    assert ntf_obj.itam_threshold2 == 0.1


def test_itam_beta1():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx, itam_beta=2.0)
    assert ntf_obj.itam_beta == 2.0


def test_itam_max_iteration1():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx, itam_max_iter=200)
    assert ntf_obj.itam_max_iter == 200


def test_itam_max_iteration2():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx, itam_max_iter=10)
    assert ntf_obj.itam_max_iter == 10


def test_itam_threshold1_1():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx, itam_threshold1=0.002)
    assert ntf_obj.itam_threshold1 == 0.002


def test_itam_threshold2_1():
    dist1 = Uniform(loc=0.0, scale=5.0)
    rx = np.array([[1.0, 0.0], [0.0, 1.0]])
    ntf_obj = Nataf(distributions=[dist1], corr_x=rx, itam_threshold2=0.3)
    assert ntf_obj.itam_threshold2 == 0.3


def test_distortion_z2x_finite_moments():
    dist1 = Lognormal(s=0.0, loc=0.0, scale=1.0)
    dist2 = Uniform(loc=0.0, scale=1.0)
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    with pytest.raises(Exception):
        assert Nataf(distributions=[dist1, dist2], corr_z=rz)


def distortion_z2x_dist_object():
    dist1 = Lognormal(s=0.0, loc=0.0, scale=1.0)
    rz = np.array([[1.0, 0.8], [0.8, 1.0]])
    with pytest.raises(Exception):
        assert Nataf.distortion_z2x(distributions=[dist1, 'Beta'], corr_z=rz)


# tests/unit_tests/eigenvalue_model.py

import numpy as np


class RunPythonModel:

    def __init__(self, samples=None, dimension=None):

        self.samples = samples
        self.dimension = dimension
        self.qoi = np.zeros_like(self.samples)
        for i in range(self.samples.shape[0]):
            p = np.array([[self.samples[i, 0]+self.samples[i, 1], -self.samples[i, 1], 0],
                          [-self.samples[i, 1], self.samples[i, 1]+self.samples[i, 2], -self.samples[i, 2]],
                          [0, -self.samples[i, 2], self.samples[i, 2]]])
            w, v = np.linalg.eig(p)
            self.qoi[i, :] = w

# tests/unit_tests/pfn.py

import numpy as np


def example1(samples=None):
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        R = samples[i, 0]
        S = samples[i, 1]
        g[i] = R - S
    return g


def example2(samples=None):
    import numpy as np
    d = 2
    beta = 3.0902
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        g[i] = -1 / np.sqrt(d) * (samples[i, 0] + samples[i, 1]) + beta
    return g


def example3(samples=None):
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        g[i] = 6.2 * samples[i, 0] - samples[i, 1] * samples[i, 2] ** 2
    return g


def example4(samples=None):
    g = np.zeros(samples.shape[0])
    for i in range(samples.shape[0]):
        g[i] = samples[i, 0] * samples[i, 1] - 80
    return g

# tests/unit_tests/pfn_models.py


import numpy as np

domain = np.linspace(0, 10, 50)


def model_quadratic(theta):
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain + inpt[1] * domain ** 2


def model_linear(theta):
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain


def model_cubic(theta):
    # this one takes one parameter vector theta and return one qoi
    inpt = np.array(theta).reshape((-1,))
    return inpt[0] * domain + inpt[1] * domain ** 2 + inpt[2] * domain ** 3

# tests/unit_tests/python_model_1Dfunction.py

def y_func(z):
    import numpy as np
    return np.sin(z)

# tests/unit_tests/python_model_function.py

def y_func(z):
    return 1/(6.2727*(abs(0.3-z[:, 0]**2-z[:, 1]**2)+0.01))

# tests/unit_tests/series.py

import numpy as np


def series(z, k=7):
    t1 = 3 + 0.1 * (z[:, 1] - z[:, 0]) ** 2 - (z[:, 1] + z[:, 0]) / np.sqrt(2)
    t2 = 3 + 0.1 * (z[:, 1] - z[:, 0]) ** 2 + (z[:, 1] + z[:, 0]) / np.sqrt(2)
    t3 = z[:, 1] - z[:, 0] + k / np.sqrt(2)
    t4 = z[:, 0] - z[:, 1] + k / np.sqrt(2)
    return min([t1, t2, t3, t4])

# tests/unit_tests/sensitivity/exponential.py

import numpy as np


def evaluate(X: np.array) -> np.array:
    r"""A non-linear function that is used to test Cramer-von Mises sensitivity index.

    .. math::
        f(x) = \exp(x_1 + 2*x_2)

    Parameters
    ----------
    X : np.array
        An `N*D` array holding values for each parameter, where `N` is the
        number of samples and `D` is the number of parameters
        (in this case, 2).

    Returns
    -------
    np.array
        [description]
    """

    Y = np.exp(X[:, 0] + 2 * X[:, 1])

    return Y


# tests/unit_tests/sensitivity/ishigami.py

"""
Auxiliary file
==============================================
"""

import numpy as np

def evaluate(X, params=[7, 0.1]):
    """Non-monotonic Ishigami-Homma three parameter test function"""

    a = params[0]
    b = params[1]

    Y = np.sin(X[:, 0]) + a * np.power(np.sin(X[:, 1]), 2) + \
        b * np.power(X[:, 2], 4) * np.sin(X[:, 0])

    return Y


# tests/unit_tests/sensitivity/pfn.py

import numpy as np


def gfun_sensitivity(samples, a_values):
    gi_xi = [(np.abs(4. * Xi - 2) + ai) / (1. + ai) for Xi, ai in zip(np.array(samples).T, a_values)]
    gfun = np.prod(np.array(gi_xi), axis=0)
    return list(gfun)


# tests/unit_tests/sensitivity/sobol_func.py

import numpy as np
import copy


def evaluate(X, a_values):

    dims = len(a_values)
    g = 1

    for i in range(dims):
        g_i = (np.abs(4 * X[:, i] - 2) + a_values[i]) / (1 + a_values[i])
        g *= g_i

    return g


def sensitivities(a_values):

    dims = len(a_values)

    Total_order = np.zeros((dims, 1))

    V_i = (3 * (1 + a_values) ** 2) ** (-1)

    total_variance = np.prod(1 + V_i) - 1

    First_order = V_i / total_variance

    for i in range(dims):

        rem_First_order = copy.deepcopy(V_i)
        rem_First_order[i] = 0
        Total_order[i] = V_i[i] * np.prod(rem_First_order + 1) / total_variance

    return First_order.reshape(-1, 1), Total_order


# tests/unit_tests/sensitivity/test_baseclass.py

"""
This module is used to test the functionalities of the baseclass.

- test_pick_and_freeze_sampling: 
    Test the `generate_pick_and_test_samples` function.
- test_bootstrap_for_vector: 
    Test the bootstrap sampling for a vector.
- test_bootstrap_for_matrix: 
    Test the bootstrap sampling for a matrix.

"""

import numpy as np
import pytest

from UQpy.run_model.RunModel import RunModel
from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.distributions import Uniform
from UQpy.distributions.collection.JointIndependent import JointIndependent
from UQpy.sensitivity.sobol import Sobol
from UQpy.sensitivity.baseclass.pickfreeze import generate_pick_freeze_samples

# Prepare
###############################################################################

# Prepare the input distribution
@pytest.fixture()
def ishigami_input_dist_object():
    """
    This function returns the input distribution for the Ishigami function.

    X1 ~ Uniform(-pi, pi)
    X2 ~ Uniform(-pi, pi)
    X3 ~ Uniform(-pi, pi)

    """
    return JointIndependent([Uniform(-np.pi, 2 * np.pi)] * 3)


@pytest.fixture()
def ishigami_model_object():
    """This function creates the Ishigami run_model_object"""
    model = PythonModel(
        model_script="ishigami.py",
        model_object_name="evaluate",
        var_names=[r"$X_1$", "$X_2$", "$X_3$"],
        delete_files=True,
        params=[7, 0.1],
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def sobol_object(ishigami_model_object, ishigami_input_dist_object):
    """This function returns the Sobol object."""

    return Sobol(ishigami_model_object, ishigami_input_dist_object)


@pytest.fixture()
def sobol_object_input_samples_small(sobol_object):
    """This creates the Sobol object."""

    SA = sobol_object

    np.random.seed(12345)  # set seed for reproducibility

    SA.n_samples = 2

    return generate_pick_freeze_samples(SA.dist_object, SA.n_samples)


# Generate N pick and free samples
@pytest.fixture()
def pick_and_freeze_samples_small():
    """
    This function returns input matrices A, B and C_i with a small number
    of samples for the Ishigami input distribution.
    This is used to test the `generate_pick_and_freeze_samples` function.

    The samples are generated as follows:

        dist_1 = JointInd([Uniform(-np.pi, 2*np.pi)]*3)

        np.random.seed(12345) #! set seed for reproducibility

        n_samples = 2
        n_vars = 3

        samples = dist_1.rvs(n_samples*2)

        # Split samples
        A_samples = samples[:n_samples, :]
        B_samples = samples[n_samples:, :]

        def _get_C_i(i, A, B):
            C_i = copy.deepcopy(B)
            C_i[:, i] = A[:, i]
            return C_i

        C_samples = np.zeros((n_vars, n_samples, n_vars))

        for i in range(3):
            C_samples[i, :, :] = _get_C_i(i, A_samples, B_samples)

        print(np.around(A_samples,3))
        print(np.around(B_samples,3))
        print(np.around(C_samples,3))

    """

    A_samples = np.array([[2.699, 0.426, 1.564], [-1.154, 0.600, 0.965]])

    B_samples = np.array([[-1.986, 2.919, 1.556], [-1.856, 0.962, 2.898]])

    C_samples = np.array(
        [
            [[2.699, 2.919, 1.556], [-1.154, 0.962, 2.898]],
            [[-1.986, 0.426, 1.556], [-1.856, 0.6, 2.898]],
            [[-1.986, 2.919, 1.564], [-1.856, 0.962, 0.965]],
        ]
    )

    return A_samples, B_samples, C_samples


@pytest.fixture()
def random_f_A():
    """This function returns an A-like vector"""

    rand_f_A = np.array([[100], [101], [102], [103], [104]])

    return rand_f_A


@pytest.fixture()
def random_f_C_i():
    """This function returns a C_i-like vector"""

    rand_f_C_i = np.array([[100, 200], [101, 201], [102, 202], [103, 203], [104, 204]])
    return rand_f_C_i


@pytest.fixture()
def manual_bootstrap_samples_f_A():
    """This function bootstraps the A-like vector using random indices"""

    # Genrated using np.random.randint(low=0, high=5, size=(5,1))
    # with np.random.seed(12345)
    # rand_indices_f_A = np.array([ [2],
    #                             [1],
    #                             [4],
    #                             [1],
    #                             [2]])

    # bootstrap_f_A = rand_f_A[rand_indices_A]
    bootstrap_sample_A = np.array([[102], [101], [104], [101], [102]])

    return bootstrap_sample_A


@pytest.fixture()
def manual_bootstrap_samples_f_C_i():
    """This function bootstraps the C_i-like vector using random indices"""

    # Genrated using np.random.randint(low=0, high=5, size=(5,2))
    # with np.random.seed(12345)
    # rand_indices_C_i = np.array([ [2, 1],
    #                               [4, 1],
    #                               [2, 1],
    #                               [1, 3],
    #                               [1, 3]])

    bootstrap_f_C_i = np.array(
        [[102, 201], [104, 201], [102, 201], [101, 203], [101, 203]]
    )

    return bootstrap_f_C_i


# Unit tests
###############################################################################


def test_pick_and_freeze_sampling(
    pick_and_freeze_samples_small, sobol_object_input_samples_small
):

    """Test the `generate_pick_and_test_samples` function."""

    # Prepare
    A_samples, B_samples, C_samples = pick_and_freeze_samples_small
    A_test, B_test, C_test_generator, _ = sobol_object_input_samples_small

    # Act
    assert np.allclose(A_samples, np.around(A_test, 3))
    assert np.allclose(B_samples, np.around(B_test, 3))

    for i in range(3):
        C_test = next(C_test_generator)
        assert np.allclose(C_samples[i, :, :], np.around(C_test, 3))


def test_bootstrap_for_vector(random_f_A, manual_bootstrap_samples_f_A):

    """Test the bootstrap sampling for a vector."""

    # Prepare
    np.random.seed(12345)  #! set seed for reproducibility

    gen_f_A = Sobol.bootstrap_sample_generator_1D(random_f_A)

    bootstrap_samples_f_A = next(gen_f_A)

    # Act
    assert np.array_equal(manual_bootstrap_samples_f_A, bootstrap_samples_f_A)


def test_bootstrap_for_matrix(random_f_C_i, manual_bootstrap_samples_f_C_i):

    """Test the bootstrap sampling for a matrix."""

    # Prepare
    np.random.seed(12345)  #! set seed for reproducibility

    gen_f_C_i = Sobol.bootstrap_sample_generator_2D(random_f_C_i)

    bootstrap_samples_C_i = next(gen_f_C_i)

    # Act
    assert np.array_equal(manual_bootstrap_samples_f_C_i, bootstrap_samples_C_i)


# tests/unit_tests/sensitivity/test_chatterjee.py

"""
This is the test module for the Chatterjee sensitivity indices. 

Here, we will use the exponential function to test the output, as in
the test module for Cramer sensitivity indices for the Chatterjee indices and 
the ishigami function as in the test module for Sobol sensitivity indices for the
Sobol indices.

The following methods are tested:
1. pick_and_freeze_estimator
2. Sobol estimate

Important
----------
The computed indices are computed using the `np.isclose` function.

Function signature:
    numpy.isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)

    Parameters:
    a, b: array_like
        Input arrays to compare.

    rtol: float
        The relative tolerance parameter.

    atol: float
        The absolute tolerance parameter.

Each element of the `diff` array is compared as follows:
diff = |a - b|
diff <= atol + rtol * abs(b)

-   relative tolerance: rtol * abs(b)
    It is the maximum allowed difference between a and b, 
    relative to the absolute value of b. 
    For example, to set a tolerance of 1%, pass rol=0.01, 
    which assures that the values are within 2 decimal places of each other.

-   absolute tolerance: atol
    When b is close to zero, the atol value is used.

"""

import numpy as np
import pytest

from UQpy.run_model.RunModel import RunModel
from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.distributions import Uniform, Normal
from UQpy.distributions.collection.JointIndependent import JointIndependent
from UQpy.sensitivity.chatterjee import Chatterjee


# Prepare
###############################################################################

# Prepare the input distribution
@pytest.fixture()
def exponential_input_dist_object():
    """
    This function returns the input distribution for the Ishigami function.

    X1 ~ Normal(0,1)
    X2 ~ Normal(0,1)

    """
    return JointIndependent([Normal(0, 1)] * 2)


@pytest.fixture()
def exponential_model_object():
    """This function creates the exponential run_model_object"""
    model = PythonModel(
        model_script="exponential.py",
        model_object_name="evaluate",
        var_names=[
            "X_1",
            "X_2",
        ],
        delete_files=True,
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def Chatterjee_object(exponential_model_object, exponential_input_dist_object):
    """This function creates the Chatterjee object"""
    return Chatterjee(exponential_model_object, exponential_input_dist_object)


@pytest.fixture()
def analytical_Chatterjee_indices():
    """This function returns the analytical Chatterjee indices.

    S1 = (6/np.pi) * np.arctan(2) - 2
    S2 = (6/np.pi) * np.arctan(np.sqrt(19)) - 2

    print(np.around(S1, 4))
    print(np.around(S2, 4))

    """

    return np.array([[0.1145], [0.5693]])


@pytest.fixture()
def numerical_Chatterjee_indices(Chatterjee_object):
    """This function returns the numerical Chatterjee indices."""

    SA = Chatterjee_object

    np.random.seed(12345)  #! set seed for reproducibility

    computed_indices = SA.run(n_samples=10_000)

    return computed_indices["chatterjee_i"]


@pytest.fixture()
def ishigami_input_dist_object():
    """
    This function returns the input distribution for the Ishigami function.

    X1 ~ Uniform(-pi, pi)
    X2 ~ Uniform(-pi, pi)
    X3 ~ Uniform(-pi, pi)

    """
    return JointIndependent([Uniform(-np.pi, 2 * np.pi)] * 3)


@pytest.fixture()
def ishigami_model_object():
    """This function creates the Ishigami run_model_object"""
    model = PythonModel(
        model_script="ishigami.py",
        model_object_name="evaluate",
        var_names=[r"$X_1$", "$X_2$", "$X_3$"],
        delete_files=True,
        params=[7, 0.1],
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def Chatterjee_object_ishigami(ishigami_model_object, ishigami_input_dist_object):
    """This function creates the Chatterjee object"""
    return Chatterjee(ishigami_model_object, ishigami_input_dist_object)


@pytest.fixture()
def numerical_Sobol_indices(Chatterjee_object_ishigami):
    """This function returns the Sobol indices."""

    SA = Chatterjee_object_ishigami

    np.random.seed(12345)

    computed_indices = SA.run(n_samples=10_000, estimate_sobol_indices=True)

    return computed_indices["sobol_i"]


@pytest.fixture()
def analytical_ishigami_Sobol_indices():
    """
    Analytical Sobol indices for the Ishigami function.

    Copy-paste the following to reproduce the given indices:

        a = 7
        b = 0.1

        V1 = 0.5*(1 + (b*np.pi**4)/5)**2
        V2 = (a**2)/8
        V3 = 0

        VT3 = (8*(b**2)*np.pi**8)/225
        VT1 = V1 + VT3
        VT2 = V2

        total_variance = V2 + (b*np.pi**4)/5 + ((b**2) * np.pi**8)/18 + 0.5

        S = np.array([V1, V2, V3])/total_variance
        S_T = np.array([VT1, VT2, VT3])/total_variance

        S = np.around(S, 4)
        S_T = np.around(S_T, 4)

    """

    S1 = 0.3139
    S2 = 0.4424
    S3 = 0

    S_T1 = 0.5576
    S_T2 = 0.4424
    S_T3 = 0.2437

    S = np.array([S1, S2, S3])
    S_T = np.array([S_T1, S_T2, S_T3])

    return S.reshape(-1, 1)


# Unit tests
###############################################################################


def test_Chatterjee_estimate(
    numerical_Chatterjee_indices, analytical_Chatterjee_indices
):
    """This function tests the Chatterjee estimate."""
    assert np.isclose(
        numerical_Chatterjee_indices, analytical_Chatterjee_indices, rtol=0, atol=1e-2
    ).all()


def test_Sobol_estimate(numerical_Sobol_indices, analytical_ishigami_Sobol_indices):
    """This function tests the Sobol estimate."""
    assert np.isclose(
        numerical_Sobol_indices, analytical_ishigami_Sobol_indices, rtol=0, atol=1e-2
    ).all()


# tests/unit_tests/sensitivity/test_cramer_von_mises.py

"""
This is the test module for Cramer sensitivity indices.

Here, we will use the an exponential function to test the output.

The following methods are tested:
1. pick_and_freeze_estimator
2. bootstrap_variance_computation

Important
----------
The computed indices are computed using the `np.isclose` function.

Function signature:
    numpy.isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)

    Parameters:
    a, b: array_like
        Input arrays to compare.

    rtol: float
        The relative tolerance parameter.

    atol: float
        The absolute tolerance parameter.

Each element of the `diff` array is compared as follows:
diff = |a - b|
diff <= atol + rtol * abs(b)

-   relative tolerance: rtol * abs(b)
    It is the maximum allowed difference between a and b, 
    relative to the absolute value of b. 
    For example, to set a tolerance of 1%, pass rol=0.01, 
    which assures that the values are within 2 decimal places of each other.

-   absolute tolerance: atol
    When b is close to zero, the atol value is used.

"""

import numpy as np
import pytest
import scipy

from UQpy.run_model.RunModel import RunModel
from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.distributions import Normal, Uniform
from UQpy.distributions.collection.JointIndependent import JointIndependent
from UQpy.sensitivity.cramer_von_mises import CramervonMises

# Prepare
###############################################################################

# Prepare the input distribution
@pytest.fixture()
def exponential_input_dist_object():
    """
    This function returns the input distribution for the Ishigami function.

    X1 ~ Normal(0,1)
    X2 ~ Normal(0,1)

    """
    return JointIndependent([Normal(0, 1)] * 2)


@pytest.fixture()
def exponential_model_object():
    """This function creates the exponential run_model_object"""
    model = PythonModel(
        model_script="exponential.py",
        model_object_name="evaluate",
        var_names=[r"$X_1$", "$X_2$"],
        delete_files=True,
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def CVM_object(exponential_model_object, exponential_input_dist_object):
    """This function returns the CVM object."""

    return CramervonMises(exponential_model_object, exponential_input_dist_object)


@pytest.fixture()
def analytical_exponential_CVM_indices():
    """This function returns the analytical Cramer-von-Mises indices.

    S1_CVM = (6/np.pi) * np.arctan(2) - 2
    S2_CVM = (6/np.pi) * np.arctan(np.sqrt(19)) - 2

    print(np.around(S1_CVM, 4))
    print(np.around(S2_CVM, 4))

    """

    return np.array([[0.1145], [0.5693]])


@pytest.fixture()
def numerical_exponential_CVM_indices(CVM_object):
    """
    This function returns the Cramer-von-Mises indices
    computed using the Pick and Freeze algorithm.

    """

    SA = CVM_object

    np.random.seed(12345)  #! set seed for reproducibility

    computed_indices = SA.run(n_samples=50_000)

    return computed_indices["CVM_i"]


@pytest.fixture()
def NUM_SAMPLES():
    """This function returns the number of samples."""

    num_bootstrap_samples = 50
    num_samples = 10_000

    return num_bootstrap_samples, num_samples


@pytest.fixture()
def bootstrap_CVM_index_variance(CVM_object, NUM_SAMPLES):
    """This function returns the variance in the computed Cramer-von-Mises index
    computed using the bootstrap algorithm."""

    #### SETUP ####
    SA = CVM_object

    np.random.seed(12345)  #! set seed for reproducibility

    confidence_level = 0.95
    delta = -scipy.stats.norm.ppf((1 - confidence_level) / 2)

    num_bootstrap_samples, n_samples = NUM_SAMPLES

    #### Compute indices ####
    computed_indices = SA.run(
        n_samples=n_samples,
        num_bootstrap_samples=num_bootstrap_samples,
        confidence_level=confidence_level,
    )

    First_order = computed_indices["CVM_i"].ravel()
    upper_bound_first_order = computed_indices["CI_CVM_i"][:, 1]

    #### Compute variance ####
    std_bootstrap_first_order = (upper_bound_first_order - First_order) / delta

    return std_bootstrap_first_order**2


@pytest.fixture()
def model_evals_CVM_index_variance():

    """
    runmodel_obj = RunModel(
                model_script='exponential.py',
                var_names=['X1', 'X2'],
                vec=True, delete_files=True)

    input_object = JointInd([Normal(0, 1)]*2)

    SA = CramervonMises(runmodel_obj, input_object)

    np.random.seed(12345)

    num_repetitions, n_samples = 1_000, 10_000

    num_vars = 2

    sample_first_order = np.zeros((num_vars, num_repetitions))

    for i in range(num_repetitions):
        CV_First_order = SA.run(n_samples=n_samples)

        sample_first_order[:, i] = CV_First_order.ravel()

    variance_first_order = np.var(sample_first_order, axis=1).reshape(-1, 1)

    print(variance_first_order)

    """

    variance_first_order = np.array([4.01099066e-05, 2.06802165e-05])

    return variance_first_order


@pytest.fixture()
def ishigami_input_dist_object():
    """
    This function returns the input distribution for the Ishigami function.

    X1 ~ Uniform(-pi, pi)
    X2 ~ Uniform(-pi, pi)
    X3 ~ Uniform(-pi, pi)

    """
    return JointIndependent([Uniform(-np.pi, 2 * np.pi)] * 3)


@pytest.fixture()
def ishigami_model_object():
    """This function creates the Ishigami run_model_object"""
    model = PythonModel(
        model_script="ishigami.py",
        model_object_name="evaluate",
        var_names=[r"$X_1$", "$X_2$", "$X_3$"],
        delete_files=True,
        params=[7, 0.1],
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def CVM_object_ishigami(ishigami_model_object, ishigami_input_dist_object):
    """This function returns the CVM object."""

    return CramervonMises(ishigami_model_object, ishigami_input_dist_object)


@pytest.fixture()
def numerical_Sobol_indices(CVM_object_ishigami):
    """
    This function returns the Sobol indices computed
    using the Pick and Freeze algorithm.
    """

    SA = CVM_object_ishigami

    np.random.seed(12345)

    computed_indices = SA.run(
        n_samples=500_000, estimate_sobol_indices=True, disable_CVM_indices=True
    )

    return computed_indices["sobol_i"], computed_indices["sobol_total_i"]


@pytest.fixture()
def analytical_ishigami_Sobol_indices():
    """
    Analytical Sobol indices for the Ishigami function.

    Copy-paste the following to reproduce the given indices:

        a = 7
        b = 0.1

        V1 = 0.5*(1 + (b*np.pi**4)/5)**2
        V2 = (a**2)/8
        V3 = 0

        VT3 = (8*(b**2)*np.pi**8)/225
        VT1 = V1 + VT3
        VT2 = V2

        total_variance = V2 + (b*np.pi**4)/5 + ((b**2) * np.pi**8)/18 + 0.5

        S = np.array([V1, V2, V3])/total_variance
        S_T = np.array([VT1, VT2, VT3])/total_variance

        S = np.around(S, 4)
        S_T = np.around(S_T, 4)

    """

    S1 = 0.3139
    S2 = 0.4424
    S3 = 0

    S_T1 = 0.5576
    S_T2 = 0.4424
    S_T3 = 0.2437

    S = np.array([S1, S2, S3])
    S_T = np.array([S_T1, S_T2, S_T3])

    return S.reshape(-1, 1), S_T.reshape(-1, 1)


# Unit tests
###############################################################################


def test_pick_and_freeze_estimator(
    numerical_exponential_CVM_indices, analytical_exponential_CVM_indices
):
    """
    This function tests the pick_and_freeze_estimator method using 50_000 samples.
    """
    S_CVM_analytical = analytical_exponential_CVM_indices
    S_CVM_numerical = numerical_exponential_CVM_indices

    assert np.isclose(S_CVM_analytical, S_CVM_numerical, rtol=0, atol=1e-2).all()


def test_bootstrap_variance_computation(
    bootstrap_CVM_index_variance, model_evals_CVM_index_variance
):
    """
    This function tests the bootstrap_variance_computation method using
    100_000 samples and 1_000 bootstrap samples.
    """
    var_first = model_evals_CVM_index_variance
    boot_var_first = bootstrap_CVM_index_variance

    assert var_first.shape == boot_var_first.shape
    assert np.isclose(boot_var_first, var_first, rtol=0, atol=1e-4).all()


def test_Sobol_estimate_computation(
    numerical_Sobol_indices, analytical_ishigami_Sobol_indices
):
    """
    This function tests the Sobol_estimate_computation method using 1_000_000 samples.
    """
    S_numerical, S_T_numerical = numerical_Sobol_indices
    S_analytical, S_T_analytical = analytical_ishigami_Sobol_indices

    assert S_analytical.shape == S_numerical.shape
    assert S_T_analytical.shape == S_T_numerical.shape
    assert np.isclose(S_numerical, S_analytical, rtol=0, atol=1e-2).all()
    assert np.isclose(S_T_numerical, S_T_analytical, rtol=0, atol=1e-2).all()


# tests/unit_tests/sensitivity/test_morris.py

from UQpy.run_model.RunModel import RunModel
from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.sensitivity.MorrisSensitivity import MorrisSensitivity
from UQpy.distributions import Uniform
import pytest


# os.chdir('./tests/Sensitivity')


@pytest.fixture
def setup():
    model = PythonModel(model_script='pfn.py', model_object_name='gfun_sensitivity', delete_files=True,
                        a_values=[0.001, 99.], var_names=['X{}'.format(i) for i in range(2)])
    runmodel_object = RunModel(model=model)
    dist_object = [Uniform(), Uniform()]
    yield runmodel_object, dist_object


def test_morris(setup):
    runmodel_object, dist_object = setup
    sens = MorrisSensitivity(runmodel_object=runmodel_object, distributions=dist_object, n_levels=8,
                             random_state=123, n_trajectories=3)
    assert round(sens.mustar_indices[1], 3) == 0.025


def test_morris_2(setup):
    runmodel_object, dist_object = setup
    sens = MorrisSensitivity(runmodel_object=runmodel_object, distributions=dist_object, n_levels=9,
                             random_state=123)
    sens.run(n_trajectories=2)
    sens.run(n_trajectories=2)
    assert round(sens.mustar_indices[1], 3) == 0.034


def test_morris_max_dispersion(setup):
    runmodel_object, dist_object = setup
    sens = MorrisSensitivity(runmodel_object=runmodel_object, distributions=dist_object, n_levels=9,
                             random_state=123, maximize_dispersion=True)
    sens.run(n_trajectories=5)
    assert round(sens.mustar_indices[1], 3) == 0.051


# tests/unit_tests/sensitivity/test_sobol.py

"""
This is the test module for Sobol sensitivity indices.

Here, we will use the Ishigami function to test the output.

The following methods are tested:
1. generate_pick_and_freeze_samples
2. pick_and_freeze_estimator (First and Total order Sobol indices)
3. pick_and_freeze_estimator (Second order Sobol indices) using [1]_.

References
----------

.. [1] Graham Glen, Kristin Isaacs, Estimating Sobol sensitivity indices using 
       correlations, Environmental Modelling & Software, Volume 37, 2012, Pages 157-166,
       ISSN 1364-8152,  https://doi.org/10.1016/j.envsoft.2012.03.014.


Important
----------
The computed indices are computed using the `np.isclose` function.

Function signature:
    numpy.isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)

    Parameters:
    a, b: array_like
        Input arrays to compare.

    rtol: float
        The relative tolerance parameter.

    atol: float
        The absolute tolerance parameter.

Each element of the `diff` array is compared as follows:
diff = |a - b|
diff <= atol + rtol * abs(b)

-   relative tolerance: rtol * abs(b)
    It is the maximum allowed difference between a and b, 
    relative to the absolute value of b. 
    For example, to set a tolerance of 1%, pass rol=0.01, 
    which assures that the values are within 2 decimal places of each other.
-   absolute tolerance: atol
    When b is close to zero, the atol value is used.

"""

import numpy as np
import pytest
import scipy

from UQpy.run_model.RunModel import RunModel
from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.distributions import Uniform
from UQpy.distributions.collection.JointIndependent import JointIndependent
from UQpy.sensitivity.sobol import Sobol

# Prepare
###############################################################################

# Prepare the input distribution
@pytest.fixture()
def ishigami_input_dist_object():
    """
    This function returns the input distribution for the Ishigami function.

    X1 ~ Uniform(-pi, pi)
    X2 ~ Uniform(-pi, pi)
    X3 ~ Uniform(-pi, pi)

    """
    return JointIndependent([Uniform(-np.pi, 2 * np.pi)] * 3)


@pytest.fixture()
def ishigami_model_object():
    """This function creates the Ishigami run_model_object"""
    model = PythonModel(
        model_script="ishigami.py",
        model_object_name="evaluate",
        var_names=[r"$X_1$", "$X_2$", "$X_3$"],
        delete_files=True,
        params=[7, 0.1],
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def sobol_object(ishigami_model_object, ishigami_input_dist_object):
    """This function returns the Sobol object."""

    return Sobol(ishigami_model_object, ishigami_input_dist_object)


@pytest.fixture()
def analytical_ishigami_Sobol_indices():
    """
    Analytical Sobol indices for the Ishigami function.

    Copy-paste the following to reproduce the given indices:

        a = 7
        b = 0.1

        V1 = 0.5*(1 + (b*np.pi**4)/5)**2
        V2 = (a**2)/8
        V3 = 0

        VT3 = (8*(b**2)*np.pi**8)/225
        VT1 = V1 + VT3
        VT2 = V2

        total_variance = V2 + (b*np.pi**4)/5 + ((b**2) * np.pi**8)/18 + 0.5

        S = np.array([V1, V2, V3])/total_variance
        S_T = np.array([VT1, VT2, VT3])/total_variance

        S = np.around(S, 4)
        S_T = np.around(S_T, 4)

    """

    S1 = 0.3139
    S2 = 0.4424
    S3 = 0

    S_T1 = 0.5576
    S_T2 = 0.4424
    S_T3 = 0.2437

    S = np.array([S1, S2, S3])
    S_T = np.array([S_T1, S_T2, S_T3])

    return S.reshape(-1, 1), S_T.reshape(-1, 1)


@pytest.fixture()
def saltelli_ishigami_Sobol_indices(sobol_object):

    SA = sobol_object

    np.random.seed(12345)  #! set seed for reproducibility

    computed_indices = SA.run(n_samples=1_000_000)

    return computed_indices["sobol_i"], computed_indices["sobol_total_i"]


@pytest.fixture()
def NUM_SAMPLES():
    """This function returns the number of samples for bootstrapping"""

    num_bootstrap_samples = 10_000
    num_samples = 100_000

    return num_bootstrap_samples, num_samples


@pytest.fixture()
def bootstrap_sobol_index_variance(sobol_object, NUM_SAMPLES):

    #### SETUP ####
    SA = sobol_object

    np.random.seed(12345)  #! set seed for reproducibility

    confidence_level = 0.95
    delta = -scipy.stats.norm.ppf((1 - confidence_level) / 2)

    num_bootstrap_samples, n_samples = NUM_SAMPLES

    #### Compute indices ####
    computed_indices = SA.run(
        n_samples=n_samples,
        num_bootstrap_samples=num_bootstrap_samples,
        confidence_level=confidence_level,
    )

    First_order = computed_indices["sobol_i"].ravel()
    Total_order = computed_indices["sobol_total_i"].ravel()
    CI_first_order = computed_indices["CI_sobol_i"]
    CI_total_order = computed_indices["CI_sobol_total_i"]

    #### Compute variance ####
    upper_bound_first_order = CI_first_order[:, 1]
    upper_bound_total_order = CI_total_order[:, 1]

    std_bootstrap_first_order = (upper_bound_first_order - First_order) / delta
    std_bootstrap_total_order = (upper_bound_total_order - Total_order) / delta

    return std_bootstrap_first_order**2, std_bootstrap_total_order**2


@pytest.fixture()
def model_eval_sobol_index_variance():

    """
    For computational efficiency, the variance of the Sobol indices
    is precomputed using model evaluations with
    NUM_SAMPLES (num_repetitions=10_000, num_samples=100_000)

    Copy-paste the following code to generate the variance
    of the Sobol indices:

        runmodel_obj = RunModel(
                model_script='ishigami.py',
                var_names=['X1', 'X2', 'X3'],
                vec=True, delete_files=True)

        input_obj = JointInd([Uniform(-np.pi, 2*np.pi)]*3)

        SA = Sobol(runmodel_obj, input_obj)

        np.random.seed(12345) # for reproducibility

        num_repetitions, n_samples = 10_000, 100_000

        num_vars = 3

        sample_first_order = np.zeros((num_vars, num_repetitions))
        sample_total_order = np.zeros((num_vars, num_repetitions))

        for i in range(num_repetitions):
            S, S_T = SA.run(n_samples=n_samples)

            sample_first_order[:, i] = S.ravel()
            sample_total_order[:, i] = S_T.ravel()

        variance_first_order = np.var(sample_first_order, axis=1, ddof=1).reshape(-1, 1)
        variance_total_order = np.var(sample_total_order, axis=1, ddof=1).reshape(-1, 1)

        print(variance_first_order)
        print(variance_total_order)

    """

    variance_first_order = np.array([1.98518409e-05, 1.69268227e-05, 2.50390610e-05])

    variance_total_order = np.array([2.82995855e-05, 2.46373399e-05, 2.59811868e-05])

    return variance_first_order, variance_total_order


@pytest.fixture()
def sobol_g_function_input_dist_object():
    """
    This function returns the input distribution object for the Sobol G function.

    X1 ~ Uniform(0, 1)
    X2 ~ Uniform(0, 1)
    X3 ~ Uniform(0, 1)
    X4 ~ Uniform(0, 1)
    X5 ~ Uniform(0, 1)
    X6 ~ Uniform(0, 1)

    """

    dist_object = JointIndependent([Uniform(0, 1)] * 6)

    return dist_object


@pytest.fixture()
def sobol_g_function_model_object():
    """This function creates the Sobol g-function model object"""

    a_vals = np.array([0.0, 0.5, 3.0, 9.0, 99.0, 99.0])

    model = PythonModel(
        model_script="sobol_func.py",
        model_object_name="evaluate",
        delete_files=True,
        a_values=a_vals,
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def sobol_object_g_func(
    sobol_g_function_input_dist_object, sobol_g_function_model_object
):
    """This function creates the Sobol object for the g-function"""

    sobol_object = Sobol(
        sobol_g_function_model_object, sobol_g_function_input_dist_object
    )

    return sobol_object


@pytest.fixture()
def analytical_sobol_g_func_second_order_indices():
    """
    This function returns the analytical second order Sobol indices for the g-function

    The values were obtained from [1]_.

    """

    S12 = 0.0869305
    S13 = 0.0122246
    S14 = 0.00195594
    S15 = 0.00001956
    S16 = 0.00001956
    S23 = 0.00543316
    S24 = 0.00086931
    S25 = 0.00000869
    S26 = 0.00000869
    S34 = 0.00012225
    S35 = 0.00000122
    S36 = 0.00000122
    S45 = 0.00000020
    S46 = 0.00000020
    S56 = 2.0e-9

    S_2 = [S12, S13, S14, S15, S16, S23, S24, S25, S26, S34, S35, S36, S45, S46, S56]

    return np.array(S_2).reshape(-1, 1)


@pytest.fixture()
def saltelli_sobol_g_function(sobol_object_g_func):

    SA = sobol_object_g_func

    np.random.seed(12345)  #! set seed for reproducibility

    # Compute Sobol indices using the pick and freeze algorithm
    # Save only second order indices
    computed_indices = SA.run(n_samples=100_000, estimate_second_order=True)

    return computed_indices["sobol_ij"]


# Unit tests
###############################################################################


def test_pick_and_freeze_estimator(
    analytical_ishigami_Sobol_indices, saltelli_ishigami_Sobol_indices
):

    """
    Test the Saltelli pick and freeze estimator using 1_000_000 samples.
    """

    # Prepare
    S_analytical, S_T_analytical = analytical_ishigami_Sobol_indices
    S_saltelli, S_T_saltelli = saltelli_ishigami_Sobol_indices

    # Act
    assert S_analytical.shape == S_saltelli.shape
    assert S_T_analytical.shape == S_T_saltelli.shape
    # Idea: Measure accuracy upto 2 decimal places -> rtol=0, atol=1e-2
    assert np.isclose(S_saltelli, S_analytical, rtol=0, atol=1e-2).all()
    assert np.isclose(S_T_saltelli, S_T_analytical, rtol=0, atol=1e-2).all()


def test_bootstrap_variance_computation(
    model_eval_sobol_index_variance, bootstrap_sobol_index_variance
):

    """Test the bootstrap variance computation."""

    # Prepare
    var_first, var_total = model_eval_sobol_index_variance
    boot_var_first, boot_var_total = bootstrap_sobol_index_variance

    # Act
    assert var_first.shape == boot_var_first.shape
    assert var_total.shape == boot_var_total.shape

    # Idea: Ensure bootstrap variance and MC variance are of same order -> rtol=0, atol=1e-4
    assert np.isclose(boot_var_first, var_first, rtol=0, atol=1e-4).all()
    assert np.isclose(boot_var_total, var_total, rtol=0, atol=1e-4).all()


def test_second_order_indices(
    analytical_sobol_g_func_second_order_indices, saltelli_sobol_g_function
):

    """Test the second order indices computation."""

    # Prepare
    S_2_analytical = analytical_sobol_g_func_second_order_indices
    S_2 = saltelli_sobol_g_function

    # Act
    # Idea: Ensure second order indices are of same order -> rtol=0, atol=1e-4
    assert np.isclose(S_2, S_2_analytical, rtol=0, atol=1e-2).all()


# tests/unit_tests/sensitivity/test_generalised_sobol.py
""""
This is the test module for the Generalised Sobol indices.

Here, we will use the toy example from [1]_, which is a multi-output problem.


References
----------

.. [1]  Gamboa F, Janon A, Klein T, Lagnoux A, others. 
        Sensitivity analysis for multidimensional and functional outputs.
        Electronic journal of statistics 2014; 8(1): 575-603.

Important
----------
The computed indices are computed using the `np.isclose` function.

Function signature:
    numpy.isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)

    Parameters:
    a, b: array_like
        Input arrays to compare.

    rtol: float
        The relative tolerance parameter.

    atol: float
        The absolute tolerance parameter.

Each element of the `diff` array is compared as follows:
diff = |a - b|
diff <= atol + rtol * abs(b)

-   relative tolerance: rtol * abs(b)
    It is the maximum allowed difference between a and b, 
    relative to the absolute value of b. 
    For example, to set a tolerance of 1%, pass rol=0.01, 
    which assures that the values are within 2 decimal places of each other.
-   absolute tolerance: atol
    When b is close to zero, the atol value is used.

"""

import numpy as np
import pytest
import scipy

from UQpy.run_model.RunModel import RunModel
from UQpy.run_model.model_execution.PythonModel import PythonModel
from UQpy.distributions import Uniform, Normal
from UQpy.distributions.collection.JointIndependent import JointIndependent
from UQpy.sensitivity.generalised_sobol import GeneralisedSobol

# Prepare
###############################################################################

# Prepare the input distribution
@pytest.fixture()
def normal_input_dist_object():
    """
    This function returns the input distribution for the toy model.

    X1 ~ Normal(0, 1)
    X2 ~ Normal(0, 1)

    """
    return JointIndependent([Normal(0, 1)] * 2)


@pytest.fixture()
def uniform_input_dist_object():
    """
    This function returns the input distribution for the toy model.

    X1 ~ Uniform(0, 1)
    X2 ~ Uniform(0, 1)

    """
    return JointIndependent([Uniform(0, 1)] * 2)


@pytest.fixture()
def toy_model_object():
    """
    This function creates the toy model.

    """
    model = PythonModel(
        model_script="multioutput.py",
        model_object_name="evaluate",
        var_names=[
            "X_1",
            "X_2",
        ],
        delete_files=True,
    )

    runmodel_obj = RunModel(model=model)

    return runmodel_obj


@pytest.fixture()
def generalised_sobol_object_normal(normal_input_dist_object, toy_model_object):
    """
    This function creates the Generalised Sobol indices object
    with normal input distribution.

    """

    return GeneralisedSobol(toy_model_object, normal_input_dist_object)


@pytest.fixture()
def generalised_sobol_object_uniform(uniform_input_dist_object, toy_model_object):
    """
    This function creates the Generalised Sobol indices object
    with uniform input distribution.

    """

    return GeneralisedSobol(toy_model_object, uniform_input_dist_object)


@pytest.fixture()
def analytical_toy_GSI_normal():
    """
    Analytical first order Generalised Sobol indices
    for the toy example with normal input distribution.
    """

    return np.array([0.2941, 0.1176]).reshape(-1, 1)


@pytest.fixture()
def analytical_toy_GSI_uniform():
    """ "
    Analytical first order Generalised Sobol indices
    for toy example with uniform input distribution.
    """

    return np.array([0.6084, 0.3566]).reshape(-1, 1)


@pytest.fixture()
def pick_and_freeze_toy_GSI_normal(generalised_sobol_object_normal):
    """ "
    Generalised first order Sobol indices computed using the Pick and Freeze
    approach for the toy example with normal input distribution.
    """

    SA = generalised_sobol_object_normal

    np.random.seed(12345)  #! set seed for reproducibility

    computed_indices = SA.run(n_samples=100_000)

    return computed_indices["gen_sobol_i"]


@pytest.fixture()
def pick_and_freeze_toy_GSI_uniform(generalised_sobol_object_uniform):
    """ "
    Generalised first order Sobol indices computed using the Pick and Freeze
    approach for the toy example with uniform input distribution.
    """

    SA = generalised_sobol_object_uniform

    np.random.seed(12345)  #! set seed for reproducibility

    computed_indices = SA.run(n_samples=100_000)

    return computed_indices["gen_sobol_i"]


@pytest.fixture()
def NUM_SAMPLES():
    """This function returns the number of samples for bootstrapping"""

    num_bootstrap_samples = 500
    num_samples = 20_000

    return num_bootstrap_samples, num_samples


@pytest.fixture()
def bootstrap_generalised_sobol_index_variance(
    generalised_sobol_object_normal, NUM_SAMPLES
):

    SA = generalised_sobol_object_normal

    np.random.seed(12345)  #! set seed for reproducibility

    num_bootstrap_samples, n_samples = NUM_SAMPLES

    confidence_level = 0.95
    delta = -scipy.stats.norm.ppf((1 - confidence_level) / 2)

    # Compute the confidence intervals

    computed_indices = SA.run(
        n_samples=n_samples,
        num_bootstrap_samples=num_bootstrap_samples,
        confidence_level=confidence_level,
    )

    gen_sobol_i = computed_indices["gen_sobol_i"].ravel()
    gen_sobol_total_i = computed_indices["gen_sobol_total_i"].ravel()
    upper_bound_first_order = computed_indices["CI_gen_sobol_i"][:, 1]
    upper_bound_total_order = computed_indices["CI_gen_sobol_total_i"][:, 1]

    std_bootstrap_first_order = (upper_bound_first_order - gen_sobol_i) / delta
    std_bootstrap_total_order = (upper_bound_total_order - gen_sobol_total_i) / delta

    return std_bootstrap_first_order**2, std_bootstrap_total_order**2


@pytest.fixture()
def model_eval_generalised_sobol_index_variance():

    """
    For computational efficiency, the variance of the generalised Sobol indices
    is precomputed using model evaluations with
    NUM_SAMPLES (num_repetitions=500, num_samples=20_000)

    Copy-paste the following code to generate the variance
    of the Sobol indices:

    runmodel_obj = RunModel(model_script='multioutput.py',
                        model_object_name='multioutput_toy',
                        vec=True, delete_files=True)

    dist_object_1 = JointInd([Normal(0, 1)]*2)

    SA = GeneralisedSobol(runmodel_obj, dist_object_1)

    np.random.seed(12345) # for reproducibility

    num_repetitions, n_samples = 500, 20_000

    num_vars = 2

    bootstrap_first_order = np.zeros((num_vars, num_bootstrap_samples))
    bootstrap_total_order = np.zeros((num_vars, num_bootstrap_samples))

    for b in range(num_repetitions):

        computed_indices = SA.run(n_samples=n_samples)

        bootstrap_first_order[:, b] = computed_indices["gen_sobol_i"].ravel()
        bootstrap_total_order[:, b] = computed_indices["gen_sobol_total_i"].ravel()

    var_bootstrap_gen_S = np.var(bootstrap_first_order, axis=1, ddof=1)
    var_bootstrap_gen_S_T = np.var(bootstrap_total_order, axis=1, ddof=1)

    print(var_bootstrap_gen_S)
    print(var_bootstrap_gen_S_T)

    """

    variance_first_order = np.array([0.00011284, 0.00012608])

    variance_total_order = np.array([0.00012448, 0.00011208])

    return variance_first_order, variance_total_order


# Unit tests
###############################################################################


def test_pick_and_freeze_estimator(
    pick_and_freeze_toy_GSI_normal,
    analytical_toy_GSI_normal,
    pick_and_freeze_toy_GSI_uniform,
    analytical_toy_GSI_uniform,
):
    """
    Test the pick and freeze estimator.

    """

    # Prepare
    N_true = analytical_toy_GSI_normal
    N_estimate = pick_and_freeze_toy_GSI_normal

    U_true = analytical_toy_GSI_uniform
    U_estimate = pick_and_freeze_toy_GSI_uniform

    # Act
    # Idea: Measure accuracy upto 2 decimal places -> rtol=0, atol=1e-2
    assert np.isclose(N_estimate, N_true, rtol=0, atol=1e-2).all()
    assert np.isclose(U_estimate, U_true, rtol=0, atol=1e-2).all()


def test_bootstrap_variance_computation(
    model_eval_generalised_sobol_index_variance,
    bootstrap_generalised_sobol_index_variance,
):

    """Test the bootstrap variance computation."""

    # Prepare
    var_first, var_total = model_eval_generalised_sobol_index_variance
    boot_var_first, boot_var_total = bootstrap_generalised_sobol_index_variance

    # Act
    assert var_first.shape == boot_var_first.shape

    # Idea: Ensure bootstrap variance and MC variance are of same order -> rtol=0, atol=1e-4
    assert np.isclose(boot_var_first, var_first, rtol=0, atol=1e-4).all()
    assert np.isclose(boot_var_total, var_total, rtol=0, atol=1e-4).all()
