Create a colorbar from an RGBA function in matplotlib

matplotlib, python

Solution

To create a colour map, you have to specify how the red/green/blue components change across a linear scale. It looks like you already have a function, `f`, that sets up the r/g/b components for you. The hard part is the 4th channel, the alpha channel. I will go through setting the alpha channel given a RGB colour map specified by your `f`.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

# some data
a = np.sort(np.random.randn(10, 10))

# use the default 'jet' colour map for showing the difference later
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.imshow(a, cmap=cm.get_cmap('jet'))
fig.savefig('map1.png')

# let's use jet and modify the alpha channel
# you would use your own colour map specified by f
my_cmap = cm.get_cmap('jet')

# this is a hack to get at the _lut array, which stores RGBA vals
my_cmap._init()

# use some made-up alphas, you would use the ones specified by f
alphas = np.abs(np.linspace(-1.0, 1.0, my_cmap.N))

# overwrite the alpha channel of the jet colour map
my_cmap._lut[:-3,-1] = alphas

# plot data with our modified colour map
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.imshow(a, cmap=my_cmap)
fig.savefig('map2.png')

Here's `map1.png`:

And here's `map2.png`:

Hope this helps.

Problem

I'm plotting an image via `imshow` with the input of a MxNx4 array `A`, a RGBA array defined for a rectangular MxN grid. This coloring was generated from `V`, a MxN array that indicates a scalar value for each one of these points. I.e. I have a function `f` that takes a scalar value and returns a RGBA tuple: `f(V) = A` I want to make a colorbar that takes as input `f,V`. Is this possible?

Original source