Pylab: сопоставьте метки с цветами

Как бы то ни было, в этом случае вы обычно делаете что-то вроде этого:

import numpy as np
import matplotlib.pyplot as plt

iris = np.recfromcsv('iris.csv')
names = set(iris['class'])

x,y = iris['sepal_length'],  iris['sepal_width']

for name in names:
    cond = iris['class'] == name
    plt.plot(x[cond], y[cond], linestyle='none', marker='o', label=name)

plt.legend(numpoints=1)
plt.show()

В том, что предложил @Yann, нет ничего плохого, но scatter лучше подходит для непрерывных данных.

Легче полагаться на цветовой цикл осей и просто вызывать график несколько раз (вы также получаете отдельных художников вместо коллекции, что хорошо для таких дискретных данных).

По умолчанию цветовой цикл для осей: синий, зеленый, красный, голубой, пурпурный, желтый, черный.

После 7 вызовов plot он будет возвращаться к этим цветам, поэтому, если у вас есть больше элементов, вам нужно задайте его вручную (или просто укажите цвет в каждом вызове plot, используя интерполированную цветовую полосу, аналогичную той, что предложил @Yann выше).

Является ли способ элегантным или нет, несколько субъективно. Я лично нахожу ваши подходы лучше, чем способ «matplotlib». Из модуля matplotlib color:

Отображение цветов обычно включает два этапа: массив данных сначала отображается в диапазоне 0-1 с использованием экземпляра Normalize или подкласса; затем это число в диапазоне 0-1 сопоставляется с цветом с использованием экземпляра подкласса Colormap.

Что я понимаю из этого в отношении вашей проблемы, так это то, что вам нужен подкласс Normalize, который принимает строки и отображает их в 0-1.

Вот пример, который наследуется от Normalize для создания подкласса TextNorm, который используется для преобразования строки в значение от 0 до 1. Эта нормализация используется для получения соответствующего цвета.

import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
import numpy as np
from numpy import ma

class TextNorm(Normalize):
    '''Map a list of text values to the float range 0-1'''

    def __init__(self, textvals, clip=False):
        self.clip = clip
        # if you want, clean text here, for duplicate, sorting, etc
        ltextvals = set(textvals)
        self.N = len(ltextvals)
        self.textmap = dict(
            [(text, float(i)/(self.N-1)) for i, text in enumerate(ltextvals)])
        self.vmin = 0
        self.vmax = 1

    def __call__(self, x, clip=None):
        #Normally this would have a lot more to do with masking
        ret = ma.asarray([self.textmap.get(xkey, -1) for xkey in x])
        return ret

    def inverse(self, value):
        return ValueError("TextNorm is not invertible")

iris = np.recfromcsv("iris.csv")
norm = TextNorm(iris.field(4))

plt.scatter(iris.field(0), iris.field(1), c=norm(iris.field(4)), cmap='RdYlGn')
plt.savefig('textvals.png')
plt.show()

Это производит:

Я выбрал цветовую карту «RdYlGn», чтобы можно было легко различать три типа точек. Я не включил функцию clip как часть __call__, хотя это возможно с некоторыми изменениями.

Традиционно вы можете проверить нормализацию метода scatter, используя ключевое слово norm, но scatter проверяет ключевое слово c, чтобы увидеть, хранит ли оно строки, и если да, то предполагается, что вы передаете цвета в качестве их строковых значений, например «Красный», «Синий» и т. д. Таким образом, вызов plt.scatter(iris.field(0), iris.field(1), c=iris.field(4), cmap='RdYlGn', norm=norm) невозможен. Вместо этого я просто использую TextNorm и «работаю» с iris.field(4), чтобы вернуть массив значений в диапазоне от 0 до 1.

Обратите внимание, что для строки, не входящей в список textvals, возвращается значение -1. Вот тут и пригодилась бы маскировка.