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:
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à:
I hope that it might help some users.
Regards,
Michaël