PythonでCSVファイルから3要素のデータを読み込み、散布図を出力させたい

Question

前提・実現したいこと

約10000件の研究データを記録したCSVファイルがあります。

X座標	Y座標	-1, 0, 1のいずれかの値(Aとする)
-16.3426246643066	12.3114566802979	1
...	...	...

これを用いて、横軸にX座標、縦軸にY座標を取った2次元の散布図に、Aの値で色分けをして点をプロットしていきたいです。

発生している問題

2次元の散布図は2変数の関係を表すものなので、どうしてもAの値で色分けしてプロットすることができません。解決方法はあるのかを教えていただきたいです。もしなければ、高さの軸にAの値を取った3次元の散布図で、Aの値で色分けする方法を教えていただきたいです。

使用しているバージョン

Python 3.8.5

追記

@can110さんの方法で以下のスクリプトで試しました。CSVファイルは3FloorLv4_result_2021-10-20-20-47-55.csvという名前です。

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# data = np.hstack([np.random.randint(0, 100, (30, 2)), np.random.randint(-1, 2, (30, 1))])
data = pd.read_csv("3FloorLv4_result_2021-10-20-20-47-55.csv", usecols=[0, 1, 2], nrows=1000, dtype="float")
df = pd.DataFrame(data, columns=list("XZS"))

colorList = {-1.0 : "r", 0.0 : "g", 1.0 : "b"}
labelName = {-1 : "Error", 0 : "Failure", 1 : "Success"}

for s in set(df["S"]) : 
    df2 = df[df["S"] == s]
    c = colorList[s]
    l = labelName[s]
    plt.scatter(df2.X, df2.Z, color=c, label=l)
    
plt.legend()
plt.show()

すると、次のエラーが出ました。

Traceback (most recent call last):
  File "c:\Users\yuulis\ML-Agents\EL_3Floor\DrawGragh.py", line 14, in <module>
    c = colorList[s]
KeyError: nan

Answer 1

力押で行けそうだなぁと思って作ってる間に綺麗なのが投稿されてたけど供養のために投稿しとこ

from re import X
import numpy as np
import matplotlib.pyplot as plt
import csv

with open('data1.csv') as f:
    reader = csv.reader(f)
    l = [row for row in reader]
	
x_=[]
y_=[]
x0=[]
y0=[]
x1=[]
y1=[]

#分岐
for i in range(len(l)):	
	if l[i][2]=="-1":			
		x_.append(float(l[i][0]))
		y_.append(float(l[i][1]))
	elif l[i][2]=="1":
		x1.append(float(l[i][0]))
		y1.append(float(l[i][1]))
	elif l[i][2]=="0":
		x0.append(float(l[i][0]))
		y0.append(float(l[i][1]))
 
# figureを生成する
fig = plt.figure()
 
# axをfigureに設定する
ax = fig.add_subplot(1,1,1)
 
# axesに散布図を設定する
ax.scatter(x_, y_, c='b')
ax.scatter(x1, y1, c='r')
ax.scatter(x0, y0, c='g')
 
# 表示する
plt.show()

Answer 2

３色しかないのであれば、単純に３回scatter描画することでも実現できます。

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

# テストデータ
data = np.hstack( [ np.random.randint(0,100, (30,2)), np.random.randint(-1,2, (30,1))])
df = pd.DataFrame(data, columns=list('XYA'))

clr = {-1:'red', 0:'green', 1:'blue'}
for a in set(df['A']):
    df2 = df[df['A'] == a]
    c = clr[a]
    plt.scatter(df2.X, df2.Y, color=c, label=c)
plt.legend()
plt.show()

Comments

[Yuulis]

回答ありがとうございます！
この方法で追記した内容のスクリプトを実行すると、エラーが出てしまったのですが、どうすればいいでしょうか。

[can110]

df["S"]に、colorListのキーに合致しないnan値が入っているのでしょう。
データを確認したうえで、fillnaで0なりキーに存在する値をセットするか
nanを含む行はdropすればよいです。

[Yuulis]

DataFrameでdataのcolumsを指定するときに、元のデータに存在していない名前をしていたことが原因で、Aのすべての値がNaNとして登録されてました...
無事解決して散布図が表示されました！ありがとうございました。

Answer 3

matplotlib で散布図 (Scatter plot) を描くの
カラーマップを指定して、値に応じてマーカーを着色
をお読みください。