python：从ScalarMappable获取整个numpy数组的颜色答案

【问题标题】：python: get colors from ScalarMappable for entire numpy arraypython：从ScalarMappable获取整个numpy数组的颜色
【发布时间】：2020-02-04 13:32:46
【问题描述】：

我有一个 4D numpy 数组中包含大量值（x、y、z 中的数千个值数千次）。对于这些值中的每一个，我都需要来自 matplotlib.cm.ScalarMappable 对象的“颜色矢量”(RGBA)。

我发现循环遍历这样的数组变得相当慢，我想知道是否有一种方法可以通过采用不同的方法来显着加快它的速度。例如，是否可以将整个 numpy 数组（大于 2D）传递给 ScalarMappable，以便以更 numpythonic 或矢量化的方式进行此操作？

我的 3D 案例示例代码：

#!/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))

第三种方法（一次传递二维数组）显示出很大的性能提升，但我想知道这是否可以进一步推动。

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

【问题讨论】：

如果您的数据始终按 0-1 缩放，您可以预先计算每个 8 位或 16 位输入值的每种颜色，并将它们存储为查找表以供将来使用，这将非常快。但是，如果您的 vmin、vmax 基于输入数据是动态的，这将不起作用。

标签： python python-3.x numpy matplotlib color-mapping

【解决方案1】：

看着help ScalarMappable.to_rgba，

def to_rgba(self, x, alpha=None, bytes=False, norm=True):
在正常情况下，x 是一维或二维的标量序列，并且将返回相应的 rgba 值的 ndarray，基于为此 ScalarMappable 设置的规范和颜色图。有一种特殊情况，用于处理已经存在的图像 rgb 或 rgba，例如可能是从图像文件中读取的。如果 x 是一个 3 维的 ndarray ...

是

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,) )

更快？

（来源：ScalarMappable to_rgba。）

【讨论】：