Сопоставить логические значения со строками

Question

Я строю график, где моя переменная x - «Mg», а моя переменная y - «Si». У меня есть третья переменная, называемая «двоичная». Если двоичный файл равен 0 или 1, как я могу покрасить нанесенную точку в красный или черный соответственно?

Мне нужно использовать функции plt.scatter и colourbar (). Я читал о colourbar, но он, кажется, генерирует непрерывный спектр цвета. Вместо этого я попытался использовать plt.colors.from_levels_and_colors, но я не совсем уверен, как правильно его использовать.

levels = [0,1]
colors = ['r','b']
cmap, norm = plt.colors.from_levels_and_colors(levels, colors)
plt.scatter(data_train['Mg'], data_train['Si'], c = data_train['binary'])
plt.show()

Кроме того, в будущем, вместо того, чтобы задавать такой вопрос на этом форуме, что я могу сделать, чтобы решить проблему самостоятельно? Сначала я пытаюсь прочитать документацию онлайн, но часто мне трудно это понять.

Answer 1

Если вы работаете с несколькими «количественными» цветами, а не с картой цветов, вам, вероятно, следует изменить c с двоичного формата на mpl-friedly. * Т.е. 1002 *

point_colors = [colors[binary] for binary in data_train['binary']]
plt.scatter(data_train['Mg'], data_train['Si'], c=point_colors)

Answer 2

np.where упрощает кодирование двоичных значений.

np.where([1, 0, 0, 1], 'yes', 'no')
# array(['yes', 'no', 'no', 'yes'], dtype='<U3')

---

colors = np.where(data_train['binary'], 'black', 'red')
plt.scatter(data_train['Mg'], data_train['Si'], c=colors)