如何根据第三列值绘制带有颜色的2D数据点

Question

Suppose I have an array with three columns:

>>>print(arr)
array(
    [[1 2 -1]
    [1 3 1]
    [3 2 -1]
    [5 2 -1]]
)

Suppose I want to turn this into a scatter plot:

plt.plot(arr[:,0], arr[:,1], 'xr')

Works fine. But, all the scattered points look like a red 'x' . Suppose I want to have a red 'x' if the third column value is -1, and a blue 'o' if the value is 1. How do I do that?

Is this possible to achieve with plt.plot() ?

Answer 1

If you want to change the color and the marker, you need to plot several scatter plots, at least one for each marker.

import matplotlib.pyplot as plt
import numpy as np

a = np.array([[1, 2, -1],
              [1, 3, 1],
              [3, 2, -1],
              [5, 2, -1]])
mapping= {-1: ("red", "x"), 1: ("blue", "o")}

for c in np.unique(a[:,2]):
    d = a[a[:,2] == c]
    plt.scatter(d[:,0], d[:,1], c=mapping[c][0], marker=mapping[c][1])
    #plt.plot(d[:,0], d[:,1], color=mapping[c][0], marker=mapping[c][1])

plt.show()

Using a plot is equally possible (see commented line in code).

Answer 2

Iterate through your list, and consider that the 3rd item of each sublist is the color. Attribute the color according to the data you have. Here, I assign the value 'r' to c if c == -1 , else, it will be equal to 'b' , meaning blue .

from matplotlib import pyplot as plt

arr = [[1, 2, -1], [1, 3, 1], [3, 2, -1], [5, 2, -1]]

for x, y, color in arr:
    c = 'r' if color == -1 else 'b'
    plt.plot(x, y, 'x' + c)
plt.show()

Answer 3

as a fellow Python beginner I have had a similar question recently and for me the best solution was to use the fantastic " seaborn " package.

The trick is to combine Pandas ' DataFrame ' and Seaborn's " FacetGrid " objects and let them do all the work. Here's how I did it for your particular example:

• Import pandas and seaborn
• Convert the array into a dataframe
• Create a seaborn facetgrid together with matplotlib scatter plot
• Use the great "Hue" parameter to get the colours right

Here is my code:

Import the necessary libraries

import seaborn as sns
import pandas as pd

Create the Pandas DataFrame

a = pd.DataFrame(
    [[1, 2, -1],
    [1, 3, 1],
    [3, 2, -1],
    [5, 2, -1]])

a.columns=['A','B','C']

Generate the seaborn FacetGrid and use Hue for the colours

sns.FacetGrid(a, hue="C", size=2).map(plt.scatter, "A", "B").add_legend()
plt.show()

Voila !

Obviously this is a trivial example, but with real data the result is great and easy to reproduce (see below)

PS: This is my very first post on stackoverflow - it feels like the beginning of something