text: #!/usr/bin/env python
  import numpy as np
  import matplotlib.pyplot as plt
  from scipy import stats
  from functools import partial
  
  import pyapprox as pya
  from pyapprox.benchmarks.sensitivity_benchmarks import *
  from pyapprox.benchmarks.surrogate_benchmarks import *
  from pyapprox.models.genz import GenzFunction
  
  from scipy.optimize import OptimizeResult
  class Benchmark(OptimizeResult):
      """
      Contains functions and results needed to implement known
      benchmarks.
  
      The quantities
  
      Attributes
      ----------
      fun : callable
          The function being analyzed
  
      variable : pya.variable
          Class containing information about each of the nvars inputs to fun
  
      jac : callable
          The jacobian of fun. (optional)
  
      hess : callable
          The Hessian of fun. (optional)
  
      hessp : callable
          Function implementing the hessian of fun multiplied by a vector. 
          (optional)
  
      mean: np.ndarray (nvars)
          The mean of the function with respect to the PDF of var
  
      variance: np.ndarray (nvars)
          The variance of the function with respect to the PDF of var
  
      main_effects : np.ndarray (nvars)
          The variance based main effect sensitivity indices
  
      total_effects : np.ndarray (nvars)
          The variance based total effect sensitivity indices
  
      sobol_indices : np.ndarray
          The variance based Sobol sensitivity indices
  
      Notes
      -----
      Use the `keys()` method to see a list of the available 
      attributes for a specific benchmark
      """
      pass
  
  def setup_sobol_g_function(nvars):
      r"""
      Setup the Sobol-G function benchmark 
  
      .. math:: f(z) = \prod_{i=1}^d\frac{\lvert 4z_i-2\rvert+a_i}{1+a_i}, \quad a_i=\frac{i-2}{2}
  
      using 
  
      >>> from pyapprox.benchmarks.benchmarks import setup_benchmark
      >>> benchmark=setup_benchmark('sobol_g',nvars=2)
      >>> print(benchmark.keys())
      dict_keys(['fun', 'mean', 'variance', 'main_effects', 'total_effects', 'variable'])
  
      Parameters
      ----------
      nvars : integer
          The number of variables of the Sobol-G function
  
      Returns
      -------
      benchmark : pya.Benchmark
         Object containing the benchmark attributes
  
      References
      ----------
      .. [Saltelli1995] `Saltelli, A., & Sobol, I. M. About the use of rank transformation in sensitivity analysis of model output. Reliability Engineering & System Safety, 50(3), 225-239, 1995. <https://doi.org/10.1016/0951-8320(95)00099-2>`_
      """
      
      univariate_variables = [stats.uniform(0,1)]*nvars
      variable=pya.IndependentMultivariateRandomVariable(univariate_variables)
      a = (np.arange(1,nvars+1)-2)/2
      mean, variance, main_effects, total_effects = \
          get_sobol_g_function_statistics(a)
      return Benchmark({'fun':partial(sobol_g_function,a),
              'mean':mean,'variance':variance,'main_effects':main_effects,
              'total_effects':total_effects,'variable':variable})
  
  def setup_ishigami_function(a,b):
      r"""
      Setup the Ishigami function benchmark 
  
      .. math:: f(z) = \sin(z_1)+a\sin^2(z_2) + bz_3^4\sin(z_0)
  
      using 
  
      >>> from pyapprox.benchmarks.benchmarks import setup_benchmark
      >>> benchmark=setup_benchmark('ishigami',a=7,b=0.1)
      >>> print(benchmark.keys())
      dict_keys(['fun', 'jac', 'hess', 'variable', 'mean', 'variance', 'main_effects', 'total_effects', 'sobol_indices'])
  
      Parameters
      ----------
      a : float
          The hyper-parameter a
  
      b : float
          The hyper-parameter b
  
      Returns
      -------
      benchmark : pya.Benchmark
         Object containing the benchmark attributes
  
      References
      ----------
      .. [Ishigami1990] `T. Ishigami and T. Homma, "An importance quantification technique in uncertainty analysis for computer models," [1990] Proceedings. First International Symposium on Uncertainty Modeling and Analysis, College Park, MD, USA, 1990, pp. 398-403 <https://doi.org/10.1109/ISUMA.1990.151285>`_
      """
      univariate_variables = [stats.uniform(-np.pi,2*np.pi)]*3
      variable=pya.IndependentMultivariateRandomVariable(univariate_variables)
      mean, variance, main_effects, total_effects, sobol_indices, \
          sobol_interaction_indices = get_ishigami_funciton_statistics()
      return Benchmark(
          {'fun':partial(ishigami_function,a=a,b=b),
           'jac':partial(ishigami_function_jacobian,a=a,b=b),
           'hess':partial(ishigami_function_hessian,a=a,b=b),
           'variable':variable,'mean':mean,'variance':variance,
           'main_effects':main_effects,'total_effects':total_effects,
           'sobol_indices':sobol_indices,
           'sobol_interaction_indices':sobol_interaction_indices})
  
  def setup_oakley_function():
      r"""
      Setup the Oakely function benchmark 
  
      .. math:: f(z) = a_1^Tz + a_2^T\sin(z) + a_3^T\cos(z) + z^TMz
  
      where :math:`z` consists of 15 I.I.D. standard Normal variables and the data :math:`a_1,a_2,a_3` and :math:`M` are defined in the function :func:`pyapprox.benchmarks.sensitivity_benchmarks.get_oakley_function_data`.
  
      >>> from pyapprox.benchmarks.benchmarks import setup_benchmark
      >>> benchmark=setup_benchmark('oakley')
      >>> print(benchmark.keys())
      dict_keys(['fun', 'variable', 'mean', 'variance', 'main_effects'])
  
      Returns
      -------
      benchmark : pya.Benchmark
         Object containing the benchmark attributes
  
      References
      ----------
      .. [OakelyOJRSB2004] `Oakley, J.E. and O'Hagan, A. (2004), Probabilistic sensitivity analysis of complex models: a Bayesian approach. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 66: 751-769. <https://doi.org/10.1111/j.1467-9868.2004.05304.x>`_
      """
      univariate_variables = [stats.norm()]*15
      variable=pya.IndependentMultivariateRandomVariable(univariate_variables)
      mean, variance, main_effects = oakley_function_statistics()
      return Benchmark(
          {'fun':oakley_function,
           'variable':variable,'mean':mean,'variance':variance,
           'main_effects':main_effects})
  
  def setup_rosenbrock_function(nvars):
      r"""
      Setup the Rosenbrock function benchmark 
  
      .. math:: f(z) = \sum_{i=1}^{d/2}\left[100(z_{2i-1}^{2}-z_{2i})^{2}+(z_{2i-1}-1)^{2}\right]
  
      This benchmark can also be used to test Bayesian inference methods. 
      Specifically this benchmarks returns the log likelihood
      
      .. math:: l(z) = -f(z)
  
      which can be used to compute the posterior distribution
      
      .. math:: \pi_{\text{post}}(\rv)=\frac{\pi(\V{y}|\rv)\pi(\rv)}{\int_{\rvdom} \pi(\V{y}|\rv)\pi(\rv)d\rv}
  
      where the prior is the tensor product of :math:`d` independent and 
      identically distributed uniform variables on :math:`[-2,2]`, i.e. 
      :math:`\pi(\rv)=\frac{1}{4^d}`, and the likelihood is given by
      
      .. math:: \pi(\V{y}|\rv)=\exp\left(l(\rv)\right)
  
      Parameters
      ----------
      nvars : integer
          The number of variables of the Rosenbrock function
  
      Returns
      -------
      benchmark : pya.Benchmark
         Object containing the benchmark attributes documented below
  
      fun : callable
  
          The rosenbrock with signature
  
          ``fun(z) -> np.ndarray``
  
          where ``z`` is a 2D np.ndarray with shape (nvars,nsamples) and the
          output is a 2D np.ndarray with shape (nsamples,1)
  
      jac : callable
          The jacobian of ``fun`` with signature
  
          ``jac(z) -> np.ndarray``
  
          where ``z`` is a 2D np.ndarray with shape (nvars,nsamples) and the
          output is a 2D np.ndarray with shape (nvars,1)
  
      hessp : callable
          Hessian of  ``fun`` times an arbitrary vector p with signature
      
          ``hessp(z, p) ->  ndarray shape (nvars,1)``
  
          where ``z`` is a 2D np.ndarray with shape (nvars,nsamples) and p is an 
          arbitraty vector with shape (nvars,1)
  
      variable : pya.IndependentMultivariateRandomVariable
          Object containing information of the joint density of the inputs z
          which is the tensor product of independent and identically distributed 
          uniform variables on :math:`[-2,2]`.
  
      mean : float
          The mean of the rosenbrock function with respect to the pdf of variable.
  
      loglike : callable
          The log likelihood of the Bayesian inference problem for inferring z
          given the uniform prior specified by variable and the negative 
          log likelihood given by the Rosenbrock function. loglike has the 
          signature
  
          ``loglike(z) -> np.ndarray``
  
          where ``z`` is a 2D np.ndarray with shape (nvars,nsamples) and the
          output is a 2D np.ndarray with shape (nsamples,1)
  
      loglike_grad : callable
          The gradient of the ``loglike`` with the signature
  
          ``loglike_grad(z) -> np.ndarray``
  
          where ``z`` is a 2D np.ndarray with shape (nvars,nsamples) and the
          output is a 2D np.ndarray with shape (nsamples,1)
  
      References
      ----------
      .. [DixonSzego1990] `Dixon, L. C. W.; Mills, D. J. "Effect of Rounding Errors on the Variable Metric Method". Journal of Optimization Theory and Applications. 80: 175–179. 1994 <https://doi.org/10.1007%2FBF02196600>`_
  
      Examples
      --------
      >>> from pyapprox.benchmarks.benchmarks import setup_benchmark
      >>> benchmark=setup_benchmark('rosenbrock',nvars=2)
      >>> print(benchmark.keys())
      dict_keys(['fun', 'jac', 'hessp', 'variable', 'mean', 'loglike', 'loglike_grad'])
      """
      univariate_variables = [stats.uniform(-2,4)]*nvars
      variable=pya.IndependentMultivariateRandomVariable(univariate_variables)
  
      benchmark = Benchmark(
          {'fun':rosenbrock_function,'jac':rosenbrock_function_jacobian,
           'hessp':rosenbrock_function_hessian_prod,'variable':variable,
           'mean':rosenbrock_function_mean(nvars)})
      benchmark.update({'loglike':lambda x: -benchmark['fun'](x),
                        'loglike_grad':lambda x: -benchmark['jac'](x)})
      return benchmark
  
  def setup_genz_function(nvars,test_name,coefficients=None):
      r"""
      Setup the Genz Benchmarks.
      
      For example, the two-dimensional oscillatory Genz problem can be defined 
      using
      
      >>> from pyapprox.benchmarks.benchmarks import setup_benchmark
      >>> benchmark=setup_benchmark('genz',nvars=2,test_name='oscillatory')
      >>> print(benchmark.keys())
      dict_keys(['fun', 'mean', 'variable'])
  
      Parameters
      ----------
      nvars : integer
          The number of variables of the Genz function
      
      test_name : string
          The test_name of the specific Genz function. See notes
          for options the string needed is given in brackets
          e.g. ('oscillatory')
  
      coefficients : tuple (ndarray (nvars), ndarray (nvars))
          The coefficients :math:`c_i` and :math:`w_i`
          If None (default) then 
          :math:`c_j = \hat{c}_j\left(\sum_{i=1}^d \hat{c}_i\right)^{-1}` where 
          :math:`\hat{c}_i=(10^{-15\left(\frac{i}{d}\right)^2)})`
  
      Returns
      -------
      benchmark : pya.Benchmark
         Object containing the benchmark attributes
  
      References
      ----------
      .. [Genz1984] `Genz, A. Testing multidimensional integration routines. In Proc. of international conference on Tools, methods and languages for scientific and engineering computation (pp. 81-94), 1984 <https://dl.acm.org/doi/10.5555/2837.2842>`_
  
      Notes
      -----
  
      Corner Peak ('corner-peak')
  
      .. math:: f(z)=\left( 1+\sum_{i=1}^d c_iz_i\right)^{-(d+1)}
  
      Oscillatory ('oscillatory')
  
      .. math:: f(z) = \cos\left(2\pi w_1 + \sum_{i=1}^d c_iz_i\right) 
  
      Gaussian Peak ('gaussian-peak')
  
      .. math:: f(z) = \exp\left( -\sum_{i=1}^d c_i^2(z_i-w_i)^2\right)
  
      Continuous ('continuous')
  
      .. math:: f(z) = \exp\left( -\sum_{i=1}^d c_i\lvert z_i-w_i\rvert\right)
      
      Product Peak ('product-peak')
  
      .. math:: f(z) = \prod_{i=1}^d \left(c_i^{-2}+(z_i-w_i)^2\right)^{-1}
  
      Discontinuous ('discontinuous')
  
      .. math:: f(z) = \begin{cases}0 & x_1>u_1 \;\mathrm{or}\; x_2>u_2\\\exp\left(\sum_{i=1}^d c_iz_i\right) & \mathrm{otherwise}\end{cases}
      
      """
      genz = GenzFunction(test_name,nvars)
      univariate_variables = [stats.uniform(0,1)]*nvars
      variable=pya.IndependentMultivariateRandomVariable(univariate_variables)
      if coefficients is None:
          genz.set_coefficients(1,'squared-exponential-decay',0)
      else:
          genz.c,genz.w = coefficients
      attributes = {'fun':genz,'mean':genz.integrate(),'variable':variable}
      if test_name=='corner-peak':
          attributes['variance']=genz.variance()
          from scipy.optimize import OptimizeResult
      return Benchmark(attributes)
  
  try:
      from pyapprox.fenics_models.advection_diffusion_wrappers import \
          setup_advection_diffusion_benchmark,\
          setup_advection_diffusion_source_inversion_benchmark,\
          setup_multi_level_advection_diffusion_benchmark
      from pyapprox.fenics_models.helmholtz_benchmarks import \
          setup_mfnets_helmholtz_benchmark
  except:
      pass
  
  def setup_benchmark(name,**kwargs):
      benchmarks = {'sobol_g':setup_sobol_g_function,
                    'ishigami':setup_ishigami_function,
                    'oakley':setup_oakley_function,
                    'rosenbrock':setup_rosenbrock_function,
                    'genz':setup_genz_function,
                    'cantilever_beam':setup_cantilever_beam_benchmark}
      try:
          # will fail if fenics is not installed and the import of the fenics
          # benchmarks fail
          fenics_benchmarks={
              'multi_index_advection_diffusion':setup_advection_diffusion_benchmark,
              'multi_index_advection_diffusion_source_inversion':setup_advection_diffusion_source_inversion_benchmark,
              'multi_level_advection_diffusion':setup_multi_level_advection_diffusion_benchmark,
              'mfnets_helmholtz':setup_mfnets_helmholtz_benchmark}
          benchmarks.update(fenics_benchmarks)
      except:
          pass
  
  
      if name not in benchmarks:
          msg = f'Benchmark "{name}" not found.\n Avaialble benchmarks are:\n'
          for key in benchmarks.keys():
              msg += f"\t{key}\n"
          raise Exception(msg)
  
      return benchmarks[name](**kwargs)
      
  
  def setup_cantilever_beam_benchmark():
      variable, design_variable = define_beam_random_variables()
      attributes = {'fun':cantilever_beam_objective,
                    'jac':cantilever_beam_objective_grad,
                    'constraint_fun':cantilever_beam_constraints,
                    'constraint_jac':cantilever_beam_constraints_jacobian,
                    'variable':variable,
                    'design_variable':design_variable,
                    'design_var_indices':np.array([4,5])}
      return Benchmark(attributes)
  