import time import numpy as np from matplotlib import pyplot as plt from reservoirpy import datasets def print_exec_time(func): def wrapper(*args, **kwargs): start = time.time() result = func(*args, **kwargs) print(f"Function {func.__name__} executed in {time.time() - start:.4f} seconds.") return result return wrapper def add_noise(data, noise_type='gaussian', intensity=0.1, **kwargs): """ 为输入数据添加噪声。 参数: data: numpy 数组, 原始数据. noise_type: str, 噪声类型，可选 'gaussian'（高斯白噪声）、'colored' 或 '1/f'（色噪声）、'impulse'（脉冲噪声）、'sine'（正弦噪声）。 intensity: float, 噪声强度. kwargs: 针对某些噪声类型的额外参数，例如: - 对于 'impulse': prob (脉冲发生概率, 默认 0.01) - 对于 'sine' : freq (正弦信号频率, 默认 5) 返回: 添加噪声后的数据. """ noise_type = noise_type.lower() # 确保噪声与数据形状匹配，处理多维数据 if data.ndim > 1: n_samples, n_dimensions = data.shape else: n_samples = len(data) n_dimensions = 1 if noise_type in ['gaussian', 'gaussian white']: noise = intensity * np.random.randn(*data.shape) return data + noise elif noise_type in ['colored', '1/f']: # 对于多维数据，为每个维度单独生成1/f噪声 if data.ndim > 1: noise = np.zeros_like(data) for dim in range(n_dimensions): f = np.fft.fftfreq(n_samples) f[0] = 1e-10 noise_complex = np.random.randn(n_samples) + 1j * np.random.randn(n_samples) factor = intensity / np.sqrt(np.abs(f)) noise[:, dim] = np.fft.ifft(noise_complex * factor).real else: f = np.fft.fftfreq(n_samples) f[0] = 1e-10 noise_complex = np.random.randn(n_samples) + 1j * np.random.randn(n_samples) factor = intensity / np.sqrt(np.abs(f)) noise = np.fft.ifft(noise_complex * factor).real return data + noise elif noise_type == 'impulse': noise = np.zeros_like(data) prob = kwargs.get('prob', 0.01) mask = np.random.rand(*data.shape) < prob impulse = intensity * (2 * np.random.rand(*data.shape) - 1) noise[mask] = impulse[mask] return data + noise elif noise_type == 'sine': t = np.arange(n_samples) freq = kwargs.get('freq', 5) # 对于多维数据，确保正弦波形状正确 if data.ndim > 1: sine_wave = intensity * np.sin(2 * np.pi * freq * t / n_samples).reshape(-1, 1) # 扩展到与数据相同的维度 sine_wave = np.tile(sine_wave, (1, n_dimensions)) else: sine_wave = intensity * np.sin(2 * np.pi * freq * t / n_samples) return data + sine_wave else: raise ValueError("未知的噪声类型。可选项：'gaussian', 'colored' (1/f), 'impulse', 'sine'") def load_data(system='lorenz', init='random', noise=None, intensity=0.1, h=0.01, n_timesteps=10000, transient=1000, normlization=True, **kwargs): """ 加载混沌系统数据. 参数: system: str, 混沌系统类型, 可选 'lorenz', 'rossler', 'multiscroll', 'kuramoto_sivashinsky'。 init: 初始值. 如果为 'random' 则随机初始化，否则期望传入数组形式的初始值. noise: None, str 或 str 列表, 指定要添加的噪声类型. 如果是列表，则混合各噪声 (取平均) 作为混合噪声. -- 可选 'gaussian'（高斯白噪声）、'colored' 或 '1/f'（色噪声）、'impulse'（脉冲噪声）、'sine'（正弦噪声） intensity: float, 噪声强度. h: float, 时间步长 (TimeDelta). n_timesteps: int, 时间步数. transient: int, 需丢弃的时间步数. normlization: bool, 是否对数据进行归一化处理. kwargs: 其他系统参数. 返回: (clean_data, noisy_data): 两个 numpy 数组, 分别为干净数据和添加噪声后的数据. 各系统默认值: - lorenz: rho=28, sigma=10, beta=8/3, h=0.03, x0=[1, 1, 1] - kuramoto_sivashinsky: N=128, M=16, h=0.25, x0=None """ system = system.lower() if init == 'random': if system in ['lorenz']: # 默认3维初始值 state = np.random.rand(3) elif system in ['kuramoto_sivashinsky']: state = np.random.rand(kwargs.get('N', 128)) else: raise ValueError("未知的混沌系统类型。") else: state = np.array(init, dtype=float) if system == 'lorenz': ''' (function) def lorenz( n_timesteps: int, rho: float = 28, sigma: float = 10, beta: float = 8 / 3, x0: list | ndarray = [1, 1, 1], h: float = 0.03, **kwargs: Any ) -> ndarray ''' sigma = kwargs.get('sigma', 10.0) rho = kwargs.get('rho', 28.0) beta = kwargs.get('beta', 8/3) clean_data = datasets.lorenz(n_timesteps=n_timesteps, h=h, sigma=sigma, rho=rho, beta=beta, x0=state) elif system == 'kuramoto_sivashinsky': ''' (function) def kuramoto_sivashinsky( n_timesteps: int, warmup: int = 0, N: int = 128, M: float = 16, x0: list | ndarray = None, h: float = 0.25 ) -> ndarray ''' clean_data = datasets.kuramoto_sivashinsky(n_timesteps=n_timesteps, h=h, **kwargs) else: raise ValueError("未知的混沌系统类型。可选项: 'lorenz', 'kuramoto_sivashinsky'") eps = 1e-8 # 定义一个极小值 normalize_01_eps = lambda data: (data - data.min(axis=0)) / (data.max(axis=0) - data.min(axis=0) + eps) if normlization: clean_data = normalize_01_eps(clean_data) print('use new 0-1 norm') #clean_data = (clean_data - clean_data.mean(axis=0)) / clean_data.std(axis=0) # Z-score 归一化 clean_data = clean_data[transient:,:] # 添加噪声 if noise is not None: if isinstance(noise, list): noise_sum = np.zeros_like(clean_data) for nt in noise: noise_sum += add_noise(np.zeros_like(clean_data), noise_type=nt, intensity=intensity, **kwargs) mixed_noise = noise_sum / len(noise) noisy_data = clean_data + mixed_noise elif isinstance(noise, str): noisy_data = add_noise(clean_data, noise_type=noise, intensity=intensity, **kwargs) else: raise ValueError("噪声参数 noise 必须为字符串或字符串列表。") else: noisy_data = clean_data.copy() return clean_data, noisy_data def get_lorenz(n_timesteps=10000, h=0.01, transient=1000, noise='gaussian', intensity=0.1, x0=[1, 1, 1], rho=28, sigma=10, beta=8/3, normlization=True): """ 生成 Lorenz 系统数据. 参数: n_timesteps: int, 时间步数. h: float, 时间步长. transient: int, 丢弃的时间步数. noise: str, 噪声类型. intensity: float, 噪声强度. x0: list, 初始值. rho, sigma, beta: float, 系统参数. kwargs: 其他参数. 返回: (clean_data, noisy_data): 两个 numpy 数组, 分别为干净数据和添加噪声后的数据. """ # 生成 Lorenz 系统数据 clean_data, noisy_data = load_data(system='lorenz', noise=noise, intensity=intensity, n_timesteps=n_timesteps, transient=transient, h=h, x0=x0, rho=rho, sigma=sigma, beta=beta, normlization=True) return clean_data, noisy_data def get_kuramoto_sivashinsky(n_timesteps=10000, h=0.25, transient=2000, noise='gaussian', intensity=0.1, N=128, M=16, x0=None, normlization=True): """ 生成 Kuramoto-Sivashinsky 系统数据. 参数: n_timesteps: int, 时间步数. h: float, 时间步长. transient: int, 丢弃的时间步数. noise: str, 噪声类型. intensity: float, 噪声强度. N: int, 系统参数. M: float, 系统参数. x0: list | ndarray, 初始值. 返回: (clean_data, noisy_data): 两个 numpy 数组, 分别为干净数据和添加噪声后的数据. """ clean_data, noisy_data = load_data(system='kuramoto_sivashinsky', noise=noise, intensity=intensity, n_timesteps=n_timesteps, transient=transient, h=h, N=N, M=M, x0=x0, normlization=True) return clean_data, noisy_data def get_wind(noise='gaussian', intensity=0.1, filepath='exp/data/scale_windspeed.txt', normlization=True): ''' 读取风速数据并添加噪声. 参数: noise: str, 噪声类型. intensity: float, 噪声强度. filepath: str, 数据文件路径. 返回: wind_data: numpy 数组, 添加噪声后的风速数据. ''' wind_data = np.loadtxt(filepath).T # shape -> [155, 157819] if normlization: wind_data = (wind_data - wind_data.min(axis=0)) / (wind_data.max(axis=0) - wind_data.min(axis=0) + 1e-8) # wind_data = (wind_data - wind_data.mean(axis=0)) / wind_data.std(axis=0) # Z-score 归一化 if noise is not None: # 添加噪声 wind_noisy = add_noise(wind_data, noise_type=noise, intensity=intensity) return wind_data, wind_noisy def visualize_data_diff(clean_data, noisy_data, prediction, warmup=0, title_prefix=""): """ 可视化数据比较：噪声数据与干净数据，预测数据与干净数据 参数: clean_data: 干净数据 noisy_data: 噪声数据 prediction: 预测数据 warmup: 起始索引 title_prefix: 标题前缀 """ # 获取数据维度 n_dims = 1 if clean_data.ndim == 1 else clean_data.shape[1] # 调整图形大小和布局 if n_dims <= 3: fig_height = 4 * n_dims fig_width = 12 n_cols = 2 else: fig_height = 3 * ((n_dims + 1) // 2) fig_width = 14 n_cols = 4 plt.figure(figsize=(fig_width, fig_height)) for i in range(n_dims): # 获取当前维度的数据 if clean_data.ndim == 1: clean_dim = clean_data[warmup:] noisy_dim = noisy_data[warmup:] pred_dim = prediction[warmup:] else: clean_dim = clean_data[warmup:, i] noisy_dim = noisy_data[warmup:, i] pred_dim = prediction[warmup:, i] # 噪声数据与干净数据对比 plt.subplot(n_dims, n_cols, n_cols*i+1) plt.plot(noisy_dim, label='noisy') plt.plot(clean_dim, label='clean') plt.title(f"{title_prefix} Dim {i+1} (Noisy vs Clean)") if i == 0 or n_dims <= 3: plt.legend() # 预测数据与干净数据对比 plt.subplot(n_dims, n_cols, n_cols*i+2) plt.plot(pred_dim, label='prediction') plt.plot(clean_dim, label='clean') plt.title(f"{title_prefix} Dim {i+1} (Prediction vs Clean)") if i == 0 or n_dims <= 3: plt.legend() plt.tight_layout() def visualize_psd_diff(clean_data, noisy_data, prediction, warmup=0, title_prefix=""): """ 可视化数据的功率谱密度比较 参数: clean_data: 干净数据 noisy_data: 噪声数据 prediction: 预测数据 warmup: 起始索引 title_prefix: 标题前缀 """ # 获取数据维度 n_dims = 1 if clean_data.ndim == 1 else clean_data.shape[1] # 调整图形大小 fig_height = 4 * min(n_dims, 5) # 限制最大高度 fig_width = 12 plt.figure(figsize=(fig_width, fig_height)) # 如果维度过多，只显示前5个维度 show_dims = min(n_dims, 5) for i in range(show_dims): plt.subplot(show_dims, 1, i+1) # 获取当前维度的数据 if clean_data.ndim == 1: clean_dim = clean_data[warmup:] noisy_dim = noisy_data[warmup:] pred_dim = prediction[warmup:] else: clean_dim = clean_data[warmup:, i] noisy_dim = noisy_data[warmup:, i] pred_dim = prediction[warmup:, i] plt.psd(clean_dim, NFFT=1024, label='clean') plt.psd(noisy_dim, NFFT=1024, label='noisy') plt.psd(pred_dim, NFFT=1024, label='prediction') plt.title(f"{title_prefix} Dim {i+1} PSD Comparison") plt.legend() plt.tight_layout() def visualize_ks_denoising(clean_data, noisy_data, prediction, warmup=0, title_prefix="", time_slice=None): """ 可视化Kuramoto-Sivashinsky系统的降噪结果 (修改版) 参数: clean_data: 干净数据，形状为[时间步数, 空间点数] noisy_data: 噪声数据，形状同上 prediction: 预测数据，形状同上 warmup: 要从数据开头移除的时间步数 title_prefix: 图形总标题的前缀 time_slice: 要显示的特定时间索引列表 (最多3个)，如果为None则自动选择 """ # 1. 数据准备 if warmup >= clean_data.shape[0]: print(f"Warning: Warmup period ({warmup}) is longer than or equal to data length ({clean_data.shape[0]}). No data to plot.") return clean = clean_data[warmup:] noisy = noisy_data[warmup:] pred = prediction[warmup:] if clean.shape[0] == 0: print("Warning: Data is empty after removing warmup period.") return # 计算误差和MSE mse_noisy = np.mean((clean - noisy) ** 2) mse_pred = np.mean((clean - pred) ** 2) diff_noisy = clean - noisy diff_pred = clean - pred # 为差异图确定对称的颜色范围 max_abs_diff = max(np.abs(diff_noisy).max(), np.abs(diff_pred).max()) # 防止 max_abs_diff 为 0 导致 vmin=vmax if max_abs_diff < 1e-9: max_abs_diff = 1e-9 diff_vmin, diff_vmax = -max_abs_diff, max_abs_diff # 为数据图确定共享的颜色范围 (可选，但有助于比较) data_vmin = min(clean.min(), noisy.min(), pred.min()) data_vmax = max(clean.max(), noisy.max(), pred.max()) if data_vmax - data_vmin < 1e-9: # Handle constant data case data_vmax = data_vmin + 1e-9 # 2. 创建图形和轴 (4x3 网格) fig, axes = plt.subplots(4, 3, figsize=(15, 17)) # 调整 figsize 以适应 4 行 fig.suptitle(f"{title_prefix} Kuramoto-Sivashinsky Denoising Results", fontsize=16) # imshow 的通用范围 extent = [0, clean.shape[0], 0, clean.shape[1]] # [时间_min, 时间_max, 空间_min, 空间_max] time_label = 'Time Step (after warmup)' space_label = 'Space' amplitude_label = 'Amplitude' difference_label = 'Difference' # --- 第 1 行: 干净, 噪声, 干净 - 噪声 --- im1 = axes[0, 0].imshow(clean.T, aspect='auto', cmap='viridis', extent=extent, vmin=data_vmin, vmax=data_vmax) fig.colorbar(im1, ax=axes[0, 0], label=amplitude_label) axes[0, 0].set_title("Clean Data") axes[0, 0].set_ylabel(space_label) im2 = axes[0, 1].imshow(noisy.T, aspect='auto', cmap='viridis', extent=extent, vmin=data_vmin, vmax=data_vmax) fig.colorbar(im2, ax=axes[0, 1], label=amplitude_label) axes[0, 1].set_title(f"Noisy Data (MSE: {mse_noisy:.4f})") # axes[0, 1].set_ylabel(space_label) # Y轴标签通常只在最左侧显示 im3 = axes[0, 2].imshow(diff_noisy.T, aspect='auto', cmap='coolwarm', extent=extent, vmin=diff_vmin, vmax=diff_vmax) fig.colorbar(im3, ax=axes[0, 2], label=difference_label) axes[0, 2].set_title("Clean - Noisy") # axes[0, 2].set_ylabel(space_label) # --- 第 2 行: 干净, 预测, 干净 - 预测 --- im4 = axes[1, 0].imshow(clean.T, aspect='auto', cmap='viridis', extent=extent, vmin=data_vmin, vmax=data_vmax) fig.colorbar(im4, ax=axes[1, 0], label=amplitude_label) axes[1, 0].set_title("Clean Data") axes[1, 0].set_ylabel(space_label) axes[1, 0].set_xlabel(time_label) # 在底部imshow行添加时间标签 im5 = axes[1, 1].imshow(pred.T, aspect='auto', cmap='viridis', extent=extent, vmin=data_vmin, vmax=data_vmax) fig.colorbar(im5, ax=axes[1, 1], label=amplitude_label) axes[1, 1].set_title(f"Prediction (MSE: {mse_pred:.4f})") # axes[1, 1].set_ylabel(space_label) axes[1, 1].set_xlabel(time_label) im6 = axes[1, 2].imshow(diff_pred.T, aspect='auto', cmap='coolwarm', extent=extent, vmin=diff_vmin, vmax=diff_vmax) fig.colorbar(im6, ax=axes[1, 2], label=difference_label) axes[1, 2].set_title("Clean - Prediction") # axes[1, 2].set_ylabel(space_label) axes[1, 2].set_xlabel(time_label) # --- 第 3 行: 时间切片 --- num_steps = clean.shape[0] if time_slice is None: # 自动选择最多3个均匀分布的时间点 if num_steps > 0: # 确保 t_points 是有效的索引，并且尽可能均匀 indices = np.linspace(0, num_steps - 1, 5, dtype=int)[1:-1] # 取中间的3个点 t_points = sorted(list(set(indices))) # 去重并排序 # 如果点数少于3个（因为数据短或linspace结果重复），补充点 while len(t_points) < 3 and len(t_points) < num_steps: if 0 not in t_points: t_points.insert(0, 0) elif num_steps - 1 not in t_points: t_points.append(num_steps - 1) else: # 如果头尾都有了，尝试在中间加 mid_ish = num_steps // 2 if mid_ish not in t_points: t_points.append(mid_ish) else: break # 无法添加更多唯一的点 t_points = sorted(list(set(t_points)))[:3] # 保持最多3个 else: t_points = [] # 如果没有时间步，则没有切片 else: # 过滤用户提供的时间切片以确保有效性，并最多取前3个 t_points = [t for t in time_slice if 0 <= t < num_steps] if len(t_points) > 3: print(f"Warning: Provided {len(time_slice)} time slices. Using the first 3 valid ones: {t_points[:3]}") t_points = t_points[:3] elif not t_points and time_slice: print(f"Warning: None of the provided time slices {time_slice} are valid for data length {num_steps}.") # 绘制时间切片图 num_slice_plots = 3 for i in range(num_slice_plots): ax_slice = axes[2, i] if i < len(t_points): t = t_points[i] ax_slice.plot(clean[t], label='Clean') ax_slice.plot(noisy[t], label='Noisy', alpha=0.7) ax_slice.plot(pred[t], label='Prediction', linestyle='--') ax_slice.set_title(f"Time slice at t={t}") ax_slice.set_xlabel(space_label) if i == 0: ax_slice.set_ylabel(amplitude_label) # 只在最左侧显示Y轴标签 ax_slice.legend() ax_slice.grid(True, linestyle='--', alpha=0.6) else: # 如果没有足够的时间点来绘制，则隐藏多余的子图 ax_slice.axis('off') # --- 第 4 行: 误差分布 & 随时间变化的 MSE --- # 误差直方图 bins = 50 ax_hist_noisy = axes[3, 0] counts_noisy, _, _ = ax_hist_noisy.hist(diff_noisy.flatten(), bins=bins, alpha=0.7, label='Noisy Error') ax_hist_noisy.set_title('Noise Error Distribution') ax_hist_noisy.set_xlabel('Error (Clean - Noisy)') ax_hist_noisy.set_ylabel('Count') ax_hist_noisy.grid(True, linestyle='--', alpha=0.6) ax_hist_pred = axes[3, 1] counts_pred, _, _ = ax_hist_pred.hist(diff_pred.flatten(), bins=bins, alpha=0.7, label='Prediction Error', color='orange') ax_hist_pred.set_title('Prediction Error Distribution') ax_hist_pred.set_xlabel('Error (Clean - Prediction)') # ax_hist_pred.set_ylabel('Count') # Y轴标签通常只在最左侧显示 # 为直方图设置共享的 y 轴限制 max_y = 0 if counts_noisy.size > 0: # 检查是否有计数 max_y = max(max_y, np.max(counts_noisy)) if counts_pred.size > 0: # 检查是否有计数 max_y = max(max_y, np.max(counts_pred)) if max_y > 0: # 仅当绘制了直方图时才设置 ylim common_ylim = (0, max_y * 1.1) # 增加 10% 的空隙 ax_hist_noisy.set_ylim(common_ylim) ax_hist_pred.set_ylim(common_ylim) else: # 处理数据为空或恒定的情况 ax_hist_noisy.set_ylim(0, 1) ax_hist_pred.set_ylim(0, 1) # 随时间变化的 MSE ax_mse = axes[3, 2] time_axis = np.arange(num_steps) # 创建时间轴 ax_mse.plot(time_axis, np.mean(diff_noisy**2, axis=1), label='Noisy MSE') ax_mse.plot(time_axis, np.mean(diff_pred**2, axis=1), label='Prediction MSE') ax_mse.set_title('MSE over Time') ax_mse.set_xlabel(time_label) ax_mse.set_ylabel('MSE') ax_mse.set_yscale('log') # MSE 通常最好用对数刻度显示 ax_mse.legend() ax_mse.grid(True, linestyle='--', alpha=0.6) # --- 最终调整 --- plt.tight_layout(rect=[0, 0.01, 1, 0.96]) # 调整布局以适应 suptitle 和 x 标签 # plt.show() # 如果需要立即显示图形，取消注释此行 if __name__ == "__main__": # 测试数据加载和噪声添加 clean_data, noisy_data = get_kuramoto_sivashinsky() print("Clean Data Shape:", clean_data.shape) print("Noisy Data Shape:", noisy_data.shape) visualize_ks_denoising(clean_data, noisy_data, noisy_data)