Пользовательское цветовое сопоставление на основе функции в matplotlib

Question

Я рисую диаграмму рассеяния, используя matplotlib в python. Я хочу, чтобы цвет точки на основе некоторой функции, например, так:

import matplotlib.pyplot as plt 

def color(x, y): 
    # based on some rules, return a color 
    if(condition): 
     return 'red' 
    else: 
     return 'blue' 

plt.scatter(index, data) #c= something? 

Я знаю функции matplotlib.from_levels_and_colors, но проблема в том, что отображение не основано на уровне значений на х или y. Существует третье значение, связанное с каждой точкой данных, которая вычисляется функцией, и именно это я хочу покрасить точки на основе.

Есть ли способ сделать это?

Answer 1

Почему бы вам просто не сделать свой массив c функцией индикатора для цветовой карты по умолчанию. Например:

import numpy as np 
import matplotlib.pyplot as plt 

x = np.arange(100) 
y = np.arange(100) 
# Colors whether or not x+y is a multiple of 5 
c = (x + y)%5 == 0 

# Use 'jet' colormap for red/blue. 
plt.scatter(x, y, c=c, cmap='jet') 

Конечно, вы могли бы использовать цветовую палитру с различными крайностями, которые будут получать отображенных в 0 и 1.

Если желаемый результат имеет более 2-х цветов, то это вполне можно передать как c массив с множеством разных значений (и его не нужно нормализовать). См. Пример here.

Answer 2

Так я и сделал. Я создал список цветов на основе значений x и y, а затем передал их функции рассеяния. Не так хорошо, как ответ wflynny, но это означает, что вы можете делать столько вычислений, сколько необходимо, чтобы придумать массив, вместо того, чтобы создавать одну функцию для этого.

import matplotlib.pyplot as plt 

colors = calculate_colors(x, y) 
plt.scatter(index, data, c=colors)