[FIXED] How to draw a two-dimensional function whose domain is a triangle?

Issue

I would like to plot a function f(x,y) = x^2 + y^2 defind on all pairs x,y such that x + y ≤ 1, x ≥ 0, y ≥ 0. Currently, I use imshow:

x = np.arange(0, 1, 0.01)
y = np.arange(0, 1, 0.01)
Z = compute_function(x, y) # evaluation of the function on the grid
imshow(Z, cmap=cm.RdBu, extent=(min(x),max(x),min(y), max(y)), origin='lower')

It looks like this, which is quite ugly:

How can I make this plot nicer, showing only the relevant triangle (at the bottom-left)?

P.S. If possible, I would like to present the relevant triangle, not as a right-angled triangle, but as a symmetric, equilateral triangle

Solution

First of all, as the name suggests, plt.imshow is meant for displaying images. For cases otherwise, you should probably use plt.pcolormesh. Regarding the masking, you can create a numpy masked array and pass that as the C array for pcolormesh and it will only plot the non-masked region.

import numpy as np
import matplotlib.pyplot as plt

plt.close("all")


x = np.arange(0, 1, 0.01)
y = np.arange(0, 1, 0.01)
X, Y = np.meshgrid(x, y)
Z = X**2 + Y**2

cond1 = X <= 0
cond2 = Y <= 0
cond3 = X + Y >= 1
cond = cond1 | cond2 | cond3
Z_masked = np.ma.array(Z, mask=cond)

fig, ax = plt.subplots()
ax.pcolormesh(X, Y, Z_masked, cmap=plt.cm.RdBu)
fig.show()

Right Triangle

In the case of the right triangle domain, we can get a sharper edge if we use plt.tripcolor and generate our triangle indices appropriately (i.e. break each square into two triangles and make sure the cut is in the right direction to match the angle of the domain cut).

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import tri

plt.close("all")

def create_triangle_indices_from_meshgrid(n, m, orientation=1):
    """Returns the indices into the flattened meshgrid which will create 
    triangles for `matplotlib.tri.Trangulation`. Provides two triangle
    orientations. One may be better than the other depending on the specific
    case.
    """
    indices = np.arange(n*m).reshape(n, m)
    p1 = indices[1:, :-1].flatten()
    p2 = indices[:-1, :-1].flatten()
    p3 = indices[:-1, 1:].flatten()
    p4 = indices[1:, 1:].flatten()

    if orientation == 1:
        triangles1 = np.stack((p1, p2, p3), axis=1)
        triangles2 = np.stack((p1, p3, p4), axis=1)
        triangles = np.vstack((triangles1, triangles2))
    elif orientation == 2:
        triangles1 = np.stack((p2, p3, p4), axis=1)
        triangles2 = np.stack((p1, p2, p4), axis=1)
        triangles = np.vstack((triangles1, triangles2))
    else:
        raise ValueError(f"Orientation number {orientation} does not exist.")

    return triangles

x = np.arange(0, 1, 0.01)
y = np.arange(0, 1, 0.01)
X, Y = np.meshgrid(x, y)
Z = X**2 + Y**2

triangle_indices = create_triangle_indices_from_meshgrid(x.size, y.size, 1)
triangles = tri.Triangulation(X.flat, Y.flat, triangles=triangle_indices)
X_triangles = X.flatten()[triangle_indices]
Y_triangles = Y.flatten()[triangle_indices]

cond1_tri = (X_triangles <= 0).any(1)
cond2_tri = (Y_triangles <= 0).any(1)
cond3_tri = (X_triangles + Y_triangles >= 1).any(1)
cond_tri = cond1_tri | cond2_tri | cond3_tri
triangles.mask = cond_tri

fig, ax = plt.subplots()
ax.tripcolor(triangles, Z.flat, cmap=plt.cm.RdBu)
fig.show()

Equilateral Triangle/Arbitrary Shape

For an equilateral triangle or an arbitrary shape, you can simply generate a fine grid of points, select only the points within your inequality, and use plt.tripcolor to plot the result. Note: for the right triangle the tri.Triangulation was used and a mask was set for the triangles. That mask masked out values, while in this case we want to select which values to keep, hence the switch in inequality direction and the & when combining conditions.

import numpy as np
import matplotlib.pyplot as plt

plt.close("all")

x = np.arange(-1, 1, 0.005)
y = np.arange(0, 1, 0.005)
X, Y = np.meshgrid(x, y)

X_flat = X.flatten()
Y_flat = Y.flatten()

cond1 = (Y_flat - np.tan(np.pi/3)*X_flat - 1 <= 0)
cond2 = (Y_flat + np.tan(np.pi/3)*X_flat - 1 <= 0)
cond3 = (Y_flat >= 0)
cond = cond1 & cond2 & cond3

X_tri = X_flat[cond]
Y_tri = Y_flat[cond]
Z_tri = X_tri**2 + Y_tri**2

fig, ax = plt.subplots()
ax.tripcolor(X_tri, Y_tri, Z_tri, cmap=plt.cm.RdBu)
ax.set_aspect(1)
fig.show()

This method is improved by a finer meshgrid or by choosing better points (i.e. choosing points more strategically).

Answered By - jared