このコードを簡素化したいです

Question

実現したいこと

グラフを出力したいディレクトリの中に、1.csv、2.csv、10.csvという3つのファイルが入っているとします。「該当のソースコード」で1.csv、2.csv、10.csvのdensity.jpgとp.jpgが正しく出力されます。
今、1.csvと10.csvだけのdensity.jpgとp.jpgを出力したときには、8行目と22行目の

files = Path(".").glob("*.csv")をfiles = Path(".").glob("*.csv")

を

files = Path(".").glob("*.csv")をfiles = Path(".").glob("1*.csv")

にすれば良いですが、8行目も22行目も直さなくてはいけないです。

今はグラフが2つですが、いくつもグラフを描きたいときは8行目と22行目のような修正をたくさんしなければいけないので、1回で済ませたいです。どのように直せば良いでしょうか。

該当のソースコード

import pandas as pd
import matplotlib.pyplot as plt
import glob
from pathlib import Path

savename ='density.jpg'
fig, ax = plt.subplots()
files = Path(".").glob("*.csv")

for file in files:
    df = pd.read_csv(file, encoding="utf-8")
    ax.plot('y[mm]', 'density[kg/m3]', data=df, marker='none', label=file.stem, linestyle='-')
    
ax.set_xlabel('$y$[mm]',fontsize=15)
ax.set_ylabel('$density[kg/m3]$',fontsize=15)
ax.grid()
plt.savefig(savename)
plt.show()

savename ='p.jpg'
fig, ax = plt.subplots()
files = Path(".").glob("*.csv")

for file in files:
    df = pd.read_csv(file, encoding="utf-8")
    ax.plot('y[mm]', 'pressure[Pa]', data=df, marker='none', label=file.stem, linestyle='-')
    
ax.set_xlabel('$y$[mm]',fontsize=15)
ax.set_ylabel('$pressure$[Pa]',fontsize=15)
ax.grid()
plt.savefig(savename)
plt.show()

試したこと

以下のようにすると、density.jpgは正しく表示されますが、p.jpgは線のない白紙のグラフが表示されました。

import pandas as pd
import matplotlib.pyplot as plt
import glob
from pathlib import Path

files = Path(".").glob("1*.csv")

savename ='density.jpg'
fig, ax = plt.subplots()

for file in files:
    df = pd.read_csv(file, encoding="utf-8")
    ax.plot('y[mm]', 'density[kg/m3]', data=df, marker='none', label=file.stem, linestyle='-')
    
ax.set_xlabel('$y$[mm]',fontsize=15)
ax.set_ylabel('$density[kg/m3]$',fontsize=15)
ax.grid()
plt.savefig(savename)
plt.show()

savename ='p.jpg'
fig, ax = plt.subplots()

for file in files:
    df = pd.read_csv(file, encoding="utf-8")
    ax.plot('y[mm]', 'pressure[Pa]', data=df, marker='none', label=file.stem, linestyle='-')
    
ax.set_xlabel('$y$[mm]',fontsize=15)
ax.set_ylabel('$pressure$[Pa]',fontsize=15)
ax.grid()
plt.savefig(savename)
plt.show()

Answer 1

重複が多いので、関数化してリファクタリングするとよいかと思います。

import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path

def plot_csv_data(y_column, ylabel, save_filename):
    fig, ax = plt.subplots()
    files = Path(".").glob("*.csv")
    
    for file in files:
        df = pd.read_csv(file, encoding="utf-8")
        ax.plot('y[mm]', y_column, data=df, marker='none', label=file.stem, linestyle='-')
    
    ax.set_xlabel('$y$[mm]', fontsize=15)
    ax.set_ylabel(f'${ylabel}$', fontsize=15)
    ax.grid()
    plt.savefig(save_filename)
    plt.show()

if __name__ == "__main__":
    plot_csv_data('density[kg/m3]', 'density[kg/m3]', 'density.jpg')
    plot_csv_data('pressure[Pa]', 'pressure[Pa]', 'p.jpg')

こうすれば、plot_csv_dataの処理を1回修正するだけで、対象ファイルを変更することができます。

Comments

[yyicp]

できました。
どうもありがとうございました。

Answer 2

他の回答にあるように、書き直すのがすっきりしそうですが、今の形を使う方法を考えるならば、
試したことの6行目を下記のようにして、filesをリストにしてしまうのが簡単です。

files = list(Path(".").glob("1*.csv"))

Path.glob() はジェネレータを返すので、一回 for を通すと中身が消費されてしまい、二回目の for では何も処理されません。
リストにしてしまえば、二回目の for でも最初から繰り返してくれます。

Comments

[yyicp]

できました。
どうもありがとうございました。

Answer 3

それぞれのグラフの属性を辞書(dict)として定義してみてはどうでしょうか。

import pandas as pd
import matplotlib.pyplot as plt
import glob
from pathlib import Path

plots = [
    {
        "files": Path(".").glob("*.csv"),
        "ycolumn": "density[kg/m3]",
        "ylabel": "$density[kg/m3]$",
        "savename": "density.jpg"
    },
    {
        "files": Path(".").glob("1*.csv"),
        "ycolumn": "pressure[Pa]",
        "ylabel": "$pressure$[Pa]",
        "savename": "p.jpg"
    },
]

for p in plots:
    fig, ax = plt.subplots()
    for file in p["files"]:
        df = pd.read_csv(file, encoding="utf-8")
        ax.plot('y[mm]', p["ycolumn"], data=df, marker='none', label=file.stem, linestyle='-')
    ax.set_xlabel('$y$[mm]',fontsize=15)
    ax.set_ylabel(p["ylabel"],fontsize=15)
    ax.grid()
    ax.legend()
    plt.savefig(p["savename"])

Comments

[yyicp]

できました。
どうもありがとうございました。