линия matplotlib с различным цветом в зависимости от другой переменной

Question

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

income = [5000, 5005, 5010, 6000, 6060, 6120, 7000]
job = [0, 0, 0, 1, 1, 1, 2]

Я пробовал что-то вроде:

plt.plot(income, c=job, cmap='RdBu')

, но это приводит к тому, что объект 'Line2D' не имеет свойство 'карта'. Я также пытался:

plt.scatter(range(0, len(income)), income, c=job, cmap='RdBu')

не дает строк, что тоже не идеально. Есть ли способ сделать фигуру, подобную приведенной ниже [созданной в Matlab] в Matplotlib?

Answer 1

Я думаю, что colormap полезна для непрерывных данных. Для дискретных лучше использовать дискретный список цветов. Таким образом, вы можете связать цвет с типом переменной:

Код:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors


income = np.array([5000, 5005, 5010, 6000, 6060, 6120, 7000])
x = np.arange(len(y))
job = np.array([0, 0, 0, 1, 1, 1, 2]).astype('int')

# iterate over zipped job and color
for j, c in zip(job, colors.TABLEAU_COLORS):
    plt.plot(x[job == j], income[job == j], 'o-', c=c)

plt.show()

Сюжет:

Я использовал TABLEAU_COLORS, но вы можете найти другой список цветов здесь , если вы sh.