Matplotlib, plot x1 and y, but color code based on x2 in python?

Question

I'm trying to create a plot of x1 against y and make the points coloured based on another variable x2.

x1 = times
y = wait
x2 = triage (from 1-5)

Every variable have the same contains data and have all the same length. The special about x2(triage), is that is numbered from 1-5 and want to colour based on those, so 1 is red, 2 is orange, 3 is yellow, 4 is green and 5 is blue.

Code UPDATED

X = dataset.iloc[:, 0:3].values
y = dataset.iloc[:, 3].values

triage = X[:, 0]
week = X[:, 1]
times = X[:, 2]
wait = y

df = pd.DataFrame(dict(times=times, wait=wait, triage=triage))
fig, ax = plt.subplots()
colors = {'1':'red', '2':'orange', '3':'yellow', '4':'green', '5':'blue'}

grouped = df.groupby('triage')
for key, group in grouped:
    group.plot(ax=ax, kind='scatter', x='times', y='wait', label=key,

triage=colors[key])

plt.show()

Managed to try something, but still does not work quite, anyone can see what could be wrong here?

Did you see my solution below? – cvanelteren Dec 23 '18 at 00:18 — cvanelteren, Dec 23 '18 at 00:18

cvanelteren · Answer 1 · 2018-12-23T00:01:15.343

First, it is important to provide a complete and verifiable example. The data dependency here does not allow for it. However, to answer your question you need to index into your conditions. I made some mock data for you to produce the results you want.

import pandas as pd, numpy as np, matplotlib.pyplot as plt
# mockdata
data = np.zeros((10, 3))
data[:, :2] = np.random.rand(len(data), 2)
data[:, -1] = np.random.randint(0, 2, size = len(data))
df = pd.DataFrame(data, columns = 'x y condition'.split(' '))
fig, ax = plt.subplots()
# index into different subsets and plot
for condition, selection in df.groupby('condition'):
    ax.plot(selection.x, selection.y, label = condition)
ax.legend(title = 'condition')
ax.set(xlabel = 'x', ylabel = 'y')
fig.show()

Matplotlib, plot x1 and y, but color code based on x2 in python?

1 Answers1