[python] yeojohnson変換の変換前後の値をプロット

import numpy as np
import pandas as pd

# ここの式をコピー
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.yeojohnson.html
def yeojohnson(lmbda, x):
    if x >= 0 and lmbda != 0:
        y = ((x + 1)**lmbda - 1) / lmbda

    elif x >= 0 and lmbda == 0:
        y = np.log(x + 1)

    elif x < 0 and lmbda != 2:
        y = -((-x + 1)**(2 - lmbda) - 1) / (2 - lmbda)

    elif x < 0 and lmbda == 2:
        y = -np.log(-x + 1)

    else:
        # 到達しないはず
        raise

    return y

# 変換する値
x_ary = np.arange(21)-10

# 2パターンの lmbda の組み合わせを試す
for lmbda_ary in [np.arange(10)*0.25, np.arange(10)*0.5]:

    # 変換して df にまとめる
    X_dic = {}
    for lmbda in lmbda_ary:
        X_dic[lmbda] = [yeojohnson(lmbda, float(x)) for x in x_ary]
    df = pd.DataFrame(X_dic, index=x_ary)

    # 描画
    df.plot(cmap='viridis', marker='.')
    plt.xlabel('元の値（x）')
    plt.ylabel('変換後の値（y）')
    plt.title('lmbda を変えた時の変換後の値の変化')
    plt.show()