How can I draw a line around the edge of the mask?

Question

I have making a heatmap with a mask. Here is a toy MWE:

import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np

images = []
vmin = 0
vmax = 80
cmap = "viridis"
size = 40
matrix = np.random.randint(vmin, vmax, size=(size,size))
np.random.seed(7)
mask = []
for _ in range(size):
    prefix_length = np.random.randint(size)
    mask.append([False]*prefix_length + [True]*(size-prefix_length))
mask = np.array(mask)
sns.heatmap(matrix, vmin=vmin, vmax=vmax, cmap="viridis", mask=mask)
plt.savefig("temp.png")
plt.show()

I want to draw a line around the edge of the mask to accentuate where it is. How can you do that?

My toy example currently looks like this:

Answer 1

Example that draws a line around the edge

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from skimage import measure

size = 10
matrix = np.random.randint(0, 80, size=(size, size))
np.random.seed(7)

mask = []
for _ in range(size):
    prefix_length = np.random.randint(size)
    mask.append([False] * prefix_length + [True] * (size - prefix_length))
mask = np.array(mask)

sns.heatmap(matrix, vmin=0, vmax=80, cmap="viridis", mask=mask)

# Find and shift the contour of the masked region
contour = measure.find_contours(image=mask, level=0.5)[0]
shift = 0.5
contour[:, 1] += shift
contour[:, 0] += shift

# Find and add additional points to improve the contour plot
points = []
for idx0 in range(len(contour)):
    pos0 = contour[idx0]
    idx1 = idx0 + 1
    if len(contour) > idx1:
        pos1 = contour[idx1]
        y1, x1 = tuple(pos1)
        y0, x0 = tuple(pos0)
        if y1 - y0 and x1 - x0:
            point = np.array([y0, x1])
            if y0 % 1:
                point = np.array([y1, x0])
            points.append((idx0, point))
points.reverse()
for idx0, point in points:
    contour = np.insert(contour, idx0 + 1, point, axis=0)

# Plot the contour
plt.plot(contour[:, 1], contour[:, 0], "r", linewidth=1.5)
plt.show()

Answer 2

• sns.heatmap is categorical, the center of each row or column of squares, corresponds to the tick location, and the edges are ± 0.5.
• The simplest option is to use .argmax to find the index location of the first instance of True in mask.
  • Given these locations, use .vlines, and .hlines, which accept a sequence of values, to draw all of the lines.
    • See How to draw vertical lines on a given plot and Plot a horizontal line on a given plot for further details and examples with these methods.
• It should be noted, this probably won't work if there are nonconsecutive True and False values in mask.

fig, ax = plt.subplots(figsize=(10, 8))
sns.heatmap(matrix, vmin=vmin, vmax=vmax, cmap="viridis", mask=mask, ax=ax)

# get the ytick locations, which corresponds to the middle of each row
y = ax.get_yticks()

# get the edges of the squares
y_min = y - 0.5
y_max = y + 0.5

# find the index of the first instance of True
row_max = mask.argmax(axis=1)

# plot the vertical lines
ax.vlines(x=row_max, ymin=y_min, ymax=y_max, colors='r', lw=1.5)

# plot the horizontal lines
ax.hlines(y=y_max[:-1], xmin=row_max[:-1], xmax=row_max[1:], colors='r', lw=1.5)

plt.show()

Answer 3

If your goal is to highlight an area, how about a fill approach using mask inversion? My approach is to create an array of colors that you want to fill and specify them in the colormap to fill. Then you can hide the color bar.

import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np

images = []
vmin = 0
vmax = 80
cmap = "viridis"
size = 40
matrix = np.random.randint(vmin, vmax, size=(size,size))
np.random.seed(7)
mask = []
for _ in range(size):
    prefix_length = np.random.randint(size)
    mask.append([False]*prefix_length + [True]*(size-prefix_length))
mask = np.array(mask)
sns.heatmap(matrix, vmin=vmin, vmax=vmax, cmap="viridis", mask=mask)
colors = ['ivory']*size*size
sns.heatmap(matrix, vmin=vmin, vmax=vmax, cmap=colors, mask=~mask, cbar=False)

plt.savefig("temp.png")
plt.show()