CODEWITHSUNDEEP
pythonmatplotlibscatter-plot
Stacey
Stacey

Reputation: 73

How to create a scatter plot for two data classes with pyplot?

I have two sets of data, x and y as ints. I need to plot both of these data points using matplotlib.pyplot.scatter. I also need to plot the first category y == 0 in one color and the second, y == 1 in a different color.

I've looked at the documentation for the scatter function, but I don't understand how to do all of this in one plot.

Sample data:

2.897534798034255,0.872359037956732,1
1.234850239781278,-0.293047584301112,1
0.238575209753427,0.129572680572429,0
-0.109757648021958,0.484048547480385,1
1.109735783200013,-0.002785328902198,0
1.572803975652908,0.098547849368397,0

x and y are defined as:

x = data[:, [0, 1]]
y = data[:, -1].astype(int)

Size of x is 2000, size of y is 1000

My attempt:

pl.scatter(x, y==0, s=3, c='r')
pl.scatter(x, y==1, s=3, c='b')
pl.show()

Upvotes: 1

Views: 7594

Answers (3)

Mr. T
Mr. T

Reputation: 12410

Not sure why you want to extract x and y first and filter later. Given that you have a lot of data and not many categories, plt.plot with markers should also be faster than plt.scatter:

import numpy as np
import matplotlib.pyplot as plt

data = np.asarray([[2.897534798034255,0.872359037956732,1],
                     [1.234850239781278,-0.293047584301112,1],
                     [0.238575209753427,0.129572680572429,0],
                     [-0.109757648021958,0.484048547480385,1],
                     [1.109735783200013,-0.002785328902198,0],
                     [1.572803975652908,0.098547849368397,0]])

colors = ["blue", "red", "green"]
labels = ["A", "B", "C"]

for i, c, l in zip(np.unique(data[:, 2]), colors, labels):   
    plt.plot(data[data[:, 2]==i][:, 0], data[data[:, 2]==i][:, 1], 
             marker="o", markersize=7, ls="None", color=c, 
             label=f"The letter {l} represents category {int(i)}")

plt.legend()
plt.show()

Sample output: enter image description here

Upvotes: 0

rpoleski
rpoleski

Reputation: 998

pyplot.scatter() accepts a list of colors, hence:

c = ['r' if yy==0 else 'b' for yy in y]
plt.scatter(x, y, c=c)

In your code, y==0 produces a mask that has only True and False values, not y values to be plotted. If x and y are numpy arrays, you can do:

mask = (y == 0)
plt.scatter(x[mask], y[mask], c='r')
mask = (y == 1)
plt.scatter(x[mask], y[mask], c='b')

Upvotes: 3

David M.
David M.

Reputation: 4588

You can do it like this:

import numpy as np
import matplotlib.pyplot as plt

data = np.array([[2.897534798034255,0.872359037956732,1],
                 [1.234850239781278,-0.293047584301112,1],
                 [0.238575209753427,0.129572680572429,0],
                 [-0.109757648021958,0.484048547480385,1],
                 [1.109735783200013,-0.002785328902198,0],
                 [1.572803975652908,0.098547849368397,0]])

x = data[:, [0, 1]]
y = data[:, -1].astype(int)

plt.scatter(x[:,0][y==0], x[:,1][y==0], s=3, c='r')
plt.scatter(x[:,0][y==1], x[:,1][y==1], s=3, c='b')
plt.show()

Although this is perhaps more readable:

x1 = data[:, 0]
x2 = data[:, 1]
y = data[:, -1].astype(int)

plt.scatter(x1[y==0], x2[y==0], s=3, c='r')
plt.scatter(x1[y==1], x2[y==1], s=3, c='b')

Output:

Scatterplot

Upvotes: 3

Related Questions