python: get colors from ScalarMappable for entire numpy array

Question

I have a large array of values packed in a 4D numpy array (thousands of values in x,y,z for thousands of times). For each of these values I need 'color vector' (RGBA) from a matplotlib.cm.ScalarMappable object.

I've discovered that looping through such an array becomes rather slow, and I'm wondering if there's a way of significantly speeding it up by taking a different approach. For example, can an entire numpy array (greater than 2D) be passed to a ScalarMappable in order for this operation to take place in a more numpythonic or vectorized way?

My example code for a 3D case:

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

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

def get_colors_1(data,x,y,z):
    colors = np.zeros( (x,y,z,4), dtype=np.float16)
    for i in range(x):
        for j in range(y):
            for k in range(z):
                colors[i,j,k,:] = m.to_rgba(data[i,j,k])
    return colors

def get_colors_2(data,x,y,z):
    colors = np.array([[[m.to_rgba(data[i,j,k]) for k in range(z)] for j in range(y)] for i in range(x)], dtype=np.float16)
    return colors

def get_colors_3(data,x,y,z):
    colors = np.zeros((x,y,z,4), dtype=np.float16)
    for i in range(x):
        colors[i,:,:,:] = m.to_rgba(data[i,:,:])
    return colors

x, y, z = 30, 20, 10
data = np.random.rand(x,y,z)
cmap = matplotlib.cm.get_cmap('jet')
norm = matplotlib.colors.PowerNorm(vmin=0.0, vmax=1.0, gamma=2.5)
m = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap)

start_time = timeit.default_timer()
colors = get_colors_1(data,x,y,z)
elapsed = timeit.default_timer() - start_time
print('time elapsed: '+str(elapsed))

start_time = timeit.default_timer()
colors = get_colors_2(data,x,y,z)
elapsed = timeit.default_timer() - start_time
print('time elapsed: '+str(elapsed))

start_time = timeit.default_timer()
colors = get_colors_3(data,x,y,z)
elapsed = timeit.default_timer() - start_time
print('time elapsed: '+str(elapsed))

The third method (passing 2D arrays at time) shows a big performance bump, but I'm wondering if this can be pushed significantly further.

time elapsed: 0.5877857000014046
time elapsed: 0.5911024999986694
time elapsed: 0.004590500000631437

Answer 1

Looking at help ScalarMappable.to_rgba,

def to_rgba(self, x, alpha=None, bytes=False, norm=True):
In the normal case, x is a 1D or 2D sequence of scalars, and the corresponding ndarray of rgba values will be returned, based on the norm and colormap set for this ScalarMappable. There is one special case, for handling images that are already rgb or rgba, such as might have been read from an image file. If x is an ndarray with 3 dimensions ...

Is

sm = ScalarMappable( norm=None, cmap=cmap )
flat = data.reshape( -1 )  # a view
rgba = sm.to_rgba( flat, bytes=True, norm=False ).reshape( data.shape + (4,) )

much faster ?

(Source: ScalarMappable to_rgba .)