   max_rank = max(rnk)
        fmean = []
        for i in range(n_mat):
            fmean.append(Grassmann.frechet_variance(data_points[i], data_points, distance_fun))

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

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

        itera = 0

        l = 0
        avg = []
        _gamma = []
        if acc:
            _gamma = Grassmann.log_map(points_grassmann=data_points, ref=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 = Grassmann.log_map(points_grassmann=data_points, ref=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 = Grassmann.exp_map(points_tangent=[step], ref=np.asarray(mean_element))

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

            if test_1 < tol:
                break

            mean_element = []
            mean_element = x[0]

            itera += 1

        # return the Karcher mean.
        return mean_element

    @staticmethod
    def _stochastic_gradient_descent(data_points, distance_fun, **kwargs):
        """
        Compute the Karcher mean using the stochastic gradient descent method.
        This method computes the Karcher mean given a set of points on the Grassmann manifold. In this regard, the
        ``stochastic_gradient_descent`` method is implemented herein. Further, this method is called by the method
        ``karcher_mean``.
        **Input:**
        * **data_points** (`list`)
            Points on the Grassmann manifold.

        * **distance_fun** (`callable`)
            Distance function.
        * **kwargs** (`dictionary`)
            Contains the keywords for the used in the optimizers to find the Karcher mean.
        **Output/Returns:**
        * **mean_element** (`list`)
            Karcher mean.
        """

        if 'tol' in kwargs.keys():
            tol = kwargs['tol']
        else:
            tol = 1e-3

        if 'maxiter' in kwargs.keys():
            maxiter = kwargs['maxiter']
        else:
            maxiter = 1000

        n_mat = len(data_points)

        rnk = []
        for i in range(n_mat):
            rnk.append(min(np.shape(data_points[i])))

        max_rank = max(rnk)

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

        index_0 = fmean.index(min(fmean))

        mean_element = data_points[index_0].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 = Grassmann.log_map(points_grassmann=[data_points[idx]], ref=np.asarray(mean_element))

                step = 2 * alpha * _gamma[0]

                X = Grassmann.exp_map(points_tangent=[step], ref=np.asarray(mean_element))

                _gamma = []
                mean_element = X[0]

                k += 1

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

            itera += 1

        return mean_element




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

import copy
from typing import Union

import numpy as np
from scipy.interpolate import LinearNDInterpolator

from UQpy.dimension_reduction.grassmann_manifold import Grassmann
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 ManifoldInterpolation:

    def __init__(self, interpolation_method: Union[Surrogate, callable, None],
                 manifold_data: list[GrassmannPoint],
                 coordinates: 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.
        :param manifold_data: Data on the Grassmann manifold.
        :param optimization_method: Optimization method for calculating the Karcher mean.
        :param coordinates: Nodes of the interpolant.
        :param distance:  Distance metric.
        """
        self.interpolation_method = interpolation_method

        self.mean = Grassmann.karcher_mean(manifold_points=manifold_data,
                                           optimization_method=optimization_method,
                                           distance=distance)

        self.tangent_points = Grassmann.log_map(manifold_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: NumpyFloatArray, ):
        """
        :param point:  Point to interpolate.
        """

        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 Grassmann.exp_map(tangent_points=[interp_point], reference_point=self.mean)


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

from UQpy.dimension_reduction.grassmann_manifold.Grassmann import Grassmann
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):
        """
        Executes proper orthogonal decomposition using the :class:`.DirectPOD` algorithm.
        """
        super().run()

    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):
        """
        Executes proper orthogonal decomposition using the :class:`.SnapshotPOD` algorithm.
        """
        super().run()

    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.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.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


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 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**.
         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 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**.
         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`
        
        """
        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.alpha = 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.jzx = None

        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=None, seed_u=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:", 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.jzx = 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.jzx = 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.jzx))

            # 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, Optional

import numpy as np
import scipy.stats as stats
from beartype import beartype

from UQpy.distributions.baseclass import DistributionContinuous1D
from UQpy.distributions.collection import JointIndependent
from UQpy.transformations.Nataf import Nataf
from UQpy.RunModel import RunModel
from UQpy.distributions.baseclass import Distribution
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 = None
        self.beta_second_order = None
        """Hasofer-Lind reliability index using the SORM method."""
        self.call = None

        self.form_object = form_object
        self._run()

    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/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 = False
        if max(eff[:, 0]) <= self.eff_stop:
            stopping_criteria_indicator = True

        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 = False
        if max(eif[:, 0]) / abs(fm) <= self.eif_stop:
            stopping_criteria_indicator = True

        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 = False
        if min(u[:, 0]) >= self.u_stop:
            stopping_criteria_indicator = True

        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 = False
        if min(u[:, 0]) >= self.weighted_u_stop:
            stopping_criteria_indicator = True

        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, List

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 *
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/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/geometric_harmonics/GeometricHarmonics.py

from UQpy.dimension_reduction.diffusion_maps.DiffusionMaps import DiffusionMaps
from UQpy.utilities import *
from UQpy.utilities.kernels import GaussianKernel


class GeometricHarmonics:
    """
    Geometric Harmonics for domain extension.
    """

    def __init__(self, eigenvectors_number: int = 5, kernel_object=GaussianKernel()):
        """

        :param eigenvectors_number: The number of eigenvectors used in the decomposition of the kernel matrix.
        :param kernel_object: Kernel used for the construction of the geometric harmonics.

        See Also
        --------

        :py:class:`.GaussianKernel`

        """

        self.kernel_object = kernel_object
        self.n_eigen_pairs = eigenvectors_number
        self.basis = None
        self.x = None
        self.y = None
        self.eigenvalues = None
        self.eigenvectors = None
        self.kwargs_kernel = None

    def fit(self, x: Numpy2DFloatArray, y: Numpy2DFloatArray, **kwargs):

        """
        Fit model with training data.

        :param x: Training points of shape :code:`(n_samples, n_features)`.
        :param y: Target function values of shape :code:`(n_samples, n_targets)`
        :param kwargs: Scale parameter of the kernel. if :code:`epsilon is None` then the parameters for the
         estimation of the kernel's scale must be provided.

        See Also
        --------

        :py:meth:`.DiffusionMaps.estimate_epsilon`

        """

        self.x = x
        self.y = y
        self.kwargs_kernel = kwargs

        if kwargs['epsilon'] is None:

            epsilon, _ = DiffusionMaps.estimate_epsilon(x, cut_off=kwargs['cut_off'], tol=kwargs['tol'],
                                                        k_nn=kwargs['k_nn'], n_partition=kwargs['n_partition'],
                                                        distance_matrix=kwargs['distance_matrix'],
                                                        random_state=kwargs['random_state'])

        else:
            epsilon = kwargs['epsilon']

        self.kernel_object.epsilon = epsilon
        kernel_matrix = self.kernel_object.calculate_kernel_matrix(points=x)

        eigenvalues, eigenvectors = DiffusionMaps.eig_solver(kernel_matrix, self.n_eigen_pairs)

        idv = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idv]
        eigenvectors = eigenvectors[:, idv]

        # Basis construction.
        basis = eigenvectors.dot(np.diag(np.reciprocal(eigenvalues))).dot(eigenvectors.T) @ y
        self.basis = basis
        self.x = x
        self.y = y
        self.eigenvalues = eigenvalues
        self.eigenvectors = eigenvectors

    def predict(self, x: Numpy2DFloatArray) -> tuple[Numpy2DFloatArray, NumpyFloatArray]:

        """
        Evaluate model for out-of-sample points.

        :param x:  Points of shape :code:`(n_samples, n_features)`
        :return: The interpolated function values of shape :code:`(n_samples, n_targets)`
        """

        if self.kwargs_kernel['epsilon'] is None:

            epsilon, _ = DiffusionMaps.estimate_epsilon(x, cut_off=self.kwargs_kernel['cut_off'],
                                                        tol=self.kwargs_kernel['tol'],
                                                        k_nn=self.kwargs_kernel['k_nn'],
                                                        n_partition=self.kwargs_kernel['n_partition'],
                                                        distance_matrix=self.kwargs_kernel['distance_matrix'],
                                                        random_state=self.kwargs_kernel['random_state'])
        else:
            epsilon = self.kwargs_kernel['epsilon']

        self.kernel_object.epsilon = epsilon
        kernel_matrix = self.kernel_object.calculate_kernel_matrix(points=x)
        y = kernel_matrix.T @ self.basis
        score = self.score(x, y, self.kwargs_kernel['score'])
        return y, score

    @staticmethod
    def score(y: Numpy2DFloatArray, y_predicted: Numpy2DFloatArray, kind: str = 'abs') -> float:

        """
        Score interpolation model with negative mean squared error metric.

        :param y: The exact function values of shape :code:`(n_test, n_targets)`
        :param y_predicted: The interpolated function values of shape :code:`(n_test, n_targets)`
        :param kind: Error metric to be used. If :code:`kind is None` then the absolute mean error will be calculated.
        :return: The mean computed error between the exact values and the predictions.
        """
        n_samples = y.shape[0]
        error = []
        for i in range(n_samples):
            # Compute the error between the exact value and the prediction.
            if kind == 'abs':
                error_value = np.linalg.norm(y_predicted[i, :] - y[i, :])
            elif kind == 'l1':
                error_value = np.linalg.norm(y_predicted[i, :] - y[i, :], 1)
            elif kind == 'max':
                error_value = np.max(abs(y_predicted[i, :] - y[i, :]))
            elif kind == 'rel':
                error_value = np.linalg.norm(y_predicted[i, :] - y[i, :]) / np.linalg.norm(y[i, :])
            else:
                raise NotImplementedError('UQpy: Not implement kind of error.')

            error.append(error_value)

        return np.mean(error)


# src/UQpy/surrogates/geometric_harmonics/Grassmannian_GeometricHarmonics.py

import numpy as np
from scipy.spatial.distance import pdist
from sklearn.neighbors import NearestNeighbors #todo: Change to scipy (using KD.Tree)


class GrassmannianGHMap:
    """
    Geometric Harmonics for domain extension.
    The class ``GeometricHarmonics`` is used in the domain extension of functions defined only on few observations.
    ``GeometricHarmonics`` is a Subclass of Similarity.
    **Input:**
    * **n_evecs** (`int`)
        The number of eigenvectors used in the eigendecomposition of the kernel matrix.
    * **kernel_method** (`callable`)
        Kernel method used in the construction of the geometric harmonics.
    **Attributes:**
    * **n_evecs** (`int`)
        The number of eigenvectors used in the eigendecomposition of the kernel matrix.
    * **X** (`list`)
        Independent variables.
    * **y** (`list`)
        Function values.
    * **basis** (`list`)
        Basis used in the domain extension.
    **Methods:**
    """

    def __init__(self,n_eigenvectors=None, n_neighbors=None, p='max',
                 epsilon_gh0=None, epsilon_gh1=None,
                 diffusion_kernel_object=ProjectionKernel()):

        # kernel_object=Gaussian()
        #self.kernel_object = kernel_object
        self.p = p
        self.diffusion_kernel_object = diffusion_kernel_object
        #self.n_evecs = n_evecs
        self.X = None
        self.y = None

        self.diffusion_object = None
        self.grassmann_object = None
        self.gh0 = None
        self.gh1 = None
        self.epsilon_gh0 = epsilon_gh0
        self.epsilon_gh1 = epsilon_gh1
        self.n_eigenvectors = n_eigenvectors
        self.n_neighbors = n_neighbors
        self.dcoords = None

    def fit(self, X, y, **kwargs):

        """
        Train the model using `fit`.
        In this method, `X` is a list of data points, `y` are the function values. `epsilon` can be
        provided, otherwise it is computed from the median of the pairwise distances of X.
        **Input:**
        * **X** (`list`)
            Input data (independent variables).
        * **y** (`list`)
            Function values.
        * **epsilon** (`float`)
            Parameter of the Gaussian kernel.
        **Output/Returns:**
        """

        if X is not None:

            self.X = X
            self.y = y
            #if not isinstance(X, list):
            #    raise TypeError('UQpy: `X` must be a list.')

            #if not isinstance(y, list):
            #    raise TypeError('UQpy: `y` must be a list.')

            self.grassmann_object = Grassmann(p=self.p)
            self.grassmann_object.fit(X=self.y)

            self.diffusion_object = GrassmannianDiffusionMaps(alpha=0.5, n_evecs=len(X), sparse=False, k_neighbors=1,
                                                              kernel_composition='prod',
                                                              kernel_object=self.diffusion_kernel_object,
                                                              p=self.p, orthogonal=False) #todo: change n_evecs!!!!

            self.diffusion_object.fit(X=self.grassmann_object)

            self.dcoords = self.diffusion_object.dcoords

            self.gh0 = GeometricHarmonics(n_evecs=self.n_eigenvectors)
            self.gh0.fit(X, self.dcoords, epsilon=self.epsilon_gh0)

            self.gh1 = GeometricHarmonics(n_evecs=self.n_eigenvectors)
            self.gh1.fit(X, self.grassmann_object.sigma, epsilon=self.epsilon_gh0)

    def predict(self, X):

        dcoords = self.dcoords
        dcoord_interp = self.gh0.predict(X, )
        sig_interp = self.gh1.predict(X, )

        nbrs = NearestNeighbors(n_neighbors=self.n_neighbors, algorithm='ball_tree').fit(dcoords)
        distances, indices = nbrs.kneighbors(dcoord_interp)
        pid = indices[0]

        psi_gr = []
        phi_gr = []
        for k in pid:
            psi_gr.append(self.grassmann_object.psi[k])
            phi_gr.append(self.grassmann_object.phi[k])

        # start = timeit.default_timer()
        # ref_psi = Gr.karcher_mean(points_grassmann=psi_gr)
        # ref_phi = Gr.karcher_mean(points_grassmann=phi_gr)
        ref_psi = psi_gr[0]
        ref_phi = phi_gr[0]
        # stop = timeit.default_timer()
        # print('Time: ', stop - start)

        gpsi = self.grassmann_object.log_map(points_grassmann=psi_gr, ref=ref_psi)
        gphi = self.grassmann_object.log_map(points_grassmann=phi_gr, ref=ref_phi)

        psi = []
        sigma = []
        phi = []
        diffc = []

        gamma_psi = []
        gamma_phi = []
        ptv = []
        n0 = np.shape(self.grassmann_object.psi[pid[0]])[0]
        n1 = np.shape(self.grassmann_object.psi[pid[0]])[1]
        n0h = np.shape(self.grassmann_object.phi[pid[0]])[0]
        n1h = np.shape(self.grassmann_object.phi[pid[0]])[1]
        for c, k in enumerate(pid):
            #n0 = np.shape(self.grassmann_object.psi[k])[0]
            #n1 = np.shape(self.grassmann_object.psi[k])[1]

            gamma_psi.append(gpsi[c].reshape(n0 * n1))
            gamma_phi.append(gphi[c].reshape(n0h * n1h))

            psi.append(self.grassmann_object.psi[k].reshape(n0 * n1))
            sigma.append(self.grassmann_object.sigma[k].reshape(len(self.grassmann_object.sigma[k])))
            phi.append(self.grassmann_object.phi[k].reshape(n0h * n1h))
            diffc.append(self.dcoords[k, :])
            ptv.append(self.X[k, :])

        #g0 = dcoord_interp
        #dg = pdist(diffc)
        epsilon_local = (np.median(pdist(diffc)) ** 2) / 4

        gh_local = GeometricHarmonics()

        gh_local.fit(diffc, gamma_psi, epsilon=epsilon_local)
        gpsipred_ = gh_local.predict(dcoord_interp, )

        gh_local.fit(diffc, gamma_phi, epsilon=epsilon_local)
        gphipred_ = gh_local.predict(dcoord_interp, )

        #gh_local.fit(diffc, sigma, epsilon=epsilon_local)
        #sigpred_ = gh_local.predict(dcoord_interp)

        gpsipred = gpsipred_[0].reshape(n0, n1)
        gphipred = gphipred_[0].reshape(n0h, n1h)
        psipred_ = self.grassmann_object.exp_map(points_tangent=[gpsipred], ref=ref_psi)[0]
        phipred_ = self.grassmann_object.exp_map(points_tangent=[gphipred], ref=ref_phi)[0]
        #sigpred_ = np.diag(sigpred_[0])

        #Stest = out[pid[0]]
        ut, st, vt = svd(self.grassmann_object.X[pid[0]], rank=self.grassmann_object.p)

        psi_ = []
        phi_ = []
        for i in range(len(ut[0, :])):
            err1 = np.linalg.norm(psipred_[:, i] - ut[:, i])
            err2 = np.linalg.norm(phipred_[:, i] - vt[:, i])

            if err1 > 1:
                psi_.append(-psipred_[:, i])
            else:
                psi_.append(psipred_[:, i])

            if err2 > 1:
                phi_.append(-phipred_[:, i])
            else:
                phi_.append(phipred_[:, i])
            # phipred_[:,i]

        psi_ = np.array(psi_).T
        phi_ = np.array(phi_).T
        sigma_ = sig_interp[0]

        y = psi_ @ np.diag(sigma_) @ phi_.T

        return y, psi_, sigma_, phi_, dcoord_interp

# src/UQpy/surrogates/geometric_harmonics/__init__.py

from UQpy.surrogates.geometric_harmonics.GeometricHarmonics import GeometricHarmonics


# src/UQpy/surrogates/kriging/Kriging.py

import logging
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 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.PolynomialBasis 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.

        :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.inv(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.

        :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.
        
        :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 defined as:

    .. math::

        k_p(x_j, x_i) = det(x_j'\cdot xj)^2

    """
    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
from UQpy.utilities.kernels import EuclideanKernel
import scipy.spatial.distance as sd


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
        """
        self.epsilon = epsilon

    def kernel_entry(self, xi, xj):
        return np.linalg.norm(xi - xj, "fro") ** 2

    def kernel_function(self, distance_pairs):
        return np.exp(-sd.squareform(distance_pairs) / (4 * self.epsilon))


# 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 defined as:

    .. math::

        k_p(x_j, x_i) = (||x_j'\cdot xj||_F)^2

    """
    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



# 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)


# 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,
            data: np.ndarray,
            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.data = data
        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: 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.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.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 = list()
        d22 = list()

        # 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(samples_number=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 = 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

            self.mcmc_objects.append(copy.deepcopy(self._sampling_class))

            # 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)
            )

            return independent_chains_cov, dependent_chains_cov
        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 = beta + 1
        beta = beta * 2

        ar = np.asarray(self.mcmc_objects[step].acceptance_rate)
        ar_mean = np.mean(ar)

        factor = 0
        for i in range(np.shape(g)[0] - 1):
            factor = factor + (1 - (i + 1) * np.shape(g)[0] / np.shape(g)[1]) * (
                1 - ar_mean
            )
        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/sampling/AdaptiveKriging.py

import logging
from beartype import beartype
from sklearn.gaussian_process import GaussianProcessRegressor

from UQpy.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:
            if 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.")
        else:
            if 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:
            if 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:
            if 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)
        else:
            if 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 callable(log_pdf_):
                if args is None:
                    args = ()
                evaluate_log_pdf = lambda x: log_pdf_(x, *args)
            else:
                raise TypeError("UQpy: log_pdf_target must be a callable")
        # pdf is provided
        elif pdf_ is not None:
            if callable(pdf_):
                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 TypeError("UQpy: pdf_target must be a callable")
        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 = add_continuous_1d + 1
                elif isinstance(distributions[i], DistributionND):
                    add_continuous_nd = 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 = list()
            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 = list()
            for j in range(nsamples):
                y = list()
                for k in range(len(self.dist_object)):
                    y.append(temp_samples[k][j])
                self.x.append(np.array(y))
        else:
            if 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)
        else:
            # If self.samples already has existing samples, append the new samples to the existing attribute.
            if 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 hasattr(self.dist_object, "cdf"):
                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
            else:
                raise ValueError("UQpy: All distributions must have a cdf method.")

        elif isinstance(self.dist_object, list) and self.list is True:
            temp_samples_u01 = list()
            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.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.samples: NumpyFloatArray = 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)
        return 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.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.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
        self.total_order_indices = None
        self.generalized_first_order_indices = None
        self.generalized_total_order_indices = None

    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.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.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 . import MorrisSensitivity


# src/UQpy/stochastic_process/BispectralRepresentation.py

import itertools
from UQpy.utilities import *


class BispectralRepresentation:
    def __init__(
        self,
        samples_number: int,
        power_spectrum: Union[list, np.ndarray],
        bispectrum: Union[list, np.ndarray],
        time_interval: Union[list, np.ndarray],
        frequency_interval: Union[list, np.ndarray],
        number_time_intervals: Union[list, np.ndarray],
        number_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 samples_number: Number of samples of the stochastic process to be simulated.
         The :meth:`run` method is automatically called if `samples_number` is provided. If `samples_number` 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 `number_frequency_intervals`.

         - For uni-variate, multi-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(number_frequency_intervals[0], ..., number_frequency_intervals[number_of_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:`(number_frequency_intervals, number_frequency_intervals)`

         - For uni-variate, multi-dimensional processes, `bispectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(number_frequency_intervals[0], ..., number_frequency_intervals[number_of_dimensions-1], number_frequency_intervals[0], ..., number_frequency_intervals[number_of_dimensions-1])`

        :param time_interval: Length of time discretizations (:math:`\Delta t`) for each dimension of size
         `number_of_dimensions`.
        :param frequency_interval: Length of frequency discretizations (:math:`\Delta \omega`) for each dimension of
         size `number_of_dimensions`.
        :param number_time_intervals: Number of time discretizations for each dimensions of size `number_of_dimensions`.
        :param number_frequency_intervals: Number of frequency discretizations for each dimension of size
         `number_of_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.nsamples = samples_number
        self.number_frequency_intervals = np.array(number_frequency_intervals)
        self.number_time_intervals = np.array(number_time_intervals)
        self.frequency_interval = np.array(frequency_interval)
        self.time_interval = np.array(time_interval)
        self.number_of_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.number_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:`(samples_number, number_of_variables, number_frequency_intervals[0], ...,
        number_frequency_intervals[number_of_dimensions-1])`"""
        self.samples: NumpyFloatArray = None
        """Generated samples.
        The shape of the samples is :code:`(samples_number, number_of_variables, number_time_intervals[0], ...,
        number_time_intervals[number_of_dimensions-1])`"""

        self.case = case

        if self.number_of_dimensions == len(self.power_spectrum.shape):
            self.case = "uni"
            self._compute_bicoherence_uni()
        else:
            self.number_of_variables: int = self.power_spectrum.shape[0]
            """Number of variables in the stochastic process."""
            self.case = "multi"

        if self.nsamples is not None:
            self.run(samples_number=self.nsamples)

    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.number_of_dimensions == 1:
            self.pure_power_sepctrum[0] = self.power_spectrum[0]
            self.pure_power_sepctrum[1] = self.power_spectrum[1]

        if self.number_of_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.number_of_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.number_frequency_intervals[i])
            for i in range(self.number_of_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.number_of_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.number_time_intervals)
        samples = samples[:, np.newaxis]
        return np.real(samples)

    def run(self, samples_number: 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 `samples_number` 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 samples_number: 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 samples_number is None:
            raise ValueError(
                "UQpy: Stochastic Process: Number of samples must be defined."
            )
        if not isinstance(samples_number, int):
            raise ValueError("UQpy: Stochastic Process: nsamples 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 :", self.number_of_dimensions
            )
            phi = (
                np.random.uniform(
                    size=np.append(
                        self.nsamples,
                        np.ones(self.number_of_dimensions, dtype=np.int32)
                        * self.number_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],
        number_time_intervals: Union[list, np.ndarray],
        number_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 number_time_intervals: The number of time discretizations.
        :param number_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.distributions = distributions
        self.frequency = np.arange(0, number_frequency_intervals) * frequency_interval
        self.time = np.arange(0, number_time_intervals) * time_interval
        self.error = percentage_error
        if (
            correlation_function_non_gaussian is None
            and power_spectrum_non_gaussian is None
        ):
            print(
                "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.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
            )
            """The inverse translated Gaussian samples from the non-Gaussian samples."""
        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 _inverse_translate_non_gaussian_samples(self):
        if hasattr(self.distributions, "cdf"):
            non_gaussian_cdf = getattr(self.distributions, "cdf")
            samples_cdf = non_gaussian_cdf(self.samples_non_gaussian)
        else:
            raise AttributeError(
                "UQpy: The marginal dist_object needs to have an inverse cdf defined."
            )
        samples_g = Normal(loc=0.0, scale=1.0).icdf(samples_cdf)
        return samples_g

    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,
        samples_number: 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 samples_number: Number of samples of the stochastic process to be simulated.
         The :meth:`run` method is automatically called if `samples_number` is provided. If `samples_number` 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:`(number_time_intervals, number_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
        if threshold:
            self.number_eigenvalues = threshold
        else:
            self.number_eigenvalues = 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.nsamples = samples_number

        self.samples: NumpyFloatArray = None
        """Array of generated samples."""
        self.xi: NumpyFloatArray = None
        """The independent gaussian random variables used in the expansion."""

        if self.nsamples is not None:
            self.run(samples_number=self.nsamples)

    def _simulate(self, xi):
        lam, phi = np.linalg.eig(self.correlation_function)
        lam = lam[: self.number_eigenvalues]
        lam = np.diag(lam)
        self.phi = np.real(phi[:, : self.number_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, samples_number: 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 `samples_number` 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 samples_number: 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 samples_number is None:
            raise ValueError(
                "UQpy: Stochastic Process: Number of samples must be defined."
            )
        if not isinstance(samples_number, int):
            raise ValueError("UQpy: Stochastic Process: nsamples 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.number_eigenvalues, self.nsamples))
        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,
        nsamples: int,
        power_spectrum: Union[list, np.ndarray, float],
        time_interval: Union[list, np.ndarray, float],
        frequency_interval: Union[list, np.ndarray, float],
        number_time_intervals: Union[list, np.ndarray, float],
        number_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 nsamples: Number of samples of the stochastic process to be simulated.
         The :meth:`run` method is automatically called if `nsamples` is provided. If `nsamples` 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 `number_frequency_intervals`.

         * For multi-variate, one-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(number_of_variables, number_of_variables, number_frequency_intervals)`.

         * For uni-variate, multi-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(number_frequency_intervals[0], ..., number_frequency_intervals[number_of_dimensions-1])`

         * For multi-variate, multi-dimensional processes, `power_spectrum` will be a :class:`list` or :class:`numpy.ndarray` of size :code:`(number_of_variables, number_of_variables, number_frequency_intervals[0],...,number_frequency_intervals[number_of_dimensions-1])`.

        :param time_interval: Length of time discretizations (:math:`\Delta t`) for each
         dimension of size `number_of_dimensions`.
        :param frequency_interval: Length of frequency discretizations
         (:math:`\Delta \omega`) for each dimension of size `number_of_dimensions`.
        :param number_time_intervals: Number of time discretizations for each dimensions of
         size `number_of_dimensions`.
        :param number_frequency_intervals: Number of frequency discretizations for each
         dimension of size `number_of_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(number_time_intervals, int)
            and isinstance(number_frequency_intervals, int)
        ):
            time_interval = [time_interval]
            frequency_interval = [frequency_interval]
            number_time_intervals = [number_time_intervals]
            number_frequency_intervals = [number_frequency_intervals]
        self.time_interval = np.array(time_interval)
        self.frequency_interval = np.array(frequency_interval)
        self.number_time_intervals = np.array(number_time_intervals)
        self.number_frequency_intervals = np.array(number_frequency_intervals)
        self.nsamples = nsamples

        # 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:`(nsamples, number_of_variables, number_time_intervals[0], ...,
        number_time_intervals[number_of_dimensions-1])`"""
        self.number_of_variables: int = None
        """Number of variables in the stochastic process."""
        self.number_of_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:`(nsamples, number_of_variables, number_frequency_intervals[0], ...,
        number_frequency_intervals[number_of_dimensions-1])`"""

        if self.number_of_dimensions == len(self.power_spectrum.shape):
            self.case = "uni"
        else:
            self.number_of_variables = self.power_spectrum.shape[0]
            self.case = "multi"

        # Run Spectral Representation Method
        if self.nsamples is not None:
            self.run(nsamples=self.nsamples)

    def run(self, nsamples):
        """
        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 `nsamples` 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 nsamples: 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 nsamples is None:
            raise ValueError(
                "UQpy: Stochastic Process: Number of samples must be defined."
            )
        if not isinstance(nsamples, 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 :", self.number_of_dimensions
            )
            phi = (
                np.random.uniform(
                    size=np.append(
                        self.nsamples,
                        np.ones(self.number_of_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 :",
                self.number_of_variables,
            )
            self.logger.info(
                "UQpy: Stochastic Process: The number of dimensions is :",
                self.number_of_dimensions,
            )
            phi = (
                np.random.uniform(
                    size=np.append(
                        self.nsamples,
                        np.append(
                            np.ones(self.number_of_dimensions, dtype=np.int32)
                            * self.number_frequency_intervals,
                            self.number_of_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.number_of_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.number_of_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.number_of_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],
            number_time_intervals: Union[list, np.ndarray, float],
            number_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 number_time_intervals: The number of time discretizations.
        :param number_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:`(number_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:`(number_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.distributions = distributions
        self.time_interval = time_interval
        self.frequency_interval = frequency_interval
        self.number_time_intervals = number_time_intervals
        self.number_frequency_intervals = number_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.number_frequency_intervals) * self.frequency_interval,
                np.arange(0, self.number_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.number_frequency_intervals) * self.frequency_interval,
                np.arange(0, self.number_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.samples_shape = samples_gaussian.shape
            self.samples_gaussian: NumpyFloatArray = samples_gaussian.flatten()[:, np.newaxis]
            """Translated non-Gaussian stochastic process from Gaussian samples."""
            self.samples_non_gaussian = self._translate_gaussian_samples().reshape(
                self.samples_shape
            )
        (
            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.number_frequency_intervals) * self.frequency_interval,
            np.arange(0, self.number_time_intervals) * self.time_interval,
        )
        """The power spectrum of the translated non-Gaussian stochastic processes."""

    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 hasattr(self.distributions, "icdf"):
            non_gaussian_icdf = getattr(self.distributions, "icdf")
            samples_non_gaussian = non_gaussian_icdf(samples_cdf)
        else:
            raise AttributeError(
                "UQpy: The marginal dist_object needs to have an inverse cdf defined."
            )
        return samples_non_gaussian

    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.geometric_harmonics import *
from UQpy.surrogates.baseclass import *

from . import polynomial_chaos, stochastic_reduced_order_models, kriging, baseclass, geometric_harmonics


# 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 ``(samples_number, dimension)``.
        :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 ``(samples_number, dimension)``."""


# 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 ``(samples_number, dimension)``.
        :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 ``(samples_number, dimension)``."""


# 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 ``(samples_number, dimension)`` 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 ``(samples_number, dimension)``
         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.dimension = 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 ``(samples_number, dimension)`` with prescribed probability distributions."""
        self.samples_z:NumpyFloatArray = samples_z
        """Standard normal random vector of shape ``(samples_number, dimension)``"""
        self.jacobian = jacobian
        self.jzx: NumpyFloatArray = None
        """The Jacobian of the transformation of shape ``(dimension, dimension)``."""
        self.jxz: NumpyFloatArray = None
        """The Jacobian of the transformation of shape ``(dimension, dimension)``."""
        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.dimension)
            """Distorted correlation matrix (:math:`\mathbf{C_X}`) of the random vector **X**."""
            self.corr_z: NumpyFloatArray = np.eye(self.dimension)
            """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.dimension))):
                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.dimension))):
                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, dimension)`.
        :param jacobian: The jacobian of the transformation of shape ``(dimension, dimension)``. 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 = list()
        itam_error2 = list()
        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 ", k)
            logger.info("UQpy: Current error, ", 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, dimension)`` 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, dimension)``, the jacobian of the transformation of
         shape ``(dimension, dimension)``.
        """

        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
        else:
            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, dimension)``
        :param jacobian: A boolean whether to return the jacobian of the transformation. Default: False
        :return: Random vector of shape ``(nsamples, dimension)`` with prescribed probability distributions,
         The jacobian of the transformation of shape ``(dimension, dimension)``.
        """
        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
        else:
            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, samples_number: PositiveInteger):
        """
        Generate realizations from the joint pdf of the random vector **X**.

        :param samples_number: Number of samples to generate.
        :return: Random vector in the parameter space of shape ``(nsamples, dimension)``.
        """
        h = cholesky(self.corr_z, lower=True)
        n = int(samples_number)
        m = np.size(self.dist_object)
        y = np.random.randn(samples_number, 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.dimension += 1
        elif isinstance(dist_object, JointIndependent):
            self.dimension += len(dist_object.marginals)
        else:
            raise TypeError(
                "UQpy: A  ``DistributionContinuous1D``  or ``JointInd`` 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/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.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 = 1024
    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 *


# tests/unit_tests/RunModel/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/RunModel/process_third_party_output_blank.py

import numpy as np



# tests/unit_tests/RunModel/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/RunModel/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/RunModel/python_model_blank.py

import numpy as np


# tests/unit_tests/RunModel/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/RunModel/python_model_function.py

import numpy as np


def sum_rvs(samples=None):
    x = np.sum(samples, axis=1)
    return x


# tests/unit_tests/RunModel/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/RunModel/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/RunModel/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/RunModel/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/RunModel/test_RunModel.py

from UQpy.sampling import MonteCarloSampling
from UQpy.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], samples_number=5, random_state=1234)
x_mcs_new = MonteCarloSampling(distributions=[d, d, d], samples_number=5, random_state=2345)
verbose_parameter = True
os.chdir('./tests/RunModel')


def test_div_zero():
    print(os.getcwd())
    with pytest.raises(TypeError):
        RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs', fmt=20, delete_files=True)


def test_fmt_1():
    with pytest.raises(TypeError):
        RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs', fmt=20, delete_files=True)


def test_fmt_2():
    with pytest.raises(ValueError):
        RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs', fmt="random_string",
                 delete_files=True)


def test_var_names():
    with pytest.raises(ValueError):
        RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs', var_names=[20],
                 delete_files=True)


def test_model_script():
    with pytest.raises(ValueError):
        RunModel(ntasks=1, model_script='random_file_name', model_object_name='SumRVs', delete_files=True)


def test_samples():
    with pytest.raises(ValueError):
        RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs', samples="samples_string",
                 delete_files=True)


def test_python_serial_workflow_class_vectorized():
    model_python_serial_class = RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs',
                                         verbose=verbose_parameter)
    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_python_serial_class = RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs',
                                         vec=False, verbose=verbose_parameter)
    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_python_serial_class = RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs',
                                         vec=False, verbose=verbose_parameter, 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_python_serial_class = RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs',
                                         vec=False, verbose=verbose_parameter, 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_python_serial_class = RunModel(ntasks=1, model_script='python_model.py', model_object_name='SumRVs',
                                         vec=False, verbose=verbose_parameter, 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_python_serial_function = RunModel(ntasks=1, model_script='python_model.py', model_object_name='sum_rvs',
                                            verbose=verbose_parameter)
    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_python_serial_function = RunModel(ntasks=1, model_script='python_model.py', model_object_name='sum_rvs',
                                            vec=False, verbose=verbose_parameter)
    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,
                                            verbose=verbose_parameter)
    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_class_no_object_name():
    model_python_serial_function = RunModel(ntasks=1, model_script='python_model_class.py', vec=False,
                                            verbose=verbose_parameter)
    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_parallel_workflow_class():
    model_python_parallel_class = RunModel(ntasks=3, model_script='python_model.py', model_object_name='SumRVs',
                                           verbose=verbose_parameter)
    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))


def test_python_parallel_workflow_function():
    model_python_parallel_function = RunModel(ntasks=3, model_script='python_model.py', model_object_name='sum_rvs',
                                              verbose=verbose_parameter)
    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))


def test_third_party_serial():
    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.py', output_object_name='read_output',
                 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)


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)


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)


def test_third_party_serial_no_output_function():
    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.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)


def test_third_party_parallel():
    names = ['var1', 'var11', 'var111']
    m = RunModel(ntasks=3, 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',
                 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)


def test_third_party_default_var_names():
    model_third_party_default_names = RunModel(ntasks=1, model_script='python_model_sum_scalar_default.py',
                                               input_template='sum_scalar_default.py', model_object_name="python",
                                               output_script='process_third_party_output.py',
                                               output_object_name='read_output',
                                               resume=False, fmt="{:>10.4f}", verbose=verbose_parameter,
                                               delete_files=True, 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)


def test_third_party_var_names():
    names = ['var1', 'var11', 'var111', 'var1111']
    with pytest.raises(ValueError):
        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.py', output_object_name='read_output',
                 resume=False, fmt="{:>10.4f}", verbose=verbose_parameter, delete_files=True, samples=x_mcs.samples)


def test_python_serial_workflow_function_object_name_error():
    with pytest.raises(ValueError):
        model = RunModel(ntasks=1, model_script='python_model.py', vec=False, verbose=verbose_parameter)
        model.run(x_mcs.samples)


def test_python_serial_workflow_function_wrong_object_name():
    with pytest.raises(ValueError):
        model = RunModel(ntasks=1, model_script='python_model.py', vec=False, verbose=verbose_parameter,
                         model_object_name="random_model_name")
        model.run(x_mcs.samples)


def test_python_serial_workflow_function_no_objects():
    with pytest.raises(ValueError):
        model = RunModel(ntasks=1, model_script='python_model_blank.py', vec=False, verbose=verbose_parameter)
        model.run(x_mcs.samples)


# 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)
    pod_output=pod_dir.run()
    actual_result=pod_output[0][0][1][1]
    expected_result=0.3054
    assert expected_result==round(actual_result,6)

def test_DirectPOD():
    pod_dir = DirectPOD(solution_snapshots=Data, modes=1)
    assert round(pod_dir.run()[0][0][1][1], 6) == 0.761704


def test_SnapshotPOD():
    pod_snap = SnapshotPOD(solution_snapshots=Data, modes=1)
    assert round(pod_snap.run()[0][0][1][1], 6) == -0.181528

def test_SnapshotPOD_listData():
    list_data=list(Data)
    pod_dir = SnapshotPOD(solution_snapshots=list_data)
    pod_output=pod_dir.run()
    actual_result=pod_output[0][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, modes=-1)


def test_DirectPOD_mode_non_integer():
    with pytest.raises(BeartypeCallHintPepParamException):
        pod_dir = DirectPOD(solution_snapshots=Data, 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, 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, modes=-1)


def test_SnapshotPOD_mode_non_integer():
    with pytest.raises(BeartypeCallHintPepParamException):
        snap = SnapshotPOD(solution_snapshots=Data, 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, 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 EuclideanDistance
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 = EuclideanDistance(metric=DistanceMetric.EUCLIDEAN)
    distance = np.matrix.round(d.compute_distance(x), 3)

    assert distance == 2.724


def test_euclidean_distance_3points():
    x = np.array([[2.0, 3.1], [2.0, 2.1], [4.0, 1.25]])
    d = EuclideanDistance(metric=DistanceMetric.EUCLIDEAN)
    distance = np.matrix.round(d.compute_distance(x), 3)

    assert np.allclose(distance, np.array([1, 2.724, 2.173]))


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


def test_distances():
    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.
    # Solutions: original space.
    Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
    Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
    Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
    Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

    # Creating a list of matrices.
    matrices = [Sol0, Sol1, Sol2, Sol3]

    manifold_projection = SvdProjection(matrices, p=sys.maxsize)

    distance_metric = GeodesicDistance()
    value = distance_metric.compute_distance(manifold_projection.psi[0], manifold_projection.psi[1])

    assert value == 1.6024416339920522


# 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.create_from_data(data=swiss_roll,
                                           alpha=0.5, n_eigenvectors=3,
                                           is_sparse=True, neighbors_number=100,
                                           kernel=GaussianKernel(epsilon=0.03))

    diff_coords, evals, evecs = dmaps.fit()

    assert round(evals[0], 9) == 1.0
    assert round(evals[1], 9) == 0.999489295
    assert round(evals[2], 9) == 0.999116766


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.create_from_data(data=X, alpha=1, n_eigenvectors=3,
                                           kernel=GaussianKernel(epsilon=0.3))

    diff_coords, evals, evecs = dmaps.fit()

    assert evals[0] == 1.0000000000000002
    assert evals[1] == 0.9964842223723996
    assert evals[2] == 0.9964453129642372


def test_diff_matrices():
    import numpy as np

    np.random.seed(111)
    npts = 1000
    pts = np.random.rand(npts, 2)

    a0 = 0
    a1 = 1
    b0 = 0
    b1 = 1

    nodes = np.zeros(np.shape(pts))

    nodes[:, 0] = pts[:, 0] * (a1 - a0) + a0
    nodes[:, 1] = pts[:, 1] * (b1 - b0) + b0

    ns = 40

    x = np.linspace(0, 1, ns)
    samples = []
    for i in range(npts):

        M = np.zeros((ns, ns))
        for k in range(ns):
            f = np.sin(0.1 * k * np.pi * nodes[i, 0] * x + 2 * np.pi * nodes[i, 1])
            M[:, k] = f

        samples.append(M)

    dmaps = DiffusionMaps.create_from_data(data=samples, alpha=0.5, n_eigenvectors=10)

    diff_coords, evals, evecs = dmaps.fit()

    assert evals[0] == 1.0000000000000004
    assert evals[1] == 0.12956887787384186
    assert evals[2] == 0.07277038589085978



# 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.interpolation.methods.LinearInterpolation import LinearInterpolation
from UQpy.dimension_reduction.grassmann_manifold.interpolation.baseclass.InterpolationMethod import InterpolationMethod
from UQpy.dimension_reduction.grassmann_manifold.projections.SvdProjection import SvdProjection
from UQpy.optimization.GradientDescent import GradientDescent
from UQpy.dimension_reduction.grassmann_manifold.ManifoldInterpolation import ManifoldInterpolation


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.
    # Solutions: original space.
    Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
    Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
    Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
    Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

    # Creating a list of solutions.
    Solutions = [Sol0, Sol1, Sol2, Sol3]

    manifold_projection = SvdProjection(Solutions, p_planes_dimensions=sys.maxsize)

    optimization_method = GradientDescent(acceleration=False, error_tolerance=1e-3, max_iterations=1000)
    interpolation = ManifoldInterpolation(LinearInterpolation())

    interpolated_solution = manifold_projection\
        .reconstruct_solution(interpolation=interpolation, coordinates=nodes,
                              point=point,
                              p_planes_dimensions=manifold_projection.p_planes_dimensions,
                              optimization_method=optimization_method,
                              distance=GeodesicDistance(),
                              element_wise=False)
    assert round(interpolated_solution[0, 0], 9) == 0.61556849


def test_user_interpolation():
    import numpy as np
    nodes = np.array([[0, 0], [1, 0], [1, 1], [0, 1]])  # node_0, node_1, node_2.
    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.
    # Solutions: original space.
    Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
    Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
    Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
    Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

    # Creating a list of solutions.
    Solutions = [Sol0, Sol1, Sol2, Sol3]

    manifold_projection = SvdProjection(Solutions, p_planes_dimensions=sys.maxsize)

    optimization_method = GradientDescent(acceleration=False, error_tolerance=1e-3, max_iterations=1000)

    class UserInterpolation(InterpolationMethod):

        def interpolate(self, coordinates, samples, point):

            """

            **Input:**

            * **coordinates** (`ndarray`)
                Coordinates of the input data points.

            * **samples** (`ndarray`)
                Matrices corresponding to the points on the Grassmann manifold.

            * **point** (`ndarray`)
                Coordinates of the point to be interpolated.

            **Output/Returns:**

            * **interp_point** (`ndarray`)
                Interpolated point.
            """

            if not isinstance(coordinates, list) and not isinstance(coordinates, np.ndarray):
                raise TypeError('UQpy: `coordinates` must be either list or ndarray.')
            else:
                coordinates = np.array(coordinates)

            if not isinstance(samples, list) and not isinstance(samples, np.ndarray):
                raise TypeError('UQpy: `samples` must be either list or ndarray.')
            else:
                samples = np.array(samples)

            if not isinstance(point, list) and not isinstance(point, np.ndarray):
                raise TypeError('UQpy: `point` must be either list or ndarray.')
            else:
                point = np.array(point)

            myInterpolator = scipy.interpolate.LinearNDInterpolator(coordinates, samples)
            interp_point = myInterpolator(point)
            interp_point = interp_point[0]

            return interp_point

    interpolation = ManifoldInterpolation(UserInterpolation())

    interpolated_solution = manifold_projection.reconstruct_solution(interpolation=interpolation, coordinates=nodes,
                                                                     point=point,
                                                                     p_planes_dimensions=manifold_projection.p_planes_dimensions,
                                                                     optimization_method=optimization_method,
                                                                     distance=GeodesicDistance(),
                                                                     element_wise=False)

    assert round(interpolated_solution[0, 0], 9) == 0.61556849



# tests/unit_tests/dimension_reduction/test_karcher.py

import copy

import numpy as np
from beartype import beartype

from UQpy.utilities.distances.grassmannian_distances import GeodesicDistance
from UQpy.utilities.GrassmannPoint import GrassmannPoint
from UQpy.optimization.baseclass.OptimizationMethod import \
    OptimizationMethod
from UQpy.dimension_reduction.grassmann_manifold.projections.SvdProjection import SvdProjection
from UQpy.dimension_reduction.grassmann_manifold.Grassmann import Grassmann
import sys
from UQpy.optimization.GradientDescent import GradientDescent


def test_karcher():
    np.random.seed(1111)
    # Solutions: original space.
    sol0 = np.dot(np.random.rand(6, 2), np.random.rand(2, 6))
    sol1 = np.dot(np.random.rand(6, 3), np.random.rand(3, 6))
    sol2 = np.dot(np.random.rand(6, 4), np.random.rand(4, 6))
    sol3 = np.dot(np.random.rand(6, 3), np.random.rand(3, 6))

    # Creating a list of matrices.
    matrices = [sol0, sol1, sol2, sol3]
    manifold_projection = SvdProjection(matrices, p=sys.maxsize)

    optimization_method = GradientDescent(acceleration=True, error_tolerance=1e-4, max_iterations=1000)
    psi_mean = Grassmann.karcher_mean(manifold_points=manifold_projection.psi,
                                      optimization_method=optimization_method,
                                      distance=GeodesicDistance())

    phi_mean = Grassmann.karcher_mean(manifold_points=manifold_projection.phi,
                                      optimization_method=optimization_method,
                                      distance=GeodesicDistance())

    assert round(psi_mean.data[0, 0], 9) == -0.399231356
    assert round(phi_mean.data[0, 0], 9) == -0.382092372


def test_user_karcher():
    import numpy as np

    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.
    # Solutions: original space.
    Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
    Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
    Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
    Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

    # Creating a list of matrices.
    matrices = [Sol0, Sol1, Sol2, Sol3]

    class UserOptimizationMethod(OptimizationMethod):

        def __init__(self, acceleration: bool = False,
                     error_tolerance: float = 1e-3,
                     max_iterations: int = 1000):
            self.max_iterations = max_iterations
            self.error_tolerance = error_tolerance
            self.acceleration = acceleration

        def optimize(self, data_points: list[GrassmannPoint], distance):
            # Number of points.
            points_number = len(data_points)
            alpha = 0.5
            rank = []
            for i in range(points_number):
                rank.append(min(np.shape(data_points[i].data)))

            from UQpy.dimension_reduction.grassmann_manifold.Grassmann import Grassmann

            max_rank = max(rank)
            fmean = []
            for i in range(points_number):
                fmean.append(
                    Grassmann.frechet_variance(reference_point=data_points[i], grassmann_points=data_points,
                                               distance=distance)
                )

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

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

            counter_iteration = 0

            l = 0
            avg = []
            _gamma = []
            if self.acceleration:
                _gamma = Grassmann.log_map(manifold_points=data_points,
                                           reference_point=mean_element)

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

            # Main loop
            while counter_iteration <= self.max_iterations:
                _gamma = Grassmann.log_map(manifold_points=data_points,
                                           reference_point=mean_element)
                avg_gamma.fill(0)

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

                test_0 = np.linalg.norm(avg_gamma, "fro")
                if test_0 < self.error_tolerance and counter_iteration == 0:
                    break

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

                x = Grassmann.exp_map(tangent_points=[step], reference_point=mean_element)

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

                if test_1 < self.error_tolerance:
                    break

                mean_element = []
                mean_element = x[0]

                counter_iteration += 1

            # return the Karcher mean.
            return mean_element

    from UQpy.utilities.distances.baseclass.GrassmannianDistance import GrassmannianDistance

    class UserDistance(GrassmannianDistance):
        @beartype
        def compute_distance(self, xi: GrassmannPoint, xj: GrassmannPoint):
            GrassmannianDistance.check_rows(xi, xj)

            rank_i = xi.data.shape[1]
            rank_j = xj.data.shape[1]

            r = np.dot(xi.data.T, xj.data)
            (ui, si, vi) = np.linalg.svd(r, full_matrices=True)
            si[np.where(si > 1)] = 1.0
            theta = np.arccos(si)
            distance = (np.sqrt(abs(rank_i - rank_j) * np.pi ** 2 / 4 + np.sum(theta ** 2)))

            return distance

    manifold_projection = SvdProjection(matrices, p=sys.maxsize)
    psi_mean = Grassmann.karcher_mean(manifold_points=manifold_projection.psi,
                                      optimization_method=GradientDescent(acceleration=True, error_tolerance=1e-4,
                                                                          max_iterations=1000),
                                      distance=GeodesicDistance())

    assert round(psi_mean.data[0, 0], 9) == -0.399231356



# 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]
    kernel = np.matrix.round(ProjectionKernel().calculate_kernel_matrix(points), 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 = np.matrix.round(BinetCauchyKernel().calculate_kernel_matrix(points), 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)
    kernel = gaussian.calculate_kernel_matrix(points)

    assert np.allclose(np.matrix.round(kernel, 4), np.array([[1., 0.2645, 1.], [0.2645, 1., 0.2645], [1, 0.2645, 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()
    kernel = gaussian.calculate_kernel_matrix(points)

    assert np.allclose(np.matrix.round(kernel, 4), np.array([[1., 0.8834, 0.7788], [0.8834, 1., 0.6937],
                                                             [0.7788, 0.6937, 1.]]))
    assert np.round(gaussian.epsilon, 4) == 3.7538


def test_kernel():
    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.
    # Solutions: original space.
    Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
    Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
    Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
    Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

    # Creating a list of solutions.
    Solutions = [Sol0, Sol1, Sol2, Sol3]
    from UQpy.dimension_reduction.grassmann_manifold.Grassmann import Grassmann
    manifold_projection = SvdProjection(Solutions, p=sys.maxsize)
    manifold = Grassmann(manifold_projected_points=manifold_projection)
    kernel = manifold.evaluate_kernel_matrix(kernel=ProjectionKernel())

    assert kernel[0, 1] == 2.740411439218967



# tests/unit_tests/dimension_reduction/test_log_exp_maps.py

import numpy as np
from UQpy.dimension_reduction.grassmann_manifold.projections.SvdProjection import SvdProjection
from UQpy.dimension_reduction.grassmann_manifold.Grassmann import Grassmann
import sys


def test_log_exp_maps():
    np.random.seed(1111)  # For reproducibility.
    # Solutions: original space.
    sol0 = np.dot(np.random.rand(6, 2), np.random.rand(2, 6))
    sol1 = np.dot(np.random.rand(6, 3), np.random.rand(3, 6))
    sol2 = np.dot(np.random.rand(6, 4), np.random.rand(4, 6))
    sol3 = np.dot(np.random.rand(6, 3), np.random.rand(3, 6))

    # Creating a list of matrices.
    matrices = [sol0, sol1, sol2, sol3]
    manifold_projection = SvdProjection(matrices, p=sys.maxsize)

    points_tangent = Grassmann.log_map(manifold_points=manifold_projection.psi,
                                       reference_point=manifold_projection.psi[0])

    assert np.round(points_tangent[0][0][0], 2) == 0.0
    assert points_tangent[1][0][0] == -0.002899719618992682
    assert points_tangent[2][0][0] == 0.012574949291454723
    assert points_tangent[3][0][0] == 0.017116995689638644

    manifold_points = Grassmann.exp_map(tangent_points=points_tangent,
                                        reference_point=manifold_projection.psi[0])

    assert np.round(manifold_points[0].data[0][0], 5) == -0.49845
    assert manifold_points[1].data[0][0] == -0.5013644537794977
    assert manifold_points[2].data[0][0] == -0.4498150778515823
    assert manifold_points[3].data[0][0] == -0.4980716121312462




# 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.py

import numpy as np

domain = np.linspace(0, 10, 50)

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


# tests/unit_tests/inference/pfn1.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

# tests/unit_tests/inference/pfn2.py

import numpy as np

domain = np.linspace(0, 10, 50)

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_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/test_bayes_model_selection.py

from UQpy.inference import *
from UQpy.RunModel import RunModel
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


def test_models():
    data_ex1 = np.loadtxt('data_ex1a.txt')

    runmodel4 = RunModel(model_script='pfn.py', model_object_name='model_linear', vec=False, var_names=['theta_0'])
    runmodel5 = RunModel(model_script='pfn1.py', model_object_name='model_quadratic', vec=False,
                         var_names=['theta_0', 'theta_1'])
    runmodel6 = RunModel(model_script='pfn2.py', model_object_name='model_cubic', vec=False,
                         var_names=['theta_0', 'theta_1', 'theta_2'])

    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],
                                    data=data_ex1,
                                    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'

    shutil.rmtree(runmodel4.model_dir)
    shutil.rmtree(runmodel5.model_dir)
    shutil.rmtree(runmodel6.model_dir)


# 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], data=data,
                                    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(data=data, 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
import shutil
import numpy as np

data = [0., 1., -1.5, -0.2]


@pytest.fixture
def setup():
    h_func = RunModel(model_script='pfn_models.py', model_object_name='model_quadratic', vec=False,
                      var_names=['theta_0', 'theta_1'], delete_files=True)
    yield h_func
    shutil.rmtree(h_func.model_dir)


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.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))

    h_func = RunModel(model_script='pfn_models.py', model_object_name='model_quadratic', vec=False,
                      var_names=['theta_0', 'theta_1'])
    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

    shutil.rmtree(h_func.model_dir)


# 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 model_i(samples):
    resistance = samples[0, 0]
    stress = samples[0, 1]
    return resistance - stress

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)

def model_k(samples):
    return samples[0, 0] * samples[0, 1] - 80

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


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/test_form.py

from UQpy.RunModel import RunModel
from UQpy.distributions import *
from UQpy.reliability import FORM
import glob
import shutil
import numpy as np
import pytest


@pytest.fixture
def setup():
    h_func = RunModel(model_script='pfn.py', model_object_name='model_i', vec=False, delete_files=True)
    yield h_func
    # shutil.rmtree(h_func.model_dir)


def test_seeds_xu_is_none(setup):
    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):
    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):
    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):
    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):
    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):
    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):
    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):
    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):
    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():
    RunModelObject = RunModel(model_script='pfn.py',
                              model_object_name="example1",
                              vec=False, ntasks=3)
    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.]))

    import shutil
    shutil.rmtree(RunModelObject.model_dir)


# tests/unit_tests/reliability/test_sorm.py

from UQpy.distributions import Normal
from UQpy.reliability import FORM, SORM
from UQpy.RunModel import RunModel
import glob
import shutil
import numpy as np
import pytest


@pytest.fixture
def setup():
    h_func = RunModel(model_script='pfn.py', model_object_name='model_k', vec=False, delete_files=True)
    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_taylor_series.py

from UQpy.distributions import Normal, JointIndependent
from UQpy.reliability import TaylorSeries
from UQpy.transformations import Nataf
import glob
import shutil
import numpy as np
import pytest
from UQpy.RunModel import RunModel


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():
    h_func = RunModel(model_script='pfn.py', model_object_name='model_i', vec=False, delete_files=True)
    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():
    h_func = RunModel(model_script='pfn.py', model_object_name='model_j', vec=False, delete_files=True)
    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')
    shutil.rmtree(h_func.model_dir)
    np.testing.assert_allclose(hessian, [[-0.00720754, 0.00477726], [0.00477726, -0.00316643]], rtol=1e-04)




# 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 import Multinomial
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

from UQpy.surrogates.kriging.Kriging import Kriging
from UQpy.sampling import MonteCarloSampling, AdaptiveKriging
from UQpy.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 Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=0)
    rmodel = RunModel(model_script='series.py', vec=False)
    regression_model = Linear()
    correlation_model = Exponential()
    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

    shutil.rmtree(rmodel.model_dir)


def test_akmcs_u():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    rmodel = RunModel(model_script='series.py', vec=False)
    regression_model = Linear()
    correlation_model = Exponential()
    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

    shutil.rmtree(rmodel.model_dir)


def test_akmcs_expected_feasibility():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    rmodel = RunModel(model_script='series.py', vec=False)
    regression_model = Linear()
    correlation_model = Exponential()
    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

    shutil.rmtree(rmodel.model_dir)


def test_akmcs_expected_improvement():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    rmodel = RunModel(model_script='series.py', vec=False)
    regression_model = Linear()
    correlation_model = Exponential()
    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

    shutil.rmtree(rmodel.model_dir)


def test_akmcs_expected_improvement_global_fit():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    rmodel = RunModel(model_script='series.py', vec=False)
    regression_model = Linear()
    correlation_model = Exponential()
    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

    shutil.rmtree(rmodel.model_dir)

def test_akmcs_samples_error():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=0)
    rmodel = RunModel(model_script='series.py', vec=False)
    regression_model = Linear()
    correlation_model = Exponential()
    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)

    shutil.rmtree(rmodel.model_dir)

def test_akmcs_u_run_from_init():
    from UQpy.surrogates.kriging.regression_models.LinearRegression import Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential

    marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
    x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)
    rmodel = RunModel(model_script='series.py', vec=False)
    regression_model = Linear()
    correlation_model = Exponential()
    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

    shutil.rmtree(rmodel.model_dir)

# 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_samples2():
    """ Check the samples attribute, when dist_object is a jointInd class object and criterion is 'correlate'."""
    expected_samples = np.array([[0.94, 0.54], [0.26, 0.08], [0.45, 0.88], [0.14, 0.7 ], [0.71, 0.4 ]])
    assert (x1b._samples.round(2) == expected_samples).all()


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.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.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 y.samples[16, 0] == 0.3637932367281488
    assert y.samples[16, 1] == 0.4676253860574614
    assert y.samples[17, 0] == 0.4245856389630922
    assert y.samples[17, 1] == 0.21730082648922827


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)
    rmodel = RunModel(model_script='python_model_function.py', vec=False)
    from UQpy.surrogates.kriging.regression_models import Linear
    from UQpy.surrogates.kriging.correlation_models import Exponential

    K = Kriging(regression_model=Linear(), correlation_model=Exponential(), 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.59254104, 0.96577043], [1.97386531, 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.98693265, 0.24237455]]))


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 Linear
    from UQpy.surrogates.kriging.correlation_models.ExponentialCorrelation import Exponential
    rmodel_ = RunModel(model_script='python_model_function.py', vec=False)
    K_ = Kriging(regression_model=Linear(), correlation_model=Exponential(), 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 Linear
    from UQpy.surrogates.kriging.correlation_models import Exponential

    K = Kriging(regression_model=Linear(), correlation_model=Exponential(), optimizations_number=20,
                correlation_model_parameters=[1, 1], optimizer=MinimizeOptimizer('l-bfgs-b'),)
    rmodel = RunModel(model_script='python_model_function.py', vec=False)
    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)
    rmodel_ = RunModel(model_script='python_model_function.py', vec=False)
    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 sts_obj.samples[2, 0] == 1.902581742436106


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 x.samples[7, 0] == 3.6928440862661223
    assert x.samples[20, 1] == 1.1555963246730931
    assert x.samples[1, 0] == 1.8393015015282757
    assert x.samples[15, 1] == 2.117727620746169


# 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/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/test_morris.py

from UQpy.sensitivity.MorrisSensitivity import MorrisSensitivity
from UQpy.RunModel import RunModel
import numpy as np
from UQpy.distributions import Uniform
import os
import pytest
import shutil

#os.chdir('./tests/Sensitivity')


@pytest.fixture
def setup():
    a_values = [0.001, 99.]
    runmodel_object = RunModel(
        model_script='pfn.py', model_object_name='gfun_sensitivity', vec=True, a_values=a_values,
        var_names=['X{}'.format(i) for i in range(2)], delete_files=True)
    dist_object = [Uniform(), Uniform()]
    yield runmodel_object, dist_object
    shutil.rmtree(runmodel_object.model_dir)


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/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,
                               number_time_intervals=nt, number_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,
                                             number_time_intervals=nt, number_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(samples_number=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,
                               number_time_intervals=nt, number_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_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 Linear, Constant
from UQpy.surrogates.kriging.correlation_models import Gaussian


samples = np.linspace(0, 5, 20).reshape(-1, 1)
values = np.cos(samples)
optimizer = MinimizeOptimizer(method="L-BFGS-B")
krig = Kriging(regression_model=Linear(), correlation_model=Gaussian(), 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=Constant(), correlation_model=Gaussian(), 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=Linear(), correlation_model=Gaussian(), optimizer=optimizer,
                                  correlation_model_parameters=[1]).regression_model
optimizer = MinimizeOptimizer(method="L-BFGS-B")
gaussian_corrleation_model = Kriging(regression_model=Linear(), correlation_model=Gaussian(), 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 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_correlation_models():
    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_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=Gaussian(), 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=Linear(), 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=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_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.PolynomialBasis import PolynomialBasis

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 = PolynomialBasis.create_total_degree_basis(dist, max_degree).polynomials
    value = polynomials[1].evaluate(x)[0]
    assert round(value, 4) == -0.2874

def test_2():
    """
    Test tp basis
    """
    polynomial_basis = PolynomialBasis.create_tensor_product_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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.first_order_indices()
    assert round(first_order_sobol[0][0], 3) == 1.0
#
def test_8():
    """
    Test Sobol indices
    """
    polynomial_basis = PolynomialBasis.create_total_degree_basis(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.total_order_indices()
    assert round(total_order_sobol[0][0], 3) == 1.0
#
def test_9():
    """
    Test Sobol indices
    """
    polynomial_basis = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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.generalized_total_order_indices()
#
def test_11():
    """
    PCE mean
    """
    polynomial_basis = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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)
    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)
    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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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)
    """

    pceLAR=polynomial_chaos.regressions.LeastAngleRegression.model_selection(pce_2)
    pce2LAR_sens=PceSensitivity(pceLAR)
    
    assert all((np.argwhere(np.round(pce2LAR_sens.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.dimension == 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)


# src/UQpy/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 collections
import datetime
import glob
import logging
import os
import pathlib
import platform
import re
import shutil
import subprocess
from beartype import beartype

from UQpy.utilities.ValidationTypes import *


class RunModel:
    # Authors:
    # B.S. Aakash, Lohit Vandanapu, Michael D.Shields
    #
    # Last
    # modified: 5 / 8 / 2020 by Michael D.Shields
    @beartype
    def __init__(
            self,
            samples: Union[list, NumpyFloatArray] = None,
            model_script: str = None,
            model_object_name: str = None,
            input_template: str = None,
            var_names: list[str] = None,
            output_script: str = None,
            output_object_name: str = None,
            ntasks: int = 1,
            cores_per_task: int = 1,
            nodes: int = 1,
            cluster: bool = False,
            resume: bool = False,
            model_dir: str = "Model_Runs",
            fmt=None,
            separator: str = ", ",
            vec: bool = True,
            delete_files: bool = False,
            **kwargs
    ):
        """
        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 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 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.

         `input_template` is not used in the Python 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 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.

         `output_object_name` is not used in the Python model workflow.

        :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.

         `nodes` is not used in the Python model workflow.

        :param cluster: Set ``cluster = True`` to run on an HPC cluster.

         :class:`.RunModel` currently supports computations on HPC clusters using the SLURM scheduler
         (https://slurm.schedmd.com). The set of model evaulations is submitted using the GNU `parallel` command with
         option `-j ntasks`. Individual model evaluations are submitted using the `srun` command with options `-N nodes`
         and `-c cores_per_node`.

         When ``cluster = False``, the `srun` command is not used, but the `parallel` command still is.

         `cluster` is not used for the Python model workflow.

        :param resume: If ``resume = True``, `GNU parallel` enables :py:mod:`UQpy` to resume execution of any model
         evaluations that failed to execute in the third-party software model workflow.

         To use this feature, execute the same call to :class:`.RunModel` that failed to complete but with ``resume = True``.
         The same set of samples must be passed to resume processing from the last successful execution of the model.

         `resume` is not used in the Python model workflow.

        :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.

        :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}".

         `fmt` is not used in the Python model workflow.

        :param separator: A string used to delimit values when printing arrays to the `template_input`.

         `separator` is not used in the Python model workflow.

        :param vec: Specifies vectorized (``vec = True``) or looped (``vec = False``) model evaluation in the
         serial Python model workflow.

         In the Python model workflow, `model_script` may be written to accept a single sample or multiple samples at a
         time. If it is written to accept a single sample, set ``vec = False`` and :class:`.RunModel` will run the model in a
         loop over the number of samples. If `model_script` is written to accept multiple samples, set ``vec = True``
         and :class:`.RunModel` will pass all of the samples to the model for vectorized computation.

         `vec` is not used in the third-party 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 kwargs: Additional inputs to the Python object specified by `model_object_name` in the Python model workflow.

         `**kwargs` is not used in the third-party model workflow.

        """
        # Check the platform and build appropriate call to Python
        if platform.system() in ["Windows"]:
            self.python_command = "python"
        else:
            self.python_command = "python3"

        self.logger = logging.getLogger(__name__)

        # Vectorized computation
        self.vec = vec

        # Format option
        self.separator = separator
        self.fmt = 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")

        self.delete_files = delete_files

        # kwargs options, used only for python runs
        self.python_kwargs = kwargs

        # Input related
        self.input_template = input_template
        self.var_names = var_names
        self.n_vars: int = 0
        """Number of variables to be passed for each model evaluation, ``n_vars = len(samples[0])``.

        Note that variables do not need to be scalars. Variables can be scalars, vectors, matrices, or tensors. When
        writing vectors, matrices, and tensors to a `input_template` they are first flattened and written in delimited
        form."""

        # Check if var_names is a list of strings
        if self.var_names is not None:
            if not self._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 = []
        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

        # Create a new directory where the model will be executed
        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)
        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
        )

        # 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'."
            )
        # Save the model object name
        self.model_object_name = model_object_name
        # Save option for resuming parallel execution
        self.resume = resume

        # Output related
        self.output_script = output_script
        self.output_object_name = output_object_name

        # Number of tasks
        self.ntasks = ntasks
        # Number of cores_per_task
        self.cores_per_task = cores_per_task
        # Number of nodes
        self.nodes = nodes
        self.template_text = ""
        self.output_module = None
        self.python_model = None

        # If running on cluster or not
        self.cluster = cluster

        # Initialize sample related variables
        self.samples: Numpy2DFloatArray = []
        """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.nexist: 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.nsim: 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)

        # Change current working directory to model run directory
        os.chdir(self.model_dir)
        self.logger.info("\nUQpy: All model evaluations will be executed from the following directory: \n"
                         + self.model_dir)

        # Number of simulations to be performed
        self.nsim = len(samples)

        # Number of variables
        self.n_vars = len(samples[0])

        # 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:

            # If there are no samples, check to ensure that len(var_names) = n_vars
            if self.input_template is not None:
                if self.var_names is not None:
                    # Check to see if self.var_names has the correct length
                    if len(self.var_names) != self.n_vars:
                        raise ValueError("\nUQpy: var_names must have the same length as the number of variables (i.e. "
                                         "len(var_names) = len(samples[0]).\n")
                else:
                    # 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_vars):
                        self.var_names.append("x%d" % i)

            self.nexist = 0
            self.samples = samples
            self.qoi_list = [None] * self.nsim

        else:  # Samples already exist in the RunModel object, append new ones
            self.nexist = len(self.samples)
            self.qoi_list.extend([None] * self.nsim)
            self.samples = np.vstack((self.samples, samples))

        # Check if there is a template input file or not and execute the appropriate function
        if self.input_template is not None:  # If there is a template input file
            # Check if it is a file and is readable
            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())

            # 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()

            # Run the serial execution or parallel execution depending on ntasks
            if self.ntasks == 1:
                self._serial_execution()
            else:
                self._parallel_execution()

        else:  # If there is no template input file supplied
            # Import the python module
            self.python_model = __import__(self.model_script[:-3])
            # Run function which checks if the python model has the model object
            self._check_python_model()

            # Run the serial execution or parallel execution depending on ntasks
            if self.ntasks == 1:
                self._serial_python_execution()
            else:
                self._parallel_python_execution()

        # Return to parent directory
        self.logger.info("\nUQpy: Returning to the parent directory:\n" + self.parent_dir)
        os.chdir(self.parent_dir)

        if self.delete_files:
            self.logger.info("UQpy: Deleting individual run files.")
            for dirname in glob.glob(os.path.join(self.model_dir, "run*")):
                shutil.rmtree(dirname)

        return None

    ####################################################################################################################

    def _serial_execution(self):
        """
        Perform serial execution of a third-party model using a template input file

        This function loops over the number of simulations, executing the model once per loop. In each loop, the
        function creates a directory for each model run, copies files to the model run directory,
        changes the current working directory to the model run directory, calls the input function, executes the model,
        calls the output function, removes the copied files and folders, and returns to the previous directory.
        """
        self.logger.info("\nUQpy: Performing serial execution of the third-party model.\n")

        # Loop over the number of simulations, executing the model once per loop
        for i in range(self.nexist, self.nexist + self.nsim):
            # Create a directory for each model run
            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)

            # Execute the model
            self._execute_serial(i)

            # Call the output function
            if self.output_script is not None:
                self._output_serial(i)

            # 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(i) + " complete.\n")
            self.logger.info("\nUQpy: Returning to the model directory:\n" + self.model_dir)

        self.logger.info("\nUQpy: Serial execution of the third-party model complete.\n")

    ####################################################################################################################

    def _parallel_execution(self):
        """
        Execute a third-party model in parallel

        This function calls the input function and generates input files for all the samples, then creates a directory
        for each model run, copies files to the model run directory, executes the model in parallel, collects output,
        removes the copied files and folders.
        """
        self.logger.info("\nUQpy: Performing parallel execution of the third-party model.\n")
        # Call the input function
        self.logger.info("\nUQpy: Creating inputs for parallel execution of the third-party model.\n")

        # Create all input files for the parallel execution and place them in the proper directories
        for i in range(self.nexist, self.nexist + self.nsim):
            # Create a directory for each model run
            work_dir = os.path.join(self.model_dir, "run_" + str(i))
            self._copy_files(work_dir=work_dir)

        self._input_parallel()

        # Execute the model
        self.logger.info("\nUQpy: Executing the third-party model in parallel.\n")

        self._execute_parallel()

        # Call the output function
        self.logger.info("\nUQpy: Collecting outputs from parallel execution of the third-party model.\n")

        for i in range(self.nexist, self.nexist + self.nsim):
            # Change current working directory to model run directory
            work_dir = os.path.join(self.model_dir, "run_" + str(i))
            self.logger.info("\nUQpy: Changing to the following directory for output processing:\n" + work_dir)
            os.chdir(work_dir)

            # Run output processing function
            if self.output_script is not None:
                self.logger.info("\nUQpy: Processing output from parallel execution of the third-party model run "
                                 + str(i) + ".\n")
                self._output_serial(i)

            # Remove the copied files and folders
            self._remove_copied_files(work_dir)

            # Change back to the upper directory
            self.logger.info("\nUQpy: Changing back to the following model directory:\n" + self.model_dir)
            os.chdir(self.model_dir)

        self.logger.info("\nUQpy: Parallel execution of the third-party model complete.\n")

    ####################################################################################################################
    def _serial_python_execution(self):
        """
        Execute a python model in serial

        This function imports the model_object from the model_script, and executes the model in series by passing the
        corresponding sample/samples along with keyword arguments, if any, as inputs to the model object.
        """
        self.logger.info("\nUQpy: Performing serial execution of a Python model.\n")

        model_object = getattr(self.python_model, self.model_object_name)
        # Run python model
        if self.vec:
            # If the Python model is vectorized to accept many samples.
            self.model_output = model_object(self.samples[self.nexist: self.nexist + self.nsim], **self.python_kwargs)
            if self.model_is_class:
                self.qoi_list[self.nexist: self.nexist + self.nsim] = list(self.model_output.qoi)
            else:
                self.qoi_list[self.nexist: self.nexist + self.nsim] = list(self.model_output)
        else:
            # If the Python model is not vectorized and accepts only a single sample.
            for i in range(self.nexist, self.nexist + self.nsim):
                sample_to_send = np.atleast_2d(self.samples[i])

                if len(self.python_kwargs) == 0:
                    self.model_output = model_object(sample_to_send)
                else:
                    self.model_output = model_object(sample_to_send, **self.python_kwargs)
                if self.model_is_class:
                    self.qoi_list[i] = self.model_output.qoi
                else:
                    self.qoi_list[i] = self.model_output

        self.logger.info("\nUQpy: Serial execution of the python model complete.\n")

    ####################################################################################################################
    def _parallel_python_execution(self):
        """
        Execute a python model in parallel

        This function imports the model object from the model script, and executes the model in parallel by passing the
        samples along with keyword arguments, if any, as inputs to the model object.
        """
        self.logger.info("\nUQpy: Performing parallel execution of the model without template input.\n")
        import multiprocessing
        import UQpy.utilities.Utilities as Utilities

        sample = []
        pool = multiprocessing.Pool(processes=self.ntasks)
        for i in range(self.nexist, self.nexist + self.nsim):
            sample_to_send = np.atleast_2d(self.samples[i])
            if len(self.python_kwargs) == 0:
                sample.append([self.model_script, self.model_object_name, sample_to_send])
            else:
                sample.append([self.model_script, self.model_object_name, sample_to_send, self.python_kwargs, ])

        results = pool.starmap(Utilities.run_parallel_python, sample)

        for i in range(self.nsim):
            if self.model_is_class:
                self.qoi_list[i + self.nexist] = results[i].qoi
            else:
                self.qoi_list[i + self.nexist] = results[i]

        pool.close()

        self.logger.info("\nUQpy: Parallel execution of the python model complete.\n")

    ####################################################################################################################
    def _input_serial(self, index):
        """
        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(index=index)
        # 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 _execute_serial(self, index):
        """
        Execute the model once using the input file of index number

        ** Input: **

        :param index: The simulation number
        :type index: int
        """
        self.model_command = [self.python_command, str(self.model_script), str(index)]
        subprocess.run(self.model_command)

    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
        self.output_module = __import__(self.output_script[:-3])
        output_object = getattr(self.output_module, self.output_object_name)
        self.model_output = output_object(index)

        if self.output_is_class:
            self.qoi_list[index] = self.model_output.qoi
        else:
            self.qoi_list[index] = self.model_output

    def _input_parallel(self):
        """
        Create all the input files required
        """
        # Loop over the number of samples and create input files in a folder in current directory
        for i in range(self.nsim):
            new_text = self._find_and_replace_var_names_with_values(index=i + self.nexist)
            folder_to_write = "run_" + str(i + self.nexist) + "/InputFiles"
            # Write the new text to the input file
            self._create_input_files(file_name=self.input_template,
                                     num=i + self.nexist,
                                     text=new_text,
                                     new_folder=folder_to_write)
            self.logger.info("\nUQpy: Created input files for run " + str(i) + " in the directory: \n"
                             + os.path.join(self.model_dir, folder_to_write))

    def _execute_parallel(self):
        """
        Build the command string and execute the model in parallel using subprocess and gnu parallel

        ** Input: **

        :param timestamp: Timestamp which is appended to the name of the input file
        :type timestamp: str
        """

        # TODO: Generalize run_string and parallel_string for any cluster
        # Check if logs folder exists, if not, create it
        if not os.path.exists("logs"):
            os.makedirs("logs")
        # If the user sets resume=True, do not delete log file. Else, delete logfile before running
        if self.resume is False:
            try:
                os.remove("logs/runtask.log")
            except OSError:
                pass
        self.parallel_string = ("parallel --delay 0.2 --joblog logs/runtask.log --resume -j " + str(self.ntasks) + " ")

        # If running on SLURM cluster
        if self.cluster:
            self.srun_string = ("srun -N" + str(self.nodes) + " -n1 -c" + str(self.cores_per_task) + " --exclusive ")
            self.model_command_string = (
                    self.parallel_string + "'(cd run_{1}" + " && " + self.srun_string + " " + self.python_command
                    + " -u " + str(self.model_script) + " {1})'  ::: {" + str(self.nexist) + ".."
                    + str(self.nexist + self.nsim - 1) + "}")
        else:  # If running locally
            self.model_command_string = (
                    self.parallel_string + " 'cd run_{1}" + " && " + self.python_command + " -u " + str(
                self.model_script)
                    + "' {1}  ::: {" + str(self.nexist) + ".." + str(self.nexist + self.nsim - 1) + "}")

        subprocess.run(self.model_command_string, shell=True)

    ####################################################################################################################
    # Helper functions

    @staticmethod
    def _create_input_files(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, index):
        """
        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
        samples = self.samples[index]

        new_text = template_text
        for j in range(self.n_vars):
            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], "samples[" + 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
            if index == 0:
                if count > 1:
                    self.logger.info(
                        "\nUQpy: Found "
                        + str(count)
                        + " instances of variable: '"
                        + var_names[j]
                        + "' in the input file.\n"
                    )
                else:
                    self.logger.info(
                        "\nUQpy: Found "
                        + str(count)
                        + " instance of variable: '"
                        + var_names[j]
                        + "' in the input file.\n"
                    )
        return new_text

    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)

    @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 _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 = []
        for name, obj in inspect.getmembers(self.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")

    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

        class_list = []
        function_list = []
        for name, obj in inspect.getmembers(self.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 _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)


# 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.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.RunModel 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

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/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])

# docs/code/dimension_reduction/diffusion_maps/plot_diffusion_maps_circle.py

"""

Circle
===================================================================

This example shows how to use the UQpy DiffusionMaps class to reveal the embedded structure of noisy data.
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the DiffusionMaps class from UQpy implemented in the DimensionReduction module.

#%%

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from UQpy.dimension_reduction import DiffusionMaps, GaussianKernel

#%% md
#
# Sample points randomly following a parametric curve and plot the 3D graphic.

#%%

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()

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(x, y, z, c='b', cmap=plt.cm.Spectral, s=8)
ax.plot(x0, y0, z0, 'r', label='parametric curve')
plt.show()

#%% md
#
# Instantiate the class `DiffusionMaps` using `alpha=1`; `n_evecs=3`, because the first eigenvector is non-informative.
# Moreover, a Gaussian is used in the kernel construction.

#%%

dmaps = DiffusionMaps.create_from_data(data=X,
                                       alpha=0.5, eigenvectors_number=3,
                                       is_sparse=True, neighbors_number=100,
                                       kernel=GaussianKernel(epsilon=0.3))

#%% md
#
# Use the method `mapping` to compute the diffusion coordinates assuming `epsilon=0.3`.

#%%

diff_coords, evals, evecs = dmaps.fit()

#%% md
#
# Plot the second and third diffusion coordinates to reveal the embedded structure of the data.

#%%

color = evecs[:, 1]
plt.scatter(diff_coords[:, 1], diff_coords[:, 2], c=color, cmap=plt.cm.Spectral, s=8)
plt.axis('equal')
plt.show()

#%% md
#
# Use the colormap to observe how the embedded structure is distributed in the original set.

#%%

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(x, y, z, c=color, cmap=plt.cm.Spectral, s=8)
plt.show()


#%% md
#
# Now, use an user defined kernel capable to handle this kind of data.

#%%

import scipy.spatial.distance as sd
import itertools
from UQpy import Kernel
class UserKernelDmaps(Kernel):
    def __init__(self, epsilon: float = None):
        self.epsilon = epsilon

    def pointwise_operator(self, point1, point2):
        pass

    def compute_default_epsilon(self, epsilon):
        if epsilon is None:
            # Compute a suitable epsilon when it is not provided by the user.
            # Compute epsilon as the median of the square of the euclidean distances
            epsilon = np.median(np.array(self.distance_pairs) ** 2)
        return epsilon

    def apply_method(self, data):
        self.distance_pairs = []
        if len(np.shape(data)) == 2:
            # Set of 1-D arrays
            self.distance_pairs = sd.pdist(data, 'euclidean')
        elif len(np.shape(data)) == 3:
            # Set of 2-D arrays
            # Check arguments: verify the consistency of input arguments.
            nargs = len(data)
            indices = range(nargs)
            pairs = list(itertools.combinations(indices, 2))

            self.distance_pairs = []
            for id_pair in range(np.shape(pairs)[0]):
                ii = pairs[id_pair][0]  # Point i
                jj = pairs[id_pair][1]  # Point j

                x0 = data[ii]
                x1 = data[jj]

                distance = np.linalg.norm(x0 - x1, 'fro')

                self.distance_pairs.append(distance)
        else:
            raise TypeError('UQpy: The size of the input data is not consistent with this method.')

        self.epsilon = self.compute_default_epsilon(self.epsilon)

        return np.exp(-sd.squareform(self.distance_pairs) ** 2 / (4 * self.epsilon))

#%% md
#
# Use the method `mapping` to compute the diffusion coordinates assuming `epsilon=0.3`.

#%%

dfm = DiffusionMaps.create_from_data(data=X,
                                     alpha=1, eigenvectors_number=3,
                                     kernel=UserKernelDmaps(epsilon=0.3))

diff_coord_new, evals, evecs = dfm.mapping()

color = evecs[:, 1]
plt.scatter(diff_coords[:, 1], diff_coords[:, 2], c=color, cmap=plt.cm.Spectral, s=8)
plt.axis('equal')
plt.show()

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(x, y, z, c=color, cmap=plt.cm.Spectral, s=8)
plt.show()

# docs/code/dimension_reduction/diffusion_maps/plot_diffusion_maps_matrices.py

"""

Matrices
===================================================================

This example shows how to use the UQpy Grassmann class to
* project points onto the Grassmann manifold;
* calculate the pairwise distances;
* compute the affinity matrix.
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the Grassmann class from UQpy implemented in the DimensionReduction module.

#%%

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.mplot3d import Axes3D
from UQpy.dimension_reduction import DiffusionMaps

#%% md
#
# Generate `npts` matrices to serve as input points.

#%%

np.random.seed(111)
npts = 1000
pts = np.random.rand(npts, 2)

a0 = 0
a1 = 1
b0 = 0
b1 = 1

nodes = np.zeros(np.shape(pts))

nodes[:, 0] = pts[:, 0] * (a1 - a0) + a0
nodes[:, 1] = pts[:, 1] * (b1 - b0) + b0

ns = 40

x = np.linspace(0, 1, ns)
samples = []
for i in range(npts):

    M = np.zeros((ns, ns))
    for k in range(ns):
        f = np.sin(0.1 * k * np.pi * nodes[i, 0] * x + 2 * np.pi * nodes[i, 1])
        M[:, k] = f

    samples.append(M)

#%% md
#
# Plot the coordinates of each point in the Euclidean space and plot two distinct points.

#%%

plt.plot(nodes[:, 0], nodes[:, 1], 'ko')
plt.plot(nodes[0, 0], nodes[0, 1], 'ro')
plt.plot(nodes[1, 0], nodes[1, 1], 'bo')
plt.text(nodes[0, 0], nodes[0, 1], '0')
plt.text(nodes[1, 0], nodes[1, 1], '1')
plt.show()

fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.imshow(samples[0])
ax1.title.set_text('Element 0')
ax2.imshow(samples[1])
ax2.title.set_text('Element 1')

#%% md
#
# Instantiate the ``UQpy`` class ``DiffusionMaps`` and compute the diffusion coordinates using the method ``mapping``.

#%%

from UQpy.dimension_reduction.kernels import Gaussian
dfm = DiffusionMaps.create_from_data(data=samples,
                                     alpha=0.5, eigenvectors_number=10,
                                     kernel=Gaussian(epsilon=0.3))
diff_coords, evals, evecs = dfm.mapping()


#%% md
#
# Plot the diffusion coordinates and the points presented above.

#%%

p0 = 1
p1 = 2
p2 = 3

fig = plt.figure(num=None, figsize=(10, 10))
ax = fig.gca(projection='3d')
ax.scatter(diff_coords[:, p0], diff_coords[:, p1], diff_coords[:, p2], c=[[0.3, 0.3, 0.3]], cmap=plt.cm.Spectral,
           s=30, marker='.')
ax.scatter(diff_coords[0, p0], diff_coords[0, p1], diff_coords[0, p2], c='r', cmap=plt.cm.Spectral, s=50)
ax.scatter(diff_coords[1, p0], diff_coords[1, p1], diff_coords[1, p2], c='b', cmap=plt.cm.Spectral, s=50)
ax.text(diff_coords[0, p0], diff_coords[0, p1], diff_coords[0, p2], '%s' % (str(0)), size=20, zorder=1, color='r')
ax.text(diff_coords[1, p0], diff_coords[1, p1], diff_coords[1, p2], '%s' % (str(1)), size=20, zorder=1, color='b')
ax.set_xlabel('diff_coords 1')
ax.set_ylabel('diff_coords 2')
ax.set_zlabel('diff_coords 3')
plt.show()

# docs/code/dimension_reduction/diffusion_maps/plot_dmaps_swiss_role.py

"""

Swiss role
=====================

This example shows how to use the :class:`.DiffusionMaps` class to reveal the embedded structure of noisy Swiss role data.
"""

#%%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.dimension_reduction import DiffusionMaps, GaussianKernel
from sklearn.datasets import make_swiss_roll

#%% md
#
# Sample points randomly following a parametric curve and plot the 3D graphic.

#%%

n = 4000  # number of samples

# generate point cloud
X, X_color = make_swiss_roll(n, random_state=3, noise=0)


#%% md
# plot the point cloud
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection="3d")
ax.scatter(X[:, 0], X[:, 1], X[:, 2], c=X_color, cmap=plt.cm.Spectral,
)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.set_title("Swiss role data")
ax.view_init(10, 70)


#%% md
#
# Case 1: Find the optimal parameter of the Gaussian kernel scale epsilon
kernel = GaussianKernel()

dmaps_object = DiffusionMaps.create_from_data(data=X, alpha=1.0, n_eigenvectors=9,
                                              optimize_parameters=True,
                                              kernel=kernel)


print('epsilon', kernel.epsilon)


#%% md
# Fit the data and calculate the embedding, the eigenvectors and eigenvalues
diff_coords, eigenvalues, eigenvectors = dmaps_object.fit()


#%% md
#
# Find the parsimonious representation of the eigenvectors. Identify the two most informative
# eigenvectors.

index, residuals = DiffusionMaps.parsimonious(eigenvectors, 2)

print('most informative eigenvectors:', index)

#%% md
#
# Plot the diffusion coordinates

DiffusionMaps._plot_eigen_pairs(eigenvectors, pair_indices=[1, 2], color=X_color, figure_size=[12, 12])
plt.show()



# docs/code/dimension_reduction/diffusion_maps/plot_grassmann_dmaps.py

"""

Grassmannian Diffusion Maps
===========================

"""
import numpy as np
from UQpy.dimension_reduction.grassmann_manifold.projections.SvdProjection import SvdProjection
from UQpy.dimension_reduction.grassmann_manifold.projections import KernelComposition
from UQpy.dimension_reduction.grassmann_manifold import Grassmann
from UQpy.utilities.kernels import ProjectionKernel
from UQpy.dimension_reduction.diffusion_maps.DiffusionMaps import DiffusionMaps
import sys
import matplotlib.pyplot as plt

#%%
#
# Create the data
N = 3000
r = 2 * np.random.rand(N)
x = np.zeros((N, 3))
c = np.zeros(N)

phi = np.linspace(0, 2*np.pi, N)
theta = np.arcsin(np.cos(phi)**2)

X = r * np.cos(phi) * np.sin(theta)
Y = r * np.sin(phi) * np.sin(theta)
Z = r * np.cos(theta)

c =X**2 + Y**2 + Z**2

fig = plt.figure()
plt.rcParams["figure.figsize"] = (5, 5)
plt.rcParams.update({'font.size': 18})
ax = fig.add_subplot(111, projection="3d")
ax.scatter(X, Y, Z, c=c, cmap=plt.cm.Spectral,
)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.set_title("data")

plt.show()

#%%

# Recast the data into a matrix
data_matrix = np.concatenate([X.reshape(-1, 1), Y.reshape(-1, 1), Z.reshape(-1, 1)], axis=1)

# Convert each data point into a 2d array with shape (1, 3)
data = [data2dArray.reshape(1, -1).T for data2dArray in data_matrix]

#%%
#
# Project each data point onto the Grassmann manifold using SVD. Use only 300 points from the origonal data.
# Select the work with maximum rank (number of planes) from each data point (p=sys.maxsize).
# Use the matrix of left eigenvectors to calculate the kernel (KernelComposition.LEFT).

Grassmann_projection = SvdProjection(data=data[::10], p=sys.maxsize, kernel_composition=KernelComposition.LEFT)
psi = Grassmann_projection.psi

#%%
#
# Plot the projected data on the Grassmann (3, 1) which is a unit sphere.

Grassmann_data = np.zeros((len(psi), 3))
for i in range(len(psi)):
    Grassmann_data[i, :] = psi[i].data.reshape(1, -1)

fig = plt.figure()
plt.rcParams["figure.figsize"] = (8, 8)
plt.rcParams.update({'font.size': 18})
ax = fig.add_subplot(111, projection="3d")

u, v = np.mgrid[0:2*np.pi:60j, 0:np.pi:60j]
x = np.cos(u)*np.sin(v)
y = np.sin(u)*np.sin(v)
z = np.cos(v)
ax.plot_surface(
    x, y, z,  rstride=1, cstride=1, color='c', alpha=0.2, linewidth=0)
ax.scatter(Grassmann_data[:, 0], Grassmann_data[:, 1], Grassmann_data[:, 2], c=c[:len(psi)], s=3, cmap=plt.cm.Spectral,
)
ax.set_title("Grassmann (3, 1)")


plt.show()

#%%
#
# Diffusion maps with Grassmann kernel
Grassmann_object = Grassmann(Grassmann_projection)
kernel = ProjectionKernel()
kernel_matrix = Grassmann_object.evaluate_kernel_matrix(kernel)

Gdmaps_UQpy = DiffusionMaps(alpha=1.0, n_eigenvectors=5, is_sparse=True, n_neighbors=250,
                            kernel_matrix=kernel_matrix)
diff_coords, evals_uqpy, evecs_uqpy = Gdmaps_UQpy.fit()

#%%
#
# Plot the first and second eigenvectors
DiffusionMaps._plot_eigen_pairs(evecs_uqpy, pair_indices=[1, 2], color=evecs_uqpy[:, 1],
                                figure_size=[5, 5], font_size=18)
DiffusionMaps._plot_eigen_pairs(evecs_uqpy, pair_indices=[1, 2], color=evecs_uqpy[:, 2],
                                figure_size=[5, 5], font_size=18)

#%%




# docs/code/dimension_reduction/grassmann/plot_grassmann_distances.py

"""

Compute distances between points
==================================

This example shows how to use the UQpy Grassmann class to compute distances
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the Grassmann class from UQpy implemented in the DimensionReduction module.

#%%

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from UQpy import SvdProjection
import sys
from UQpy import GrassmannDistance, RiemannianDistance

from UQpy.dimension_reduction import Grassmann

#%% md
#
# Generate four random matrices with reduced rank corresponding to the different samples. The samples are stored in
# `matrices`.

#%%

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.
# Solutions: original space.
Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

# Creating a list of solutions.
matrices = [Sol0, Sol1, Sol2, Sol3]

# Plot the solutions
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(Sol0)
ax2.title.set_text('Matrix 1')
ax2.imshow(Sol1)
ax3.title.set_text('Matrix 2')
ax3.imshow(Sol2)
ax4.title.set_text('Matrix 3')
ax4.imshow(Sol3)
plt.show()

#%% md
#
# Instatiate the UQpy class Grassmann considering the `"grassmann_distance"` as the a definition of distance on the
# manifold.

#%%

manifold_projection = SvdProjection(matrices, p_planes_dimensions=sys.maxsize)

#%% md
#
# Compute the pairwise distances for $\Psi$ and $\Phi$, the left and right -singular eigenvectors, respectively, of
# singular value decomposition of each solution.

#%%

pairwise_distance = Grassmann.calculate_pairwise_distances(distance_method=GrassmannDistance(),
                                                           points_grassmann=manifold_projection.psi)
print(pairwise_distance)

#%% md
#
# Compute the distance between 2 points.

#%%

distance_metric = GrassmannDistance()
distance01 = distance_metric.compute_distance(manifold_projection.psi[0], manifold_projection.psi[1])
print(distance01)

#%% md
#
# Compute the pairwise distances for $\Psi$ and $\Phi$, the left and right -singular eigenvectors, respectively, of
# singular value decomposition of each solution. In this case, use an user defined function `my_distance`.

#%%

class UserDistance(RiemannianDistance):

    def compute_distance(self, x0, x1):
        """
            Estimate the user distance.

            **Input:**

            * **x0** (`list` or `ndarray`)
                Point on the grassman manifold.

            * **x1** (`list` or `ndarray`)
                Point on the grassman manifold.

            **Output/Returns:**

            * **distance** (`float`)
                Procrustes distance between x0 and x1.
            """

        if not isinstance(x0, list) and not isinstance(x0, np.ndarray):
            raise TypeError('UQpy: x0 must be either list or numpy.ndarray.')
        else:
            x0 = np.array(x0)

        if not isinstance(x1, list) and not isinstance(x1, np.ndarray):
            raise TypeError('UQpy: x1 must be either list or numpy.ndarray.')
        else:
            x1 = np.array(x1)

        l = min(np.shape(x0))
        k = min(np.shape(x1))
        rank = min(l, k)

        r = np.dot(x0.T, x1)
        # (ui, si, vi) = svd(r, rank)

        ui, si, vi = np.linalg.svd(r, full_matrices=True, hermitian=False)  # Compute the SVD of matrix
        si = np.diag(si)  # Transform the array si into a diagonal matrix containing the singular values
        vi = vi.T  # Transpose of vi

        u = ui[:, :rank]
        s = si[:rank, :rank]
        v = vi[:, :rank]

        index = np.where(si > 1)
        si[index] = 1.0
        theta = np.arccos(si)
        theta = np.sin(theta / 2) ** 2
        distance = np.sqrt(abs(k - l) + 2 * np.sum(theta))

        return distance

pairwise_distance_psi = \
    Grassmann.calculate_pairwise_distances(distance_method=UserDistance(),
                                           points_grassmann=manifold_projection.psi)
pairwise_distance_phi = \
    Grassmann.calculate_pairwise_distances(distance_method=UserDistance(),
                                           points_grassmann=manifold_projection.phi)
print(pairwise_distance_psi)
print(pairwise_distance_phi)

# docs/code/dimension_reduction/grassmann/plot_grassmann_interpolation.py

"""

Interpolation
==================================

This example shows how to use the UQpy Grassmann class to perform interpolation on the Grassmann manifold
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the Grassmann class from UQpy implemented in the DimensionReduction module.

#%%

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from UQpy.dimension_reduction.grassman.manifold_projections.SvdProjection import SvdProjection
import sys
from UQpy.dimension_reduction import Grassmann
from UQpy.dimension_reduction.grassman.optimization_methods.GradientDescent import GradientDescent
from UQpy.dimension_reduction.distances.grassmanian.GrassmanDistance import GrassmannDistance
from UQpy import Interpolation
from UQpy import LinearInterpolation
from UQpy import InterpolationMethod

#%% md
#
# Generate the initial samples located at the vertices of a triangle. The coordinates of each vertix are stored in
# `nodes` and `point` contain the point to be interpolated on the tangent space.

#%%

nodes = np.array([[0, 0], [1, 0], [1, 1], [0, 1]])  # node_0, node_1, node_2.
point = np.array([0.1, 0.1])  # Point to interpolate.

plot_ = nodes[0:]
Xplot = plot_.T[0].tolist()
Xplot.append(plot_[0][0])
Yplot = plot_.T[1].tolist()
Yplot.append(plot_[0][1])
plt.plot(Xplot, Yplot)
plt.plot(nodes[0][0], nodes[0][1], 'ro')
plt.plot(nodes[1][0], nodes[1][1], 'ro')
plt.plot(nodes[2][0], nodes[2][1], 'ro')
plt.plot(nodes[3][0], nodes[3][1], 'ro')
plt.plot(point[0], point[1], 'bo')

dt = 0.015
plt.text(nodes[0][0] + dt, nodes[0][1] + dt, '0')
plt.text(nodes[1][0] + dt, nodes[1][1] + dt, '1')
plt.text(nodes[2][0] + dt, nodes[2][1] + dt, '2')
plt.text(nodes[3][0] + dt, nodes[3][1] + dt, '2')
plt.text(point[0] + dt, point[1] + dt, 'point')
plt.show()
plt.close()

#%% md
#
# Generate three random matrices with reduced rank corresponding to the different samples. The samples are stored in
# `Solutions`.

#%%

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.
# Solutions: original space.
Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

# Creating a list of solutions.
Solutions = [Sol0, Sol1, Sol2, Sol3]

# Plot the solutions
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Solution 0')
ax1.imshow(Sol0)
ax2.title.set_text('Solution 1')
ax2.imshow(Sol1)
ax3.title.set_text('Solution 2')
ax3.imshow(Sol2)
ax4.title.set_text('Solution 3')
ax4.imshow(Sol3)
plt.show()

#%% md
#
# Firs, let's perform the interpolation step-by-step using an object of the ``UQpy.Kriging`` class to interpolate.
# Further, instatiate the ``UQpy`` class ``Grassmann`` considering the `grassmann_distance` for the distance,
# `gradient_descent` to estimate the Karcher mean.

#%%

from UQpy.surrogates import Kriging
from UQpy.surrogates.kriging.regression_models.Linear import Linear
from UQpy.surrogates.kriging.correlation_models.Exponential import Exponential

Krig = Kriging(regression_model=Linear(),
               correlation_model=Exponential(),
               correlation_model_parameters=[1.0, 1.0],
               optimizations_number=1)

manifold_projection = SvdProjection(Solutions, p_planes_dimensions=sys.maxsize)

#%% md
#
# Compute the Karcher mean for $\Psi$ and $\Phi$, the left and right -singular eigenvectors, respectively, of singular
# value decomposition of each solution.

#%%

optimization_method = GradientDescent(acceleration=True, error_tolerance=1e-4, max_iterations=1000)
karcher_psi = Grassmann.karcher_mean(points_grassmann=manifold_projection.psi,
                                  p_planes_dimensions=manifold_projection.p_planes_dimensions,
                                  optimization_method=optimization_method,
                                  distance=GrassmannDistance())
karcher_phi = Grassmann.karcher_mean(points_grassmann=manifold_projection.phi,
                                     p_planes_dimensions=manifold_projection.p_planes_dimensions,
                                     optimization_method=optimization_method,
                                     distance=GrassmannDistance())

#%% md
#
# Rearrange the singular values $\Sigma$ of each solution as a diagonal matrix.

#%%

sigma_m = []
for i in range(len(manifold_projection.sigma)):
    sigma_m.append(np.diag(manifold_projection.sigma[i]))

#%% md
#
# Select the Karcher mean as a reference point for $\Psi$ and $\Phi$ and map those points on the manifold to the
# tangent space generated at the reference point.

#%%

gammaPsi = Grassmann.log_map(manifold_projection.psi, reference_point=karcher_psi)
gammaPhi = Grassmann.log_map(manifold_projection.phi, reference_point=karcher_phi)

#%% md
#
# Perform the standard linear interpolation of `point` on the tangent space for $\Psi$, $\Phi$, and $\Sigma$. The
# interpolated points are given by $\tilde{\Psi}$, $\tilde{\Phi}$, and $\tilde{\Sigma}$. Thus, the interpolated
# solution is given by $\tilde{\mathrm{X}}=\tilde{\Psi}\tilde{\Phi}\tilde{\Sigma}$.

#%%

interpolation = Interpolation(interpolation_method=Krig)

interpPsi = interpolation.interpolate_sample(coordinates=nodes, samples=gammaPsi, point=point)
interpPhi = interpolation.interpolate_sample(coordinates=nodes, samples=gammaPhi, point=point)
interpS = interpolation.interpolate_sample(coordinates=nodes, samples=sigma_m, point=point)

PsiTilde = Grassmann.exp_map([interpPsi], reference_point=karcher_psi)
PhiTilde = Grassmann.exp_map([interpPhi], reference_point=karcher_phi)

PsiTilde = np.array(PsiTilde[0])
PhiTilde = np.array(PhiTilde[0])
SolTilde = np.dot(np.dot(PsiTilde, interpS), PhiTilde.T)

#%% md
#
# Print the interpolated solution and compare to the given solutions associated to the different vertices of the
# triangle.

#%%

print(Sol0)
print(" ")
print(Sol1)
print(" ")
print(Sol2)
print("-------------------------")
print(SolTilde)

plot_ = nodes[0:]
Xplot = plot_.T[0].tolist()
Xplot.append(plot_[0][0])
Yplot = plot_.T[1].tolist()
Yplot.append(plot_[0][1])
plt.plot(Xplot, Yplot)
plt.plot(nodes[0][0], nodes[0][1], 'ro')
plt.plot(nodes[1][0], nodes[1][1], 'ro')
plt.plot(nodes[2][0], nodes[2][1], 'ro')
plt.plot(nodes[3][0], nodes[3][1], 'ro')
plt.plot(point[0], point[1], 'bo')

dt = 0.015
plt.text(nodes[0][0] + dt, nodes[0][1] + dt, '0')
plt.text(nodes[1][0] + dt, nodes[1][1] + dt, '1')
plt.text(nodes[2][0] + dt, nodes[2][1] + dt, '2')
plt.text(nodes[3][0] + dt, nodes[3][1] + dt, '3')
plt.text(point[0] + dt, point[1] + dt, 'p')
plt.show()
plt.close()

# Plot the solutions
fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(1, 5)
ax1.title.set_text('Solution 0')
ax1.imshow(Sol0)
ax2.title.set_text('Solution 1')
ax2.imshow(Sol1)
ax3.title.set_text('Solution 2')
ax3.imshow(Sol2)
ax4.title.set_text('Solution 3')
ax4.imshow(Sol3)
ax5.title.set_text('Interpolated')
ax5.imshow(SolTilde)
plt.show()

#%% md
#
# All the operations above are implemented in the method `interpolate`, which is used below. in this case, the
# interpolation is performed on the entries of the input matrices.

#%%

X = manifold_projection.reconstruct_solution(interpolation=interpolation,
                                             coordinates=nodes,
                                             point=point,
                                             p_planes_dimensions=manifold_projection.p_planes_dimensions,
                                             distance=GrassmannDistance(),
                                             element_wise=True,
                                             optimization_method=GradientDescent())

# Plot the solutions
fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(1, 5)
ax1.title.set_text('Solution 0')
ax1.imshow(Sol0)
ax2.title.set_text('Solution 1')
ax2.imshow(Sol1)
ax3.title.set_text('Solution 2')
ax3.imshow(Sol2)
ax4.title.set_text('Solution 3')
ax4.imshow(Sol3)
ax5.title.set_text('Interpolated')
ax5.imshow(X)
plt.show()

#%% md
#
# Now, let's use an object of sklearn.gaussiann_process. (To run this example, you have to install the scikit learn
# toolbox in advance.)

#%%

from sklearn.gaussian_process import GaussianProcessRegressor
gp = GaussianProcessRegressor()

# Instantiate the method again with: interp_object=gp.
interpolation = Interpolation(interpolation_method=gp)

X = manifold_projection.reconstruct_solution(coordinates=nodes, point=[point], element_wise=True,
                                             interpolation=interpolation,
                                             optimization_method=GradientDescent(),
                                             distance=GrassmannDistance(),
                                             p_planes_dimensions=manifold_projection.p_planes_dimensions)

#%% md
#
# Plot the solution.

#%%

# Plot the solutions
fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(1, 5)
ax1.title.set_text('Solution 0')
ax1.imshow(Sol0)
ax2.title.set_text('Solution 1')
ax2.imshow(Sol1)
ax3.title.set_text('Solution 2')
ax3.imshow(Sol2)
ax4.title.set_text('Solution 3')
ax4.imshow(Sol3)
ax5.title.set_text('Interpolated')
ax5.imshow(X)
plt.show()


# docs/code/dimension_reduction/grassmann/plot_grassmann_karcher.py

"""

Karcher
==================================

This example shows how to use the UQpy Grassmann class to use the logarithmic map and the exponential map
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the Grassmann class from UQpy implemented in the DimensionReduction module.

#%%

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import sys
from UQpy.dimension_reduction.grassman.manifold_projections.SvdProjection import SvdProjection
from UQpy.dimension_reduction import Grassmann

#%% md
#
# Generate four random matrices with reduced rank corresponding to the different samples. The samples are stored in
# `matrices`.

#%%

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.
# Solutions: original space.
Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

# Creating a list of matrices.
matrices = [Sol0, Sol1, Sol2, Sol3]

# Plot the matrices
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(Sol0)
ax2.title.set_text('Matrix 1')
ax2.imshow(Sol1)
ax3.title.set_text('Matrix 2')
ax3.imshow(Sol2)
ax4.title.set_text('Matrix 3')
ax4.imshow(Sol3)
plt.show()

#%% md
#
# Instatiate the UQpy class Grassmann.

#%%

manifold_projection = SvdProjection(matrices, p_planes_dimensions=sys.maxsize)

# Plot the points on the Grassmann manifold defined by the left singular eigenvectors.
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(manifold_projection.psi[0])
ax2.title.set_text('Matrix 1')
ax2.imshow(manifold_projection.psi[0])
ax3.title.set_text('Matrix 2')
ax3.imshow(manifold_projection.psi[0])
ax4.title.set_text('Matrix 3')
ax4.imshow(manifold_projection.psi[0])
plt.show()

#%% md
#
# Compute the Karcher mean.

#%%

from UQpy.dimension_reduction.grassman.optimization_methods.GradientDescent import GradientDescent
from UQpy.dimension_reduction.distances.grassmanian.GrassmanDistance import GrassmannDistance

optimization_method = GradientDescent(acceleration=True, error_tolerance=1e-4, max_iterations=1000)
distance_method = GrassmannDistance()
karcher_psi = Grassmann.karcher_mean(points_grassmann=manifold_projection.psi,
                                     p_planes_dimensions=manifold_projection.p_planes_dimensions,
                                     distance=distance_method,
                                     optimization_method=optimization_method)
karcher_phi = Grassmann.karcher_mean(points_grassmann=manifold_projection.phi,
                                     p_planes_dimensions=manifold_projection.p_planes_dimensions,
                                     distance=distance_method,
                                     optimization_method=optimization_method)

#%% md
#
# Project $\Psi$, the left singular eigenvectors, on the tangent space centered at the Karcher mean.

#%%

points_tangent = Grassmann.log_map(points_grassmann=manifold_projection.psi,
                                   reference_point=karcher_psi)

print(points_tangent[0])
print(points_tangent[1])
print(points_tangent[2])
print(points_tangent[3])

# Plot the matrices
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(points_tangent[0])
ax2.title.set_text('Matrix 1')
ax2.imshow(points_tangent[1])
ax3.title.set_text('Matrix 2')
ax3.imshow(points_tangent[2])
ax4.title.set_text('Matrix 3')
ax4.imshow(points_tangent[3])
plt.show()

#%% md
#
# Map the points back to the Grassmann manifold.

#%%

points_grassmann = Grassmann.exp_map(points_tangent=points_tangent,
                                     reference_point=manifold_projection.psi[0])

print(points_grassmann[0])
print(points_grassmann[1])
print(points_grassmann[2])
print(points_grassmann[3])

# Plot the matrices
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(points_grassmann[0])
ax2.title.set_text('Matrix 1')
ax2.imshow(points_grassmann[1])
ax3.title.set_text('Matrix 2')
ax3.imshow(points_grassmann[2])
ax4.title.set_text('Matrix 3')
ax4.imshow(points_grassmann[3])
plt.show()

# docs/code/dimension_reduction/grassmann/plot_grassmann_kernel.py

"""

Kernel
==================================

This example shows how to use the UQpy Grassmann class to compute kernels
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the Grassmann class from UQpy implemented in the DimensionReduction module.

#%%

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from UQpy.dimension_reduction.grassman.manifold_projections.SvdProjection import SvdProjection
from UQpy.dimension_reduction import Grassmann
import sys
from UQpy import ProjectionKernel
from UQpy import KernelComposition
from UQpy import OrthoMatrixPoints
from UQpy import Kernel

#%% md
#
# Generate four random matrices with reduced rank corresponding to the different samples. The samples are stored in
# `matrices`.

#%%
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.
# Solutions: original space.
Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

# Creating a list of solutions.
matrices = [Sol0, Sol1, Sol2, Sol3]

# Plot the solutions
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(Sol0)
ax2.title.set_text('Matrix 1')
ax2.imshow(Sol1)
ax3.title.set_text('Matrix 2')
ax3.imshow(Sol2)
ax4.title.set_text('Matrix 3')
ax4.imshow(Sol3)
plt.show()

#%% md
#
# Instatiate the UQpy class Grassmann considering the `projection_kernel` as the a kernel definition on the Grassmann
# manifold.

#%%

manifold_projection = SvdProjection(matrices, p_planes_dimensions=sys.maxsize,
                                    kernel_composition=KernelComposition.LEFT)
manifold = Grassmann(manifold_projected_points=manifold_projection)

#%% md
#
# Compute the kernels for $\Psi$ and $\Phi$, the left and right -singular eigenvectors, respectively, of singular value
# decomposition of each solution.

#%%

kernel_psi = manifold.evaluate_kernel_matrix(kernel=ProjectionKernel())

manifold_projection.kernel_composition = KernelComposition.RIGHT
kernel_phi = manifold.evaluate_kernel_matrix(kernel=ProjectionKernel())

fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.title.set_text('kernel_psi')
ax1.imshow(kernel_psi)
ax2.title.set_text('kernel_phi')
ax2.imshow(kernel_phi)
plt.show()

#%% md
#
# Compute the kernel only for 2 points.

#%%

grassmann_points = OrthoMatrixPoints(input_points=[manifold_projection.psi[0],
                                                   manifold_projection.psi[1],
                                                   manifold_projection.psi[2]],
                                     p_planes_dimensions=manifold_projection.p_planes_dimensions)
manifold = Grassmann(manifold_projected_points=grassmann_points)
kernel_01 = manifold.evaluate_kernel_matrix(kernel=ProjectionKernel())

fig = plt.figure()
plt.imshow(kernel_01)
plt.show()

#%% md
#
# Compute the kernels for $\Psi$ and $\Phi$, the left and right -singular eigenvectors, respectively, of singular value
# decomposition of each solution. In this case, use an user defined function `my_kernel`.

#%%

class UserKernel(Kernel):

    def apply_method(self, data):
        data.evaluate_matrix(self, self.kernel_operator)

    def pointwise_operator(self, x0, x1):
        if not isinstance(x0, list) and not isinstance(x0, np.ndarray):
            raise TypeError('UQpy: x0 must be either list or numpy.ndarray.')
        else:
            x0 = np.array(x0)

        if not isinstance(x1, list) and not isinstance(x1, np.ndarray):
            raise TypeError('UQpy: x1 must be either list or numpy.ndarray.')
        else:
            x1 = np.array(x1)

        r = np.dot(x0.T, x1)
        det = np.linalg.det(r)
        ker = det * det
        return ker

manifold_projection.kernel_composition = KernelComposition.LEFT
kernel_user_psi = manifold.evaluate_kernel_matrix(kernel=UserKernel())

manifold_projection.kernel_composition = KernelComposition.RIGHT
kernel_user_phi = manifold.evaluate_kernel_matrix(kernel=UserKernel())


fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.title.set_text('kernel_psi')
ax1.imshow(kernel_user_psi)
ax2.title.set_text('kernel_phi')
ax2.imshow(kernel_user_phi)
plt.show()


# docs/code/dimension_reduction/grassmann/plot_grassmann_log_exp.py

"""

Mappings to & from tangent space
==================================

This example shows how to use the UQpy Grassmann class to use the logarithmic map and the exponential map
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the Grassmann class from UQpy implemented in the DimensionReduction module.

#%%

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from UQpy.dimension_reduction.grassman.manifold_projections.SvdProjection import SvdProjection
from UQpy.dimension_reduction import Grassmann
import sys

#%% md
#
# Generate four random matrices with reduced rank corresponding to the different samples. The samples are stored in
# `matrices`.

#%%

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.
# Solutions: original space.
Sol0 = np.dot(np.random.rand(D1, r0), np.random.rand(r0, D1))
Sol1 = np.dot(np.random.rand(D1, r1), np.random.rand(r1, D1))
Sol2 = np.dot(np.random.rand(D1, r2), np.random.rand(r2, D1))
Sol3 = np.dot(np.random.rand(D1, r3), np.random.rand(r3, D1))

# Creating a list of matrices.
matrices = [Sol0, Sol1, Sol2, Sol3]

# Plot the matrices
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(Sol0)
ax2.title.set_text('Matrix 1')
ax2.imshow(Sol1)
ax3.title.set_text('Matrix 2')
ax3.imshow(Sol2)
ax4.title.set_text('Matrix 3')
ax4.imshow(Sol3)
plt.show()

#%% md
#
# Instantiate the UQpy class Grassmann.

#%%

manifold_projection = SvdProjection(matrices, p_planes_dimensions=sys.maxsize)

# Plot the points on the Grassmann manifold defined by the left singular eigenvectors.
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(manifold_projection.psi[0])
ax2.title.set_text('Matrix 1')
ax2.imshow(manifold_projection.psi[0])
ax3.title.set_text('Matrix 2')
ax3.imshow(manifold_projection.psi[0])
ax4.title.set_text('Matrix 3')
ax4.imshow(manifold_projection.psi[0])
plt.show()

#%% md
#
# Project $\Psi$, the left singular eigenvectors, on the tangent space centered at $\Psi_0$.

#%%

points_tangent = Grassmann.log_map(points_grassmann=manifold_projection.psi,
                                   reference_point=manifold_projection.psi[0])

print(points_tangent[0])
print(points_tangent[1])
print(points_tangent[2])
print(points_tangent[3])

# Plot the matrices
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(points_tangent[0])
ax2.title.set_text('Matrix 1')
ax2.imshow(points_tangent[1])
ax3.title.set_text('Matrix 2')
ax3.imshow(points_tangent[2])
ax4.title.set_text('Matrix 3')
ax4.imshow(points_tangent[3])
plt.show()

#%% md
#
# Map the points back to the Grassmann manifold.

#%%

points_grassmann = Grassmann.exp_map(points_tangent=points_tangent,
                                     reference_point=manifold_projection.psi[0])

print(points_grassmann[0])
print(points_grassmann[1])
print(points_grassmann[2])
print(points_grassmann[3])

# Plot the matrices
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
ax1.title.set_text('Matrix 0')
ax1.imshow(points_grassmann[0])
ax2.title.set_text('Matrix 1')
ax2.imshow(points_grassmann[1])
ax3.title.set_text('Matrix 2')
ax3.imshow(points_grassmann[2])
ax4.title.set_text('Matrix 3')
ax4.imshow(points_grassmann[3])
plt.show()

# docs/code/dimension_reduction/grassmann/user_distance.py

"""

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


"""
import numpy as np

def my_distance(x0, x1):

    """
    Estimate the user distance.

    **Input:**

    * **x0** (`list` or `ndarray`)
        Point on the grassmann_manifold manifold.

    * **x1** (`list` or `ndarray`)
        Point on the grassmann_manifold manifold.

    **Output/Returns:**

    * **distance** (`float`)
        Procrustes distance between x0 and x1.
    """

    if not isinstance(x0, list) and not isinstance(x0, np.ndarray):
        raise TypeError('UQpy: x0 must be either list or numpy.ndarray.')
    else:
        x0 = np.array(x0)

    if not isinstance(x1, list) and not isinstance(x1, np.ndarray):
        raise TypeError('UQpy: x1 must be either list or numpy.ndarray.')
    else:
        x1 = np.array(x1)

    l = min(np.shape(x0))
    k = min(np.shape(x1))
    rank = min(l, k)

    r = np.dot(x0.T, x1)
    #(ui, si, vi) = svd(r, rank)
    
    ui, si, vi = np.linalg.svd(r, full_matrices=True,hermitian=False)  # Compute the SVD of matrix
    si = np.diag(si)  # Transform the array si into a diagonal matrix containing the singular values
    vi = vi.T  # Transpose of vi

    u = ui[:, :rank]
    s = si[:rank, :rank]
    v = vi[:, :rank]
    
    index = np.where(si > 1)
    si[index] = 1.0
    theta = np.arccos(si)
    theta = np.sin(theta / 2) ** 2
    distance = np.sqrt(abs(k - l) + 2 * np.sum(theta))

    return distance

# docs/code/dimension_reduction/grassmann/user_interp.py

"""

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


"""
import numpy as np
from scipy.interpolate import LinearNDInterpolator

def my_interp(coordinates, samples, point):

    """

    **Input:**

    * **coordinates** (`ndarray`)
        Coordinates of the input second_order_tensor points.

    * **samples** (`ndarray`)
        Matrices corresponding to the points on the grassmann_manifold manifold.

    * **point** (`ndarray`)
        Coordinates of the point to be interpolated.

    **Output/Returns:**

    * **interp_point** (`ndarray`)
        Interpolated point.
    """

    if not isinstance(coordinates, list) and not isinstance(coordinates, np.ndarray):
        raise TypeError('UQpy: `coordinates` must be either list or ndarray.')
    else:
        coordinates = np.array(coordinates)

    if not isinstance(samples, list) and not isinstance(samples, np.ndarray):
        raise TypeError('UQpy: `samples` must be either list or ndarray.')
    else:
        samples = np.array(samples)

    if not isinstance(point, list) and not isinstance(point, np.ndarray):
        raise TypeError('UQpy: `point` must be either list or ndarray.')
    else:
        point = np.array(point)

    myInterpolator = LinearNDInterpolator(coordinates, samples)
    interp_point = myInterpolator(point)
    interp_point = interp_point[0]

    return interp_point



# docs/code/dimension_reduction/grassmann/user_karcher.py

"""

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


"""
import numpy as np
from UQpy.dimension_reduction import Grassmann

def my_karcher(data_points, distance_fun, kwargs):

        """


        **Input:**

        * **data_points** (`list`)
            Points on the grassmann_manifold manifold.
            
        * **distance_fun** (`callable`)
            Distance function.

        * **kwargs** (`dictionary`)
            Contains the keywords for the used in the optimizers to find the Karcher mean.

        **Output/Returns:**

        * **mean_element** (`list`)
            Karcher mean.
        """

        if 'tol' in kwargs.keys():
            tol = kwargs['tol']
        else:
            tol = 1e-3

        if 'maxiter' in kwargs.keys():
            maxiter = kwargs['maxiter']
        else:
            maxiter = 1000

        n_mat = len(data_points)

        rnk = []
        for i in range(n_mat):
            rnk.append(min(np.shape(data_points[i])))

        max_rank = max(rnk)

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

        index_0 = fmean.index(min(fmean))

        mean_element = data_points[index_0].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 = Grassmann.log_map(points_grassmann=[data_points[idx]], reference_point=np.asarray(mean_element))

                step = 2 * alpha * _gamma[0]

                X = Grassmann.exp_map(points_tangent=[step], reference_point=np.asarray(mean_element))

                _gamma = []
                mean_element = X[0]

                k += 1

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

            itera += 1

        return mean_element



# docs/code/dimension_reduction/grassmann/user_kernel.py

"""

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


"""
import numpy as np

def my_kernel(x0, x1):

    """
    User defined kernel.

    **Input:**

    * **x0** (`list` or `ndarray`)
        Point on the grassmann_manifold manifold.

    * **x1** (`list` or `ndarray`)
        Point on the grassmann_manifold manifold.

    **Output/Returns:**

    * **distance** (`float`)
        Kernel value for x0 and x1.
    """

    if not isinstance(x0, list) and not isinstance(x0, np.ndarray):
        raise TypeError('UQpy: x0 must be either list or numpy.ndarray.')
    else:
        x0 = np.array(x0)

    if not isinstance(x1, list) and not isinstance(x1, np.ndarray):
        raise TypeError('UQpy: x1 must be either list or numpy.ndarray.')
    else:
        x1 = np.array(x1)

    r = np.dot(x0.T, x1)
    det = np.linalg.det(r)
    ker = det * det
    return ker

# docs/code/dimension_reduction/grassmann/user_kernel_diffusion.py

"""

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


"""
import numpy as np
import itertools
import scipy.spatial.distance as sd

def my_kernel_diffusion(data):

    epsilon=0.3
    # Compute the pairwise distances.
    if len(np.shape(data)) == 2:
        # Set of 1-D arrays
        distance_pairs = sd.pdist(data, 'euclidean')
    elif len(np.shape(data)) == 3:
        # Set of 2-D arrays
        # Check arguments: verify the consistency of input arguments.
        nargs = len(data)
        indices = range(nargs)
        pairs = list(itertools.combinations(indices, 2))

        distance_pairs = []
        for id_pair in range(np.shape(pairs)[0]):
            ii = pairs[id_pair][0]  # Point i
            jj = pairs[id_pair][1]  # Point j

            x0 = data[ii]
            x1 = data[jj]

            distance = np.linalg.norm(x0 - x1, 'fro')

            distance_pairs.append(distance)
    else:
        raise TypeError('UQpy: The size of the input second_order_tensor is not consistent with this method.')
        
    kernel_matrix = np.exp(-sd.squareform(distance_pairs) ** 2 / (4 * epsilon))
    
    return kernel_matrix

# docs/code/dimension_reduction/pod/DiffusionEquation.py

"""

Diffusion Equation File
========================


"""
import numpy as np
def diffusion(w, h, dx, dy, D, Tcool, Thot, r, cx, cy, nsteps):
    nx, ny = int(w/dx), int(h/dy)

    dx2, dy2 = dx*dx, dy*dy
    dt = dx2 * dy2 / (2 * D * (dx2 + dy2))

    u0 = Tcool * np.ones((nx, ny))
    u = u0.copy()

    r2 = r**2
    for i in range(nx):
        for j in range(ny):
            p2 = (i*dx-cx)**2 + (j*dy-cy)**2
            if p2 < r2:
                u0[i,j] = Thot

    def do_timestep(u0, u):
        # Propagate with forward-difference in time, central-difference in space
        u[1:-1, 1:-1] = u0[1:-1, 1:-1] + D * dt * (
              (u0[2:, 1:-1] - 2*u0[1:-1, 1:-1] + u0[:-2, 1:-1])/dx2
              + (u0[1:-1, 2:] - 2*u0[1:-1, 1:-1] + u0[1:-1, :-2])/dy2 )

        u0 = u.copy()
        return u0, u

    u_f = []
    for m in range(nsteps):
        u0, u = do_timestep(u0, u)
        u_f.append(u0)
    return u_f

# docs/code/dimension_reduction/pod/plot_pod_1.py

"""

POD on data
=====================

In this example, the Direct Proper Orthogonal Decomposition (POD) method is used to decompose a set of data and extract
basis functions which can be used for the reconstruction of the solution.
"""

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy, matplotlib, and we also import the
# POD class from UQpy.

# %%

from UQpy.dimension_reduction import DirectPOD
import time
import numpy as np

# %% md
#
# Input dataset in the form of a second order tensor consists of three snapshot solutions. Using the POD method we
# reconstruct the input dataset by keeping a predefined number of modes.

# %%

Data = np.zeros((3, 5, 3))

Data[:, :, 0] = [
    [0.9073, 1.7842, 2.1236, 1.1323, 1.6545],
    [0.8924, 1.7753, -0.6631, 0.5654, 2.1235],
    [2.1488, 4.2495, 1.8260, 0.3423, 4.9801]]

Data[:, :, 1] = [
    [0.7158, 1.6970, -0.0740, 5.478, 1.0987],
    [-0.4898, -1.5077, 1.9103, 6.7121, 0.5334],
    [0.3054, 0.3207, 2.1335, 1.1082, 5.5435]]

Data[:, :, 2] = [
    [-0.3698, 0.0151, 1.4429, 2.5463, 6.9871],
    [2.4288, 4.0337, -1.7495, -0.5012, 1.7654],
    [2.3753, 4.7146, -0.2716, 1.6543, 0.9121]]

# %% md
#
# The Direct POD method is used to compute the reconstructed solutions and reduced-order solutions in the spatial
# dimension of the data. Full reconstruction is achieved when the number of modes chosen, equals the number of
# dimensions.

# %%

start_time = time.time()

pod = DirectPOD(solution_snapshots=Data, modes=1)
Data_reconstr = pod.run()

# %% md
#
# Print the reconstructed dataset.

# %%

print('Reconstructed snapshot no.1:')
print(Data_reconstr[:, :, 0])

if np.allclose(Data, Data_reconstr) == True:
    print('Input data and reconstructed data are identical.')

elapsed_time = time.time() - start_time
time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
print('Elapsed time: ', round(elapsed_time, 8))


# docs/code/dimension_reduction/pod/plot_pod_diffusion.py

"""

POD on Diffusion Equation 2D
============================

In this example, the diffusion equation is solved and then methods from the POD class are used to decompose the output
solutions/dataset and extract its basis functions which can be used for the reconstruction of the solution.

2D Diffusion equation
---------------------
> $\displaystyle \frac{\partial U}{\partial t} = D \bigg(\frac{\partial^2 U}{\partial x^2} + \frac{\partial^2 U}{\partial y^2}\bigg)$
> <br>
> <br>
> where $D$ is the diffusion coefficient. $U$ describes the behavior of the particles in Brownian motion, resulting from their random movements and collisions.
<img src="plate_and_disc.png" alt="plate_and_disc.png" height="160" width="160" align=right>

### Problem description:

>  - A 2D metal plate is initially at temperature $T_{cool}$.
>  - A disc of a specified size inside the plate is at temperature $T_{hot}$.
>  - Suppose that the edges of the plate are held fixed at $T_{cool}$.
>  - The diffusion equation is applied to follow the evolution of the temperature of the plate.
"""

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy, matplotlib, and we also import the
# POD class from UQpy.

# %%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.dimension_reduction import DirectPOD, SnapshotPOD, HigherOrderSVD
from DiffusionEquation import diffusion
import time

# %% md
#
# The diffusion equation is solved by calling the 'diffusion' function and a dataset (list) is obtained. To run this function the following need to be specified:
# >  - w, h - Plate size, mm
# >  - dx, dy - Intervals in x-, y- directions, mm
# >  - D - Thermal diffusivity, mm2.s-1
# >  - Tcool, Thot - Plate and disc temperature
# >  - r, cx, cy - Initial conditions - ring of inner radius r, width dr centred at (cx,cy) (mm)
# >  - nsteps - Number of time steps

# %%

w = h = 5.
dx, dy = 0.1, 0.1
D = 5
Tcool, Thot = 400, 700
r, cx, cy = 1.5, 2.5, 2.5
nsteps = 500

Data = diffusion(w, h, dx, dy, D, Tcool, Thot, r, cx, cy, nsteps)

# %% md
#
# The Direct POD method is used to reconstruct the data for different values of spatial modes. Full reconstruction is
# achieved when the number of modes chosen equals the number of dimensions. Since the dataset for this problem is large,
# the Snapshot POD is recommended as the decomposition will be performed much faster.

# %%

start_time = time.time()

n_modes = [1, 2, 3, 4, 40, 50]
frame = 40

Data_modes = []

for n_mode in n_modes:
    pod = DirectPOD(solution_snapshots=Data, modes=n_mode)
    Data_reconstr, Data_reduced = pod.run()
    Data_modes.append(Data_reconstr[:, :, frame])
    del Data_reconstr

elapsed_time = time.time() - start_time
time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
print('Elapsed time: ', elapsed_time)

# %% md
#
# Comparison of input and reduced solution.

# %%

# Plot input solution
plt.figure(figsize=(16, 2.5))
c = plt.imshow(Data[frame], cmap=plt.get_cmap('coolwarm'), vmin=Tcool, vmax=Thot)
plt.colorbar(c)
plt.title('Input solution', fontweight="bold", size=15)
plt.show()

# Plot reduced solution
fig = plt.figure(figsize=(10, 7))
fig.subplots_adjust(hspace=0.5, wspace=0.4)
for i in range(len(n_modes)):
    ax = fig.add_subplot(2, 3, i + 1)
    im = ax.imshow(Data_modes[i], cmap=plt.get_cmap('coolwarm'), vmin=Tcool, vmax=Thot)
    ax.set_axis_off()
    ax.set_title('Mode {}'.format(n_modes[i]), size=15)

fig.subplots_adjust(right=0.85)
cbar_ax = fig.add_axes([0.9, 0.15, 0.03, 0.7])
cbar_ax.set_xlabel('$T$ / K', labelpad=20)
fig.colorbar(im, cax=cbar_ax)
fig.suptitle('Reconstructed solution', fontweight="bold", size=15)
plt.show()


# docs/code/distributions/continuous_1d/plot_distribution_continuous_1D.py

"""

Distribution Continuous 1D Example
==================================

This examples shows the use of a univariate continuous distributions class. In particular:
"""

#%% md
#
# - How to define one of the univariate distributions supported by UQpy
# - How to plot the pdf and log_pdf of the distribution
# - How to modify the parameters of the distribution
# - How to extract the moments of the distribution
# - How to draw random samples from the distribution

#%%

#%% md
#
# Initially we have to import the necessary modules.

#%%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Lognormal

#%% md
#
# Example of a univariate lognormal distribution
# ----------------------------------------------
# In order to define the lognormal distribution, the user must provide its three parameters, *s*, *loc* and *scale*.
# Printing the dist object created after the Lognormal distribution initialization, will display as output the class
# of the distribution as can be seen below. The provided parameters of the distribution can be retrieved using the
# *parameters* attribute available to all distributions.


#%%

dist = Lognormal(s=1., loc=0., scale=np.exp(5))
print(dist)
print(dist.parameters)

#%% md
#
# Plot the pdf of the distribution.
# ------------------------------------
#
# The user must provide x as a ndarray of shape (nsamples, 1) or (nsamples,) - the former if preferred.
# The result of pdf or log_pdf will be a 1D array (nsamples, ).

#%%

x = np.linspace(0.01, 1000, 1000).reshape((-1, 1))  # Use reshape to provide a 2D array (1000, 1)
fig, ax = plt.subplots(ncols=2, figsize=(15, 4))
ax[0].plot(x, dist.pdf(x))  # Do not give params
ax[0].set_xlabel('x')
ax[0].set_ylabel('pdf(x)')
ax[0].set_title('pdf of lognormal distribution')

ax[1].plot(x, dist.log_pdf(x))
ax[1].set_xlabel('x')
ax[1].set_ylabel('log pdf(x)')
ax[1].set_title('Log pdf of lognormal distribution')
plt.show()

print('size of input x:')
print(x.shape)
print('size of dist.pdf(x):')
print(dist.pdf(x).shape)

#%% md
#
# Modify one of the parameters of the distribution.
# -------------------------------------------------
#
# Use the *update_parameters* method. The user must provide as input to the *update_parameters* method the name,
# as well as the updated value of the specified parameter. Note that in case the user provides a non-existing parameter
# name, the *update_parameters* method will raise an exception.

#%%

dist.update_parameters(loc=100.)
print(dist.parameters)

#%% md
#
# Plot the pdf and log_pdf functions of the lognormal distribution with the updated parameters.

#%%

x = np.linspace(0.01, 1000, 1000).reshape((-1, 1))  # Use reshape to provide a 2D array (1000, 1)
fig, ax = plt.subplots(ncols=2, figsize=(15, 4))
ax[0].plot(x, dist.pdf(x))  # Do not give params
ax[0].set_xlabel('x')
ax[0].set_ylabel('pdf(x)')
ax[0].set_title('pdf of lognormal distribution')

ax[1].plot(x, dist.log_pdf(x))
ax[1].set_xlabel('x')
ax[1].set_ylabel('log pdf(x)')
ax[1].set_title('Log pdf of lognormal distribution')
plt.show()

#%% md
#
# Print the mean, standard deviation, skewness, and kurtosis of the distribution.
# -------------------------------------------------------------------------------
# Using the moments method existing in all univariate distributions, the user can retrieve the available
# moments. The order in which the moments are extracted can be seen in the moments_list variable.

#%%

moments_list = ['mean', 'variance', 'skewness', 'kurtosis']
m = dist.moments()
print('Moments with inherited parameters:')
for i, moment in enumerate(moments_list):
    print(moment + ' = {0:.2f}'.format(m[i]))

# %% md
#
# Generate 5000 random samples from the lognormal distribution.
# -------------------------------------------------------------
#
# The number of samples is provided as nsamples (default 1).
# The user can fix the seed of the pseudo-random generator via input random_state.
#
# Important: the output of rvs is a (nsamples, 1) ndarray.

# %%

y1 = dist.rvs(nsamples=5000)
print('Shape of output provided by rvs is (nsamples, dimension), i.e. here:')
print(y1.shape)
plt.hist(y1[:, 0], bins=50)
plt.xlabel('x')
plt.ylabel('count')
plt.show()

# docs/code/distributions/continuous_1d/plot_distribution_normal_fitting.py

"""

Distribution fitting
==================================

This examples showcases the calculation of a distributions parameters based on fitting available data
"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

from UQpy.distributions.collection.Normal import Normal

#%% md
#
# Define a Normal distribution and use the fit method.
# ------------------------------------------------------
#
# Parameters to be learnt should be instantiated as None.
# Note that the fit method of each distribution returns a dictionary containing as keys the names of parameters
# and its values are their data fitted values.

#%%

normal1 = Normal(loc=None, scale=None)
fitted_parameters1 = normal1.fit(data=[-4, 2, 2, 1])
print(fitted_parameters1)

normal2 = Normal(loc=0., scale=None)
fitted_parameters2 = normal2.fit(data=[-4, 2, 2, 1])
print(fitted_parameters2)



# docs/code/distributions/discrete_1d/plot_distribution_discrete_1D.py

"""

Distribution Discrete 1D example
=================================

This examples shows the use of the univariate discrete distributions class. In particular:
"""
#%% md
#
# - How to define one of the univariate discrete distributions supported by UQpy
# - How to extract the moments of the distribution
# - How to draw random samples from the distribution

#%%

#%% md
#
# Import the necessary modules.

#%%

import matplotlib.pyplot as plt
from UQpy.distributions.collection.Binomial import Binomial

#%% md
#
# Example of a 1D discrete distribution
# --------------------------------------
#
# Define a univariate binomial distribution.
# By using the __bases__ attribute we can verify that the Binomial distribution extends the DistributionDiscrete1D
# baseclass, while in order to define the Binomial distribution, two parameters are required, namely, *n* and *p*.

#%%

print(Binomial.__bases__)
dist = Binomial(n=5, p=0.4)


#%% md
#
# Print the mean, standard deviation, skewness, and kurtosis of the distribution.
# -------------------------------------------------------------------------------
# Using the moments method existing in all univariate distributions, the user can retrieve the available
# moments. The order in which the moments are extracted can be seen in the moments_list variable.

#%%

moments_list = ['mean', 'variance', 'skewness', 'kurtosis']
m = dist.moments()
print('Moments with inherited parameters:')
for i, moment in enumerate(moments_list):
    print(moment + ' = {0:.2f}'.format(m[i]))

# %% md
#
# Generate 5000 random samples from the binomial distribution.
# -------------------------------------------------------------
#
# The number of samples is provided as nsamples (default 1).
# The user can fix the seed of the pseudo random generator via input random_state.
#
# Important: the output of rvs is a (nsamples, 1) ndarray.

# %%

y1 = dist.rvs(nsamples=5000)
print('Shape of output provided by rvs is (nsamples, dimension), i.e. here:')
print(y1.shape)
plt.hist(y1[:, 0], bins=50)
plt.xlabel('x')
plt.ylabel('count')
plt.show()








# docs/code/distributions/multivariate/plot_joint_independent.py

"""

Multivariate from independent marginals
===============================================

This examples shows the use of the multivariate normal distribution class. In particular:
"""

#%% md
#
# - How to define one multivariate distribution from independent marginals supported by UQpy
# - How to plot the pdf of the distribution
# - How to extract the moments of the distribution
# - How to draw random samples from the distribution

#%%

#%% md
#
# Import the necessary modules.

#%%


from UQpy.distributions.collection import Normal, Lognormal, JointIndependent
import numpy as np
import matplotlib.pyplot as plt


#%% md
#
# Define a multivariate distribution from independent univariate marginals.
# --------------------------------------------------------------------------
#
# In order to define a JointIndependent distribution, a list of marginal distributions is initially created.
# These marginals are then provided as input to the JointIndependent class of UQpy.
# Retrieving the multivariate distribution parameters is achieved using the get_parameters method.

#%%

marginals = [Normal(loc=2., scale=2.),
             Lognormal(s=1., loc=0., scale=np.exp(5))]

dist = JointIndependent(marginals=marginals)

print(dist.get_parameters())

#%% md
#
# Sample the distribution  and plot the pdf of the distribution.
# ---------------------------------------------------------------
# In a similar manner to univariate distributions, samples can be drawn from an JointIndependent distribution
# using its rvs method.

#%%

data = dist.rvs(nsamples=1000)
fig, ax = plt.subplots(ncols=2, figsize=(10, 4))
ax[0].scatter(data[:, 0], data[:, 1], alpha=0.2)
ax[0].set_xlabel('dimension 1')
ax[0].set_ylabel('dimension 2')
ax[0].set_title('random samples')
ax[0].set_ylim([0, 500])
ax[0].set_xlim([-2, 6])

x = np.arange(-2.0, 6.0, 0.2)
y = np.arange(0.01, 500, 1)
X, Y = np.meshgrid(x, y)
Z = dist.pdf(x=np.concatenate([X.reshape((-1, 1)), Y.reshape((-1, 1))], axis=1))
CS = ax[1].contour(X, Y, Z.reshape(X.shape))
ax[1].clabel(CS, inline=1, fontsize=10)
ax[1].set_xlabel('dimension 1')
ax[1].set_ylabel('dimension 2')
ax[1].set_title('Contour plot of pdf')
ax[0].set_ylim([0, 500])
ax[0].set_xlim([-2, 6])
plt.show()

#%% md
#
# Print the moments of the distribution
# -------------------------------------------------------------------------------
# Providing a single or multiple consecutive initials of the following four distributions moments
# 'm': mean, 'v': variance, 's': skewness, 'k': kurtosis
# allows the user to obtains the respective moments from all underlying univariate distributions.
# In the following examples providing the string 'mv' to the moments function, returns the respective means and
# variances.

#%%

print(dist.moments())
print(dist.moments(moments2return='mv'))


#%% md
#
# Modify the parameters of the distribution.
# -------------------------------------------------
#
# Use the update_parameters method.

#%%
print(dist)
print()
print('Parameters of the marginals:')
print([m.parameters for m in marginals])
print('Parameters of the joint:')
print(dist.get_parameters())
print()
print('Update the location parameter of the second marginal and scale parameter of first marginal...')
dist.update_parameters(loc_1=1., scale_0=3., )
print('Parameters of the marginals:')
print([m.parameters for m in marginals])
print('Parameters of the joint:')
print(dist.get_parameters())

# docs/code/distributions/multivariate/plot_multivariate_copulas.py

"""

Multivariate from independent marginals and copula
==================================================

"""

#%% md
#
# - How to define α bivariate distribution from independent marginals and change its structure based on a copula supported by UQpy
# - How to plot the pdf of the distribution
# - How to modify the parameters of the distribution

#%%

#%% md
#
# Import the necessary modules.

#%%

import numpy as np
import matplotlib.pyplot as plt

#%% md
#
# Example of a multivariate distribution from joint independent marginals
# ------------------------------------------------------------------------


#%%
from UQpy.distributions import Normal, JointIndependent
from UQpy.distributions import Gumbel, JointCopula

#%% md
#
# Define a Copula
# ---------------
# The definition of bivariate distribution with a copula, is similar to defining a multivariate distribution from
# independent marginals. In both cases a list of marginals needs to be defined. In case of 

#%%
marginals = [Normal(loc=0., scale=1), Normal(loc=0., scale=1)]
copula = Gumbel(theta=3.)

# dist_1 is a multivariate normal with independent marginals
dist_1 = JointIndependent(marginals)
print('Does the distribution with independent marginals have an rvs method?')
print(hasattr(dist_1, 'rvs'))

# dist_2 exhibits dependence between the two dimensions, defined using a gumbel copula
dist_2 = JointCopula(marginals=marginals, copula=copula)
print('Does the distribution with copula have an rvs method?')
print(hasattr(dist_2, 'rvs'))

#%% md
#
# Plot the pdf of the distribution before and after the copula
# -------------------------------------------------------------
#

#%%

fig, ax = plt.subplots(ncols=2, figsize=(10, 4))

x = np.arange(-3, 3, 0.1)
y = np.arange(-3, 3, 0.1)
X, Y = np.meshgrid(x, y)
Z = dist_1.pdf(x=np.concatenate([X.reshape((-1, 1)), Y.reshape((-1, 1))], axis=1))
CS = ax[0].contour(X, Y, Z.reshape(X.shape))
ax[0].clabel(CS, inline=1, fontsize=10)
ax[0].set_title('Contour plot of pdf - independent normals')

x = np.arange(-3, 3, 0.1)
y = np.arange(-3, 3, 0.1)
X, Y = np.meshgrid(x, y)
Z = dist_2.pdf(x=np.concatenate([X.reshape((-1, 1)), Y.reshape((-1, 1))], axis=1))
CS = ax[1].contour(X, Y, Z.reshape(X.shape))
ax[1].clabel(CS, inline=1, fontsize=10)
ax[1].set_title('Contour plot of pdf - normals with Gumbel copula')
plt.show()



#%% md
#
# Modify the parameters of the multivariate copula.
# -------------------------------------------------
#
# Use the update_parameters method.

#%%
print(dist_2.copula.parameters)
dist_2.update_parameters(theta_c=2.)
print(dist_2.copula.parameters)

# docs/code/distributions/multivariate/plot_multivariate_normal.py

"""

Multivariate normal distribution
==================================

This examples shows the use of the multivariate normal distributions class. In particular:
"""

#%% md
#
# - How to define one of the univariate distributions supported by UQpy
# - How to plot the pdf of the distribution
# - How to extract the moments of the distribution
# - How to draw random samples from the distribution

#%%

#%% md
#
# Initially we have to import the necessary modules.

#%%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions.collection.MultivariateNormal import MultivariateNormal

#%% md
#
# Example of a multivariate normal distribution
# ---------------------------------------------
# Note that multivariate normal distribution can facilitate any number of dimensions.
# In order to define it two arguments are necessary, specifically, a list containing the
# mean values for each one of the dimensions and a covariance matrix with shape
# (ndimensions, ndimensions)

#%%

print(MultivariateNormal.__bases__)
dist = MultivariateNormal(mean=[1., 2.], cov=[[4., -0.2], [-0.2, 1.]])

#%% md
#
# Plot the two-dimensional pdf of the distribution.
# -------------------------------------------------
#

#%%

fig, ax = plt.subplots(ncols=1, figsize=(10, 4))
x = np.arange(-6.0, 6.0, 0.1)
y = np.arange(-6.0, 6.0, 0.1)
X, Y = np.meshgrid(x, y)
Z = dist.pdf(x=np.concatenate([X.reshape((-1, 1)), Y.reshape((-1, 1))], axis=1))
CS = ax.contour(X, Y, Z.reshape(X.shape))
ax.clabel(CS, inline=1, fontsize=10)
ax.set_xlabel('dimension 1')
ax.set_ylabel('dimension 2')
ax.set_title('Contour plot of pdf')
ax.set_xlim([-4, 6])
ax.set_ylim([-2, 6])
plt.show()

#%% md
#
# Print the multivariate moments of the distribution.
# ---------------------------------------------------
# Providing a single or multiple consecutive initials of the following four distributions moments
# 'm': mean, 'v': variance, 's': skewness, 'k': kurtosis
# allows the user to obtains the respective moments from all underlying univariate distributions.
# In the following examples providing the string 'mv' to the moments function, returns the respective means and
# variances.

#%%
print(dist.moments())
print(dist.moments(moments2return='mv'))

# %% md
#
# Generate 1000 random samples from the binomial distribution.
# -------------------------------------------------------------
#
# Important: the output of rvs is a (nsamples, 1) ndarray.

# %%

data = dist.rvs(nsamples=1000)

fig, ax = plt.subplots(ncols=1, figsize=(10, 4))
ax.scatter(data[:, 0], data[:, 1], alpha=0.2)
ax.set_xlabel('dimension 1')
ax.set_ylabel('dimension 2')
ax.set_title('random samples')
ax.set_xlim([-4, 6])
ax.set_ylim([-2, 6])



# docs/code/distributions/user_defined/plot_user_defined.py

"""

User-defined Rosenbrock distribution
=====================================

This examples shows the use of the multivariate normal distributions class. In particular:
"""


import numpy as np
from UQpy.distributions import DistributionND
import matplotlib.pyplot as plt

#%% md
#
# Example with a custom distribution
# ----------------------------------
# In order to define a new distribution, the user must extend the one of the abstract base classes
# *DistributionContinuous1D*, *DistributionDiscrete1D* or *DistributionND*. For the purpose of this example a
# new multivariate Rosenbrock distribution is defined. Note that three methods are implemented, namely, the __init__
# which allows the user to define custom distribution arguments, as well as the pdf and log_pdf of this new
# distribution. Note that it is required for the user to call the **__init__** method of the baseclass by providing all
# arguments names and values required in the custom distribution initializer. In our case, the parameter name **p** and
# its values are provided to the *super().__init__* method.

#%%


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']

#%% md
#
# Initialize custom distribution
# ----------------------------------
# Given the newly defined distribution class, a Rosenbrock distribution object can be defined by providing to the
# initializer the user-defined argument p. Since the Rosenbrock distribution extends the *DistributionND* class, methods
# and attributes of the baseclass are already available. For instance, the user can retrieve the already defined
# parameters using the *parameters* attribute and subsequently update them using the *update_parameters* method.

#%%

dist = Rosenbrock(p=20)
print(dist.parameters)
dist.update_parameters(p=40)
print(dist.parameters)

dist = Rosenbrock(p=20)
print(dist.parameters)

#%% md
#
# Plot pdf of the user defined distribution
# ------------------------------------------
#

#%%

fig, ax = plt.subplots(figsize=(8, 5))
x = np.arange(-5, 8, 0.1)
y = np.arange(-5, 50, 0.1)
X, Y = np.meshgrid(x, y)
Z = dist.pdf(x=np.concatenate([X.reshape((-1, 1)), Y.reshape((-1, 1))], axis=1))
CS = ax.contour(X, Y, Z.reshape(X.shape))
ax.clabel(CS, inline=1, fontsize=10)
ax.set_title('Contour plot of custom pdf - Rosenbrock')
plt.show()

#%% md
#
# Check if the user-defined distribution has rvs, pdf and update_parameters method
# ----------------------------------------------------------------------------------
#

#%%

print('Does the rosenbrock distribution have an rvs method?')
print(hasattr(dist, 'rvs'))

print('Does the rosenbrock distribution have an pdf method?')
print(hasattr(dist, 'pdf'))

print('Does the rosenbrock distribution have an update_parameters method?')
print(hasattr(dist, 'update_parameters'))

# docs/code/inference/bayes_model_selection/bayes_model_selection.py

"""

Bayes Model Selection - Regression Models
=============================================================================

In the following we present an example for which the posterior pdf of the parameters, evidences and model probabilities
can be computed analytically. We drop the :math:`m_{j}` subscript when referring to model parameters for simplicity. Three
models are considered (the domain :math:`x` is fixed and consists in 50 equally spaced points):

.. math:: m_{linear}: \quad y = θ_{0} x + \epsilon

.. math:: m_{quadratic}: \quad y = θ_{0} x + θ_{1} x^2 + \epsilon

.. math:: m_{cubic}: \quad y = θ_{0} x + θ_{1} x^2+ θ_{2} x^3 + \epsilon

All three models can be written in a compact form as :math:`y=X θ + \epsilon`, where :math:`X` contains the
necessary powers of :math:`x`. For all three models, the prior is chosen to be Gaussian,
:math:`p(θ) = N(\cdot, θ_{prior}, \Sigma_{prior})`, and so is the noise
:math:`\epsilon \sim N(\cdot; 0, \sigma_{n}^{2} I)`.
"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

import numpy as np
import matplotlib.pyplot as plt
from UQpy.sampling.mcmc import MetropolisHastings
from UQpy.inference import BayesModelSelection, BayesParameterEstimation, ComputationalModel
from UQpy.RunModel import RunModel  # required to run the quadratic model
from UQpy.distributions import Normal, JointIndependent
from scipy.stats import multivariate_normal, norm


#%% md
#
# Generate data from a quadratic function

#%%

param_true = np.array([1.0, 2.0]).reshape(1, -1)
var_n = 1
error_covariance = var_n * np.eye(50)
print(param_true.shape)

z = RunModel(samples=param_true, model_script='local_pfn_models.py', model_object_name='model_quadratic', vec=False,
             var_names=['theta_1', 'theta_2'])
data_clean = z.qoi_list[0].reshape((-1,))
data = data_clean + Normal(scale=np.sqrt(var_n)).rvs(nsamples=data_clean.size, random_state=456).reshape((-1,))

#%% md
#
# Define the models, compute the true values of the evidence.
#
# For all three models, a Gaussian prior is chosen for the parameters, with mean and covariance matrix of the
# appropriate dimensions. Each model is given prior probability :math:`P(m_{j}) = 1/3`.

#%%

model_names = ['model_linear', 'model_quadratic', 'model_cubic']
model_n_params = [1, 2, 3]
model_prior_means = [[0.], [0., 0.], [0., 0., 0.]]
model_prior_stds = [[10.], [1., 1.], [1., 2., 0.25]]


evidences = []
model_posterior_means = []
model_posterior_stds = []
for n, model in enumerate(model_names):
    # compute matrix X
    X = np.linspace(0, 10, 50).reshape((-1, 1))
    if n == 1:  # quadratic model
        X = np.concatenate([X, X ** 2], axis=1)
    if n == 2:  # cubic model
        X = np.concatenate([X, X ** 2, X ** 3], axis=1)

    # compute posterior pdf
    m_prior = np.array(model_prior_means[n]).reshape((-1, 1))
    S_prior = np.diag(np.array(model_prior_stds[n]) ** 2)
    S_posterior = np.linalg.inv(1 / var_n * np.matmul(X.T, X) + np.linalg.inv(S_prior))
    m_posterior = np.matmul(S_posterior,
                            1 / var_n * np.matmul(X.T, data.reshape((-1, 1))) + np.matmul(np.linalg.inv(S_prior),
                                                                                          m_prior))
    m_prior = m_prior.reshape((-1,))
    m_posterior = m_posterior.reshape((-1,))
    model_posterior_means.append(list(m_posterior))
    model_posterior_stds.append(list(np.sqrt(np.diag(S_posterior))))
    print('posterior mean and covariance for ' + model)
    print(m_posterior, S_posterior)

    # compute evidence, evaluate the formula at the posterior mean
    like_theta = multivariate_normal.pdf(data, mean=np.matmul(X, m_posterior).reshape((-1,)), cov=error_covariance)
    prior_theta = multivariate_normal.pdf(m_posterior, mean=m_prior, cov=S_prior)
    posterior_theta = multivariate_normal.pdf(m_posterior, mean=m_posterior, cov=S_posterior)
    evidence = like_theta * prior_theta / posterior_theta
    evidences.append(evidence)
    print('evidence for ' + model + '= {}\n'.format(evidence))

# compute the posterior probability of each model
tmp = [1 / 3 * evidence for evidence in evidences]
model_posterior_probas = [p / sum(tmp) for p in tmp]

print('posterior probabilities of all three models')
print(model_posterior_probas)

#%% md
#
# Define the models for use in UQpy

#%%

candidate_models = []
for n, model_name in enumerate(model_names):
    run_model = RunModel(model_script='local_pfn_models.py', model_object_name=model_name, vec=False)
    prior = JointIndependent([Normal(loc=m, scale=std) for m, std in
                              zip(model_prior_means[n], model_prior_stds[n])])
    model = ComputationalModel(n_parameters=model_n_params[n],
                               runmodel_object=run_model, prior=prior,
                               error_covariance=error_covariance,
                               name=model_name)
    candidate_models.append(model)


#%% md
#
# Run MCMC for one model

#%%

# Quadratic model
sampling = MetropolisHastings(args_target=(data, ),
                              log_pdf_target=candidate_models[1].evaluate_log_posterior,
                              jump=10, burn_length=100,
                              proposal=JointIndependent([Normal(scale=0.1), ] * 2),
                              seed=[0., 0.], random_state=123)

bayesMCMC = BayesParameterEstimation(sampling_class=sampling,
                                     inference_model=candidate_models[1],
                                     data=data,
                                     nsamples=3500)

# plot prior, true posterior and estimated posterior
fig, ax = plt.subplots(1, 2, figsize=(16, 5))
for n_p in range(2):
    domain_plot = np.linspace(-0.5, 3, 200)
    ax[n_p].plot(domain_plot, norm.pdf(domain_plot, loc=model_prior_means[1][n_p], scale=model_prior_stds[1][n_p]),
                 label='prior', color='green', linestyle='--')
    ax[n_p].plot(domain_plot, norm.pdf(domain_plot, loc=model_posterior_means[1][n_p],
                                       scale=model_posterior_stds[1][n_p]),
                 label='true posterior', color='red', linestyle='-')
    ax[n_p].hist(bayesMCMC.sampler.samples[:, n_p], density=True, bins=30, label='estimated posterior MCMC')
    ax[n_p].legend()
    ax[n_p].set_title('MCMC for quadratic model')
plt.show()


#%% md
#
# Run Bayesian Model Selection for all three models

#%%

proposals = [Normal(scale=0.1),
             JointIndependent([Normal(scale=0.1), Normal(scale=0.1)]),
             JointIndependent([Normal(scale=0.15), Normal(scale=0.1), Normal(scale=0.05)])]
nsamples = [2000, 6000, 14000]
nburn = [500, 2000, 4000]
jump = [5, 10, 25]


sampling_inputs=list()
estimators = []
for i in range(3):
    sampling = MetropolisHastings(args_target=(data, ),
                                  log_pdf_target=candidate_models[i].evaluate_log_posterior,
                                  jump=jump[i],
                                  burn_length=nburn[i],
                                  proposal=proposals[i],
                                  seed=model_prior_means[i],
                                  random_state=123)
    estimators.append(BayesParameterEstimation(inference_model=candidate_models[i], data=data,
                                                  sampling_class=sampling))

selection = BayesModelSelection(parameter_estimators=estimators,
                                data=data,
                                prior_probabilities=[1. / 3., 1. / 3., 1. / 3.],
                                nsamples=nsamples)

sorted_indices = np.argsort(selection.probabilities)[::-1]
print('Sorted models:')
print([selection.candidate_models[i].name for i in sorted_indices])
print('Evidence of sorted models:')
print([selection.evidences[i] for i in sorted_indices])
print('Posterior probabilities of sorted models:')
print([selection.probabilities[i] for i in sorted_indices])

#%% md
#
# As of version 2, the implementation of BayesModelSelection in UQpy uses the method of the harmonic mean to compute
# the models' evidence. This method is known to behave quite poorly, in particular it yeidls estimates with large
# variance. In the problem above, this implementation does not consistently detects that the quadratic model has the
# highest model probability. Future versions of UQpy will integrate more advanced methods for the estimation of the
# evidence.

#%%

for i, (m, be) in enumerate(zip(selection.candidate_models, selection.bayes_estimators)):
    # plot prior, true posterior and estimated posterior
    print('Posterior parameters for model ' + m.name)
    fig, ax = plt.subplots(1, 3, figsize=(16, 5))
    for n_p in range(m.n_parameters):
        domain_plot = np.linspace(min(be.sampler.samples[:, n_p]), max(be.sampler.samples[:, n_p]), 200)
        ax[n_p].plot(domain_plot, norm.pdf(domain_plot, loc=model_prior_means[i][n_p],
                                           scale=model_prior_stds[i][n_p]),
                     label='prior', color='green', linestyle='--')
        ax[n_p].plot(domain_plot, norm.pdf(domain_plot, loc=model_posterior_means[i][n_p],
                                           scale=model_posterior_stds[i][n_p]),
                     label='true posterior', color='red', linestyle='-')
        ax[n_p].hist(be.sampler.samples[:, n_p], density=True, bins=30, label='estimated posterior MCMC')
        ax[n_p].legend()
    plt.show()

shutil.rmtree(z.model_dir)
shutil.rmtree(candidate_models[0].runmodel_object.model_dir)
shutil.rmtree(candidate_models[1].runmodel_object.model_dir)
shutil.rmtree(candidate_models[2].runmodel_object.model_dir)




# docs/code/inference/bayes_model_selection/local_pfn_models.py

"""

Auxiliary file
==============================================

"""
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

# docs/code/inference/bayes_parameter_estimation/local_pfn_models.py

"""

Auxiliary file
==============================================

"""
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

# docs/code/inference/bayes_parameter_estimation/plot_bayes_parameter_IS_probability.py

"""

Parameter estimation using Importance Sampling - Probability Model
=============================================================================

In the following we learn the mean and covariance of a univariate gaussian distribution from data.

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.inference import DistributionModel, BayesParameterEstimation
from sklearn.neighbors import KernelDensity  # for the plots
from UQpy.distributions import JointIndependent, Uniform, Lognormal, Normal

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)

#%% md
#
# First, for the sake of this example, we generate fake data from a gaussian distribution with mean 10 and standard
# deviation 1.

#%%

# 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
data = np.random.normal(mu, sigma, 100).reshape((-1, 1))
np.random.seed()

# plot the data and true distribution
count, bins, ignored = plt.hist(data, 30, density=True)
plt.plot(bins, 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(- (bins - mu) ** 2 / (2 * sigma ** 2)),
         linewidth=2, color='r')
plt.title('data as histogram and true distribution to be estimated')
plt.show()


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)

#%% md
#
# Learn the unknown parameters using IS. If no proposal is given, the samples are sampled from the prior.

#%%

from UQpy.sampling import ImportanceSampling
sampling = ImportanceSampling()
bayes_estimator = BayesParameterEstimation(sampling_class=sampling,
                                               inference_model=candidate_model,
                                               data=data,
                                               nsamples=50000)


s_prior = bayes_estimator.sampler.samples
bayes_estimator.sampler.resample()
s_posterior = bayes_estimator.sampler.unweighted_samples

# print results
fig, ax = plt.subplots(1, 2, figsize=(10, 4))

domain = np.linspace(0, 15, 200)[:, np.newaxis]
pdf_ = pdf_from_kde(domain, s_posterior[:, 0])
ax[0].plot(domain, pdf_, label='posterior')
pdf_ = pdf_from_kde(domain, s_prior[:, 0])
ax[0].plot(domain, pdf_, label='prior')
ax[0].legend()
ax[0].set_title('theta 1')

domain = np.linspace(0, 2, 200)[:, np.newaxis]
pdf_ = pdf_from_kde(domain, s_posterior[:, 1])
ax[1].plot(domain, pdf_, label='posterior')
pdf_ = pdf_from_kde(domain, s_prior[:, 1])
ax[1].plot(domain, pdf_, label='prior')
ax[1].legend()
ax[1].set_title('theta 1')

plt.show()



# docs/code/inference/bayes_parameter_estimation/plot_bayes_parameter_IS_regression.py

"""

Parameter estimation using Importance Sampling - Regression Model
=============================================================================

Here a model is defined that is of the form

.. math:: y=f(θ) + \epsilon

where :math:`f` consists in running RunModel. In particular, here :math:`f(θ)= θ_{0} x + θ_{1} x^{2}` is
a regression model.

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

import numpy as np
import matplotlib.pyplot as plt

from UQpy.sampling.ImportanceSampling import ImportanceSampling
from UQpy.inference import DistributionModel, BayesParameterEstimation, ComputationalModel
from UQpy.RunModel import RunModel  # required to run the quadratic model
from sklearn.neighbors import KernelDensity  # for the plots
from UQpy.distributions import JointIndependent, Uniform, Lognormal, Normal

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)

#%% md
#
# First we generate synthetic data, and add some noise to it.

#%%

param_true = np.array([1.0, 2.0]).reshape((1, -1))
print('Shape of true parameter vector: {}'.format(param_true.shape))

h_func = RunModel(model_script='local_pfn_models.py', model_object_name='model_quadratic', vec=False,
                  var_names=['theta_0', 'theta_1'])
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=50).reshape((50,))
data_3 = data_clean + noise
print('Shape of data: {}'.format(data_3.shape))


inference_model = ComputationalModel(n_parameters=2, runmodel_object=h_func, error_covariance=error_covariance)

sampling = ImportanceSampling(proposal=JointIndependent([Normal(scale=2, )] * 2))
bayes_estimator =\
    BayesParameterEstimation(inference_model=inference_model,
                             data=data_3,
                             sampling_class=sampling,
                             nsamples=5000)

s = bayes_estimator.sampler.samples
w = bayes_estimator.sampler.weights
print(sum(w))

# print results
fig, ax = plt.subplots(1, 2)
for i in range(2):
    ax[i].hist(x=s[:, i], weights=None, density=True, range=(-4, 4), bins=20, color='blue', alpha=0.4,
               label='prior')
    ax[i].hist(x=s[:, i], weights=w, density=True, range=(-4, 4), bins=20, color='orange', alpha=0.7,
               label='posterior')
    ax[i].legend()
    ax[i].set_title('theta {}'.format(i + 1))

plt.show()

shutil.rmtree(h_func.model_dir)


# docs/code/inference/bayes_parameter_estimation/plot_bayes_parameter_MCMC_probability.py

"""

Parameter estimation using MCMC - Probability Model
=============================================================================

In the following we learn the mean and covariance of a univariate gaussian distribution from data.

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

import numpy as np
import matplotlib.pyplot as plt

from UQpy.inference import DistributionModel, BayesParameterEstimation
from sklearn.neighbors import KernelDensity  # for the plots
from UQpy.distributions import JointIndependent, Uniform, Lognormal, Normal

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)

#%% md
#
# First, for the sake of this example, we generate fake data from a gaussian distribution with mean 10 and standard
# deviation 1.

#%%

np.random.seed(100)
mu, sigma = 10, 1  # true mean and standard deviation
data_1 = np.random.normal(mu, sigma, 100).reshape((-1, 1))
np.random.seed()

# plot the data and true distribution
count, bins, ignored = plt.hist(data_1, 30, density=True)
plt.plot(bins, 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(- (bins - mu) ** 2 / (2 * sigma ** 2)),
         linewidth=2, color='r')
plt.title('data as histogram and true distribution to be estimated')
plt.show()


#%% md
#
# In a Bayesian setting, the definition of a prior pdf is a key point. The prior for the parameters must be defined in
# the model. Note that if no prior is given, an improper, uninformative, prior is chosen, :math:`p(\theta)=1` for all
# :math:`\theta`.

#%%

p0 = Uniform(loc=0., scale=15)
p1 = Lognormal(s=1., loc=0., scale=1.)
prior = JointIndependent(marginals=[p0, p1])

candidate_model = DistributionModel(distributions=Normal(loc=None, scale=None), n_parameters=2, prior=prior)

# Learn the unknown parameters using MCMC
from UQpy.sampling import MetropolisHastings

mh1 = MetropolisHastings(jump=10, burn_length=10, seed=[1.0, 0.2], random_state=123)

bayes_estimator = BayesParameterEstimation(inference_model=candidate_model,
                                           data=data_1,
                                           nsamples=500,
                                           sampling_class=mh1)

# print results
s = bayes_estimator.sampler.samples
plt.scatter(s[:, 0], s[:, 1])
plt.scatter(10, 1, marker='+', label='true parameter')
plt.title('MCMC samples')
plt.legend()
plt.show()

fig, ax = plt.subplots(1, 2, figsize=(10, 4))

domain = np.linspace(0, 15, 200)[:, np.newaxis]
pdf = pdf_from_kde(domain, s[:, 0])
ax[0].plot(domain, p0.pdf(domain), label='prior')
ax[0].plot(domain, pdf, label='posterior')
ax[0].set_title('posterior pdf of theta=mu')
ax[0].legend()

domain = np.linspace(0, 2, 200)[:, np.newaxis]
pdf = pdf_from_kde(domain, s[:, 1])
ax[1].plot(domain, p1.pdf(domain), label='prior')
ax[1].plot(domain, pdf, label='posterior')
ax[1].set_title('posterior pdf of theta=sigma')
ax[1].legend()

plt.show()

# docs/code/inference/bayes_parameter_estimation/plot_bayes_parameter_MCMC_regression.py

"""

Parameter estimation using MCMC - Regression Model
=============================================================================

Here a model is defined that is of the form

.. math:: y=f(θ) + \epsilon

where :math:`f` consists in running RunModel. In particular, here :math:`f(θ)=θ_{0} x + θ_{1} x^{2}` is a
regression model.

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

import numpy as np
import matplotlib.pyplot as plt

from UQpy.sampling.mcmc.MetropolisHastings import MetropolisHastings
from UQpy.inference.inference_models.ComputationalModel import ComputationalModel
from UQpy.RunModel import RunModel
from UQpy.inference import DistributionModel, BayesParameterEstimation
from sklearn.neighbors import KernelDensity  # for the plots
from UQpy.distributions import JointIndependent, Uniform, Lognormal, Normal

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)

#%% md
#
# First we generate synthetic data, and add some noise to it.

#%%

# Generate data
param_true = np.array([1.0, 2.0]).reshape((1, -1))
print('Shape of true parameter vector: {}'.format(param_true.shape))

h_func = RunModel(model_script='local_pfn_models.py', model_object_name='model_quadratic', vec=False,
                  var_names=['theta_0', 'theta_1'])
h_func.run(samples=param_true)
data_clean = np.array(h_func.qoi_list[0])
print(data_clean.shape)

# Add noise, use a RandomState for reproducible results
error_covariance = 1.
noise = Normal(loc=0., scale=np.sqrt(error_covariance)).rvs(nsamples=50, random_state=123).reshape((50,))
data_3 = data_clean + noise
print('Shape of data: {}'.format(data_3.shape))
print(data_3[:4])

p0 = Normal()
p1 = Normal()
prior = JointIndependent(marginals=[p0, p1])

inference_model = ComputationalModel(n_parameters=2, runmodel_object=h_func, error_covariance=error_covariance,
                                     prior=prior)

proposal = JointIndependent([Normal(scale=0.1), Normal(scale=0.05)])

mh1 = MetropolisHastings(jump=10, burn_length=0, proposal=proposal, seed=[0.5, 2.5],
                         random_state=456)

bayes_estimator = BayesParameterEstimation(inference_model=inference_model,
                                           data=data_3,
                                           sampling_class=mh1,
                                           nsamples=500)

s = bayes_estimator.sampler.samples
plt.scatter(s[:, 0], s[:, 1])
plt.scatter(1.0, 2.0, label='true value')
plt.title('MCMC samples')
plt.show()

fig, ax = plt.subplots(1, 2, figsize=(10, 4))

domain = np.linspace(-4, 4, 200)[:, np.newaxis]
pdf_ = pdf_from_kde(domain, s[:, 0])
ax[0].plot(domain, pdf_, label='posterior')
ax[0].plot(domain, p0.pdf(domain), label='prior')
ax[0].set_title('posterior pdf of theta_{1}')

domain = np.linspace(-4, 4, 200)[:, np.newaxis]
pdf_ = pdf_from_kde(domain, s[:, 1])
ax[1].plot(domain, pdf_, label='posterior')
ax[1].plot(domain, p1.pdf(domain), label='prior')
ax[1].set_title('posterior pdf of theta_{2}')

plt.show()

print(bayes_estimator.sampler.samples[:4])

shutil.rmtree(h_func.model_dir)


# docs/code/inference/info_model_selection/local_pfn_models.py

"""

Auxiliary file
==============================================

"""
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

# docs/code/inference/info_model_selection/plot_selection_distributions.py

"""

Selection between univariate distributions
==============================================

Here a model is defined that is of the form

.. math:: y=f(θ) + \epsilon

where f consists in running RunModel. In particular, here :math:`f(θ)=θ_{0} x + θ_{1} x^{2}` is a regression model.


"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

from UQpy.inference import DistributionModel, InformationModelSelection, MLE
import numpy as np
import matplotlib.pyplot as plt
from UQpy.distributions import Gamma, Exponential, ChiSquare

#%% md
#
# Generate data using a gamma distribution.

#%%

data = Gamma(a=2, loc=0, scale=2).rvs(nsamples=500, random_state=12)
print(data.shape)

#%% md
#
# Define the models to be compared, then call InfoModelSelection to perform model selection. By default,
# InfoModelSelection returns its outputs, fitted parameters, value of the chosen criteria, model probabilities and so
# on, in a sorted order, i.e., starting with the most probable model. However, if setting sorted_outputs=False, the
# class output attributes are given in the same order as the candidate_models.

#%%


# Define the models to be compared, for each model one must create an instance of the model class

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)

#%% md
#
# Perform model selection using different information criteria

#%%

from UQpy.inference import BIC, AIC, AICc

criteria = [BIC(), AIC(), AICc()]
sorted_names =[]
criterion_value = []
param =[]
for criterion in criteria:
    selector = InformationModelSelection(parameter_estimators=[mle1, mle2, mle3], criterion=criterion,
                                         n_optimizations=[5]*3)
    selector.sort_models()
    print('Sorted model using ' + str(criterion) + ' criterion: ' + ', '.join(
        m.name for m in selector.candidate_models))
    if isinstance(criterion, BIC):
        criterion_value = selector.criterion_values
        sorted_names = [m.name for m in selector.candidate_models]
        param = [m.mle for m in selector.parameter_estimators]

width = 0.5
ind = np.arange(len(sorted_names))
p1 = plt.bar(ind, criterion_value, width=width)
# p2 = plt.bar(ind, criterion_value-data_fit_value, bottom=data_fit_value, width = width)

plt.ylabel('BIC criterion')
plt.title('Model selection using BIC criterion: model fit vs. Ockam razor')
plt.xticks(ind, sorted_names)

plt.show()

print('Shape parameter of the gamma distribution: {}'.format(param[sorted_names.index('gamma')][0]))
print('DoF of the chisquare distribution: {}'.format(param[sorted_names.index('chi-square')][0]))


#%% md
#
# Note that here both the chisquare and gamma are capable of explaining the data, with :math:`a = \nu/2`, :math:`a` is
# gamma's shape parameter and :math:`\nu` is the number of DOFs in chi-square distribution.


# docs/code/inference/info_model_selection/plot_selection_regression_model.py

"""

Selection between regression models
==============================================

Here candidate models are defined as

.. math:: y=f(θ) + \epsilon

where :math:`f` consists in running RunModel. The three models considered are:

.. math:: f(θ)=θ_{0} x

.. math:: f(θ)=θ_{0} x + θ_{1} x^{2}

.. math:: f(θ)=θ_{0} x + θ_{1} x^{2} + θ_{2} x^{3}


"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

from UQpy.inference import DistributionModel, InformationModelSelection, MLE
from UQpy.RunModel import RunModel
import numpy as np
from UQpy.inference import BIC, AIC, AICc
import matplotlib.pyplot as plt
from UQpy.distributions import Gamma, Exponential, ChiSquare, Normal
from UQpy.inference import ComputationalModel


#%% md
#
# First we generate synthetic data using the quadratic model, and add some noise to it.

#%%

param_true = np.array([1.0, 2.0]).reshape((1, -1))
print('Shape of true parameter vector: {}'.format(param_true.shape))

h_func = RunModel(model_script='local_pfn_models.py', model_object_name='model_quadratic', vec=False,
                  var_names=['theta_0', 'theta_1'])
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=50).reshape((50,))
data_1 = data_clean + noise
print('Shape of data: {}'.format(data_1.shape))

shutil.rmtree(h_func.model_dir)

#%% md
#
# Create instances of the Model class for three models: linear, quadratic and cubic

#%%

names = ['linear', 'quadratic', 'cubic']
estimators = []
for i in range(3):
    h_func = RunModel(model_script='local_pfn_models.py', model_object_name='model_' + names[i], vec=False,
                      var_names=['theta_{}'.format(j) for j in range(i + 1)])
    M = ComputationalModel(runmodel_object=h_func, n_parameters=i + 1,
                           name=names[i], error_covariance=error_covariance)
    estimators.append(MLE(inference_model=M, data=data_1))

#%% md
#
# Apart from the data, candidate models and method (BIC, AIC...), InfoModelSelection also takes as inputs lists of
# inputs to the maximum likelihood class (iter_optim, method_optim, ...). Those inputs should be lists of length
# len(candidate_models).

#%%

from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer
optimizer = MinimizeOptimizer(method='nelder-mead')
selector = InformationModelSelection(parameter_estimators=estimators, criterion=BIC(), n_optimizations=[1]*3)
selector.sort_models()
print('Sorted models: ', [m.name for m in selector.candidate_models])
print('Values of criterion: ', selector.criterion_values)
print('Values of data fit:', [cr - pe for (cr, pe) in zip(selector.criterion_values, selector.penalty_terms)])
print('Values of penalty term (complexity):', selector.penalty_terms)
print('Values of model probabilities:', selector.probabilities)

#%% md
#
# Plot the results

#%%

domain = np.linspace(0, 10, 50)
fig, ax = plt.subplots(figsize=(8, 6))

for i, (model, estim) in enumerate(zip(selector.candidate_models, selector.parameter_estimators)):
    model.runmodel_object.run(samples=estim.mle.reshape((1, -1)), append_samples=False)
    y = model.runmodel_object.qoi_list[-1].reshape((-1,))
    ax.plot(domain, y, label=selector.candidate_models[i].name)
    shutil.rmtree(model.runmodel_object.model_dir)

plt.plot(domain, data_1, linestyle='none', marker='.', label='data')
plt.xlabel('x')
plt.ylabel('y')

plt.legend()
plt.show()



#%% md
#
# For this case, one can observe that both the quadratic and cubic model are capable of explaining the data. The cubic
# model is penalized due to its higher complexity (penalty_term) and thus the quadratic model is preferred.


# docs/code/inference/mle/local_pfn_models.py

"""

Auxiliary file
==============================================

"""
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

# docs/code/inference/mle/plot_complex_probability_model.py

"""

Complex probability distribution model
==============================================

Here we define a bivariate probability model, with a dependence structure defined using a gumbel copula. The goal of
inference is to learn the paremeters of the Gaussian marginals and the copula parameter, i.e., the model has 5 unknown
parameters.


"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

import matplotlib.pyplot as plt
from UQpy.inference import DistributionModel, MLE
from UQpy.distributions import Normal
from UQpy.inference import MinimizeOptimizer
from UQpy.distributions import JointIndependent, JointCopula, Gumbel
from UQpy.sampling import ImportanceSampling

#%% md
#
# First data is generated from a true model. A distribution with copulas does not possess a fit method, thus sampling is
# performed using importance sampling/resampling.

#%%

# dist_true exhibits dependence between the two dimensions, defined using a gumbel copula
dist_true = JointCopula(marginals=[Normal(), Normal()], copula=Gumbel(theta=2.))

# generate data using importance sampling: sample from a bivariate gaussian without copula, then weight samples
u = ImportanceSampling(proposal = JointIndependent(marginals=[Normal(), Normal()]),
                       log_pdf_target = dist_true.log_pdf,
                       nsamples=500)
print(u.samples.shape)
print(u.weights.shape)
# Resample to obtain 5,000 data points
u.resample(nsamples=5000)
data_2 = u.unweighted_samples
print('Shape of data: {}'.format(data_2.shape))

fig, ax = plt.subplots()
ax.scatter(data_2[:, 0], data_2[:, 1], alpha=0.2)
ax.set_title('Data points from true bivariate normal with gumbel dependency structure')
plt.show()

#%% md
#
# To define a model for inference, the user must create a custom file, here bivariate_normal_gumbel.py, to compute the
# log_pdf of the distribution, given a bivariate data matrix and a parameter vector of length 5. Note that for any
# probability model that is not one of the simple univariate pdfs supported by UQpy, such a custom file will be
# necessary.

#%%

d_guess = JointCopula(marginals=[Normal(loc=None, scale=None), Normal(loc=None, scale=None)],
                      copula=Gumbel(theta=None))
print(d_guess.get_parameters())
candidate_model = DistributionModel(n_parameters=5, distributions=d_guess)
print(candidate_model.list_params)

#%% md
#
# When calling MLEstimation, the function minimize from the scipy.optimize package is used by default. The user can
# define bounds for the optimization, a seed, the algorithm to be used, and set the algorithm to perform several
# optimization iterations, starting at a different random seed every time.

#%%

optimizer = MinimizeOptimizer(bounds=[[-5, 5], [0, 10], [-5, 5], [0, 10], [1.1, 4]], method="SLSQP")
ml_estimator = MLE(inference_model=candidate_model, data=data_2, optimizer=optimizer)

ml_estimator = MLE(inference_model=candidate_model, data=data_2, optimizer=optimizer,
                   initial_parameters=[1., 1., 1., 1., 4.])

print('ML estimates of the mean={0:.3f} and std. dev={1:.3f} of 1st marginal (true: 0.0, 1.0)'.
      format(ml_estimator.mle[0], ml_estimator.mle[1]))
print('ML estimates of the mean={0:.3f} and std. dev={1:.3f} of 2nd marginal (true: 0.0, 1.0)'.
      format(ml_estimator.mle[2], ml_estimator.mle[3]))
print('ML estimates of the copula parameter={0:.3f} (true: 2.0)'.format(ml_estimator.mle[4]))

#%% md
#
# Again, some known parameters can be fixed during learning.

#%%

d_guess = JointCopula(marginals=[Normal(loc=None, scale=None), Normal(loc=0., scale=1.)],
                      copula=Gumbel(theta=None))
candidate_model = DistributionModel(n_parameters=3, distributions=d_guess)

optimizer = MinimizeOptimizer(bounds=[[-5, 5], [0, 10], [1.1, 4]],
                              method="SLSQP")
ml_estimator = MLE(inference_model=candidate_model, data=data_2, optimizer=optimizer,
                   initial_parameters=[1., 1., 4.])

print('ML estimates of the mean={0:.3f} and std. dev={1:.3f} of 1st marginal (true: 0.0, 1.0)'.
      format(ml_estimator.mle[0], ml_estimator.mle[1]))
print('ML estimates of the copula parameter={0:.3f} (true: 2.0)'.format(ml_estimator.mle[2]))

# docs/code/inference/mle/plot_learn_distribution_model.py

"""

Simple probability distribution model
==============================================

In the following we learn the mean and covariance of a univariate gaussian distribution from data.
"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

import numpy as np
import matplotlib.pyplot as plt
from UQpy.inference import DistributionModel, MLE
from UQpy.distributions import Normal
from UQpy.inference import MinimizeOptimizer

#%% md
#
# First, for the sake of this example, we generate fake data from a gaussian distribution with mean 0 and
# standard deviation 1.

#%%

mu, sigma = 0, 0.1  # true mean and standard deviation
data_1 = np.random.normal(mu, sigma, 1000).reshape((-1, 1))
print('Shape of data vector: {}'.format(data_1.shape))

count, bins, ignored = plt.hist(data_1, 30, density=True)
plt.plot(bins, 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(- (bins - mu) ** 2 / (2 * sigma ** 2)),
         linewidth=2, color='r')
plt.title('Histogram of the data')
plt.show()

#%% md
#
# Create an instance of the class Model. The user must define the number of parameters to be estimated, in this case 2
# (mean and standard deviation), and set those parameters to be learnt as None when instantiating the Distribution
# object. For maximum likelihood estimation, no prior pdf is required.

#%%

# set parameters to be learnt as None
dist = Normal(loc=None, scale=None)
candidate_model = DistributionModel(n_parameters=2, distributions=dist)

ml_estimator = MLE(inference_model=candidate_model, data=data_1, n_optimizations=3)
print('ML estimates of the mean={0:.3f} (true=0.) and std. dev={1:.3f} (true=0.1)'.format(
    ml_estimator.mle[0], ml_estimator.mle[1]))

#%% md
#
# We can also fix one of the parameters and learn the remaining one

#%%

d = Normal(loc=0., scale=None)
candidate_model = DistributionModel(n_parameters=1, distributions=d)

optimizer = MinimizeOptimizer(bounds=[[0.0001, 2.]])
ml_estimator = MLE(inference_model=candidate_model, data=data_1,
                   n_optimizations=1)
print('ML estimates of the std. dev={0:.3f} (true=0.1)'.format(ml_estimator.mle[0]))

# docs/code/inference/mle/plot_regression_model.py

"""

Regression model
==============================================

Here a model is defined that is of the form

.. math:: y=f(\theta) + \epsilon

where f consists in running RunModel. In particular, here $f(\theta)=\theta_{0} x + \theta_{1} x^{2}$ is a regression model.


"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

import numpy as np
from UQpy.inference import ComputationalModel, MLE
from UQpy.distributions import Normal
from UQpy.inference import MinimizeOptimizer
from UQpy.RunModel import RunModel

#%% md
#
# First we generate synthetic data, and add some noise to it.

#%%

# Generate data

param_true = np.array([1.0, 2.0]).reshape((1, -1))
print('Shape of true parameter vector: {}'.format(param_true.shape))

h_func = RunModel(model_script='local_pfn_models.py', model_object_name='model_quadratic', vec=False,
                  var_names=['theta_0', 'theta_1'])
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=50).reshape((50,))
data_3 = data_clean + noise
print('Shape of data: {}'.format(data_3.shape))


#%% md
#
# Then we create an instance of the Model class, using model_type='python', and we perform maximum likelihood estimation
# of the two parameters.

#%%

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)
print('fitted parameters: theta_0={0:.3f} (true=1.), and theta_1={1:.3f} (true=2.)'.format(
    ml_estimator.mle[0], ml_estimator.mle[1]))

shutil.rmtree(h_func.model_dir)

# docs/code/reliability/subset_simulation/local_Resonance_pfn.py

"""

Resonance Auxiliary File
======================================================================


"""
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)

# docs/code/reliability/subset_simulation/local_Rosenbrock.py

"""

Rosenbrock Distribution Auxiliary File
======================================================================

"""
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']


# docs/code/reliability/subset_simulation/local_Rosenbrock_pfn.py

"""

Rosenbrock Distribution Auxiliary File 2
======================================================================

"""
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]

# docs/code/reliability/subset_simulation/local_pfn.py

"""

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

"""
import numpy as np


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


# docs/code/reliability/subset_simulation/plot_subset_linear_performance_function.py

"""

Correlated Normal Random Variables with a Linear Performance Function
======================================================================

This example runs subset simulation for a linear performance function having a specified reliability index
:math:`β` and standard normal random variables with specified correlation using the affine invariate ensemble
"stretch" sampler and the conventional Modified Metropolis Hastings algorithm.

For more details, refer to:

- Shields, M.D. and Giovanis, D.G. and Sundar, V.S. "Subset Simulation for problems with strongly non-Gaussian, highly anisotropic, and degenerate distributions," Computers and Structures. (In Review)
"""

# %% md
#
# 1. Import the necessary libraries

# %%
import shutil

from UQpy.reliability import SubsetSimulation
import matplotlib.pyplot as plt
from UQpy.sampling import ModifiedMetropolisHastings, Stretch
import time
import numpy as np
from UQpy.distributions import MultivariateNormal
from UQpy.RunModel import RunModel
import scipy.stats as stats

# %% md
#
# 2. Define the reliability index, correlation, dimension, algorithm, and the number of trials to run.

# %%

# Specified Reliability Index
beta = 4

# Specified Correlation
rho = 0.5

# Dimension
dim = 2

# Specify the number of trials to run
ntrials = 1

# Define the correlation matrix
C = np.ones((dim, dim)) * rho
np.fill_diagonal(C, 1)
print(C)

# Print information related to the true probability of failure
e, v = np.linalg.eig(np.asarray(C))
print(e)
print(v)
beff = np.sqrt(np.max(e)) * beta
print(beff)
pf_true = stats.norm.cdf(-beta)
print(pf_true)

# %% md
#
# 3. Execute subset simulation with the MMH algorithm.

# %%

pf = np.zeros((ntrials, 1))
cov1 = np.zeros((ntrials, 1))
cov2 = np.zeros((ntrials, 1))
for i in range(ntrials):
    model = RunModel(model_script='local_pfn.py', model_object_name="RunPythonModel", ntasks=1, b_eff=beff, d=dim)
    dist = MultivariateNormal(mean=np.zeros((dim)), cov=C)
    x = dist.rvs(nsamples=1000, random_state=349857)
    sampling = ModifiedMetropolisHastings(dimension=dim, log_pdf_target=dist.log_pdf,
                                          n_chains=100, random_state=342985)

    x_ss = SubsetSimulation(sampling=sampling, runmodel_object=model, samples_init=x, conditional_probability=0.1,
                            nsamples_per_subset=1000)
    shutil.rmtree(model.model_dir)
    pf[i] = x_ss.failure_probability
    cov1[i] = x_ss.independent_chains_CoV
    cov2[i] = x_ss.dependent_chains_CoV

# %% md
#
# 4. Plot samples from each conditional level and print the subset simulation results.

# %%

for i in range(len(x_ss.performance_function_per_level)):
    plt.scatter(x_ss.samples[i][:, 0], x_ss.samples[i][:, 1], marker='o')
plt.show()

print('Mean Pf:', np.mean(pf))
pf[pf == 0] = 1e-100
print('Mean beta: ', -np.mean(stats.norm.ppf(pf)))
print('CoV: ', stats.variation(pf))
print('CoV beta: ', np.absolute(stats.variation(stats.norm.ppf(pf))))
print('CoV log10: ', np.absolute(stats.variation(np.log10(pf))))
print(np.mean(cov1))
print(np.mean(cov2))

# %% md
#
# 5. Plot histograms of the failure probabilities and the reliability indices from subset simulation.

# %%
plt.figure()
plt.hist(pf)
plt.show()

plt.figure()
beta = -stats.norm.ppf(pf)
plt.hist(beta)
plt.show()

# %% md
#
# 6. Execute subset simulation with the Stretch algorithm.

# %%

pf = np.zeros((ntrials, 1))
cov1 = np.zeros((ntrials, 1))
cov2 = np.zeros((ntrials, 1))
for i in range(ntrials):
    model = RunModel(model_script='local_pfn.py', model_object_name="RunPythonModel", ntasks=1, b_eff=beff, d=dim)
    dist = MultivariateNormal(mean=np.zeros((dim)), cov=C)
    x = dist.rvs(nsamples=1000, random_state=349857)
    sampling = Stretch(dimension=dim, log_pdf_target=dist.log_pdf, n_chains=100, random_state=342985)

    x_ss = SubsetSimulation(sampling=sampling, runmodel_object=model, samples_init=x, conditional_probability=0.1,
                            nsamples_per_subset=1000)
    shutil.rmtree(model.model_dir)
    pf[i] = x_ss.failure_probability
    cov1[i] = x_ss.independent_chains_CoV
    cov2[i] = x_ss.dependent_chains_CoV

# %% md
#
# 7. Plot samples from each conditional level and print the subset simulation results.

# %%

for i in range(len(x_ss.performance_function_per_level)):
    plt.scatter(x_ss.samples[i][:, 0], x_ss.samples[i][:, 1], marker='o')
plt.show()

print('Mean Pf:', np.mean(pf))
pf[pf == 0] = 1e-100
print('Mean beta: ', -np.mean(stats.norm.ppf(pf)))
print('CoV: ', stats.variation(pf))
print('CoV beta: ', np.absolute(stats.variation(stats.norm.ppf(pf))))
print('CoV log10: ', np.absolute(stats.variation(np.log10(pf))))
print(np.mean(cov1))
print(np.mean(cov2))

# %% md
#
# 8. Plot histograms of the failure probabilities and the reliability indices from subset simulation.

# %%
plt.figure()
plt.hist(pf)
plt.show()

plt.figure()
beta = -stats.norm.ppf(pf)
plt.hist(beta)
plt.show()


# docs/code/reliability/subset_simulation/plot_subset_resonance.py

"""

Subset Simulation for a SDOF resonance problem
======================================================================

Here, the Modified Metropolis Hastings Algorithm for MCMC is compared to the Affine Invariant Ensemble sampler for
MCMC in subset simulation to estimate probability of failure of a single degree of freedom system subjected to
harmonic excitation with a given frequency.

Problem Definition
-------------------
A stochastic single degree of freedom system having stiffness :math:`k\sim N(\mu_k,\sigma_k)` and mass
:math:`m\sim N(\mu_m,\sigma_m)` is excited by a sinusoidal load with frequency :math:`\omega \ rad/sec`.
The system is undamped and has equation of motion given by:

.. math:: `m\ddot{u}+ku=sin(\omega t)`

Resonance occurs when the natural frequency of the system :math:`\omega_n=\sqrt{\dfrac{k}{m}}=\omega \ rad/sec`.
To avoid resonance, we consider failure of the system to be associated with the natural frequency being within a
threshold, :math:`\epsilon` of the excitation frequency $\omega$. That is, failure of the system occurs when
:math:`\omega-\epsilon\le\sqrt{\dfrac{k}{m}}\le\omega+\epsilon`.

"""

# %% md
#
# 1. Import the necessary libraries

# %%
import shutil

from UQpy.reliability import SubsetSimulation
from UQpy.RunModel import RunModel
from UQpy.sampling import Stretch, ModifiedMetropolisHastings, MonteCarloSampling
import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt
import time
from local_Resonance_pfn import RunPythonModel
from UQpy.distributions import Normal, JointIndependent, MultivariateNormal

# %% md
#
# Set Parameters
#
# :math:`\omega=5`
#
# :math:`\epsilon = 0.0001`
#
# :math:`\mu_k=125`
#
# :math:`\sigma_k=20`
#
# :math:`\mu_m=5`
#
# :math:`\sigma_m=1`

# %%

omega = 6
epsilon = 0.0001
mu_m = 5
sigma_m = 1
mu_k = 125
sigma_k = 20
m = np.linspace(mu_m - 3 * sigma_m, mu_m + 3 * sigma_m, 101)
k_hi = (omega - epsilon) ** 2 * m
k_lo = (omega + epsilon) ** 2 * m

# %% md
#
# Plot the failure domain

# %%

x = np.linspace(2, 8, 1000)
y = np.linspace(25, 225, 1000)

X, Y = np.meshgrid(x, y)
Z = np.zeros((1000, 1000))

d1 = Normal(loc=5, scale=1)
d2 = Normal(loc=125, scale=20)

dist = JointIndependent(marginals=[d1, d2])

for i in range(len(x)):
    Z[i, :] = dist.pdf(np.append(np.atleast_2d(X[i, :]), np.atleast_2d(Y[i, :]), 0).T)

fig, ax = plt.subplots()
CS = ax.contour(X, Y, Z, 15)
plt.plot(m, k_hi, 'k')
plt.plot(m, k_lo, 'k')
# plt.fill_between(m,k_lo,k_hi)
plt.xlim([mu_m - 3 * sigma_m, mu_m + 3 * sigma_m])
plt.ylim([mu_k - 3 * sigma_k, mu_k + 3 * sigma_k])
plt.xlabel(r'Mass ($m$)')
plt.ylabel(r'Stiffness ($k$)')
plt.grid(True)
plt.tight_layout()
plt.show()

fig, ax = plt.subplots()
CS = ax.contour(X, Y, Z, 15)
plt.plot(m, k_hi, 'k')
plt.plot(m, k_lo, 'k')
# plt.fill_between(m,k_lo,k_hi)
plt.xlim([3.5, 4.5])
plt.ylim([130, 150])
plt.xlabel(r'Mass ($m$)')
plt.ylabel(r'Stiffness ($k$)')
plt.grid(True)
plt.tight_layout()
plt.show()

model = RunModel(model_script='local_Resonance_pfn.py', model_object_name="RunPythonModel", ntasks=1)

# %% md
#
# Monte Carlo Simulation

# %%

x_mcs = MonteCarloSampling(distributions=[d1, d2])
x_mcs.run(nsamples=1000000)

model.run(samples=x_mcs.samples)

shutil.rmtree(model.model_dir)
A = np.asarray(model.qoi_list) < 0
pf = np.shape(np.asarray(model.qoi_list)[np.asarray(model.qoi_list) < 0])[0] / 1000000
print(pf)

ntrials = 1
pf_stretch = np.zeros((ntrials, 1))
cov1_stretch = np.zeros((ntrials, 1))
cov2_stretch = np.zeros((ntrials, 1))
m = np.ones(2)
m[0] = 5
m[1] = 125
C = np.eye(2)
C[0, 0] = 1
C[1, 1] = 20 ** 2

for i in range(ntrials):
    model = RunModel(model_script='local_Resonance_pfn.py', model_object_name="RunPythonModel", ntasks=1)
    dist = MultivariateNormal(mean=m, cov=C)
    xx = dist.rvs(nsamples=1000, random_state=123)
    xx1 = dist.rvs(nsamples=100, random_state=123)

    sampling=Stretch(dimension=2, n_chains=100, log_pdf_target=dist.log_pdf)
    x_ss_stretch = SubsetSimulation(sampling=sampling, runmodel_object=model, conditional_probability=0.1,
                                    nsamples_per_subset=1000, samples_init=xx, )
    shutil.rmtree(model.model_dir)
    pf_stretch[i] = x_ss_stretch.failure_probability
    cov1_stretch[i] = x_ss_stretch.independent_chains_CoV
    cov2_stretch[i] = x_ss_stretch.dependent_chains_CoV

print(pf_stretch)
print(np.mean(pf_stretch, 0))
b_stretch = -stats.norm.ppf(pf_stretch)
print(np.mean(b_stretch, 0))
print(stats.variation(b_stretch))


for i in range(len(x_ss_stretch.performance_function_per_level)):
    plt.scatter(x_ss_stretch.samples[i][:, 0], x_ss_stretch.samples[i][:, 1], marker='o')

plt.xlim([mu_m - 3 * sigma_m, mu_m + 3 * sigma_m])
plt.ylim([mu_k - 3 * sigma_k, mu_k + 3 * sigma_k])
plt.xlabel(r'Mass ($m$)')
plt.ylabel(r'Stiffness ($k$)')
plt.grid(True)
plt.tight_layout()
plt.show()


ntrials = 100
pf_mmh = np.zeros((ntrials, 1))
cov1_mmh = np.zeros((ntrials, 1))
cov2_mmh = np.zeros((ntrials, 1))
m = np.ones(2)
m[0] = 5
m[1] = 125
C = np.eye(2)
C[0, 0] = 1
C[1, 1] = 20 ** 2

for i in range(ntrials):
    model = RunModel(model_script='local_Resonance_pfn.py', model_object_name="RunPythonModel", ntasks=1)
    dist = MultivariateNormal(mean=m, cov=C)
    xx = dist.rvs(nsamples=1000, random_state=123)
    xx1 = dist.rvs(nsamples=100, random_state=123)

    sampling = ModifiedMetropolisHastings(dimension=2, n_chains=100, log_pdf_target=dist.log_pdf)
    x_ss_mmh = SubsetSimulation(sampling=sampling, runmodel_object=model, conditional_probability=0.1,
                                nsamples_per_subset=1000, samples_init=xx)

    shutil.rmtree(model.model_dir)
    pf_mmh[i] = x_ss_mmh.failure_probability
    cov1_mmh[i] = x_ss_mmh.independent_chains_CoV
    cov2_mmh[i] = x_ss_mmh.dependent_chains_CoV

pf_mmh[pf_mmh == 0] = 1e-100
print(np.mean(pf_mmh, 0))
b_mmh = -stats.norm.ppf(pf_mmh)
print(np.mean(b_mmh, 0))
print(stats.variation(b_mmh))


for i in range(len(x_ss_mmh.performance_function_per_level)):
    plt.scatter(x_ss_mmh.samples[i][:, 0], x_ss_mmh.samples[i][:, 1], marker='o')

plt.xlim([mu_m - 3 * sigma_m, mu_m + 3 * sigma_m])
plt.ylim([mu_k - 3 * sigma_k, mu_k + 3 * sigma_k])
plt.xlabel(r'Mass ($m$)')
plt.ylabel(r'Stiffness ($k$)')
plt.grid(True)
plt.tight_layout()
plt.show()

# docs/code/reliability/subset_simulation/plot_subset_rosenbrock.py

"""

Rosenbrock performance function
======================================================================


"""

# %% md
#
# 1. Import the necessary libraries

# %%

import shutil

from UQpy.reliability import SubsetSimulation
import matplotlib.pyplot as plt
from UQpy.sampling import ModifiedMetropolisHastings, Stretch
# First import the file that contains the newly defined Rosenbrock distribution
from local_Rosenbrock import Rosenbrock
import time
import numpy as np
# Import this newly defined Rosenbrock distribution into the Distributions module
from UQpy.distributions import Normal
from UQpy.RunModel import RunModel
import scipy.stats as stats

# %% md
#
# MMH Initial Samples

# %%

model = RunModel(model_script='local_Rosenbrock_pfn.py', model_object_name="RunPythonModel", ntasks=1)
dist = Rosenbrock(p=100.)
dist_prop1 = Normal(loc=0, scale=1)
dist_prop2 = Normal(loc=0, scale=10)

x = stats.norm.rvs(loc=0, scale=1, size=(100, 2), random_state=83276)

mcmc_init1 = ModifiedMetropolisHastings(dimension=2, log_pdf_target=dist.log_pdf, seed=x.tolist(),
                                                           burn_length=1000, proposal=[dist_prop1, dist_prop2],
                                                           random_state=8765)
mcmc_init1.run(10000)

sampling=Stretch(log_pdf_target=dist.log_pdf, dimension=2, n_chains=1000, random_state=38546)
x_ss_MMH = SubsetSimulation(sampling=sampling, runmodel_object=model, conditional_probability=0.1,
                            nsamples_per_subset=10000, samples_init=mcmc_init1.samples)

for i in range(len(x_ss_MMH.performance_function_per_level)):
    plt.scatter(x_ss_MMH.samples[i][:, 0], x_ss_MMH.samples[i][:, 1], marker='o')
plt.grid(True)
plt.xlabel(r'$X_1$')
plt.ylabel(r'$X_2$')
plt.yticks(np.arange(-20, 180, step=20))
plt.xlim((-10, 15))
plt.tight_layout()
plt.show()

shutil.rmtree(model.model_dir)

print(x_ss_MMH.failure_probability)


# %% md
#
# Stretch Initial Samples

# %%

model = RunModel(model_script='local_Rosenbrock_pfn.py', model_object_name="RunPythonModel", ntasks=1)
dist = Rosenbrock(p=100.)

x = stats.norm.rvs(loc=0, scale=1, size=(100, 2), random_state=83276)

mcmc_init2 = Stretch(dimension=2, log_pdf_target=dist.log_pdf, seed=x.tolist(),
                                                burn_length=1000, random_state=8765)
mcmc_init2.run(10000)

sampling=Stretch(log_pdf_target=dist.log_pdf, dimension=2, n_chains=1000, random_state=83456)
x_ss_Stretch = SubsetSimulation(sampling=sampling, runmodel_object=model, conditional_probability=0.1,
                                nsamples_per_subset=10000, samples_init=mcmc_init2.samples)

for i in range(len(x_ss_Stretch.performance_function_per_level)):
    plt.scatter(x_ss_Stretch.samples[i][:, 0], x_ss_Stretch.samples[i][:, 1], marker='o')
plt.grid(True)
plt.xlabel(r'$X_1$')
plt.ylabel(r'$X_2$')
plt.yticks(np.arange(-20, 180, step=20))
plt.xlim((-10, 15))
plt.tight_layout()
plt.show()

print(x_ss_Stretch.failure_probability)


plt.figure()
plt.plot(mcmc_init2.samples[:, 0], mcmc_init2.samples[:, 1], 'o')
plt.plot(mcmc_init1.samples[:, 0], mcmc_init1.samples[:, 1], 'x')
plt.grid(True)
plt.xlabel(r'$X_1$')
plt.ylabel(r'$X_2$')
plt.yticks(np.arange(-20, 180, step=20))
plt.xlim((-10, 15))
plt.tight_layout()
plt.show()
shutil.rmtree(model.model_dir)

# docs/code/reliability/taylor/local_pfn.py

"""

Auxiliary file
==============================================

"""
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

# docs/code/reliability/taylor/plot_FORM_linear function_2d.py

"""

FORM - Linear function - Two Random variables
==============================================

In the second example we have the limit state to be a linear function of two (:math:`d=2`) independent Gaussian random
variables
"""

#%% md
#
# :math:`g(U) = -\frac{1}{\sqrt{d}}\sum_{i=1}^{d} u_i + \beta`
#
# The probability of failure in this case is :math:`P(F) ≈ 10^{−3}` for :math:`β = 3.0902`
#
# Initially we have to import the necessary modules.

#%%

import shutil

import numpy as np
import matplotlib.pyplot as plt
from UQpy.RunModel import RunModel
from UQpy.distributions import Normal
from UQpy.sampling import MonteCarloSampling
from UQpy.reliability import FORM, SORM
from UQpy.reliability import SORM
from UQpy.distributions import Lognormal, Gamma


dist1 = Normal(loc=0., scale=1.)
dist2 = Normal(loc=0., scale=1.)
RunModelObject2 = RunModel(model_script='local_pfn.py', model_object_name="example2", vec=False, ntasks=3)
Z = FORM(distributions=[dist1, dist2], runmodel_object=RunModelObject2)
Z.run()

# print results
print('Design point in standard normal space: %s' % Z.DesignPoint_U)
print('Design point in original space: %s' % Z.DesignPoint_X)
print('Hasofer-Lind reliability index: %s' % Z.beta)
print('FORM probability of failure: %s' % Z.failure_probability)

shutil.rmtree(RunModelObject2.model_dir)


# docs/code/reliability/taylor/plot_FORM_linear_function_3d.py

"""

FORM - Linear function - Three Random variables
==============================================

In the third example we have the limit state to be a linear function of three (:math:`d`=3) independent Gaussian random
variables

.. math:: g(X_1, X_2, X_3) = 6.2X_1 -X_2X_3^2

.. math:: \mu_{X_1}=20, \mu_{X_2}=5, \mu_{X_3}=4

.. math:: \sigma_{X_1}=3.5, \sigma_{X_2}=0.8, \sigma_{X_3}=0.4

The probability of failure in this case is :math:`P_f ≈ 0.079` for :math:`\beta = 1.413`
"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

import shutil

import numpy as np
import matplotlib.pyplot as plt
from UQpy.RunModel import RunModel
from UQpy.distributions import Normal
from UQpy.sampling import MonteCarloSampling
from UQpy.reliability import FORM, SORM
from UQpy.reliability import SORM
from UQpy.distributions import Lognormal, Gamma


dist1 = Normal(loc=20., scale=3.5)
dist2 = Normal(loc=5., scale=0.8)
dist3 = Normal(loc=4., scale=0.4)
RunModelObject3 = RunModel(model_script='local_pfn.py', model_object_name="example3", vec=False, ntasks=3)
Z0 = FORM(distributions=[dist1, dist2, dist3], runmodel_object=RunModelObject3)
Z0.run()

print('Design point in standard normal space: %s' % Z0.DesignPoint_U)
print('Design point in original space: %s' % Z0.DesignPoint_X)
print('Hasofer-Lind reliability index: %s' % Z0.beta)
print('FORM probability of failure: %s' % Z0.failure_probability)

shutil.rmtree(RunModelObject3.model_dir)

# docs/code/reliability/taylor/plot_FORM_nonlinear_function.py

"""

FORM - Nonlinear function - Two Random variables
==============================================

In the fourth example we have the limit state to be a nonlinear function of two (:math:`d=2`) random variables

.. math:: g(X_1, X_2) = X_1X_1 - 80

where :math:`X_1` follows a normal distribution with mean :math:`\mu_{X_1}=20` and standard deviation
:math:`\sigma_{X_1}=7` and :math:`X_2` follows a lognormal distribution with mean :math:`\mu_{X_2}=7` and
standard deviation :math:`\sigma_{X_2}=1.4`.
"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

import shutil

import numpy as np
from UQpy.RunModel import RunModel
from UQpy.distributions import Normal
from UQpy.reliability import FORM
from UQpy.reliability import SORM
from UQpy.distributions import Lognormal

m0 = 7
v0 = 1.4
mu = np.log(m0) - np.log(np.sqrt(1 + (v0 / m0) ** 2))
scale = np.exp(mu)
s = np.sqrt(np.log(1 + (v0 / m0) ** 2))
loc_ = 0.0

dist1 = Normal(loc=20., scale=2)
dist2 = Lognormal(s=s, loc=0.0, scale=scale)
RunModelObject4 = RunModel(model_script='local_pfn.py', model_object_name="example4", vec=False, ntasks=3)
form = FORM(distributions=[dist1, dist2], runmodel_object=RunModelObject4)
form.run()
Q0 = SORM(form_object=form)


# print results
print('Design point: %s' % Q0.DesignPoint_X)
print('Design point: %s' % Q0.DesignPoint_U)
print('beta: %s' % Q0.beta)
print('FORM probability of failure: %s' % Q0.Pf_form)
print('SORM probability of failure: %s' % Q0.failure_probability)

shutil.rmtree(RunModelObject4.model_dir)

# docs/code/reliability/taylor/plot_FORM_structural_reliability.py

"""

FORM - Structural Reliability
==============================================

The benchmark problem is a simple structural reliability problem
defined in a two-dimensional parameter space consisting of a resistance :math:`R` and a stress :math:`S`. The failure
happens when the stress is higher than the resistance, leading to the following limit-state function:

.. math:: \textbf{X}=\{R, S\}

.. math:: g(\textbf{X}) = R - S

The two random variables are independent  and  distributed
according to:

.. math:: R \sim N(200, 20)

.. math:: S \sim N(150, 10)
"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

import numpy as np
import matplotlib.pyplot as plt
from UQpy.RunModel import RunModel
from UQpy.distributions import Normal
from UQpy.sampling import MonteCarloSampling
from UQpy.reliability import FORM, SORM
from UQpy.reliability import SORM
from UQpy.distributions import Lognormal, Gamma

RunModelObject = RunModel(model_script='local_pfn.py', model_object_name="example1", vec=False, ntasks=3)
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
print('Design point in standard normal space: %s' % Q.DesignPoint_U)
print('Design point in original space: %s' % Q.DesignPoint_X)
print('Hasofer-Lind reliability index: %s' % Q.beta)
print('FORM probability of failure: %s' % Q.failure_probability)
print(Q.dg_u_record)


# Supporting function
def multivariate_gaussian(pos, mu, Sigma):
    n = mu.shape[0]
    Sigma_det = np.linalg.det(Sigma)
    Sigma_inv = np.linalg.inv(Sigma)
    N = np.sqrt((2 * np.pi) ** n * Sigma_det)
    fac = np.einsum('...k,kl,...l->...', pos - mu, Sigma_inv, pos - mu)
    return np.exp(-fac / 2) / N

N = 60
XX = np.linspace(150, 250, N)
YX = np.linspace(120, 180, N)
XX, YX = np.meshgrid(XX, YX)

XU = np.linspace(-3, 3, N)
YU = np.linspace(-3, 3, N)
XU, YU = np.meshgrid(XU, YU)

# Mean vector and covariance matrix in the original space
parameters = [[200, 20], [150, 10]]
mu_X = np.array([parameters[0][0], parameters[1][0]])
Sigma_X = np.array([[parameters[0][1] ** 2, 0.0], [0.0, parameters[1][1] ** 2]])

# Mean vector and covariance matrix in the standard normal space
mu_U = np.array([0., 0.])
Sigma_U = np.array([[1., 0.0], [0.0, 1]])

# Pack X and Y into a single 3-dimensional array for the original space
posX = np.empty(XX.shape + (2,))
posX[:, :, 0] = XX
posX[:, :, 1] = YX
ZX = multivariate_gaussian(posX, mu_X, Sigma_X)

# Pack X and Y into a single 3-dimensional array for the standard normal space
posU = np.empty(XU.shape + (2,))
posU[:, :, 0] = XU
posU[:, :, 1] = YU
ZU = multivariate_gaussian(posU, mu_U, Sigma_U)

# Figure 4a
plt.figure()
plt.rcParams["figure.figsize"] = (12, 12)
plt.rcParams.update({'font.size': 22})
plt.plot(parameters[0][0], parameters[1][0], 'r.')
plt.plot([0, 200], [0, 200], 'k', Linewidth=5)
plt.plot(Q.DesignPoint_X[0][0], Q.DesignPoint_X[0][1], 'bo', markersize=12)
plt.contour(XX, YX, ZX, levels=20)
plt.xlabel(r'$X_1$')
plt.ylabel(r'$X_2$')
plt.text(170, 182, '$X_1 - X_2=0$',
         rotation=45,
         horizontalalignment='center',
         verticalalignment='top',
         multialignment='center')
plt.ylim([120, 200])
plt.xlim([130, 240])
plt.grid()
plt.title('Original space')
plt.axes().set_aspect('equal', 'box')
plt.show()

# Figure 4b
plt.figure()
plt.rcParams["figure.figsize"] = (12, 12)
plt.rcParams.update({'font.size': 22})
plt.plot([0, Q.DesignPoint_U[0][0]], [0, Q.DesignPoint_U[0][1]], 'b', Linewidth=2)
plt.plot([0, -3], [5, -1], 'k', Linewidth=5)
plt.plot(Q.DesignPoint_U[0][0], Q.DesignPoint_U[0][1], 'bo', markersize=12)
plt.contour(XU, YU, ZU, levels=20)
plt.axhline(0, color='black')
plt.axvline(0, color='black')
plt.plot(0, 0, 'r.')

plt.xlabel(r'$U_1$')
plt.ylabel(r'$U_2$')
plt.text(-1.0, 1.1, '$U^\star$=({:1.2f}, {:1.2f})'.format(-2.0, 1.0),
         rotation=0,
         horizontalalignment='center',
         verticalalignment='top',
         multialignment='center')

plt.text(-2.1, 2.05, '$20U_1 - 10U_2 + 50=0$',
         rotation=63,
         horizontalalignment='center',
         verticalalignment='top',
         multialignment='center')

plt.text(-1.5, 0.7, r'$\overrightarrow{\beta}$',
         rotation=0,
         horizontalalignment='center',
         verticalalignment='top',
         multialignment='center')

plt.text(0.02, -0.2, '({:1.1f}, {:1.1f})'.format(0.0, 0.0))
plt.ylim([-1, 3])
plt.xlim([-3.5, 2])
plt.grid()
plt.title('Standard Normal space')
plt.axes().set_aspect('equal', 'box')
plt.show()

shutil.rmtree(RunModelObject.model_dir)

# docs/code/sampling/adaptive_kriging/local_BraninHoo.py

"""

Branin-Hoo File
==================================


"""
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

# docs/code/sampling/adaptive_kriging/plot_adaptive_kriging_branin_hoo.py

"""

Expected Improvement - Branin Hoo
============================================

In this example, Monte Carlo Sampling is used to generate samples from Uniform distribution and new samples are
generated adaptively, using EIF (Expected Improvement Function) as the learning criteria.
"""

# %% md
#
# Branin-Hoo function
# --------------------
# Decription:
#
# >  - Dimensions: 2
# >  - This function is usually evaluated on the square $x_1 \in [-5, 10], \ x_2 \in [0, 15]$
# >  - The function has two local minima and one global minimum
# >  - Reference: Forrester, A., Sobester, A., & Keane, A. (2008). Engineering design via surrogate modelling: a practical guide. Wiley.
#
# > $\displaystyle f(x) = a(x_2-bx_1^2 + cx_1 -r)^2 + s(1-t)\cos(x_1) + s + 5x_1$
# > <br>
# > <br>
# > where the recommended values of a, b, c, r, s and t are: $a = 1,\ b = 5.1/(4\pi^2),\ c = 5/\pi, \ r = 6, \ s = 10, \ t = 1/(8\pi)$
#
# <img src="branin.png" alt="branin.png" height="350" width="400" align=left>

# %%

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy, matplotlib and other necessary
# library for plots, but also need to import the MCS, AKMCS, Kriging and RunModel class from UQpy.

# %%
import shutil

import numpy as np
from matplotlib import pyplot as plt

from UQpy.surrogates import Kriging
from UQpy.sampling import MonteCarloSampling, AdaptiveKriging
from UQpy.RunModel import RunModel
from UQpy.distributions import Uniform
from local_BraninHoo import function
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import time
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

# %% md
#
# Using UQpy MCS class to generate samples for two random variables, which are uniformly distributed

# %%

marginals = [Uniform(loc=-5, scale=15), Uniform(loc=0, scale=15)]
x = MonteCarloSampling(distributions=marginals, nsamples=20)

# %% md
#
# RunModel class is used to define an object to evaluate the model at sample points.

# %%

rmodel = RunModel(model_script='local_BraninHoo.py', vec=False)

# %% md
#
# Krig class defines an object to generate an surrogate model for a given set of data.

# %%

from UQpy.surrogates.kriging.regression_models import Linear
from UQpy.surrogates.kriging.correlation_models import Exponential
optimizer = MinimizeOptimizer(method="L-BFGS-B")
K = Kriging(regression_model=Linear(), correlation_model=Exponential(),optimizer=optimizer,
            correlation_model_parameters=[1, 1], optimizations_number=10)

# %% md
#
# Choose an appropriate learning function.

# %%

from UQpy.sampling.adaptive_kriging_functions.ExpectedImprovement import ExpectedImprovement

# %% md
#
# AKMCS class is used to generate new sample using 'U-function' as active learning function.

# %%

start_time = time.time()
learning_function= ExpectedImprovement()
a = AdaptiveKriging(runmodel_object=rmodel, samples=x.samples, surrogate=K,
                    learning_nsamples=10 ** 3, n_add=1,
                    learning_function=learning_function, distributions=marginals)
a.run(nsamples=50)
elapsed_time = time.time() - start_time

# %% md
#
# Visualize initial and new samples on top of the Branin-Hoo surface.

# %%

num = 200
xlist = np.linspace(-6, 11, num)
ylist = np.linspace(-1, 16, num)
X, Y = np.meshgrid(xlist, ylist)

Z = np.zeros((num, num))
for i in range(num):
    for j in range(num):
        tem = np.array([[X[i, j], Y[i, j]]])
        Z[i, j] = function(tem)

shutil.rmtree(rmodel.model_dir)

fig, ax = plt.subplots(1, 1)
cp = ax.contourf(X, Y, Z, 10)
plt.xlabel('x1')
plt.ylabel('x2')
fig.colorbar(cp)
nd = x.nsamples
plt.scatter(a.samples[nd:, 0], a.samples[nd:, 1], color='pink', label='New samples')
plt.scatter(x.samples[:nd, 0], x.samples[:nd, 1], color='Red', label='Initial samples')
plt.title('Branin-Hoo function');
plt.legend()






# docs/code/sampling/adaptive_kriging/plot_adaptive_kriging_normal.py

"""

U-Function & User-defined learning function
============================================

In this example, Monte Carlo Sampling is used to generate samples from Normal distribution and new samples are generated
adaptively, using U-function as the learning criteria .
"""

# %% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy, matplotlib and other necessary
# library for plots, but also need to import the MCS, AKMCS, Kriging and RunModel class from UQpy.

# %%
import shutil

from UQpy.surrogates import Kriging
from UQpy.sampling import MonteCarloSampling, AdaptiveKriging
from UQpy.RunModel import RunModel
from UQpy.distributions import Normal
from local_series import series
import numpy as np
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.pyplot as plt
import time
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer


# %% md
#
# Using UQpy MCS class to generate samples for two random variables, which are normally distributed with mean 0 and
# variance 1.

# %%

marginals = [Normal(loc=0., scale=4.), Normal(loc=0., scale=4.)]
x = MonteCarloSampling(distributions=marginals, nsamples=20, random_state=1)

# %% md
#
# RunModel class is used to define an object to evaluate the model at sample points.

# %%

rmodel = RunModel(model_script='local_series.py', vec=False)


# %% md
#
# Krig class defines an object to generate an surrogate model for a given set of data.

# %%

from UQpy.surrogates.kriging.regression_models import Linear
from UQpy.surrogates.kriging.correlation_models import Exponential
optimizer = MinimizeOptimizer(method="L-BFGS-B")
K = Kriging(regression_model=Linear(), correlation_model=Exponential(), optimizer=optimizer,
            correlation_model_parameters=[1, 1], optimizations_number=10)

# %% md
#
# This example works for all three learning function based on reliability analysis.
#
# AKMCS class is used to generate new sample using 'U-function' as active learning function.

# %%

from UQpy.sampling.adaptive_kriging_functions import *
start_time = time.time()
learning_function = WeightedUFunction(weighted_u_stop=2)
a = AdaptiveKriging(runmodel_object=rmodel, surrogate=K, learning_nsamples=10 ** 3, n_add=1,
                    learning_function=learning_function, distributions=marginals, random_state=2)
a.run(nsamples=100, samples=x.samples)

elapsed_time = time.time() - start_time


time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
g = a.surrogate.predict(a.learning_set, False)
n_ = a.learning_set.shape[0] + len(a.qoi)
pf = (sum(g < 0) + sum(np.array(a.qoi) < 0)) / n_
print('Time: ', elapsed_time)
print('Function evaluation: ', a.samples.shape[0])
print('Probability of failure: ', pf)

# %% md
#
# This figure shows the location of new samples generated using active learning function.

# %%

num = 50
x1 = np.linspace(-7, 7, num)
x2 = np.linspace(-7, 7, num)
x1v, x2v = np.meshgrid(x1, x2)
y = np.zeros([num, num])
y_act = np.zeros([num, num])
mse = np.zeros([num, num])
for i in range(num):
    for j in range(num):
        xa = marginals[0].cdf(np.atleast_2d(x1v[i, j]))
        ya = marginals[1].cdf(np.atleast_2d(x2v[i, j]))
        y[i, j] = a.surrogate.predict(np.hstack([xa, ya]))
        y_act[i, j] = series(np.array([[x1v[i, j], x2v[i, j]]]))

fig, ax = plt.subplots()
kr_a = ax.contour(x1v, x2v, y_act, levels=[0], colors='Black')

# Plot for scattered data
nd = x.nsamples
ID1 = ax.scatter(a.samples[nd:, 0], a.samples[nd:, 1], color='Grey', label='New samples')
ID = ax.scatter(x.samples[:nd, 0], x.samples[:nd, 1], color='Red', label='Initial samples')
plt.legend(handles=[ID1, ID])
plt.show()

# %% md
#
# User-defined Learning function
# ------------------------------

# %%

class UserLearningFunction(LearningFunction):

    def __init__(self, u_stop: int = 2):
        self.u_stop = u_stop

    def evaluate_function(self, distributions, n_add, surrogate, population, qoi=None, samples=None):
        # AKMS class use these inputs to compute the learning function

        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]

        indicator = False
        if min(u[:, 0]) >= self.u_stop:
            indicator = True

        return population[rows, :], u[rows, 0], indicator

# %% md
#
# Creating new instances of Kriging and RunModel class.

# %%

K1 = Kriging(regression_model=Linear(), correlation_model=Exponential(), optimizer=optimizer,
             correlation_model_parameters=[1, 1], optimizations_number=10)
rmodel1 = RunModel(model_script='local_series.py', vec=False)

# %% md
#
# Executing AKMCS with the user-defined learning function.

# %%

start_time = time.time()
ak = AdaptiveKriging(runmodel_object=rmodel1, samples=x.samples, surrogate=K1, learning_nsamples=10 ** 3,
                     n_add=1, learning_function=UserLearningFunction(), distributions=marginals, random_state=3)
ak.run(nsamples=100)


elapsed_time = time.time() - start_time

time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
g = ak.surrogate.predict(ak.learning_set, False)
n_ = ak.learning_set.shape[0] + len(ak.qoi)
pf = (sum(g < 0) + sum(np.array(ak.qoi) < 0)) / n_
print('Time: ', elapsed_time)
print('Function evaluation: ', ak.samples.shape[0])
print('Probability of failure: ', pf)

# %% md
#
# This figure shows the location of new samples generated using active learning function.

# %%

fig1, ax1 = plt.subplots()
kr_a = ax1.contour(x1v, x2v, y_act, levels=[0], colors='Black')

# Plot for scattered data
ID1 = ax1.scatter(ak.samples[nd:, 0], ak.samples[nd:, 1], color='Grey', label='New samples')
ID = ax1.scatter(x.samples[:nd, 0], x.samples[:nd, 1], color='Red', label='Initial samples')
plt.legend(handles=[ID1, ID])
plt.show()

# %% md
#
# Monte Carlo Simulation
# -----------------------
# Probability of failure and covariance is estimated using Monte Carlo Simulation. 10,000 samples are generated
# randomly using MCS class and model is evaluated at all samples.

# %%

start_time = time.time()

# Code
b = MonteCarloSampling(distributions=marginals, nsamples=10 ** 4, random_state=4)
r1model = RunModel(model_script='local_series.py', vec=False)
r1model.run(samples=b.samples)


gx = np.array(r1model.qoi_list)
pf_mcs = np.sum(np.array(gx) < 0) / b.nsamples
cov_pf_mcs = np.sqrt((1 - pf_mcs) / (pf_mcs * b.nsamples))
elapsed_time = time.time() - start_time
time.strftime("%H:%M:%S", time.gmtime(elapsed_time))

shutil.rmtree(r1model.model_dir)
shutil.rmtree(rmodel1.model_dir)
shutil.rmtree(rmodel.model_dir)

# %% md
#
# Results from Monte Carlo Simulation.

# %%

print('Time: ', elapsed_time)
print('Function evaluation: ', b.nsamples)
print('Probability of failure: ', pf_mcs)



# docs/code/sampling/importance_sampling/plot_importance_sampling_diagnostics.py

"""

Diagnostics for Importance Sampling
============================================

"""

#%%
#
# This notebook illustrates the use of some simple diagnostics for Importance Sampling. For IS, in extreme settings
# only a few samples may have a significant weight, yielding poor approximations of the target pdf :math:`p(x)`. A popular
# diagnostics is the Effective Sample Size (ESS), which is theoretically defined as the number of independent samples
# generated directly form the target distribution that are required to obtain an estimator with same variance as the
# one obtained from IS. Heuristically, ESS approximates how many i.i.d. samples, drawn from the target, are equivalent
# to :math:`N` weighted samples drawn from the IS or MCMC approximation.
#
# An approximation of the ESS is given by [1]:
#
# .. math::  ESS = \frac{1}{\sum \tilde{w}^2}
#
# where :math:`\tilde{w}` are the normalized weights.
#
# [1] *Sequential Monte Carlo Methods in Practice*, A. Doucet, N. de Freitas, and N. Gordon, 2001, Springer, New York

from UQpy.distributions import Uniform, JointIndependent
from UQpy.sampling import ImportanceSampling
import matplotlib.pyplot as plt
import numpy as np

def log_Rosenbrock(x, param):
    return (-(100 * (x[:, 1] - x[:, 0] ** 2) ** 2 + (1 - x[:, 0]) ** 2) / param)

proposal = JointIndependent([Uniform(loc=-8, scale=16), Uniform(loc=-10, scale=60)])
print(proposal.get_parameters())

w = ImportanceSampling(log_pdf_target = log_Rosenbrock, args_target = (20,),
                       proposal = proposal, nsamples=5000)

#%% md
#
# Look at distribution of weights
# -------------------------------

#%%


print('max_weight = {}, min_weight = {} \n'.format(max(w.weights), min(w.weights)))

fig, ax = plt.subplots(figsize=(8, 3))
ax.scatter(w.weights, np.zeros((np.size(w.weights),)), s=w.weights * 600, marker='o')
ax.set_xlabel('weights')
ax.set_title('Normalized weights out of importance sampling')
ax.tick_params(which='both', left=False, labelleft=False)  # labels along the bottom edge are off
plt.show()

#%% md
#
# Compute the effective sample size
# ----------------------------------

#%%

effective_sample_size = 1. / np.sum(w.weights ** 2, axis=0)

print('Effective sample size is ne={}, out of a total number of samples={} \n'.
      format(effective_sample_size, np.size(w.weights)))


# docs/code/sampling/importance_sampling/plot_importance_sampling_rosenbrock.py

"""

Rosenbrock distribution
============================================

"""

#%% md
#
# For importance sampling, the function must be written in a way that it can
# evaluate multiple samples at once.

#%%

from UQpy.distributions import Uniform, JointIndependent
from UQpy.sampling import ImportanceSampling
import time
import matplotlib.pyplot as plt
import numpy as np

def log_Rosenbrock(x, param):
    return (-(100 * (x[:, 1] - x[:, 0] ** 2) ** 2 + (1 - x[:, 0]) ** 2) / param)

proposal = JointIndependent([Uniform(loc=-8, scale=16), Uniform(loc=-10, scale=60)])
print(proposal.get_parameters())

#%% md
#
# Run IS
# -------

#%%

t4 = time.time()


w = ImportanceSampling(log_pdf_target=log_Rosenbrock, args_target=(20,), proposal=proposal, nsamples=10000)

t_IS = time.time() - t4
print(t_IS)

w.resample(nsamples=1000)
plt.plot(w.unweighted_samples[:, 0], w.unweighted_samples[:, 1], 'gs', alpha=0.2)
print(w.unweighted_samples.shape)
plt.legend(['IS'])
plt.show()


#%% md
#
# Run IS by adding samples: call the run method in a loop (one can also look at diagnostics)
# -------------------------------------------------------

#%%

t4 = time.time()

w = ImportanceSampling(log_pdf_target = log_Rosenbrock, args_target = (20,), proposal = proposal)
for nsamples in [5000, 5000, 5000, 5000]:
    w.run(nsamples)
    print(w.samples.shape)
    # IS_diagnostics(weights=w.weights, graphics=False)
t_IS = time.time() - t4
print(t_IS)

w.resample(nsamples=1000)
plt.plot(w.unweighted_samples[:, 0], w.unweighted_samples[:, 1], 'gs', alpha=0.2)
plt.legend(['IS'])
plt.show()

#%% md
#
# Another example: sampling from a bivariate with copula dependence. Giving a random state enforces that results are
# the same for repeatability.

#%%

from UQpy.distributions import Normal, Gumbel, JointCopula
dist_true = JointCopula(marginals=[Normal(), Normal()], copula=Gumbel(theta=2.))
proposal1 = JointIndependent(marginals=[Normal(), Normal()])

sampler = ImportanceSampling(log_pdf_target = dist_true.log_pdf,proposal = proposal1, random_state = 123,
                             nsamples=500)
print(sampler.samples.shape)
print(sampler.weights.shape)
print(np.round(sampler.samples[-5:], 4))




# docs/code/sampling/latin_hypercube/plot_latin_hypercube_simple.py

"""

Latin Hypercube Sampling
==================================

This example shows the use of the Latin Hypercube sampling class. In particular:
"""

# %% md
#
# - How to define the Latin Hypercube sampling method supported by UQpy
# - How to use different sampling criteria
# - How to plot the 2D samples

# %%

# %% md
#
# Initially we have to import the necessary modules.

# %%

from UQpy.sampling import LatinHypercubeSampling
import matplotlib.pyplot as plt
from UQpy.distributions import Uniform
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria import *

# %% md
#
# Define Latin Hypercube sampling class
# ----------------------------------------------
# In order to initialize the LatinHypercube sampling class, the user needs to define a list of distributions
# for each one of the parameters that need to be sampled.
#
# Apart from the distributions list, the number of samples  *nsamples* to be drawn is required. The *random_state*
# parameter defines the seed of the random generator.
#
# Finally, the design criterion can be defined by the user. The default case is the :class:`.Random`.
# For more details on the various criteria you can refer to the documentation of the criteria
# :class:`.Random`, :class:`.Centered`, :class:`.Maximin`, :class:`.MinCorrelation`

# %%

dist1 = Uniform(loc=0., scale=1.)
dist2 = Uniform(loc=0., scale=1.)

lhs_random = LatinHypercubeSampling(distributions=[dist1, dist2], nsamples=5,
                                    random_state=np.random.RandomState(789),
                                    criterion=Random())

lhs_centered = LatinHypercubeSampling(distributions=[dist1, dist2], criterion=Centered(),
                                      random_state=np.random.RandomState(789),
                                      nsamples=5)

lhs_maximin = LatinHypercubeSampling(distributions=[dist1, dist2],
                                     random_state=np.random.RandomState(789),
                                     criterion=MaxiMin(metric=DistanceMetric.CHEBYSHEV),
                                     nsamples=5)

lhs_mincorrelate = LatinHypercubeSampling(distributions=[dist1, dist2],
                                          random_state=np.random.RandomState(789),
                                          criterion=MinCorrelation(iterations=100),
                                          nsamples=5)

# %% md
#
# Plot the generated samples
# ------------------------------------
# The *samples* attribute of the latin hypercube class is a numpy array of with shape (nsamples, len(distributions))
# Both *samples* and *samplesU01* are populated at the same time since the Latin Hypercube samples are initially drawn
# in the unit hypercube, thus in contrast to Monte Carlo sampling no transformation is required. Using the
# :py:meth:`run` method to generate samples replaces the previously created ones.

# %%

# plot the samples
fig, axs = plt.subplots(2, 2)
fig.subplots_adjust(hspace=0.5)
axs[0, 0].set_title('Random-LHS design')
axs[0, 0].scatter(lhs_random._samples[:, 0], lhs_random._samples[:, 1])
axs[0, 0].set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[0, 0].set_xticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[0, 0].yaxis.grid(True)
axs[0, 0].xaxis.grid(True)


axs[0, 1].set_title('Centered-LHS design')
axs[0, 1].scatter(lhs_centered._samples[:, 0], lhs_centered._samples[:, 1])
axs[0, 1].set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[0, 1].set_xticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[0, 1].yaxis.grid(True)
axs[0, 1].xaxis.grid(True)


axs[1, 0].set_title('Maximin-LHS design')
axs[1, 0].scatter(lhs_maximin._samples[:, 0], lhs_maximin._samples[:, 1])
axs[1, 0].set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[1, 0].set_xticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[1, 0].yaxis.grid(True)
axs[1, 0].xaxis.grid(True)


axs[1, 1].set_title('MinCorrelation-LHS design')
axs[1, 1].scatter(lhs_random._samples[:, 0], lhs_random._samples[:, 1])
axs[1, 1].set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[1, 1].set_xticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
axs[1, 1].yaxis.grid(True)
axs[1, 1].xaxis.grid(True)

plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()


# docs/code/sampling/latin_hypercube/plot_latin_hypercube_user_criterion.py

"""

User Design Criterion
==================================

This example illustrates use of user-defined latin hypercube design criteria. In particular:
"""

# %% md
#
# - How to define the Monte Carlo sampling method supported by UQpy
# - How to append new samples to the existing ones
# - How to transform existing samples to the unit hypercube
# - How to plot histograms for each one of the sampled parameters and 2D scatter of the produced samples

# %%

# %% md
#
# Initially we have to import the necessary modules.

# %%

from UQpy.sampling import LatinHypercubeSampling
import matplotlib.pyplot as plt
from UQpy.distributions import Uniform
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria import *
from UQpy.sampling.stratified_sampling.latin_hypercube_criteria import Criterion


# %% md
#
# Create user-defined criterion
# ----------------------------------------------
# In order to create a user defined criterion, a concrete implementation of the :class:`.Criterion` abstract class must
# be created. The *generate_samples* method must be implemented, which receives as input the random_state, the randomly
# generated latin hypercube samples.

# %%


class UserCriterion(Criterion):
    def generate_samples(self, random_state):
        lhs_samples = np.zeros_like(self.samples)
        for j in range(self.samples.shape[1]):
            order = np.random.permutation(self.samples.shape[0])
            lhs_samples[:, j] = self.samples[order, j]
        return lhs_samples


# %% md
#
# Define Latin Hypercube sampling
# ----------------------------------------------
# In order to initialize the LatinHypercube sampling class, the user needs to define a list of distributions
# for each one of the parameters that need to be sampled.
#
# Apart from the distributions list, the number of samples  *nsamples* to be drawn is required. The *random_state*
# parameter defines the seed of the random generator.
#
# Finally, the design criterion can be defined by the user. The default case is the :class:`.Random`.
# For more details on the various criteria you can refer to the documentation of the criteria
# :class:`.Random`, :class:`.Centered`, :class:`.Maximin`, :class:`.MinCorrelation`
#
# In the case of user-defined criteria an instantiation of the UserCriterion class is provided instead of the
# built-in criteria.


# %%

dist1 = Uniform(loc=0., scale=1.)
dist2 = Uniform(loc=0., scale=1.)

lhs_user_defined = LatinHypercubeSampling(distributions=[dist1, dist2], nsamples=5,
                                          criterion=UserCriterion())
print(lhs_user_defined._samples)

# %% md
#
# Plot the samples
# ------------------------------------
#
# The samples generated using the LatinHypercube sampling method can be retrieved using the *samples* attribute. This
# attribute is a numpy.ndarray

# %%

# plot the samples
fig, ax = plt.subplots()
plt.title('LHS sampling - User Criterion')
plt.scatter(lhs_user_defined._samples[:, 0], lhs_user_defined._samples[:, 1], marker='o')
ax.yaxis.grid(True)
ax.xaxis.grid(True)
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()


# docs/code/sampling/mcmc/plot_mcmc_algorithm_comparison.py

"""

Comparison of various MCMC aglorithms
============================================================

This script illustrates performance of various MCMC algorithms currently integrated in UQpy:
- Metropolis Hastings (MH)
- Modified Metropolis Hastings (MMH)
- Affine Invariant with Stretch moves (Stretch)
- Adaptive Metropolis with delayed rejection (DRAM)
"""

# %% md
#
# Import the necessary libraries.

# %%

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

from UQpy.sampling import MetropolisHastings, Stretch, ModifiedMetropolisHastings, DREAM, DRAM

# %% md
# Affine invariant with Stretch moves
# -----------------------------------
# This algorithm requires as seed a few samples near the region of interest. Here MH is first run to obtain few samples,
# used as seed within the Stretch algorithm.

# %%

def log_Rosenbrock(x):
    return (-(100 * (x[:, 1] - x[:, 0] ** 2) ** 2 + (1 - x[:, 0]) ** 2) / 20)

x = MetropolisHastings(dimension=2, burn_length=0, jump=10, n_chains=1, log_pdf_target=log_Rosenbrock,
                       nsamples=5000)
print(x.samples.shape)
plt.figure()
plt.plot(x.samples[:, 0], x.samples[:, 1], 'o')
plt.show()

x = Stretch(burn_length=0, jump=10, log_pdf_target=log_Rosenbrock, seed=x.samples[:10].tolist(), scale=2.,
            nsamples=5000)
print(x.samples.shape)

plt.figure()
plt.plot(x.samples[:, 0], x.samples[:, 1], 'o')
plt.show()

# %% md
# DREAM algorithm: compare with MH (inputs parameters are set as their default values)
# -----------------------------------

# %%

# Define a function to sample seed uniformly distributed in the 2d space ([-20, 20], [-4, 4])
from UQpy.distributions import Uniform, JointIndependent
prior_sample = lambda nsamples: np.array([[-2, -2]]) + np.array([[4, 4]]) * JointIndependent(
    [Uniform(), Uniform()]).rvs(nsamples=nsamples)

fig, ax = plt.subplots(ncols=2, figsize=(12, 4))
seed = prior_sample(nsamples=7)

x = MetropolisHastings(dimension=2, burn_length=500, jump=50, seed=seed.tolist(),
                       log_pdf_target=log_Rosenbrock, nsamples=1000)
ax[0].plot(x.samples[:, 0], x.samples[:, 1], 'o')

x = DREAM(dimension = 2, burn_length = 500, jump = 50, seed = seed.tolist(), log_pdf_target = log_Rosenbrock,
          nsamples=1000)
ax[1].plot(x.samples[:, 0], x.samples[:, 1], 'o')

plt.show()


# %% md
# DRAM algorithm
# -----------------------------------

# %%

fig, ax = plt.subplots(ncols=2, figsize=(12, 4))
t = time.time()
seed = prior_sample(nsamples=1)

x = MetropolisHastings(dimension = 2, burn_length = 500, jump = 10, seed = seed.tolist(), log_pdf_target = log_Rosenbrock,
                       nsamples=1000)

ax[0].plot(x.samples[:, 0], x.samples[:, 1], 'o')
ax[0].set_xlabel('x1')
ax[0].set_ylabel('x2')
ax[0].set_title('algorithm: MH')
print('time to run MH' + ': {}s'.format(time.time() - t))

x = DRAM(dimension=2, burn_length=500, jump=10, seed=seed.tolist(), log_pdf_target=log_Rosenbrock,
         save_covariance=True, nsamples=1000)

ax[1].plot(x.samples[:, 0], x.samples[:, 1], 'o')
ax[1].set_xlabel('x1')
ax[1].set_ylabel('x2')
ax[1].set_title('algorithm: DRAM' )
print('time to run DRAM'  + ': {}s'.format(time.time() - t))

plt.show()

# look at the covariance adaptivity
fig, ax = plt.subplots()
adaptive_covariance = np.array(x.adaptive_covariance)
for i in range(2):
    ax.plot(np.sqrt(adaptive_covariance[:, 0, i, i]), label='dimension {}'.format(i))
ax.set_title('Adaptive proposal std. dev. in both dimensions')
ax.legend()
plt.show()

# %% md
# MMH: target pdf is given as a joint pdf
# -----------------------------------
# The target pdf should be a 1 dimensional distribution or set of 1d distributions.

# %%

from UQpy.distributions import Normal
proposal = [Normal(), Normal()]
proposal_is_symmetric = [False, False]

x = ModifiedMetropolisHastings(dimension=2, burn_length=500, jump=50, log_pdf_target=log_Rosenbrock,
                               proposal=proposal, proposal_is_symmetric=proposal_is_symmetric, n_chains=1,
                               nsamples=500)

fig, ax = plt.subplots()
ax.plot(x.samples[:, 0], x.samples[:, 1], linestyle='none', marker='.')

# %% md
# MMH: target pdf is given as a couple of independent marginals
# -----------------------------------

# %%

log_pdf_target = [Normal(loc=0., scale=5.).log_pdf, Normal(loc=0., scale=20.).log_pdf]

proposal = [Normal(), Normal()]
proposal_is_symmetric = [True, True]

x = ModifiedMetropolisHastings(dimension = 2, burn_length = 100, jump = 10, log_pdf_target = log_pdf_target,
                               proposal = proposal, proposal_is_symmetric = proposal_is_symmetric, n_chains = 1,
                               nsamples=1000)

fig, ax = plt.subplots()
ax.plot(x.samples[:, 0], x.samples[:, 1], linestyle='none', marker='.')
plt.show()
print(x.samples.shape)

# %% md
# Use random_state to provide repeated results
# -----------------------------------

# %%

from UQpy.distributions import Normal, Gumbel, JointCopula, JointIndependent, Uniform
seed = Uniform().rvs(nsamples=2 * 10).reshape((10, 2))
dist_true = JointCopula(marginals=[Normal(), Normal()], copula=Gumbel(theta=2.))
proposal = JointIndependent(marginals=[Normal(scale=0.2), Normal(scale=0.2)])

for _ in range(3):
    sampler = ModifiedMetropolisHastings(log_pdf_target=dist_true.log_pdf, proposal=proposal, seed=[0., 0.],
                                         random_state=123, nsamples=500)
    print(sampler.samples.shape)
    print(np.round(sampler.samples[-5:], 4))

plt.plot(sampler.samples[:, 0], sampler.samples[:, 1], linestyle='none', marker='+')


# docs/code/sampling/mcmc/plot_mcmc_diagnostics.py

"""

Diagnostics for MCMC methods
==================================

This notebook illustrates the use of a few diagnostics for MCMC. The example consists in learning the parameters of a
regression model via Bayesian inference.
"""

#%% md
#
# Import the necessary libraries.

#%%

from UQpy.sampling import MetropolisHastings
from UQpy.distributions import Normal, JointIndependent
import numpy as np
import matplotlib.pyplot as plt

#%% md
# Example setup
# --------------
# First generate data then define the target log pdf = log likelihood $p(data \vert x)$.

#%%

param_true = [1., 2.]
domain = np.linspace(0, 10, 50)
data_clean = param_true[0] * domain + param_true[1] * domain ** 2
rstate = np.random.RandomState(123)
data_noisy = data_clean + rstate.randn(*data_clean.shape)


from scipy.stats import norm
def log_target(x, data, x_domain):
    log_target_value = np.zeros(x.shape[0])
    for i, xx in enumerate(x):
        h_xx = xx[0] * x_domain + xx[1] * x_domain ** 2
        log_target_value[i] = np.sum([norm.logpdf(hxi - datai) for hxi, datai in zip(h_xx, data)])
    return log_target_value


proposal = JointIndependent([Normal(scale=0.1), Normal(scale=0.05)])

sampler = MetropolisHastings(nsamples=500, dimension=2, log_pdf_target=log_target,
                             burn_length=10, jump=10, n_chains=1,
                             args_target=(data_noisy, domain), proposal=proposal)

print(sampler.samples.shape)
samples = sampler.samples

plt.plot(samples[:, 0], samples[:, 1], 'o', alpha=0.5)
plt.plot(1., 2., marker='x', color='orange')
plt.show()


#%% md
# Look at acceptance rate of the chain(s)
# --------------
#

#%%

print(sampler.acceptance_rate)


#%% md
# This acceptance rate is quite low, showing that the proposal should probably be refined !
#

#%%

#%% md
# Graphics
# ---------

#%%

fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(15, 7))
for j in range(samples.shape[1]):
    ax[j, 0].plot(np.arange(samples.shape[0]), samples[:, j])
    ax[j, 0].set_title('chain - parameter # {}'.format(j+1))
    ax[j, 1].plot(np.arange(samples.shape[0]), np.cumsum(samples[:, j])/np.arange(samples.shape[0]))
    ax[j, 1].set_title('parameter convergence')
    ax[j, 2].acorr(samples[:, j] - np.mean(samples[:, j]), maxlags=40, normed=True)
    ax[j, 2].set_title('correlation between samples')
plt.show()

#%% md
# Compute effective sample size (ESS)
# -----------------------------------
# In MCMC, the $ESS$ has been used to derive termination rules, based on the quality of the estimation. In brief, the simulation stops when the computational uncertainty on a chosen quantity (for instance, the expected value of RV $x$) is small compared to its posterior uncertainty [1]. Mathematically, this allows computation of the $ESS$ and a minimum value $ESS_{min}$ (qualitatively, an acceptable approximation is reached if $ESS > ESS_{min}$). These quantities can be computed by looking at each marginal density (is $x$ is a multivariate random variable), or by looking at the joint. The reader is referred to [1, 2] for more details.
#
# [1] *A practical sequential stopping rule for high-dimensional MCMC and its application to spatial-temporal Bayesian models*, Gong and Flegal, 2014. <br>
# [2] *Multivariate Output Analysis for Markov Chain Monte Carlo*, Vats and Flegal, 2015

#%%

def compute_marginal_ESS(samples, eps_ESS=0.05, alpha_ESS=0.05):
    # Computation of ESS and min ESS, look at each marginal distribution separately
    nsamples, dim = samples.shape

    bn = np.ceil(nsamples ** (1 / 2))  # nb of samples per bin
    an = int(np.ceil(nsamples / bn))  # nb of bins
    idx = np.array_split(np.arange(nsamples), an)

    means_subdivisions = np.empty((an, samples.shape[1]))
    for i, idx_i in enumerate(idx):
        x_sub = samples[idx_i, :]
        means_subdivisions[i, :] = np.mean(x_sub, axis=0)
    Omega = np.cov(samples.T)
    Sigma = np.cov(means_subdivisions.T)

    marginal_ESS = np.empty((dim,))
    min_marginal_ESS = np.empty((dim,))
    for j in range(dim):
        marginal_ESS[j] = nsamples * Omega[j, j] / Sigma[j, j]
        min_marginal_ESS[j] = 4 * norm.ppf(alpha_ESS / 2) ** 2 / eps_ESS ** 2

    return marginal_ESS, min_marginal_ESS


def compute_joint_ESS(samples, eps_ESS=0.05, alpha_ESS=0.05):
    # Computation of ESS and min ESS, look at the joint distribution separately
    nsamples, dim = samples.shape

    bn = np.ceil(nsamples ** (1 / 2))  # nb of samples per bin
    an = int(np.ceil(nsamples / bn))  # nb of bins
    idx = np.array_split(np.arange(nsamples), an)

    means_subdivisions = np.empty((an, samples.shape[1]))
    for i, idx_i in enumerate(idx):
        x_sub = samples[idx_i, :]
        means_subdivisions[i, :] = np.mean(x_sub, axis=0)
    Omega = np.cov(samples.T)
    Sigma = np.cov(means_subdivisions.T)

    from scipy.stats import chi2
    from scipy.special import gamma
    joint_ESS = nsamples * np.linalg.det(Omega) ** (1 / dim) / np.linalg.det(Sigma) ** (1 / dim)
    chi2_value = chi2.ppf(1 - alpha_ESS, df=dim)
    min_joint_ESS = 2 ** (2 / dim) * np.pi / (dim * gamma(dim / 2)) ** (
            2 / dim) * chi2_value / eps_ESS ** 2

    return joint_ESS, min_joint_ESS

marginal_ESS, min_marginal_ESS = compute_marginal_ESS(samples, eps_ESS=0.05, alpha_ESS=0.05)

print('Univariate Effective Sample Size in each dimension:')
for j in range(2):
    print('Dimension {}: ESS = {}, minimum ESS recommended = {}'.
          format(j+1, marginal_ESS[j], min_marginal_ESS[j]))

    joint_ESS, min_joint_ESS = compute_joint_ESS(samples, eps_ESS=0.05, alpha_ESS=0.05)

    print('\nMultivariate Effective Sample Size:')
    print('Multivariate ESS = {}, minimum ESS recommended = {}'.format(joint_ESS, min_joint_ESS))

#%% md
# Results above show that the chains should probably be run for a longer amount of time in order to obtain reliable
# results in terms of the expected value of the parameters to be learnt.


# docs/code/sampling/mcmc/plot_mcmc_metropolis_hastings.py

"""

Sampling Rosenbrock distribution using Metropolis-Hastings
============================================================

In this example, the Metropolis-Hastings is employed to generate samples from a Rosenbrock distribution. The method illustrates various aspects of the UQpy MCMC class:
- various ways of defining the target pdf to sample from,
- definition of input parameters required by the algorithm (proposal_type and proposal_scale for MH),
- running several chains in parallel,
- call diagnostics functions.
"""

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

# %%

from UQpy.sampling import MetropolisHastings
import numpy as np
import matplotlib.pyplot as plt
import time


# %% md
#
# Explore various ways of defining the target pdf
# -----------------------------------------------
# Define the Rosenbrock probability density function up to a scale factor. Here the pdf is defined directly in the python script <br>
# - define the Rosenbrock probability density function up to a scale factor, this function only takes as input parameter the point x where to compute the pdf, <br>
# - define a pdf function that also takes as argument a set of parameters params,
# - define a function that computes the log pdf up to a constant.
# Alternatively, the pdf can be defined in an external file that defines a distribution and its pdf or log_pdf methods (Rosenbrock.py)

# %%

def rosenbrock_no_params(x):
    return np.exp(-(100 * (x[:, 1] - x[:, 0] ** 2) ** 2 + (1 - x[:, 0]) ** 2) / 20)


def log_rosenbrock_with_param(x, p):
    return (-(100 * (x[:, 1] - x[:, 0] ** 2) ** 2 + (1 - x[:, 0]) ** 2) / p)


x = MetropolisHastings(dimension=2, pdf_target=rosenbrock_no_params, burn_length=500, jump=50,
                       n_chains=1, nsamples=500)
print(x.samples.shape)
plt.figure()
plt.plot(x.samples[:, 0], x.samples[:, 1], 'o', alpha=0.5)
plt.show()

plt.figure()
x = MetropolisHastings(dimension=2, pdf_target=log_rosenbrock_with_param, burn_length=500,
                       jump=50, n_chains=1, args_target=(20,),
                       nsamples=500)
plt.plot(x.samples[:, 0], x.samples[:, 1], 'o')
plt.show()

# %% md
#
# In the following, a custom Rosenbrock distribution is defined and its log_pdf method is used.

# %%

from UQpy.distributions import DistributionND


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']


log_pdf_target = Rosenbrock(p=20.).log_pdf

x = MetropolisHastings(dimension=2, pdf_target=log_pdf_target, burn_length=500, jump=50,
                       n_chains=1, nsamples=500)

# %% md
#
# In the following, we show that if nburn is set to 0, the first sample is always the seed.

# %%

seed = [[0., -1.], [1., 1.]]

x = MetropolisHastings(dimension=2, log_pdf_target=log_pdf_target, burn_length=0, jump=50,
                       seed=seed, concatenate_chains=False,
                       nsamples=500)
print(x.samples.shape)
plt.plot(x.samples[:, 0, 0], x.samples[:, 0, 1], 'o', alpha=0.5)
plt.plot(x.samples[:, 1, 0], x.samples[:, 1, 1], 'o', alpha=0.5)
plt.show()

print(seed)
print(x.samples[0, :, :])

# %% md
#
# The algorithm-specific parameters for MH are proposal and proposal_is_symmetric
# -------------------------------------------------------------------------------
# The default proposal is standard normal (symmetric).

# %%

# Define a few proposals to try out
from UQpy.distributions import JointIndependent, Normal, Uniform

proposals = [JointIndependent([Normal(), Normal()]),
             JointIndependent([Uniform(loc=-0.5, scale=1.5), Uniform(loc=-0.5, scale=1.5)]),
             Normal()]

proposals_is_symmetric = [True, False, False]

fig, ax = plt.subplots(ncols=3, figsize=(16, 4))
for i, (proposal, symm) in enumerate(zip(proposals, proposals_is_symmetric)):
    print(i)
    try:
        x = MetropolisHastings(dimension=2, burn_length=500, jump=100, log_pdf_target=log_pdf_target,
                               proposal=proposal, proposal_is_symmetric=symm, n_chains=1,
                               nsamples=1000)
        ax[i].plot(x.samples[:, 0], x.samples[:, 1], 'o')
    except ValueError as e:
        print(e)
        print('This last call fails because the proposal is in dimension 1, while the target distribution is'
              ' in dimension 2')
plt.show()

# %% md
#
# Run several chains in parallel
# -------------------------------------------------------------------------------
# The user can provide the total number of samples nsamples, or the number of samples per chain nsamples_per_chain.

# %%

x = MetropolisHastings(dimension=2, log_pdf_target=log_pdf_target, jump=1000, burn_length=500,
                       seed=[[0., 0.], [1., 1.]], concatenate_chains=False,
                       nsamples=100)
plt.plot(x.samples[:, 0, 0], x.samples[:, 0, 1], 'o', label='chain 1', alpha=0.5)
plt.plot(x.samples[:, 1, 0], x.samples[:, 1, 1], 'o', label='chain 2', alpha=0.5)
print(x.samples.shape)
plt.legend()
plt.show()

x = MetropolisHastings(dimension=2, log_pdf_target=log_pdf_target, jump=1000, burn_length=500,
                       seed=[[0., 0.], [1., 1.]], concatenate_chains=False,
                       nsamples_per_chain=100)
plt.plot(x.samples[:, 0, 0], x.samples[:, 0, 1], 'o', label='chain 1', alpha=0.5)
plt.plot(x.samples[:, 1, 0], x.samples[:, 1, 1], 'o', label='chain 2', alpha=0.5)
print(x.samples.shape)
plt.legend()
plt.show()

# %% md
#
# Initialize without nsamples... then call run
# -------------------------------------------------------------------------------

# %%

t = time.time()

x = MetropolisHastings(dimension=2, log_pdf_target=log_pdf_target, jump=1000,
                       burn_length=500, seed=[[0., 0.], [1., 1.]], concatenate_chains=False)
print('Elapsed time for initialization: {} s'.format(time.time() - t))

t = time.time()
x.run(nsamples=100)
print('Elapsed time for running MCMC: {} s'.format(time.time() - t))
print('nburn, jump at first run: {}, {}'.format(x.burn_length, x.jump))
print('total nb of samples: {}'.format(x.samples.shape[0]))

plt.plot(x.samples[:, 0, 0], x.samples[:, 0, 1], 'o', label='chain 1')
plt.plot(x.samples[:, 1, 0], x.samples[:, 1, 1], 'o', label='chain 2')
plt.legend()
plt.show()

# %% md
#
# Run another example with a bivariate distributon with copula dependence - use random_state to always have the same outputs
# -------------------------------------------------------------------------------

# %%

from UQpy.distributions import Normal, Gumbel, JointCopula, JointIndependent

dist_true = JointCopula(marginals=[Normal(), Normal()], copula=Gumbel(theta=2.))
proposal = JointIndependent(marginals=[Normal(scale=0.2), Normal(scale=0.2)])

sampler = MetropolisHastings(dimension=2, log_pdf_target=dist_true.log_pdf, proposal=proposal,
                             seed=[0., 0.], random_state=123,
                             nsamples=500)
print(sampler.samples.shape)
print(np.round(sampler.samples[-5:], 4))

plt.plot(sampler.samples[:, 0], sampler.samples[:, 1], linestyle='none', marker='+')


# docs/code/sampling/monte_carlo/plot_monte_carlo.py

"""

Monte Carlo sampling with UQpy
==================================
"""

# %% md
#
# We'll be using UQpy's Monte Carlo sampling 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.sampling import MonteCarloSampling
import matplotlib.pyplot as plt

from numpy.random import RandomState

# %% md
#
# Step-by-step: continuous univariate distribution
# -------------------------------------------------
# First, we import UQpy's normal distribution class.

# %%
from UQpy.distributions import Normal

# %% md
#
# We'll start by constructing two identical standard normal distributions :code:`normal1` and :code:`normal2`

# %%

normal1 = normal2 = Normal()

# %% md
#
# Next, we'll construct a :code:`MonteCarloSampling` object :code:`mc` to generate random samples following those
# distributions. Here, we specify an optional initial number of samples, :code:`nsamples` to be generated at the
# object's construction. For teh purposes of this demonstration, we also supply a random seed :code:`random_state`.
#
# We access the generated samples via the :code:`samples` attribute.

# %%

mc = MonteCarloSampling(distributions=[normal1, normal2],
                        nsamples=5,
                        random_state=RandomState(123))

mc.samples

# %% md
#
# To generate more samples on :code:`mc` after construction, we call :code:`mc.run` and once again specify
# :code:`nsamples`.

# %%

mc.run(nsamples=2, random_state=RandomState(23))

mc.samples

# %% md
#
# We can transform the samples onto the unit hypercube via applying the probability integral transformation on the
# samples to yield similar samples from the uniform distribution. We call :code:`mc.transform_u01`, from which results
# are stored in the :code:`samplesU01` attribute.

# %%

mc.transform_u01()

mc.samplesU01

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

# %%

fig, ax = plt.subplots()
plt.title('Samples')

plt.scatter(x=mc.samples[:, 0],
            y=mc.samples[:, 1],
            marker='o')

plt.setp(ax, xlim=(-1.7, 1.7), ylim=(-2.6, 2.6))
ax.yaxis.grid(True)
ax.xaxis.grid(True)

# %% md
#
# As well, we can visualize each distribution's sample densities via histograms.

# %%

fig, ax = plt.subplots(1, 2)

for i in (0, 1):
    ax[i].set_title('Distribution ' + str(i + 1))
    ax[i].hist(mc.samples[:, i])
    ax[i].yaxis.grid(True)
    ax[i].xaxis.grid(True)

plt.setp(ax, xlim=(-3, 3), ylim=(0, 2));

# %% md
#
# Additional Examples
# -------------------------------------------------
# Continuous multivariate distribution
# """"""""""""""""""""""""""""""""""""""""
# We'll use UQpy's multivariate normal class.

# %%

from UQpy.distributions import MultivariateNormal

# %% md
#
# And we construct a multivariate normal distribution :code:`mvnormal` specifying parameters :code:`mean` with a vector
# of mean values and :code:`cov` with a covariance matrix.

# %%

mvnormal = MultivariateNormal(mean=[1, 2],
                              cov=[[4, -0.9],
                                   [-0.9, 1]])

# %% md
#
# With this distribution, we construct a :code:`MonteCarloSampling` object :code:`mvmc` and generate five samples on
# construction.

# %%

mvmc = MonteCarloSampling(distributions=mvnormal,
                          nsamples=5,
                          random_state=RandomState(456))

mvmc.samples

# %% md
#
# Mixing a multivariate and a univariate continuous distribution
# """"""""""""""""""""""""""""""""""""""""
# Here, we use one of our normal distributions and our multivariate normal distribution. Notice how each distribution
# has its own bundle (array) of samples per run of sampling—even when that bundle contains a single value.

# %%

mixedmc = MonteCarloSampling(distributions=[normal1, mvnormal],
                             nsamples=5,
                             random_state=RandomState(789))

mixedmc.samples

# %% md
#
# Mixing a continuous and a discrete distribution
# """""""""""""""""""""""""""""""""""""""""""""""""
# We'll use UQpy's binomial distribution class for our discrete distribution.

# %%


from UQpy.distributions import Binomial

# %% md
#
# With that, we'll construct a :code:`binomial` distribution binomial with five trials :code:`n` and a 40% probability
# :code:`p` of success per trial.

# %%

binomial = Binomial(n=5, p=0.4)

# %% md
#
# And we construct a :code:`MonteCarloSampling` object :code:`cdmv` with five initial samples using our binomial
# distribution and one of our normal distributions.

# %%

cdmv = MonteCarloSampling(distributions=[binomial, normal1],
                          nsamples=5,
                          random_state=RandomState(333))

cdmv.samples


# docs/code/sampling/refined_stratified_sampling/plot_refined_stratified_rectangular_gradient.py

"""

Rectangular Refined Stratified Sampling - Gradient Enhanced Refinement
===================================================================

In this example, Stratified sampling is used to generate samples from Uniform distribution and sample expansion is done
adaptively using Refined Stratified Sampling.
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy, matplotlib and other necessary
# library for plots, but also need to import the STS, RSS and Krig class from UQpy.

#%%
import shutil

from UQpy.sampling import TrueStratifiedSampling, RefinedStratifiedSampling
from UQpy.surrogates import Kriging
from UQpy.RunModel import RunModel
from UQpy.distributions import Uniform
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits import mplot3d
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.patches as patches
import numpy as np
from UQpy.sampling.stratified_sampling.strata import RectangularStrata
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

#%% 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=[4, 4])

#%% md
#
# Run stratified sampling.

#%%

x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1, random_state=1)
initial_samples=x.samples.copy()

#%% md
#
# This plot shows the samples and strata generated by the STS class.

#%%

fig1 = x.strata_object.plot_2d()
plt.title("STS samples U(0,1) and space stratification")
plt.plot(x.samples[:16, 0], x.samples[:16, 1], 'ro')
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.show()

#%% md
#
# RunModel class is used to estimate the function value at sample points generated using STS class.

#%%

rmodel = RunModel(model_script='local_python_model_function.py', vec=False)

#%% md
#
# This figure shows the actual function defined in python model script.

#%%

rmodel1 = RunModel(model_script='local_python_model_function.py', vec=False)
rmodel1.run(samples=x.samples)
num = 50
x1 = np.linspace(0, 1, num)
x2 = np.linspace(0, 1, num)
x1v, x2v = np.meshgrid(x1, x2)
y_act = np.zeros([num, num])
r1model = RunModel(model_script='local_python_model_function.py')
for i in range(num):
    for j in range(num):
        r1model.run(samples=np.array([[x1v[i, j], x2v[i, j]]]), append_samples=False)
        y_act[i, j] = r1model.qoi_list[0]

fig1 = plt.figure()
ax1 = fig1.gca(projection='3d')
# Plot for estimated values
surf = ax1.plot_surface(x1v, x2v, y_act, cmap=cm.coolwarm, linewidth=0, antialiased=False)
# Customize the z axis.
ax1.set_zlim(-1, 15)
ax1.zaxis.set_major_locator(LinearLocator(10))
ax1.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
#
# Krig class generated a surrogate model using STS samples and function value at those points.

#%%

from UQpy.surrogates.kriging.regression_models import Linear
from UQpy.surrogates.kriging.correlation_models import Exponential
K = Kriging(regression_model=Linear(), correlation_model=Exponential(), optimizer=MinimizeOptimizer(method="L-BFGS-B"),
            optimizations_number=20, correlation_model_parameters=[1, 1])
K.fit(samples=x.samples, values=rmodel1.qoi_list)
print(K.correlation_model_parameters)

#%% md
#
# This figure shows the surrogate model generated using Krig class from initial samples.

#%%

num = 25
x1 = np.linspace(0, 1, num)
x2 = np.linspace(0, 1, num)
x1v, x2v = np.meshgrid(x1, x2)
y = np.zeros([num, num])
for i in range(num):
    for j in range(num):
        y[i, j] = K.predict(np.array([x1v[i, j], x2v[i, j]]))

fig2 = plt.figure()
ax2 = fig2.gca(projection='3d')
# Plot for estimated values
kr = ax2.plot_wireframe(x1v, x2v, y, color='Green', label='Kriging interpolate')

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

#%% md
#
# A RSS class object is initiated by using the STS, RunModel and Krig object.

#%%

from UQpy.sampling import GradientEnhancedRefinement
refinement = GradientEnhancedRefinement(strata=strata, runmodel_object=rmodel, surrogate=K)
z = RefinedStratifiedSampling(stratified_sampling=x, refinement_algorithm=refinement, random_state=2)

#%% md
#
# After initiating the RSS class object, new samples are generated using the RSS.sample method.

#%%

z.run(nsamples=50)

#%% md
#
# This figure shows the final samples generated using RSS class, where red dots shows the initial samples.

#%%

fig3 = strata.plot_2d()
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.plot(initial_samples[:, 0], initial_samples[:, 1], 'ro')
plt.plot(z.samplesU01[:, 0], z.samplesU01[:, 1], 'gx')
plt.show()

#%% md
#
# Krig class is used to generate a surrogate model using final samples from RSS class.

#%%

K2 = Kriging(regression_model=Linear(), correlation_model=Exponential(), optimizer=MinimizeOptimizer(method="L-BFGS-B"),
             correlation_model_parameters=K.correlation_model_parameters.tolist())
K2.fit(samples=z.samples, values=rmodel.qoi_list)

#%% md
#
# This figure shows the final surrogate model, generated using 200 samples.

#%%

y = np.zeros([num, num])
for i in range(num):
    for j in range(num):
        y[i, j] = K2.predict(np.array([x1v[i, j], x2v[i, j]]))

plt.clf()
fig4 = plt.figure()
a4 = fig4.gca(projection='3d')
# Plot for estimated values
kr = a4.plot_wireframe(x1v, x2v, y, color='Green', label='Kriging interpolate')

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

shutil.rmtree(rmodel.model_dir)
shutil.rmtree(r1model.model_dir)
shutil.rmtree(rmodel1.model_dir)



# docs/code/sampling/refined_stratified_sampling/plot_refined_stratified_rectangular_random.py

"""

Rectangular Refined Stratified Sampling - Random Refinement
============================================================

In this example, the stratified sampling method is employed to generate samples from an exponential distribution using
Voronoi stratification.
"""

#%% md
#
# In this example, Stratified sampling is used to generate samples from Uniform probability distribution and sample are
# added using adaptive approach Refined Stratified Sampling.
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the STS and RSS class from UQpy.

#%%

from UQpy.sampling.stratified_sampling.strata import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling, RefinedStratifiedSampling, RandomRefinement
from UQpy.distributions import Uniform
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

#%% 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=[4, 4])

#%% md
#
# Run stratified sampling.

#%%

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

#%% md
#
# Using UQpy RSS class to expand samples generated by STS class. In this example, two new samples are generated using
# rectangular stratification .

#%%

refinement = RandomRefinement(strata=strata)
y = RefinedStratifiedSampling(stratified_sampling=x, refinement_algorithm=refinement,
                              nsamples=18, samples_per_iteration=2, random_state=2)

#%% md
#
# In the first figure shown below, samples generated from STS class are plotted. Second figure, shows the new
# stratification corresponding to two generated samples using RSS class.

#%%

fig1 = strata.plot_2d()
plt.plot(x.samples[:16, 0], x.samples[:16, 1], 'ro')
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.show()

fig2 = strata.plot_2d()
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.plot(y.samples[:16, 0], y.samples[:16, 1], 'ro')
plt.plot(y.samples[16:, 0], y.samples[16:, 1], 'yo')
plt.show()

#%% md
#
# Further, RSS class is used to adaptively increase the sample size. In this example, samples are added in strata with
# maximum weights associated with it and strata's are cut randomly along the maximum width.

#%%

y.run(nsamples=100)

#%% md
#
# In the figure shown below, all samples generated from STS and RSS class are plotted.

#%%

fig3 = strata.plot_2d()
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.plot(y.samplesU01[:16, 0], y.samplesU01[:16, 1], 'ro')
plt.plot(y.samplesU01[16:, 0], y.samplesU01[16:, 1], 'go')
plt.show()

# docs/code/sampling/refined_stratified_sampling/plot_refined_stratified_voronoi_gradient.py

"""

Voronoi Refined Stratified Sampling - Gradient Enhanced Refinement
===================================================================

In this example, Stratified sampling is used to generate samples from Uniform distribution and sample expansion is done
adaptively using Refined Stratified Sampling.
"""

#%% md
#
# Import the necessary libraries. Here we import standard libraries such as numpy, matplotlib and other necessary
# library for plots, but also need to import the STS, RSS and Krig class from UQpy.

#%%
import shutil

from UQpy.sampling import TrueStratifiedSampling, RefinedStratifiedSampling
from UQpy.sampling import VoronoiStrata
from UQpy.surrogates import Kriging
from UQpy.RunModel import RunModel
from UQpy.distributions import Uniform
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits import mplot3d
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.patches as patches
import numpy as np
# from scipy.spatial import Delaunay
from scipy.spatial import voronoi_plot_2d
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import Matern

#%% md
#
# Create a distribution object.

#%%

marginals = [Uniform(loc=0., scale=1.), Uniform(loc=0., scale=1.)]
strata = VoronoiStrata(seeds_number=16, dimension=2)

#%% md
#
# Using UQpy STS 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
#
# This plot shows the samples and stratas generated by the STS class.

#%%

fig = voronoi_plot_2d(x.strata_object.voronoi)
plt.title('Stratified Samples (U(0,1)) - Voronoi Stratification')
plt.plot(x.samples[:, 0], x.samples[:, 1], 'dm')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()

#%% md
#
# RunModel class is used to define an object to evaluate the model at sample points.

#%%

rmodel = RunModel(model_script='local_python_model_function.py')

#%% md
#
# This figure shows the actual function defined in python model script.

#%%

rmodel1 = RunModel(model_script='local_python_model_function.py')
rmodel1.run(samples=x.samples)

num = 50
x1 = np.linspace(0, 1, num)
x2 = np.linspace(0, 1, num)
x1v, x2v = np.meshgrid(x1, x2)
y_act = np.zeros([num, num])
r1 = RunModel(model_script='local_python_model_function.py')
for i in range(num):
    for j in range(num):
        r1.run(samples=np.array([[x1v[i, j], x2v[i, j]]]))
        y_act[i, j] = r1.qoi_list[-1]

fig1 = plt.figure()
ax = fig1.gca(projection='3d')
# Plot for estimated values
surf = ax.plot_surface(x1v, x2v, y_act, cmap=cm.coolwarm, linewidth=0, antialiased=False)
# Customize the z axis.
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
#
# Scikit-learn Gaussian Process Regrssor is used to generated a surrogate model using STS samples and function value at
# those points. User can also run the same example script using Kriging class from UQpy library.

#%%

k1 = 1.0 * Matern(length_scale=1.0, length_scale_bounds=(1e-1, 10.0), nu=1.5)
K = GaussianProcessRegressor(kernel=k1, n_restarts_optimizer=5)

#%% md
#
# This figure shows the surrogate model generated using Krig class based on initial samples. Note that, user don't have
# to fit the surrogate model before executing RSS class, it is done here to show the 3-D plot.

#%%

K.fit(x.samples, rmodel1.qoi_list)

num = 25
x1 = np.linspace(0, 1, num)
x2 = np.linspace(0, 1, num)
x1v, x2v = np.meshgrid(x1, x2)
y = np.zeros([num, num])
for i in range(num):
    for j in range(num):
        y[i, j] = K.predict(np.array([[x1v[i, j], x2v[i, j]]]))

fig2 = plt.figure()
ax2 = fig2.gca(projection='3d')
# Plot for estimated values
kr = ax2.plot_wireframe(x1v, x2v, y, color='Green', label='Kriging interpolate')

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

#%% md
#
# Using UQpy RSS class to expand samples generated by STS class. In this example, meta specifies the method used to
# estimate the gradient and Voronoi cells are used for stratification. Krig class is used with 'Gaussian' correlation
# model.

#%%

from UQpy.sampling import GradientEnhancedRefinement
refinement = GradientEnhancedRefinement(strata=strata, runmodel_object=rmodel, surrogate=K)

z = RefinedStratifiedSampling(stratified_sampling=x, random_state=2, refinement_algorithm = refinement)

#%% md
#
# After initiating the RSS class object, new samples are generated using the RSS.sample method.

#%%

z.run(nsamples=50)

#%% md
#
# This figure shows the final samples generated using RSS class, where red dots shows the initial samples.

#%%

fig = voronoi_plot_2d(refinement.strata.voronoi)
plt.title('Gradient Enhanced RSS - Voronoi Stratification')
plt.plot(z.samplesU01[:, 0], z.samplesU01[:, 1], 'dm')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()

#%% md
#
# This figure shows the final surrogate model, generated using 200 samples.

#%%

y = np.zeros([num, num])
for i in range(num):
    for j in range(num):
        y[i, j] = refinement.surrogate.predict(np.array([[x1v[i, j], x2v[i, j]]]))

plt.clf()
fig4 = plt.figure()
a2 = fig4.gca(projection='3d')
# Plot for estimated values
kr = a2.plot_wireframe(x1v, x2v, y, color='Green', label='Kriging interpolate')

# Plot for scattered data
ID = a2.scatter3D(z.samples[:, 0], z.samples[:, 1], refinement.runmodel_object.qoi_list, color='Red',
                  label='Input data')
plt.legend(handles=[kr, ID])
plt.show()

shutil.rmtree(rmodel.model_dir)
shutil.rmtree(rmodel1.model_dir)
shutil.rmtree(r1.model_dir)


# docs/code/sampling/refined_stratified_sampling/plot_refined_stratified_voronoi_random.py

"""

Voronoi Refined Stratified Sampling - Random Refinement
============================================================

In this example, Stratified sampling is used to generate samples from Uniform probability distribution and sample are
added using adaptive approach Refined Stratified Sampling.
"""

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

#%%

from UQpy.sampling.stratified_sampling.strata import VoronoiStrata
from UQpy.sampling import TrueStratifiedSampling, RefinedStratifiedSampling, RandomRefinement
from UQpy.distributions import Uniform
import matplotlib.pyplot as plt

#%% md
#
# Create a distribution object.

#%%

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

strata = VoronoiStrata(seeds_number=16, dimension=2)

#%% md
#
# Run stratified sampling.

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

#%% md
#
# Plot the resulting stratified samples and the boundaries of the strata in the U(0,1) space.

#%%

from scipy.spatial import voronoi_plot_2d
fig = voronoi_plot_2d(strata.voronoi)
plt.title('Stratified Samples (U(0,1)) - Voronoi Stratification')
plt.plot(x.samples[:, 0], x.samples[:, 1], 'dm')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()

#%% md
#
# Using UQpy RSS class to expand samples generated by STS class. In this example, two new samples are generated inside
# cells with maximum weight associated with it.

#%%

refinement = RandomRefinement(strata=strata)
y = RefinedStratifiedSampling(stratified_sampling=x, nsamples=18, refinement_algorithm=refinement,
                              samples_per_iteration=2, random_state=2)

#%% md
#
# The figure shows the voronoi tesselation of initial samples and samples generated using RSS class. The RSS class
# creates a Delaunay triangulation using existing samples and generate a new sample inside the triangle with maximum
# volume.

#%%

fig2 = voronoi_plot_2d(strata.voronoi)
plt.title('Refined Stratified Samples - Voronoi Stratification')
plt.plot(x.samples[:16, 0], x.samples[:16, 1], 'dm')
plt.scatter(y.samplesU01[16:18, 0], y.samplesU01[16:18, 1], marker='s', color='black')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()

#%% md
#
# Further, RSS class is used to adaptively increase the sample size. In this example, samples are randomly added in
# cell with maximum weights associated with it and new sample generated using Simplex class.

#%%

y.run(nsamples=50)

#%% md
#
# In the figure shown below, all samples generated from STS and RSS class are plotted.

#%%

fig2 = voronoi_plot_2d(y.refinement_algorithm.strata.voronoi)
plt.title('Refined Stratified Samples - Voronoi Stratification')
plt.plot(x.samples[:16, 0], x.samples[:16, 1], 'dm')
plt.scatter(y.samplesU01[16:, 0], y.samplesU01[16:, 1], marker='s', color='black')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()

# docs/code/sampling/simplex/plot_simplex.py

"""

Simplex Sampling
==================================

This example shows the use of the Monte Carlo sampling class. In particular:
"""

#%% md
#
# - How to define the Simplex sampling method supported by UQpy
# - How to draw samples using Simplex sampling
# - How to plot simplices as well as their corresponding samples.

#%%

#%% md
#
# Initially we have to import the necessary modules.

#%%

from UQpy.sampling import SimplexSampling
import numpy as np
import matplotlib.pyplot as plt

#%% md
#
# Define Simplices
# ----------------------------------------------
# Create an array of 3 points in 2-D, which will be the vertex coordinates of the simplex.

#%%

vertex = np.array([[0, 0], [0.5, 1], [1, 0]])

#%% md
#
# Add new samples to existing sampling object
# ----------------------------------------------
# Use Simplex class to generate uniformly distributed samples. This class needs two input parameters,
# i.e. nodes and nsamples. Nodes are the vertex coordinates of simplex and nsamples is the number of new samples to be
# generated. In this example, we are generating ten new samples inside our simplex.

#%%

x = SimplexSampling(nodes=vertex, nsamples=10, random_state=1)

#%% md
#
# Plot the samples of the simplex
# ------------------------------------
# A schematic illustration of the 2-D simplex and new samples generated using Simplex class are presented below.

#%%

plt.plot(np.array([0, 0.5, 1, 0]), np.array([0, 1, 0, 0]), color='blue')
plt.scatter(x.samples[:, 0], x.samples[:, 1], color='red')
plt.show()

#%% md
#
# User can also define a Simplex object using vertices and generate samples using 'run' method.

#%%

y = SimplexSampling(nodes=vertex)
y.run(nsamples=5)



# docs/code/sampling/true_stratified_sampling/plot_true_stratified_delaunay.py

"""

Delaunay Stratified Sampling
==================================

In this example, the stratified sampling method is employed to generate samples from an exponential distribution
using Delaunay stratification.
"""

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

#%%

from UQpy.sampling.stratified_sampling.TrueStratifiedSampling import TrueStratifiedSampling
from UQpy.sampling.stratified_sampling.strata import DelaunayStrata
from UQpy.distributions import Exponential
import numpy as np
import matplotlib.pyplot as plt


#%% md
#
# Run STS for 25 samples.
#
# - 2 dimensions
# - Five strata in each dimension
# - Exponential distribution with location parameter = 1 and scale parameter = 1.
#
# Create a distribution object.

#%%

marginals = [Exponential(loc=1., scale=1.), Exponential(loc=1., scale=1.)]

#%% md
#
# Equal number of samples in each stratum
# ---------------------------------------
#
# Create strata object using VoronoiStrata class.

#%%

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)


#%% md
#
# Figure shows the stratification of domain using randomly generated seed points. Notice that DelaunayStrata class
# include the corners of [0, 1]^dimension hypercube before constructing Delaunay Triangulation. In this plot, orange
# points are the seed points and left corner is also included in the delaunay construction

#%%

plt.triplot(strata_obj.delaunay.points[:, 0], strata_obj.delaunay.points[:, 1], strata_obj.delaunay.simplices)
plt.plot(seeds[:, 0], seeds[:, 1], 'or')
plt.show()

#%% md
#
# Run stratified sampling

#%%

sts_obj = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj)
sts_obj.run(nsamples_per_stratum=2)

#%% md
#
# Plot the resulting stratified samples and the boundaries of the strata in the U(0,1) space.

#%%

plt.triplot(strata_obj.delaunay.points[:, 0], strata_obj.delaunay.points[:, 1], strata_obj.delaunay.simplices)
plt.plot(seeds[:, 0], seeds[:, 1], 'or')
plt.plot(sts_obj.samplesU01[:, 0], sts_obj.samplesU01[:, 1], 'dm')
plt.title('Stratified Sample - U(0,1)')
plt.show()

sts_obj.weights


#%% md
#
# Proportional Sampling
# ----------------------
# Delaunay class can generate samples proportional to volume of each stratum.

#%%

sts_obj = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj)
sts_obj.run(nsamples=10)


#%% md
#
# It can be noticed that new sample in each stratum is proportional to volume

#%%

print('Volume: ', sts_obj.strata_object.volume)
print('Number of samples in each stratum: ', sts_obj.nsamples_per_stratum)

plt.triplot(strata_obj.delaunay.points[:, 0], strata_obj.delaunay.points[:, 1], strata_obj.delaunay.simplices)
plt.plot(seeds[:, 0], seeds[:, 1], 'or')
plt.plot(sts_obj.samplesU01[:, 0], sts_obj.samplesU01[:, 1], 'dm')
plt.show()

# docs/code/sampling/true_stratified_sampling/plot_true_stratified_rectangular.py

"""

Rectangular Stratified Sampling
==================================

In this example, the stratified sampling method is employed to generate samples from an exponential distribution.
"""

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

#%%

from UQpy.sampling.stratified_sampling.TrueStratifiedSampling import TrueStratifiedSampling
from UQpy.distributions import Exponential
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import expon
from UQpy.sampling.stratified_sampling.strata import RectangularStrata


#%% md
#
# Run STS for 25 samples.
#
# - 2 dimensions
# - Five strata in each dimension
# - Exponential distribution with location parameter = 1 and scale parameter = 1.
#
# Create a distribution object.

#%%

marginals = [Exponential(loc=1., scale=1.), Exponential(loc=1., scale=1.)]

#%% md
#
# Create strata with equal volume
# --------------------------------
# Create a strata object using RectangularStrata class.

#%%

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

#%% md
#
# Generate samples using RectangularSTS class, one sample is generate inside each stratum.

#%%

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


#%% md
#
# Plot the resulting stratified samples and the boundaries of the strata in the U(0,1) space.

#%%

fig = strata.plot_2d()
plt.title('Stratified Sample - U(0,1)')
plt.scatter(x_sts.samplesU01[:, 0], x_sts.samplesU01[:, 1], color='r')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()


print(x_sts.weights)

#%% md
#
# Plot the resulting stratified exponential samples and the boundaries of the strata in the exponential space.

#%%

fig, ax = plt.subplots()
plt.title('Stratified Sample - Exponential')
plt.scatter(x_sts.samples[:, 0], x_sts.samples[:, 1])
ax.set_yticks([1.0, expon.ppf(0.2, 1, 1), expon.ppf(0.4, 1, 1), expon.ppf(0.6, 1, 1), expon.ppf(0.8, 1, 1),
               expon.ppf(0.99, 1, 1)])
ax.set_xticks([1.0, expon.ppf(0.2, 1, 1), expon.ppf(0.4, 1, 1), expon.ppf(0.6, 1, 1), expon.ppf(0.8, 1, 1),
               expon.ppf(0.99, 1, 1)])
ax.yaxis.grid(True)
ax.xaxis.grid(True)
plt.ylim(1, expon.ppf(0.99, 1, 1))
plt.xlim(1, expon.ppf(0.99, 1, 1))
plt.show()

print(x_sts.samples)


#%% md
#
# Create stratification using seeds and widths
# --------------------------------
# Strata object can be initiated by defining seeds and widths of the strata.

#%%

seeds = np.array([[0, 0], [0.4, 0], [0, 0.5], [0.4, 0.5]])
widths = np.array([[0.4, 0.5], [0.6, 0.5], [0.4, 0.5], [0.6, 0.5]])
strata_obj = RectangularStrata(seeds=seeds, widths=widths)


#%% md
#
# Generate samples using RectangularSTS class. User can control the number of samples generated inside each stratum.
# In this illustration, 10 samples are generated such that nsamples_per_stratum governs the number of sa

#%%

sts_obj = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj, random_state=20)
sts_obj.run(nsamples_per_stratum=[1, 2, 3, 4])

#%% md
#
# Plot show the strata and samples generated in each stratum.

#%%

fig = strata_obj.plot_2d()
plt.title('Stratified Sample - U(0,1)')
plt.scatter(sts_obj.samplesU01[:, 0], sts_obj.samplesU01[:, 1], color='r')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()

#%% md
#
# Probability weights corresponding to each samples computed using Stratified sampling.

#%%

sts_obj.weights

#%% md
#
# Create stratification using input file
# --------------------------------
# Strata object can be defined using a input file, which contains the seeds and widths of each stratum.

#%%

strata_obj1 = RectangularStrata(input_file='strata.txt')

#%% md
#
# Generate samples inside eaach stratum using RectangularSTS class.

#%%

sts_obj1 = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj1, nsamples_per_stratum=1)

fig = strata_obj1.plot_2d()
plt.title('Stratified Sample - U(0,1)')
plt.scatter(sts_obj1.samplesU01[:, 0], sts_obj1.samplesU01[:, 1], color='r')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()

#%% md
#
# Proportional sampling
# ----------------------
# RectangularSTS class can generate samples proportional to volume of each stratum.

#%%

strata_obj.random_state = 24
sts_obj2 = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj)
sts_obj2.run(nsamples=10)

#%% md
#
# It can be noticed that new sample in each stratum is proportional to volume.

#%%

print('Volume: ', sts_obj2.strata_object.volume)
print('Number of samples in each stratum: ', sts_obj2.nsamples_per_stratum)


fig = strata_obj.plot_2d()
plt.title('Stratified Sample - U(0,1)')
plt.scatter(sts_obj.samplesU01[:, 0], sts_obj.samplesU01[:, 1], color='r')
plt.ylim(0, 1)
plt.xlim(0, 1)
plt.show()


# docs/code/sampling/true_stratified_sampling/plot_true_stratified_voronoi.py

"""

Voronoi Stratified Sampling
==================================

In this example, the stratified sampling method is employed to generate samples from an exponential distribution using
Voronoi stratification.
"""

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

#%%

from UQpy.sampling.stratified_sampling.TrueStratifiedSampling import TrueStratifiedSampling
from UQpy.sampling.stratified_sampling.strata import VoronoiStrata
from UQpy.distributions import Exponential
import matplotlib.pyplot as plt
from scipy.spatial import voronoi_plot_2d

#%% md
#
# Run STS for 25 samples.
#
# - 2 dimensions
# - Five strata in each dimension
# - Exponential distribution with location parameter = 1 and scale parameter = 1.
#
# Create a distribution object.

#%%

marginals = [Exponential(loc=1., scale=1.), Exponential(loc=1., scale=1.)]

#%% md
#
# Equal number of samples in each stratum
# ----------------------------------------
# Create strata object using VoronoiStrata class.

#%%

strata_obj = VoronoiStrata(seeds_number=8, dimension=2)
sts_vor_obj = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj,  random_state=3)

#%% md
#
# Figure shows the stratification of domain using randomly generated seed points

#%%

strata_obj

voronoi_plot_2d(strata_obj.voronoi)
plt.title('Stratified Sample - U(0,1)')
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.show()

#%% md
#
# Run stratified sampling

#%%

sts_vor_obj.run(nsamples_per_stratum=3)

#%% md
#
# Plot the resulting stratified samples and the boundaries of the strata in the U(0,1) space.

#%%

voronoi_plot_2d(strata_obj.voronoi)
plt.title('Stratified Sample - U(0,1)')
plt.plot(sts_vor_obj.samplesU01[:, 0], sts_vor_obj.samplesU01[:, 1], 'dm')
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.show()

sts_vor_obj.weights

#%% md
#
# Proportional Sampling
# ---------------------
# VoronoiSTS class can generate samples proportional to volume of each stratum.

#%%

sts_vor_obj1 = TrueStratifiedSampling(distributions=marginals, strata_object=strata_obj, random_state=1)
sts_vor_obj1.run(nsamples=10)

#%% md
#
# It can be noticed that new sample in each stratum is proportional to volume

#%%

print('Volume: ', sts_vor_obj1.strata_object.volume)
print('Number of samples in each stratum: ', sts_vor_obj1.nsamples_per_stratum)


voronoi_plot_2d(strata_obj.voronoi)
plt.title('Stratified Sample - U(0,1)')
plt.plot(sts_vor_obj1.samplesU01[:, 0], sts_vor_obj1.samplesU01[:, 1], 'dm')
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.show()

# docs/code/sensitivity/morris/local_pfn.py

"""

Auxiliary file
==============================================

"""

import numpy as np
import sys


class RunPythonModel:

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

        self.samples = samples
        self.dimension = dimension
        self.qoi = [0]*self.samples.shape[0]

        P = 750
        for i in range(self.samples.shape[0]):
            self.qoi[i] = 1800-np.maximum(self.samples[i,0],self.samples[i,1])-2*np.sqrt(2)/3*P


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)


def fun2_sensitivity(samples):
    fun_vals = 0.01 * samples[:, 0] + 1. * samples[:, 1] + 0.4 * samples[:, 2] ** 2 + samples[:, 3] * samples[:, 4]
    return list(fun_vals)



# docs/code/sensitivity/morris/plot_12_dimensional_gfunction.py

"""

12-dimensional g-function
==============================================

To be compared with results from "An effective screening design for sensitivity analysis of large models",
Campolongo et al, 2007

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

from UQpy.RunModel import RunModel
from UQpy.distributions import Uniform
from UQpy.sensitivity import MorrisSensitivity
import numpy as np
import matplotlib.pyplot as plt

#%% md
#
# Set-up problem with g-function.

#%%

a_values = [0.001, 89.9, 5.54, 42.10, 0.78, 1.26, 0.04, 0.79, 74.51, 4.32, 82.51, 41.62]
na = len(a_values)

var_names = ['X{}'.format(i) for i in range(na)]
runmodel_object = RunModel(
    model_script='local_pfn.py', model_object_name='gfun_sensitivity', var_names=var_names, vec=True, a_values=a_values)

dist_object = [Uniform(), ] * na

sens = MorrisSensitivity(runmodel_object=runmodel_object,
                         distributions=dist_object,
                         n_levels=20,
                         maximize_dispersion=True)
sens.run(n_trajectories=10)


print(['a{}={}'.format(i + 1, ai) for i, ai in enumerate(a_values)])

fig, ax = plt.subplots()
ax.scatter(sens.mustar_indices, sens.sigma_indices)
for i, (mu, sig) in enumerate(zip(sens.mustar_indices, sens.sigma_indices)):
    ax.text(x=mu + 0.01, y=sig + 0.01, s='X{}'.format(i + 1))
ax.set_xlabel(r'$\mu^{\star}$', fontsize=14)
ax.set_ylabel(r'$\sigma$', fontsize=14)
ax.set_title('Morris sensitivity indices', fontsize=16)
plt.show()

shutil.rmtree(runmodel_object.model_dir)

# docs/code/sensitivity/morris/plot_morris_2d_gfunction.py

"""

2-dimensional g-function
==============================================

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

from UQpy.RunModel import RunModel
from UQpy.distributions import Uniform
from UQpy.sensitivity import MorrisSensitivity
import numpy as np
import matplotlib.pyplot as plt

#%% md
#
# Set-up problem with g-function.

#%%

a_values = [0.001, 99.]
na = len(a_values)

var_names = ['X{}'.format(i) for i in range(na)]
runmodel_object = RunModel(
    model_script='local_pfn.py', model_object_name='gfun_sensitivity', var_names=var_names, vec=True,
    a_values=a_values)

dist_object = [Uniform(), ] * na

#%% md
#
# First plot contour of the function, clearly X2 has little influence on the function compared to X1.

#%%

x = np.arange(0, 1, 0.02)
y = np.arange(0, 1, 0.02)
xx, yy = np.meshgrid(x, y, sparse=False)
runmodel_object.run(samples=np.vstack([xx.reshape((-1,)), yy.reshape((-1,))]).T, append_samples=False)
h = plt.contourf(x, y, np.array(runmodel_object.qoi_list).reshape(xx.shape))
plt.xlabel('X1')
plt.ylabel('X2')
plt.show()

sens = MorrisSensitivity(runmodel_object=runmodel_object,
                         distributions=dist_object,
                         n_levels=9)
sens.run(n_trajectories=10)


fig, ax = plt.subplots(figsize=(4, 4))
for trajectory in sens.trajectories_physical_space[:5]:
    ax.plot(trajectory[:, 0], trajectory[:, 1], marker='x', linewidth=3)
ax.set_xlim([0, 1])
ax.set_ylim([0, 1])
ax.set_xticks(np.arange(0., 1. + 0.5 / (sens.n_levels - 1), 1. / (sens.n_levels - 1)), minor=True)
ax.set_xticks([0., 1.], minor=False)
ax.set_xticklabels([0., 1.])
ax.set_yticks(np.arange(0., 1. + 0.5 / (sens.n_levels - 1), 1. / (sens.n_levels - 1)), minor=True)
ax.set_yticks([0., 1.], minor=False)
ax.set_yticklabels([0., 1.])
ax.grid(True, which='both')
ax.set_xlabel(r'$X_{1}$', fontsize=14)
ax.set_ylabel(r'$X_{2}$', fontsize=14)
plt.show()

print(['a{}={}'.format(i + 1, ai) for i, ai in enumerate(a_values)])

fig, ax = plt.subplots()
ax.scatter(sens.mustar_indices, sens.sigma_indices)
for i, (mu, sig) in enumerate(zip(sens.mustar_indices, sens.sigma_indices)):
    ax.text(x=mu + 0.01, y=sig + 0.01, s='X{}'.format(i + 1))
ax.set_xlabel(r'$\mu^{\star}$', fontsize=14)
ax.set_ylabel(r'$\sigma$', fontsize=14)
ax.set_title('Morris sensitivity indices', fontsize=16)
plt.show()

shutil.rmtree(runmodel_object.model_dir)

# docs/code/sensitivity/morris/plot_morris_nonlinearities.py

"""

Function with nonlinearities / parameter dependencies
=================================================================

.. math:: Y = h(X) = 0.01 X_{1} + 1.0 X_{2} + 0.4 X_{3}^{2} + X_{4} X_{5}

ranking of input parameters:

- :math:`X_{1}` is non-influential
- :math:`X_{2}` is influential, linear/additive effect (expect large :math:`\mu^{\star}` and small :math:`\sigma`)
- :math:`X_{3}` is somewhat influential, nonlinear effect,
- :math:`X_{4}, X_{5}` are influential with dependence

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%
import shutil

from UQpy.RunModel import RunModel
from UQpy.distributions import Uniform
from UQpy.sensitivity import MorrisSensitivity
import numpy as np
import matplotlib.pyplot as plt

#%% md
#
# Set-up problem with g-function.

#%%

var_names = ['X{}'.format(i) for i in range(5)]
runmodel_object = RunModel(
    model_script='local_pfn.py', model_object_name='fun2_sensitivity', var_names=var_names, vec=True)

dist_object = [Uniform(), ] * 5


sens = MorrisSensitivity(runmodel_object=runmodel_object,
                         distributions=dist_object,
                         n_levels=20, maximize_dispersion=True)
sens.run(n_trajectories=10)


fig, ax = plt.subplots(figsize=(5, 3.5))
ax.scatter(sens.mustar_indices, sens.sigma_indices, s=60)
for i, (mu, sig) in enumerate(zip(sens.mustar_indices, sens.sigma_indices)):
    ax.text(x=mu + 0.01, y=sig + 0.01, s='X{}'.format(i + 1), fontsize=14)
ax.set_xlabel(r'$\mu^{\star}$', fontsize=18)
ax.set_ylabel(r'$\sigma$', fontsize=18)
# ax.set_title('Morris sensitivity indices', fontsize=16)
plt.show()

shutil.rmtree(runmodel_object.model_dir)


# docs/code/stochastic_processes/bispectral/plot_bispectral_1d.py

"""

BiSpectral Representation Method - 1D
=================================================================

In this example, the BiSpectral Representation Method is used to generate stochastic processes from a prescribed Power
Spectrum and associated Bispectrum. This example illustrates how to use the BSRM class for one dimensional 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 BSRM class from the StochasticProcesses module of UQpy.

#%%

from UQpy.stochastic_process import BispectralRepresentation
import numpy as np
from scipy.stats import skew
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

# Input parameters
T = 60  # Time(1 / T = dw)
nt = 1200  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nf = 600  # 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
#
# Defining the Power Spectral Density(S)

#%%

S = 32 * 1 / np.sqrt(2 * np.pi) * np.exp(-1 / 2 * f ** 2)

#%% md
#
# 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

#%% md
#
# Defining the Bispectral Density(B)

#%%

B_Complex = B_Real + 1j * B_Imag
B_Ampl = np.absolute(B_Complex)

#%% 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')

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

fig, ax = plt.subplots()
plt.title('Realisation of the BiSpectral 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)*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(S)*df*2)**(3/2))

# docs/code/stochastic_processes/bispectral/plot_bispectral_nd.py

"""

Bispectral Representation Method - ND
=================================================================

In this example, the BiSpectral Representation Method is used to generate stochastic processes from a prescribed Power
Spectrum and associated Bispectrum. This example illustrates how to use the BSRM class for 'n' dimensional 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 BSRM class from the StochasticProcesses module of UQpy.

#%%

from UQpy.stochastic_process import BispectralRepresentation
import numpy as np
from scipy.stats import skew
import matplotlib.pyplot as plt
from copy import deepcopy
plt.style.use('seaborn')

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

#%%

n_sim = 1000  # Num of samples
dim = 2
n = 1  # Num of dimensions

T = 20  # Time(1 / T = dw)
nt = 128  # Num.of Discretized Time
F = 1 / T * nt / 2  # Frequency.(Hz)
nf = 64  # 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)

f_list = [f for _ in range(dim)]
F_P = np.array(np.meshgrid(*f_list, indexing='ij'))
P = 20 / np.sqrt(2 * np.pi) * np.exp(-1/2 * np.linalg.norm(F_P, axis=0) ** 2)

F_B = np.meshgrid(*[*f_list, *f_list])
b = 40 / (2 * np.pi) * np.exp(2 * (-1/2 * np.linalg.norm(F_B, axis=0) ** 2))
B_Real = deepcopy(b)
B_Imag = deepcopy(b)

B_Real[0, :, :, :] = 0
B_Real[:, 0, :, :] = 0
B_Real[:, :, 0, :] = 0
B_Real[:, :, :, 0] = 0
B_Imag[0, :, :, :] = 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 KLE 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 KLE class from the StochasticProcesses 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(samples_number=n_sim, correlation_function=R, time_interval=dt)
samples = KLE_Object.samples

fig, ax = plt.subplots()
plt.title('Realisation of the Karhunen Louve 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/plot_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/plot_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 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 = 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/plot_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 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 SRM 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 generatino 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/plot_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 SRM 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 generatino 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/plot_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 Translate 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 Translate class from the StochasticProcesses 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 1 to 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,
                               number_time_intervals=nt, number_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/local_python_model_1Dfunction.py

"""

Auxiliary file 1
======================================================================

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



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

"""

User defined Regression & Correlation
======================================================================

In this example, Kriging is used to generate a surrogate model for a given data. In this data, sample points are
generated using STS class and functional value at sample points are estimated using a model defined in python script
s('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 STS, RunModel and Krig class from UQpy.

# %%
import shutil

from UQpy.surrogates import Kriging
from UQpy.sampling import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling
from UQpy.RunModel import RunModel
from UQpy.surrogates.kriging.regression_models.baseclass import Regression
from UQpy.surrogates.kriging.correlation_models.baseclass import Correlation
from UQpy.distributions import Gamma
import numpy as np
import matplotlib.pyplot as plt

# %% md
#
# Create a distribution object.

# %%

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

# %% md
#
# Create a strata object.

# %%

strata = RectangularStrata(strata_number=[20])

# %% md
#
# Using UQpy STS 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'.

# %%

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


# %% md
#
# A regression model is defined, this function return the basis function and its jacobian.

# %%

class UserRegression(Regression):
    def r(self, s):
        fx = np.concatenate((np.ones([np.size(s, 0), 1]), s), 1)
        jf_b = np.zeros([np.size(s, 0), np.size(s, 1), np.size(s, 1)])
        np.einsum('jii->ji', jf_b)[:] = 1
        jf = np.concatenate((np.zeros([np.size(s, 0), np.size(s, 1), 1]), jf_b), 2)
        return fx, jf


# %% md
#
# A user-defined correlation model is created, which returns covariance matrix and its derivatives.

# %%

class UserCorrelation(Correlation):
    def c(self, x, s, params, dt=False, dx=False):
        x, s = np.atleast_2d(x), np.atleast_2d(s)
        # Create stack matrix, where each block is x_i with all s
        stack = - np.tile(np.swapaxes(np.atleast_3d(x), 1, 2), (1, np.size(s, 0), 1)) + np.tile(s,
                                                                                                (np.size(x, 0), 1,
                                                                                                 1))
        rx = np.exp(np.sum(-params * (stack ** 2), axis=2))
        if dt:
            drdt = -(stack ** 2) * np.transpose(np.tile(rx, (np.size(x, 1), 1, 1)), (1, 2, 0))
            return rx, drdt
        if dx:
            drdx = 2 * params * stack * np.transpose(np.tile(rx, (np.size(x, 1), 1, 1)), (1, 2, 0))
            return rx, drdx
        return rx


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

# %%

regression_model = UserRegression()
correlation_model = UserCorrelation()

from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

optimizer = MinimizeOptimizer(method="L-BFGS-B")

K = Kriging(regression_model=regression_model, optimizer=optimizer,
            correlation_model=correlation_model,
            optimizations_number=10,
            correlation_model_parameters=[10],
            random_state=2)
K.fit(samples=x.samples, values=rmodel.qoi_list)
print(K.correlation_model_parameters)


# %% md
#
# Kriging surrogate is used to compute the response surface and its gradient.

# %%

num = 1000
x1 = np.linspace(min(x.samples), max(x.samples), num)
y, mse = K.predict(x1.reshape([num, 1]), return_std=True)
y_grad = K.jacobian(x1.reshape([num, 1]))

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

# %%

rmodel.run(samples=x1, append_samples=False)
shutil.rmtree(rmodel.model_dir)

# %% 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(x1, 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 * mse,
                         (y + 1.9600 * mse)[::-1]]),
         alpha=.5, fc='y', ec='None', label='95% CI')
# plt.legend(loc='lower right')
# Shrink current axis by 20%
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/gpr/plot_gpr_function_2D.py

"""

Kriging of a custom 2D function
======================================================================

In this example, Kriging is used to generate a surrogate model for a given data. In this data, sample points are
generated using STS 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 STS, RunModel and Krig class from UQpy.

# %%
import shutil

from UQpy.surrogates import Kriging
from UQpy.surrogates.kriging.regression_models import Constant
from UQpy.surrogates.kriging.correlation_models import Exponential
from UQpy.sampling import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling
from UQpy.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

# %% 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 STS 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'.

# %%

rmodel = RunModel(model_script='local_python_model_function.py', vec=False)
rmodel.run(samples=x.samples)
shutil.rmtree(rmodel.model_dir)

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

# %%

regression_model = Constant()
correlation_model = Exponential()

from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

optimizer = MinimizeOptimizer(method="L-BFGS-B")
K = Kriging(regression_model=regression_model, optimizer=optimizer,
            correlation_model=correlation_model,
            optimizations_number=20,
            correlation_model_parameters=[1, 1])
K.fit(samples=x.samples, values=rmodel.qoi_list)
print(K.correlation_model_parameters)


# %% 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])
r2model = RunModel(model_script='local_python_model_function.py')
r2model.run(samples=np.concatenate([x1gv, x2gv], 1))
y_act = np.array(r2model.qoi_list).reshape(x1g.shape[0], x1g.shape[1])

shutil.rmtree(r2model.model_dir)

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 wirefram 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_sine.py

"""

Kriging of a sinusoidal function
======================================================================

In this example, Kriging is used to generate a surrogate model for a given data. In this data, sample points are
generated using STS 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 STS, RunModel and Krig class from UQpy.

# %%
import shutil

from UQpy.surrogates import Kriging
from UQpy.sampling.stratified_sampling.strata import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling
from UQpy.RunModel import RunModel
from UQpy.distributions import Gamma
import numpy as np
import matplotlib.pyplot as plt

# %% md
#
# Create a distribution object.

# %%

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

# %% md
#
# Create a strata 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'.

# %%

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

from UQpy.surrogates.kriging.regression_models import Linear
from UQpy.surrogates.kriging.correlation_models import Gaussian
from UQpy.utilities.MinimizeOptimizer import MinimizeOptimizer

optimizer = MinimizeOptimizer(method="L-BFGS-B")
K = Kriging(regression_model=Linear(), correlation_model=Gaussian(), optimizer=optimizer,
            optimizations_number=20, correlation_model_parameters=[1], random_state=2)
K.fit(samples=x.samples, values=rmodel.qoi_list)
print(K.correlation_model_parameters)


# %% md
#
# Kriging surrogate is used to compute the response surface and its gradient.

# %%

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]))

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

# %%

rmodel.run(samples=x1, append_samples=False)
shutil.rmtree(rmodel.model_dir)

# %% 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(x1, 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/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

<img src="Example_Camel_function.png" alt="Drawing" style="width: 400px;"/>

.. math:: f(\textbf{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(\textbf{x}^*)=-1.0316,` at :math:`\textbf{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 = PolynomialBasis.create_total_degree_basis(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 = 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 using Ridge regression.

# %%

polynomial_basis = PolynomialBasis.create_total_degree_basis(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 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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 $P$ and all multivariate polynomial have a total-degree
# (sum of degrees of corresponding univariate polynomials) at most equal to $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 = PolynomialBasis.create_total_degree_basis(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 $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)
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 $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 $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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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_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.

<img src="Example_RobotArm_function.png" alt="Drawing" style="width: 200px;"/>

**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:`\theta_i` inputs.

**Function:**

.. math:: f(\textbf{x}) = \Big( \sum_{i=1}^{4}L_i cos\Big(\sum_{j=1}^{i} \theta_j\Big) \Big)^{2} + \Big( \sum_{i=1}^{4}L_i sin\Big(\sum_{j=1}^{i} \theta_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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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_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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 = PolynomialBasis.create_total_degree_basis(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 $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 = PolynomialBasis.create_total_degree_basis(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)
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 = PolynomialBasis.create_total_degree_basis(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 $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 $q$ in interval (0,1). The set of multi indices $\boldsymbol{\alpha}$ is reduced as follows:
# \begin{equation}
#         {\boldsymbol{\alpha}}\in \mathbb{N}^{N} :
#         || \boldsymbol{\alpha}||_q
#         \equiv
#         \Big( \sum_{i=1}^{N} \alpha_i^q \Big)^{1/q}
#         \leq P
# \end{equation}
#
# Note that $q=1$ leads to full TD set.

# %%

print('Size of the full set of PCE basis:', PolynomialBasis.create_total_degree_basis(joint,P).polynomials_number)
q=0.8
polynomial_basis_hyperbolic = PolynomialBasis.create_total_degree_basis(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.

<img src="Example_Sphere_function.png" alt="Drawing" style="width: 400px;"/>

.. math:: f(\textbf{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 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 = PolynomialBasis.create_total_degree_basis(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 = 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 with Ridge regression.

# %%

polynomial_basis = PolynomialBasis.create_total_degree_basis(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 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. SROM is created for all three random variables and distribution of eigenvalues are identified
using SROM.
"""

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

# %%
import shutil

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

# %% md
#
# Create a distribution object for Gamma distribution with shape, shift and scale parameters as 2, 1 and 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 STS class to generate samples for two random variables having Gamma distribution.

# %%

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


# %% md
#
# Run 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 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 STS class. This model defines the stiffness
# matrix corresponding to each sample and estimate the eigenvalues of the matrix.

# %%

model = RunModel(model_script='local_eigenvalue_model.py')
model.run(samples=y.samples)
r_srom = model.qoi_list

# %% md
#
# MCS 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 MCS class.

# %%

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

# %% md
#
# Plot the distribution of each eigenvalue, estimated using SROM and MCS weights.

# %%

# Plot SROM and MCS approximation for first eigenvalue
r = np.array(r_srom)[:, 0]
r_mcs = np.array(r_mcs)
com = np.append(np.atleast_2d(r).T, 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(r_srom)[:, 1]
com = np.append(np.atleast_2d(r).T, 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(r_srom)[:, 2]
com = np.append(np.atleast_2d(r).T, 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 STS and 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 Gamma distribution with shape, shift and scale parameters as 2, 1 and 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 STS class to generate samples for two random variables having Gamma distribution.

#%%

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

#%% md
#
# Run SROM using the defined Gamma distribution. Here we use the following parameters.
#
# - Gamma distribution with shape, shift and scale parameters as 2, 1 and 3.
# - First and second order moments about origin are 6 and 54.
# - Notice that pdf_target references the 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 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_weigths are not defined and default value is [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 STS and 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 Gamma distribution with shape, shift and scale parameters as 2, 1 and 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 STS class to generate samples for two random variables having Gamma distribution.

# %%

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

# %% md
#
# Run SROM using the defined Gamma distribution. Here we use the following parameters.
#
# - Gamma distribution with shape, shift and scale parameters as 2, 1 and 3.
# - First and second order moments about origin are 6 and 54.
# - Notice that pdf_target references the 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 weights_distribution, weights_moments and
# 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, weights_distribution is modified by SROM class. First, it defines an array of size 2×16 with all
# elements equal to 1 and then multiply first column by 0.4 and second column by 0.5 . Similarly, weights_moments and
# 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
# s'float list' type.


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

"""

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

"""

#%% md
#
# Initially we have to import the necessary modules.

#%%

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

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]])

nataf_obj = Nataf(distributions=[dist1, dist2], corr_x=Rx)
print('Cz', nataf_obj.corr_z)
corr_x = Nataf.distortion_z2x(distributions=[dist1, dist2], corr_z=nataf_obj.corr_z)
print('Cx', corr_x)

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]])
nataf_obj = Nataf(distributions=[dist1, dist2], corr_x=Rx)


samples_x = nataf_obj.rvs(1000)
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()


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

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()

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

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()


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

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()


dist1 = Gamma(4.0, loc=0.0, scale=1.0)
dist2 = Lognormal(s=2., loc=0., scale=np.exp(1))
N = 20
w3 = np.zeros(N)
rho = np.linspace(-0.999, 0.999, N)
for i in range(10):
    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(w3, rho)
plt.plot(rho, rho)
plt.show()


#%% md
#
# Calculate correlation distortion for the transformation of two random variables from normal to lognormal.
# Compute the correlation distortion at various values of Gaussian correlation.

#%%

N = 20
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()


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

nataf_obj1 = Nataf(distributions=[dist1, dist2])

samples_x = nataf_obj1.rvs(1000)
plt.figure()
plt.title('Random samples')
plt.scatter(samples_x[:, 0], samples_x[:, 1])
plt.grid(True)
plt.xlabel('$X_1$')
plt.ylabel('$X_2$')
plt.show()

nataf_obj1.run(samples_x=samples_x, jacobian=True)

plt.figure()
plt.title('Random samples')
plt.scatter(nataf_obj1.samples_z[:, 0], nataf_obj1.samples_z[:, 1])
plt.grid(True)
plt.xlabel('$X_1$')
plt.ylabel('$X_2$')
plt.show()

# 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"
]
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/stochastic_processes/bispectral',
                      '../code/stochastic_processes/karhunen_loeve',
                      '../code/stochastic_processes/spectral',
                      '../code/stochastic_processes/translation',
                      '../code/reliability/taylor',
                      '../code/reliability/subset_simulation',
                      '../code/surrogates/srom',
                      '../code/surrogates/gpr',
                      '../code/surrogates/pce'],  # 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/stochastic_processes/bispectral',
                     'auto_examples/stochastic_processes/karhunen_loeve',
                     'auto_examples/stochastic_processes/spectral',
                     'auto_examples/stochastic_processes/translation',
                     'auto_examples/reliability/taylor',
                     'auto_examples/reliability/subset_simulation',
                     'auto_examples/surrogates/srom',
                     'auto_examples/surrogates/gpr',
                     'auto_examples/surrogates/pce',],  # 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),
}


# 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',
    ],
)


# 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 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.

# %%

# %% 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.RunModel import RunModel
from UQpy.distributions import Normal
import matplotlib.pyplot as plt
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)