import pywt
import matplotlib.pyplot as plt
import numpy as np

# 信号を定義
t = np.linspace(0, 1, 400)
signal = np.sin(50 * 2 * np.pi * t) * (t < 0.5) + np.sin(5 * 2 * np.pi * t) * (t >= 0.5)

# CWTで解析
coeffs, freqs = pywt.cwt(signal, scales=np.arange(1, 100), wavelet='mexh')

# プロット
plt.imshow(np.abs(coeffs), extent=[0, 1, 1, 100], cmap='jet', aspect='auto')
plt.title("CWT（連続ウェーブレット変換）")
plt.xlabel("Time")
plt.ylabel("Scale")
plt.colorbar(label='Coefficient magnitude')
plt.show()