python: получить цвета из ScalarMappable для всего массива numpy - PullRequest
Новая
       32

python: получить цвета из ScalarMappable для всего массива numpy

3 голосов
/ 04 февраля 2020

У меня есть большой массив значений, упакованный в массив 4D numpy (тысячи значений в x, y, z тысячи раз). Для каждого из этих значений мне нужен «цветовой вектор» (RGBA) от объекта matplotlib.cm.ScalarMappable.

Я обнаружил, что цикл через такой массив становится довольно медленным, и мне интересно, есть ли способ значительно ускорить его, приняв другой подход. Например, может ли весь массив numpy (больше, чем 2D) быть передан в ScalarMappable, чтобы эта операция выполнялась более numpythoni c или векторизованным способом?

Мой пример кода для трехмерного случая: 

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