2차원 플롯의 경우, 표현할 수 있는 데이터의 종류는 x축, y축과 색을 사용한 최대 3종류.
MATLAB 편 ( htps : // m / 유리 cst / ms / 1594 또는 cb930bd4 6d63 )에 이어 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 와 ColorBar 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()