Contour plots using Matplotlib

[efekhari27]

Hi Michaël,

Thanks for this script, the levels computation is very handy in this case!
When dealing propagating uncertainties (i.e., central tendency estimation, rare even estimation) it can be useful to plot both a function and a joint distribution. Here is a small class that I wrote starting from your code:

import numpy as np
from matplotlib import cm
import matplotlib.pyplot as plt
import openturns as ot 

class Draw2DFunctions:
    def __init__(self, lowerbound, upperbound):
        self.grid_size = 100
        mesher = ot.IntervalMesher([self.grid_size-1] * 2)
        interval = ot.Interval(lowerbound, upperbound)
        mesh = mesher.build(interval)
        self.nodes = mesh.getVertices()
        self.X0, self.X1 = np.array(self.nodes).T.reshape(2, self.grid_size, self.grid_size)

    def draw_2D_contour(self, title, function=None, distribution=None, colorbar=cm.coolwarm, nb_isocurves=8, contour_values=True, opt_function_levels=True):
        fig = plt.figure(figsize=(5, 4))
        if distribution is not None:
            Zpdf = np.array(distribution.computePDF(self.nodes)).reshape(self.grid_size, self.grid_size)
            contours = plt.contour(self.X0, self.X1, Zpdf, nb_isocurves, colors='black', alpha=0.6)
            if contour_values:
                plt.clabel(contours, inline=True, fontsize=8)
        if function is not None:
            Z = np.array(function(self.nodes)).reshape(self.grid_size, self.grid_size)
            if opt_function_levels: 
                Z_sample = ot.Sample.BuildFromPoint(Z.flatten())
                # Compute levels
                nb_levels = 10
                epsilon_grid = 5e-2
                regular_grid = np.linspace(epsilon_grid, 1.0 - epsilon_grid, nb_levels)
                computed_levels = np.array(Z_sample.computeQuantile(regular_grid)).flatten()
            else: 
                computed_levels = 20
            plt.contourf(self.X0, self.X1, Z, levels=computed_levels, cmap=colorbar)
            plt.colorbar()
        plt.title(title, fontsize=20)
        plt.xlabel("$x_0$", fontsize=20)
        plt.ylabel("$x_1$", fontsize=20)
        return fig
#
rosenbrock = ot.SymbolicFunction(["x1", "x2"], ["(1 - x1)^2 + 100 * (x2 - x1^2)^2"])
funky_distribution = ot.ComposedDistribution([ot.Normal(0., 0.7)] * 2, ot.ClaytonCopula(2.))
d = Draw2DFunctions([-2.] * 2, [2.] * 2)
d.draw_2D_contour("Rosenbrock function", function=rosenbrock, distribution=funky_distribution, colorbar=cm.plasma)
plt.show()

Best,
Elias