Python カテゴリ別にプロットする散布図を作成する（カテゴリデータが行方向に塊りである場合）「plt.scatter」

'21/05/26更新：コードの可読性を多少良くしました。
　本記事では、例えば、下図のように「t=1, t=2, t=3」をカテゴリとした散布図を作成する雛形コードを載せました。

　本プログラムの仕様について説明します。下図のようなcsvファイルがあって、Time_1, Time_2, …というカテゴリで行方向に塊りで分けられています。これらの塊りの行数は同じとは限りません。プロットしたいデータは、横軸に「x」,縦軸に「stress」です。複数のcsvファイルに対して、全て順番に処理する仕様です。csvファイルの数だけグラフ画像ファイルを生成します。

処理の過程で、一旦下図のように、例えば「label」と言うカテゴリ変数の列を作成します。この列を用いることで、冒頭のようなグラフに凡例別にプロットすることが出来ます。

■本プログラム

import os, glob
import re
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm
plt.rcParams['font.size'] = 16 # グラフの基本フォントサイズの設定

def initialize(my_save_dir):
    if os.path.isdir(my_save_dir):
        path_list = [my_save_dir, '*.jpg']
        path_list = [dir.replace('./', '') for dir in path_list]
        my_path = os.path.join(*path_list)
        print(my_path)
        delete_file_list = glob.glob(my_path)
        for file in delete_file_list:
            os.remove(file)
    else:
        os.mkdir(my_save_dir)


def main():
    initialize(my_save_dir)
    for my_file in my_files:
        print(my_file)
        base_filename = os.path.basename(my_file)
        
        df = pd.read_csv(my_file, sep=',', header=(my_header-1), encoding='utf_8')
        
        # 列名をリストで取得する
        column_names = df.columns
        #print(column_names)
        
        # 抽出する行区間を判別するため、指定列をリストで取得する
        x_list = df[my_columns[0]].values.tolist()
        
        # リスト内にある指定要素の複数インデックスをリストで返す
        index_start_list = [i for i, x in enumerate(x_list) if x == my_index_start]
        #print(index_start_list)
        index_end_list = [i for i, x in enumerate(x_list) if x == my_index_end]
        #print(index_end_list)

        target_str = '_'
        dict_df = {}
        for start, end in zip(index_start_list, index_end_list):
            # 開始行の1つ上の行のラベルを取得する
            df_label_row = df[start - 1: start]
            string_buf = df_label_row['x'].values[0]
            idx = string_buf.find(target_str)
            label = string_buf[idx + 1:]
            
            # 開始行と終了行を指定して抽出する
            df_extract = df[start: end + 1]
            dict_df[int(label)] = df_extract

        # 辞書をソート
        data_list = sorted(dict_df.items(), key=lambda x:x[0])
        data_list

        # カテゴリ分けのため、ラベル列を生成する
        df_list = []
        for i, taple in enumerate(data_list):
            if i % 1 == 0: # 間引きたい場合に2以上を指定
                df_buf = taple[1].copy()
                df_buf[my_label] = taple[0]
                df_list.append(df_buf)
        df_list

        # ラベル付きデータフレームを作成
        DF = pd.concat(df_list)
        DF

        # 型をfloat型へ変換する
        for my_col in column_names:
            DF[my_col] = DF[my_col].astype('float')

        # ラベルはカテゴリ変数へ変換する
        DF[my_label] = DF[my_label].astype('category')
        # 型を確認する
        print(DF.dtypes)

        # csvファイルに保存
        #DF.to_csv('new_' + base_filename)

        # ひとつの散布図にカテゴリ別にプロットする
        colors = cm.rainbow(np.linspace(0, 1, len(DF[my_label].unique())))
        for i, legend in enumerate(DF[my_label].unique()):
            plt.scatter(DF.loc[DF[my_label] == legend, my_columns[0]],
                        DF.loc[DF[my_label] == legend, my_columns[1]],
                        color = colors[i],
                        label = 't=' + str(legend),
                       )
        plt.grid(which='both')
        plt.title(base_filename)
        plt.xlabel(my_columns[0])
        plt.ylabel(my_columns[1])
        plt.legend(loc='best')
        plt.tick_params()
        plt.savefig(my_save_dir + '/' + base_filename + "_" \
                    + my_columns[0] + "_" + my_columns[1] + ".jpg")
        plt.close()

if __name__ == '__main__':
    my_files = glob.glob("./Csv/*.csv")
    my_files = sorted(my_files, key=lambda x:int((re.search(r"[0-9]+", x)).group(0)))
    my_header = 1
    my_index_start = '3'
    my_index_end = '17'
    my_columns = ("x", "stress")
    my_label = 'label'
    my_Xscale = "linear" # log or linear
    my_Yscale = "linear" # log or linear
    my_y_range = () #(50, 900) # レンジの範囲を指定しない場合は空()
    my_save_dir = './Jpg'
    
    main()

以上

<広告>

リンク