Matplotlibでcolorbarいろいろ

二次元プロットの場合、表現できるデータの種類はx軸、y軸と色を使った最大三種類。
MATLAB編（ https://qiita.com/Yuricst/items/1594cacb930bd4fa6d63 ）に引き続き、PythonのMatplotlibでplot内でcolorbarの追加についてメモ。
（最適解ではないにしも、取り合えずこの手法で間に合う、といった感じ。随時更新予定。）

散布図（plt.scatter）の場合

scatterの場合は結構シンプルで、scatter()を使う際に色のパラメーターとして使用したいarrayをcで指定、さらにカラーマップをcmapで指定するだけ。
注意点としては、colorbarを追加するために使用するfig.colorbar()ではmappable objectが必要になるため、scatter自体をオブジェクトとしてassignしないといけない（下記例ではim0としている）。subplot使用時、色のパラメーターが同じなら、全サブプロットをオブジェクトとしてassignしなくても問題ない。

# import modules
#import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
#from mpl_toolkits.mplot3d import Axes3D
#from matplotlib.colors import ListedColormap, BoundaryNorm

# assign color array
cc = transfer_df['Jacobi_LPO']

plt.rcParams["font.size"] = 16
fig, axs = plt.subplots(1, 3, figsize=(15, 6))
im0 = axs[0].scatter(transfer_df['loi_cost']*1000, transfer_df['incl'], c=cc, cmap='plasma', s=15, marker='x')
axs[0].set_xlabel('LOI cost [m/sec]')
axs[0].set_ylabel('Inclination [deg]')
axs[0].grid(True)

axs[1].scatter(transfer_df['loi_cost']*1000, transfer_df['raan'], c=cc, cmap='plasma', s=15, marker='x')
axs[1].set_xlabel('LOI cost [m/sec]')
axs[1].set_ylabel('RAAN [deg]')
axs[1].grid(True)

axs[2].scatter(transfer_df['loi_cost']*1000, -transfer_df['tof_EM_SOI']*Tstar_EM/(60*60*24), c=cc, cmap='plasma', s=15, marker='x')
axs[2].set_xlabel('LOI cost [m/sec]')
axs[2].set_ylabel('Tof til lunar SOI [days]')
axs[2].grid(True)


fig.colorbar(im0, label='LPO Jacobi')
plt.suptitle(f'Moon SOI leg, beta = {transfer_df.iloc[0,:]['beta0']}')
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.show()

軌跡（plt.plot）の場合

plotの場合、少しプロセスが複雑になる。大まかにいうと、
- プロットしたいデータをpointsに保存
- pointsとカラーバーをLineCollectionでアサイン
- lineにLineCollectionを追加
といった感じ（なのかな？）。

ということで、まず軌跡一つ一つに色をアサインしてLineCollectionを作る関数を書く。

def get_lc_traj_singleColor(xs, ys, c, vmin, vmax, cmap, lw=0.8):
    """
    Get line collection object a trajectory with a single color based on a colormap defined by vmin ~ vmax
    For plotting, run:
        line = ax.add_collection(lc)
        fig.colorbar(line, ax=ax, label="Colorbar label")

    Args:
        xs (np.array): array-like object of x-coordinates of the trajectory

    Returns:
        (obj): line collection object
    """
    # generate segments
    points = np.array([ xs , ys ]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    plt_color = c * np.ones((len(xs), ))
    # create color bar
    norm = plt.Normalize( vmin, vmax )
    lc = LineCollection(segments, cmap=cmap, norm=norm)
    # Set the values used for colormapping
    lc.set_array( plt_color )
    lc.set_linewidth(lw)
    return lc

で、これをプロットを作る際に、軌跡毎に上の関数を呼ぶ。

# plot transfers inside lunar SOI
plt.rcParams["font.size"] = 20
fig, axs = plt.subplots(1, 1, figsize=(18, 10))

for idxTraj in tqdm(range(len(proplst))):
    # create line collection object for each trajectory
    lc = get_lc_traj_singleColor(xs=proplst[idxTraj]["xs"], ys=proplst[idxTraj]["ys"], 
                            c=es[idxTraj], vmin=min(es), vmax=max(es), cmap='viridis', lw=1.2)
    line = axs.add_collection(lc)

# display colorbar
fig.colorbar(line, ax=axs, label="Eccentricity")

axs.axis('equal')
axs.grid(True)
axs.set_xlabel('x, km')
axs.set_ylabel('y, km')
axs.set_title('Simulated LEOs')
#axs.legend(loc='lower right')

plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.show()