Contour plots using Matplotlib

Hello!
I sometimes have to create contour plots e.g. when performing optimization using OpenTURNS. If g is a Function with 2 inputs and 1 output, the draw() method of any function can plot contours (i.e. iso-values plots) of the function. This is shown e.g. in the next example:

https://openturns.github.io/openturns/latest/auto_graphs/plot_graphs_loglikelihood_contour.html

For example, the iso-values of Rosenbrock’s function are:

image

This is straightforward to produce:

import openturns as ot
rosenbrock = ot.SymbolicFunction(
    ["x1", "x2"],
    ["(1 - x1)^2 + 100 * (x2 - x1^2)^2"]
    )
graph = rosenbrock.draw([-2.0] * 2, [2.0] * 2, [100] * 2)

But we may sometimes want to fill the plot with colors so that the value of the function is better displayed. The contourf function from Matplotlib does the job, but we need to provide the required input arguments: the goal of this topic is to share the function I came up with, which involves both OpenTURNS and Numpy. The function is mainly based on the numpy.meshgrid function, along with appropriate flatten, transpose and reshape function calls. Its goal is to be fast and avoid a double for loop. To achieve this goal, I provide the appropriate Sample to the function, so that all function values are produced in a single call.

def ComputeMeshGrid2DFunctionValue(
    function, interval, numberOfXCells=100, numberOfYCells=100
):
    """
    Compute the value of the function on a mesh. 
    
    This is handy for contour plots.

    Parameters
    ----------
    function : ot.Function(2, 1)
        The function, which input has dimension 2 and output has dimension 1.
    interval : ot.Interval(2)
        A 2D interval.
    numberOfXCells : int, optional
        The number of cells on the X axis. The default is 100.
    numberOfYCells : int, optional
        The number of cells on the Y axis. The default is 100.

    Raises
    ------
    ValueError
        If the input dimension is not equal to 2 or the output dimension 
        is not equal to 1.

    Returns
    -------
    xArray : np.array(numberOfYCells, numberOfXCells)
        The input X array.
    yArray : np.array(numberOfYCells, numberOfXCells)
        The input Y array.
    zArray : np.array(numberOfYCells, numberOfXCells)
        The output array.

    """
    if interval.getDimension() != 2:
        raise ValueError(
            "Dimension of interval is equal "
            "to %d, which is not equal to 2" % (interval.getDimension())
        )
    if function.getInputDimension() != 2:
        raise ValueError(
            "Input dimension of function is equal "
            "to %d, which is not equal to 2" % (function.getInputDimension())
        )
    if function.getOutputDimension() != 1:
        raise ValueError(
            "Output dimension of function is equal "
            "to %d, which is not equal to 1" % (function.getOutputDimension())
        )
    # Create the input mesh
    lowerBound = interval.getLowerBound()
    upperBound = interval.getUpperBound()
    xList = np.linspace(lowerBound[0], upperBound[0], numberOfXCells)
    yList = np.linspace(lowerBound[1], upperBound[1], numberOfYCells)
    xArray, yArray = np.meshgrid(xList, yList)
    # Convert to Sample
    xArrayFlat = xArray.flatten()
    yArrayFlat = yArray.flatten()
    size = numberOfXCells * numberOfYCells
    inputSample = ot.Sample(size, 2)
    inputSample[:, 0] = ot.Sample.BuildFromPoint(xArrayFlat)
    inputSample[:, 1] = ot.Sample.BuildFromPoint(yArrayFlat)
    # Evaluate the function
    zArray = function(inputSample)
    zArray = np.array(zArray)
    zArray = zArray.transpose()
    zArray = np.reshape(zArray, (numberOfYCells, numberOfXCells))
    return xArray, yArray, zArray

The next script plots the contour. It first computes appropriate levels corresponding to the contour plot. This is done using quantiles of the function. Then I use the reversed Viridis color map to create the contour plot using contourf. One of the tricky points is to set up the vmin and vmax parameters, which sets the threshold down and up to which the colors are truncated.

import matplotlib.pyplot as plt
import openturns as ot
import numpy as np
# Compute value of standard deviation on a grid
interval = ot.Interval([-2.0] * 2, [2.0] * 2)
xArray, yArray, zArray = ComputeMeshGrid2DFunctionValue(
    rosenbrock, interval, numberOfXCells=100, numberOfYCells=100
)
# Compute levels
numberOfLevels = 15
zSample = ot.Sample.BuildFromPoint(zArray.flatten())
epsilon_grid = 5.e-2
regularGrid = np.linspace(epsilon_grid, 1.0 - epsilon_grid, numberOfLevels)
levelsSample = zSample.computeQuantile(regularGrid)
levels = [point[0] for point in levelsSample]  # Convert into a list
# Plot with customized levels
_ = plt.figure()
fig, ax = plt.subplots(1, 1)
cmap = mpl.cm.viridis_r
cp = ax.contourf(
    xArray, yArray, zArray, levels, cmap=cmap, 
    vmin = 10.0, vmax=500.0
    )
_ = fig.colorbar(cp)
_ = ax.set_title("Rosenbrock")
_ = ax.set_xlabel(r"$x_1$")
_ = ax.set_ylabel(r"$x_2$")

Et voilà:

image

I hope that it might help some users.

Regards,

Michaël

2 Likes

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

1 Like

Hi Elias,
I was asked :slight_smile: to see if OpenTURNS was able to produce the same plot as in Matplotlib’s tricontourf. This functions takes a sample of x, y and z, triangulate the input domain (x, y) and plot a colored contour of the data.
Surely, OpenTURNS cannot produce filled contour plots. Anyway, I wanted to see what can be done with DatabaseFunction. This function takes an input / output sample pair and defines a function. Given any new input (x,y), the function searches for the closest input data point and returns the corresponding database output. This is done using DatabaseEvaluation. In turn, this class uses NearestNeighbourAlgorithm to solve the optimization problem. Depending on the dimension, efficient algorithms are provided.

Using the same example as in Tricontourf, I came up with the following script. It produces an input sample with size 100 uniform in the cube [-3, 3] \times [-3, 3]. Then we use the draw() method of this function. Finally, we add the Cloud of points on the same graph.

import openturns as ot
import openturns.viewer as otv
import matplotlib.pyplot as plt
import matplotlib as mpl
import pylab as pl
import numpy as np

func = ot.SymbolicFunction(
    ['x', 'y'],
    ['(1 - x / 2 + x^5 + y^3) * exp(-x^2 - y^2)']
)

dimension = 2
distribution = ot.ComposedDistribution([ot.Uniform(-3.0, 3.0)] * dimension)
sampleSize = 100
inputSample = distribution.getSample(sampleSize)
outputSample = func(inputSample)
database = ot.DatabaseFunction(inputSample, outputSample)
numberOfPointsPerDimension = 50
graph = database.draw([-3.0] * dimension, [3.0] * dimension, [numberOfPointsPerDimension] * dimension)
graph.add(ot.Cloud(inputSample))
graph.setTitle("Database Function")
graph.setXTitle("$x_1$")
graph.setYTitle("$x_2$")
view = otv.View(graph, figure_kw={"figsize": (6.0, 4.0)}, legend_kw={"bbox_to_anchor":(1.0, 1.0), "loc":"upper left"})
plt.subplots_adjust(right = 0.7)

This produces the next figure.

databaseContour-basic
Figure 1. Contour plot of a DatabaseFunction using OpenTURNS’s draw() method.

We can see that, on the one hand, where there are many points, the levels are lines with a locally high curvature. On the other hand, where there is no point, the levels are almost locally straight lines. These are effects of the search for the closest point.

I then wanted to see what produces Matplotlib’s contour on this function.

# Compute value of standard deviation on a grid
interval = ot.Interval([-3.0] * 2, [3.0] * 2)
xArray, yArray, zArray = ComputeMeshGrid2DFunctionValue(
    database, interval, numberOfXCells=100, numberOfYCells=100
)
# Compute levels
numberOfLevels = 10
zSample = ot.Sample.BuildFromPoint(zArray.flatten())
epsilon_grid = 1.e-3
regularGrid = np.linspace(epsilon_grid, 1.0 - epsilon_grid, numberOfLevels)
levelsSample = zSample.computeQuantile(regularGrid)
levels = [point[0] for point in levelsSample] # Convert into a list
# Plot with customized levels
fig, ax = plt.subplots(1, 1)
cmap = mpl.cm.viridis_r
cp = ax.contourf(xArray, yArray, zArray, levels, cmap=cmap, vmin = -0.4, vmax = 0.7)
_ = fig.colorbar(cp)
_ = ax.set_title("Database Function")
_ = ax.set_xlabel(r"$x_1$")
_ = ax.set_ylabel(r"$x_2$")
pl.plot(inputSample[:, 0], inputSample[:, 1], "o")

databaseContour-filled
Figure 2. Contour plot of a DatabaseFunction using Matplotlib’s contourf() function.

This plot shows perhaps more clearly that each area of constant function value has limits defined as a polygon in the neighborhood of a single isolated input point.

Best regards,

Michaël

1 Like