matplotlib - Draw a heatmap/pixelmap with ability to edit individual pixel colours (different colormaps by row)

Question

I'm trying to draw a heat map/pixelmap representation of a matrix using matplotlib. I currently have the following code which gives me the pixelmap as required (adapted from Heatmap in matplotlib with pcolor?):

import matplotlib.pyplot as plt
import numpy as np

column_labels = list('ABCD')
row_labels = list('0123')

data = np.array([[0,1,2,0],
                     [1,0,1,1],
                     [1,2,0,0],
                     [0,0,0,1]])

fig, ax = plt.subplots()
heatmap = ax.pcolor(data, cmap=plt.cm.Blues)

# put the major ticks at the middle of each cell
ax.set_xticks(np.arange(data.shape[0])+0.5, minor=False)
ax.set_yticks(np.arange(data.shape[1])+0.5, minor=False)

# want a more natural, table-like display
ax.invert_yaxis()
ax.xaxis.tick_top()

ax.set_xticklabels(row_labels, minor=False)
ax.set_yticklabels(column_labels, minor=False)

ax.yaxis.grid(True, which='minor', linestyle='-', color='k', linewidth = 0.3, alpha = 0.5)
ax.xaxis.grid(True, which='minor', linestyle='-', color='k', linewidth = 0.3, alpha = 0.5)

# Set the location of the minor ticks to the edge of pixels for the x grid
minor_locator = AutoMinorLocator(2)
ax.xaxis.set_minor_locator(minor_locator)

# Lets turn off the actual minor tick marks though
for tickmark in ax.xaxis.get_minor_ticks():
    tickmark.tick1On = tickmark.tick2On = False

# Set the location of the minor ticks to the edge of pixels for the y grid
minor_locator = AutoMinorLocator(2)
ax.yaxis.set_minor_locator(minor_locator)

# Lets turn off the actual minor tick marks though
for tickmark in ax.yaxis.get_minor_ticks():
    tickmark.tick1On = tickmark.tick2On = False

plt.show()

Which gives the following plot:

However I would like to extend this such that on mouse click I can highlight a 'row' in the pixelmap in green, e.g. if the user selected row 'C' I would have (I appreciate the green highlight is not clear for pixels with a 0 value):

I know how to deal with the mouse events but I'm not sure how to modify the colour of a single row in the pixelmap. It would also help if I could set labels for individual pixels of the pixel map to be retrieved on mouse click, as opposed to using the mouse x/y location to index the label lists.

score 0 · Accepted Answer · edited May 23 '17 at 12:07

I have figured out my own problem, with help from this question: Plotting of 2D data : heatmap with different colormaps.

The code is below and the comments should explain the steps taken clearly.

import matplotlib.pyplot as plt
import numpy as np
from numpy.ma import masked_array
import matplotlib.cm as cm
from matplotlib.ticker import AutoMinorLocator

column_labels = list('ABCD')
row_labels = list('0123')
data = np.array([[0,1,2,0],
                 [1,0,1,1],
                 [1,2,0,0],
                 [0,0,0,1]])

fig, ax = plt.subplots()

# List to keep track of handles for each pixel row
pixelrows = []

# Lets create a normalizer for the whole data array
norm = plt.Normalize(vmin = np.min(data), vmax = np.max(data))

# Let's loop through and plot each pixel row
for i, row in enumerate(data):
    # First create a mask to ignore all others rows than the current
    zerosarray = np.ones_like(data)
    zerosarray[i, :] = 0

    plotarray = masked_array(data, mask=zerosarray)

    # If we are not on the 3rd row down let's use the red colormap
    if i != 2:
        pixelrows.append(ax.matshow(plotarray, norm=norm, cmap=cm.Reds))

    # Otherwise if we are at the 3rd row use the green colormap
    else:
        pixelrows.append(ax.matshow(plotarray, norm=norm, cmap=cm.Greens))

# put the major ticks at the middle of each cell
ax.set_xticks(np.arange(data.shape[0]), minor=False)
ax.set_yticks(np.arange(data.shape[1]), minor=False)

# want a more natural, table-like display
ax.xaxis.tick_top()

ax.set_xticklabels(row_labels, minor=False)
ax.set_yticklabels(column_labels, minor=False)

ax.yaxis.grid(True, which='minor', linestyle='-', color='k', linewidth = 0.3, alpha = 0.5)
ax.xaxis.grid(True, which='minor', linestyle='-', color='k', linewidth = 0.3, alpha = 0.5)

# Set the location of the minor ticks to the edge of pixels for the x grid
minor_locator = AutoMinorLocator(2)
ax.xaxis.set_minor_locator(minor_locator)

# Lets turn of the actual minor tick marks though
for tickmark in ax.xaxis.get_minor_ticks():
    tickmark.tick1On = tickmark.tick2On = False

# Set the location of the minor ticks to the edge of pixels for the y grid
minor_locator = AutoMinorLocator(2)
ax.yaxis.set_minor_locator(minor_locator)

# Lets turn of the actual minor tick marks though
for tickmark in ax.yaxis.get_minor_ticks():
    tickmark.tick1On = tickmark.tick2On = False

plt.show()

matplotlib - Draw a heatmap/pixelmap with ability to edit individual pixel colours (different colormaps by row)

1 Answers1