"""
Polynomial chaos exploitation
=============================

The goal of this example is to create a polynomial chaos expansion using the
:class:`~openturns.FunctionalChaosAlgorithm` class and see the different methods
of the :class:`~openturns.FunctionalChaosResult` class.
In order to understand the different results, let us review some notations associated to the
polynomial chaos expansion.

We first introduce the physical input, output and model:

- the random vector :math:`\\vect{X}` is the physical input random vector,
- the output random vector :math:`\\vect{Y}` is the output of the physical model,
- the physical model :math:`g` is a function of the physical input:

.. math:: \\vect{Y} = g(\\vect{X}).

Then we introduce the iso-probabilistic transformation:

- the random vector :math:`\\vect{Z}` is the standardized input random vector,
- the transformation :math:`T` is the iso-probabilistic transformation mapping
  from the physical input to the standardized input:

.. math:: \\vect{Z} = T(\\vect{X}),

- the standard distribution :math:`\\mu` of the standardized random vector :math:`\\vect{Z}`.

We expand the model on the multivariate basis:

- the full (i.e. unselected) multivariate basis is :math:`\\left(\\Psi_k\\right)_{k \\in \\mathbb{N}}`,
- the composition of each polynomial of the truncated multivariate basis :math:`\\Psi_k`,
- the full set of coefficients of the polynomial expansion is the set :math:`(\\vect{a}_k)_{k \\in \\mathbb{N}}`,
- the composed model :math:`h` is a function of the standardized random vector :math:`\\vect{Z}`:

.. math:: \\vect{Y} = h(\\vect{Z}) = g \\circ T^{-1}(\\vect{Z})

Then we use a model selection method (e.g. from the :class:`~openturns.LARS` class):

- the set :math:`\\mathcal{K} \\subseteq \\mathbb{N}` is the set of indices corresponding to the
  result of the selection method,
- the truncated (i.e. selected) multivariate basis is :math:`\\left(\\Psi_k\\right)_{k \\in \\mathcal{K}}`,
- the selected set of coefficients of the polynomial expansion is the set :math:`(\\vect{a}_k)_{k \\in \\mathcal{K}}`,
- the composed meta model :math:`\\hat{h}` is the function
  of the standardized variables based on the truncated
  multivariate basis :math:`\\left(\\Psi_k\\right)_{k \\in \\mathcal{K}}`.
- the meta model is a function of the physical input:

.. math:: \\vect{Y} = \\hat{g} (\\vect{X}) = \\hat{h} \\circ T(\\vect{X}).


Based on the previous definitions, the composed model is:

.. math::

    h(\\vect{Z}) =  \\sum_{k \\in \\mathbb{N}} \\vect{a}_k \\Psi_k(\\vect{Z}),

the composed metamodel is:

.. math::

    \\hat{h}(\\vect{Z}) = \\sum_{k \\in \\mathcal{K}} \\vect{a}_k \\Psi_k(\\vect{Z}),

and the metamodel is:

.. math::

    \\hat{g}(\\vect{X}) = \\sum_{k \\in \\mathcal{K}} \\vect{a}_k \\Psi_k \\circ T(\\vect{X}).

The three previous mathematical functions are implemented as instances of the
:class:`~openturns.Function` class.
"""

# %%
# Create the polynomial chaos expansion
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# %%
import openturns as ot


# %%
# Prepare some input and output samples.
# We define a model which has two inputs and two outputs.
# Then we define a Normal input random vector with independent marginals,
# and we generate a sample from the input random vector.
# Finally, we evaluate the output sample from the model.
dimension = 2
input_names = ["x1", "x2"]
formulas = ["cos(x1 + x2)", "(x2 + 1) * exp(x1 - 2 * x2)"]
model = ot.SymbolicFunction(input_names, formulas)
distribution = ot.Normal(dimension)
x = distribution.getSample(30)
y = model(x)

# %%
# Create a functional chaos algorithm.
algo = ot.FunctionalChaosAlgorithm(x, y)

# %%
# The previous constructor is the simplest since it has only two input arguments.
# In this case, the algorithm has to retrieve the distribution from the `x`
# sample, which can be difficult, especially if the sample size is small.
# Notice, however, that the input distribution is known in our simple case.
# This is why we update the previous script and give the input distribution as a third
# input argument of the constructor.
algo = ot.FunctionalChaosAlgorithm(x, y, distribution)
algo.run()

# %%
# Get the result.
result = algo.getResult()
result

# %%
# Get the coefficients
# ~~~~~~~~~~~~~~~~~~~~

# %%
# In the next cells, we observe several methods of the `result` object.
# First, we get the polynomial chaos coefficients.
result.getCoefficients()

# %%
# The coefficients of the i-th output marginal.
i = 1
result.getCoefficients()[i]

# %%
# Get the indices of the selected polynomials.
subsetK = result.getIndices()
subsetK

# %%
# Get the composition of the polynomials
# of the truncated multivariate basis.
for i in range(subsetK.getSize()):
    print(
        "Polynomial number ",
        i,
        " in truncated basis <-> polynomial number ",
        subsetK[i],
        " = ",
        ot.LinearEnumerateFunction(dimension)(subsetK[i]),
        " in complete basis",
    )

# %%
# Get the basis
# ~~~~~~~~~~~~~

# %%
# Get the multivariate basis
# as a collection of :class:`~openturns.Function` objects.
reduced = result.getReducedBasis()

# %%
# Get the orthogonal basis.
orthgBasis = result.getOrthogonalBasis()

# %%
# Get the standard distribution :math:`\mu` of the standardized
# random vector :math:`\vect{Z}`.
orthgBasis.getMeasure()

# %%
# Get the metamodel
# ~~~~~~~~~~~~~~~~~

# %%
# Get the composed meta model :math:`\hat{h}` which is
# the model of the standardized random vector :math:`\vect{Z}`
# within the reduced polynomials basis.
result.getComposedMetaModel()

# %%
# Get the meta model :math:`\hat{g}` which is the composed meta model composed with the
# iso probabilistic transformation :math:`T`.
result.getMetaModel()

# %%
# Get the projection method
# ~~~~~~~~~~~~~~~~~~~~~~~~~

# %%
# Get the projection strategy: this is the method to estimate
# the coefficients, i.e. regression or integration.
algo.getProjectionStrategy()
