# 第1部
# 第4章：並列処理の魔法と、消えた順番の謎
# 【体験ドリル】なぜ正弦波なのか？ 数学的背景の視覚化
# 佐藤　聖

import numpy as np
import matplotlib.pyplot as plt
import japanize_matplotlib  # 日本語表示を可能にする

def visualize_pe():
    # パラメータ設定：文章の長さ(100単語)、情報の切り口（128次元）
    seq_len, d_model = 100, 128
    
    # 位置（縦軸）と次元（横軸）の計算
    pos = np.arange(seq_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
    
    # 位置エンコーディング行列の生成
    pe = np.zeros((seq_len, d_model))
    pe[:, 0::2] = np.sin(pos * div_term)  # 偶数次元にはサイン波
    pe[:, 1::2] = np.cos(pos * div_term)  # 奇数次元にはコサイン波

    # グラフの描画
    plt.figure(figsize=(10, 6))
    plt.imshow(pe, cmap='RdBu', aspect='auto')
    
    # 日本語のラベル設定
    plt.title("位置エンコーディング（各行が単語の「住所」を表す固有の模様）")
    plt.xlabel("次元（情報の切り口）")
    plt.ylabel("文中の位置（単語の順番）")
    
    # カラーバーと表示
    cbar = plt.colorbar()
    cbar.set_label('加算される数値の強さ')
    plt.show()

visualize_pe()
