ml-resources/MambaClimateNet

Mirror of ZHangang1/MambaClimateNet from ModelScope. All credit to the original authors. Not an official release from them. No license was declared at the source; all rights remain with the original authors.

MambaClimateNet

MambaClimateNet 是一种基于选择性状态空间模型（Selective State Space Model, SSM）的高效气候时间序列预测网络。该模型以 Mamba 模块替代传统自注意力结构，实现线性复杂度的长序列建模，并在此基础上引入局部动态增强模块以捕捉长期气候变化特征，从而在全局趋势建模与局部突变响应之间取得平衡。模型在输入阶段采用倒置式嵌入结构沿变量维度进行线性映射，实现多气候要素间的特征融合；多层状态空间编码器结合非线性变换与归一化机制，进一步强化时序模式学习。通过时间与位置嵌入的联合设计，MambaClimateNet 能在较低算力条件下稳定、高效地捕捉气候序列的周期性与季节性规律，实现精准的多要素气候预测。

技术架构

MambaClimateNet 模型的架构包括四个关键处理阶段(如图 1 所示)：首先，倒置变量嵌入层（Inverted Variate Embedding Layer） 将多维气候变量的时间序列进行维度翻转，把每个气候因子的完整历史演变映射为独立的潜在特征Token，从而将单一的时间依赖问题转化为不同气候变量（如温度、湿度、风场等）之间的耦合关系建模。随后，前置特征增强模块（Pre-processing FFN） 利用卷积前馈网络对潜在的气候特征进行预处理，在进入状态空间之前提取特征内部的局部非线性波动。紧接着，设计了一个双向状态空间融合模块（Bidirectional SSM Fusion Module），利用两个并行工作的 Mamba 块分别处理正向演变和逆向溯源序列，通过双向扫描机制全方位捕捉气候系统中的长期记忆性与复杂的动力学依赖。最后，线性解码投影器（Linear Decoding Projector） 将编码器输出的高维隐层表征直接映射至预测视窗长度，高效生成未来的气候状态序列与趋势。

图 1：MambaClimateNet 模型架构图

关键组件详解

模型整体架构由以下四个核心阶段组成：

1. 倒置变量嵌入 (Inverted Variate Embedding)

为了解决多变量气候数据的异质性问题，模型并未对时间步进行切片，而是采用了倒置嵌入策略：

• 维度变换：输入张量 $\mathbf{X} \in \mathbb{R}^{B \times L \times N}$ 被重塑为 $\mathbf{X}^{T} \in \mathbb{R}^{B \times N \times L}$，其中 $N$ 个气候变量被视为独立的语义 Token。
• 特征映射：通过线性层将完整的回顾窗口长度 $L$ 映射为隐层维度 $D_{model}$。此设计保证了模型能够完整保留单一变量的历史物理特性，为后续的变量间相关性建模奠定基础。

2. 卷积特征预处理 (Convolutional Feature Pre-processing)

在进入状态空间之前，潜在特征首先通过一个基于 Conv1d 的前馈网络（FeedForward Network）进行增强。

• 该模块通过局部卷积操作，提取特征空间内的短程局部依赖（Local Dependencies）。
• 引入非线性激活函数，增强了模型对气候突变信号的表征能力，平滑了输入噪声。

3. 双向 Mamba 融合机制 (Bidirectional Mamba Fusion)

这是 MambaClimateNet 的核心推理引擎。考虑到气候系统不仅存在顺向的时间演变，还存在逆向的因果追溯，模型设计了双向 SSM 扫描机制：

• **正向扫描 (Forward Path)**：利用 Mamba 模块处理标准变量序列，捕捉顺向的系统动力学特征。
• **逆向扫描 (Backward Path)**：通过对序列进行 Flip 操作，利用独立的 Mamba 模块捕捉逆向依赖，随后再次翻转复原。
• 特征融合：将双向路径输出进行加权融合（Element-wise Summation），构建出具有全上下文感知（Full-context Awareness）的气候状态表征。

4. 线性解码投影 (Linear Decoding Projector)

解码阶段摒弃了复杂的自回归生成，采用高效的线性投影器（Linear Projector）。

• 直接将编码器输出的高维隐层特征 $\mathbb{R}^{B \times N \times D}$ 映射回预测视窗长度 $\mathbb{R}^{B \times N \times T}$。
• 配合 RevIN (Reversible Instance Normalization) 技术，有效缓解了气候数据中的分布漂移（Distribution Shift）问题，确保预测结果的稳健性。

运行条件

🧩 环境要求

项目	配置说明
GPU	NVIDIA RTX 3090 × 1（24 GB 显存）
CPU	10 vCPU
内存	60 GB
操作系统	Ubuntu 22.04
开发框架	PyTorch 2.1.0
Python	3.9

📦 核心依赖

库名称	版本	功能说明
PyTorch	2.8.0+cu128	深度学习框架
mamba-ssm	2.2.6.post3	选择性状态空间模块（核心组件）
numpy	2.0.2	数值计算与矩阵运算
pandas	2.3.3	数据加载与时间序列预处理
scikit-learn	1.6.1	数据标准化与评估指标计算
matplotlib	3.9.4	可视化工具
einops	0.8.1	张量重排与模型结构定义
transformers	4.57.3	可扩展 Transformer 模块接口
tqdm	4.67.1	训练进度可视化

运行说明

如何配置环境并运行 MambaClimateNet

1️⃣ 安装依赖

运行以下命令以安装所有必要的 Python 依赖：

# 创建并激活新环境
conda create -n mamba_climate python=3.9
conda activate mamba_climate

# 安装 Python 依赖：
pip install -r requirements.txt

2️⃣ 准备数据

确保你的项目目录组织如下：

├── dataset/
│   ├── China.csv          # 中国气象数据
│   └── Global.csv         # 全球气象数据
├── scripts/
│   └── multivariate_forecasting/
│       └── MambaClimateNet_mu.sh   # 多步预测批量脚本
├── results/               # 模型结果与可视化输出
├── run.py                 # 主程序入口
├── requirements.txt       # 依赖包列表
└── README.md              # 项目说明文档

数据集介绍 (Dataset Description)

1.1 概述 (Overview)

为了全面评估模型在不同时间粒度和不同气候条件下的时空预测能力，本研究采用了两个经过严格筛选和处理的气象数据集：中国区域日级气象数据集 (China-Daily) 和 **全球多站点高频气象数据集 (Global-Hourly)**。原始数据来源于 Kaggle 开源平台，经过清洗、去噪及标准化处理后，构建了适用于深度学习模型训练的标准时间序列基准。

1.2. 数据集详情 (Details)

数据集 A：中国区域日级气象数据集 (China Meteorological Dataset) 该数据集主要用于评估模型对长期气候趋势和长序列依赖关系 (Long-term Dependencies) 的捕捉能力。

• 时间跨度：2020年1月1日至2023年12月（日级分辨率）。
• 特征维度：包含8个关键气象指标，涵盖温度、降水、风速及日照时长。具体包括：
  • 温度指标：最高温 (temp_max)、最低温 (temp_min)、平均温 (temp_mean)。
  • 动力指标：最大风速 (windspeed_max)、阵风峰值 (windgusts_max)。
  • 其他指标：降水量 (precip)、日照时长 (sunshine_duration)。
• 数据特性：该数据集具有显著的季节性波动特征（如图 2 所示），适合验证模型在处理季节性周期变化（Seasonality）方面的有效性。

  
  图 2：中国气象数据趋势分析面板 (2020-2023年)

数据集 B：全球多站点高频气象数据集 (Global Weather Dataset) 该数据集用于评估模型在高频采样和多变量耦合场景下的细粒度预测性能。

• 数据结构：包含多个不同地理区域（标记为 Country 0 至 Country 4）的混合气象记录。
• 时间粒度：亚小时/小时级（High-frequency），能够捕捉气象要素的瞬时突变。
• 特征维度：包含10个多维特征，侧重于大气物理状态的描述。具体包括：
  • 基础气象：平均气温、最大风速、气压 (pressure)、降水。
  • 环境感知：相对湿度 (humidity)、云量 (cloud)、体感温度 (feels_like)、紫外线指数 (uv_index)。
• 数据特性：该数据展现了更复杂的高频波动和非线性特征（如图 3 所示），要求模型具备较强的局部特征提取能力和对突发气象事件的响应能力。

  
  图 3：全球气象数据趋势分析面板 (2025年)

3. 数据预处理 (Data Preprocessing) 为了消除原始数据中的噪声并提高模型的收敛速度，我们对原始 Kaggle 数据进行了以下预处理步骤：

4. **缺失值处理 (Imputation)**：对于时间序列中的个别缺失值，采用线性插值法（Linear Interpolation）进行填充；对于连续缺失，采用基于时间周期的均值填充。

5. **异常值剔除 (Outlier Removal)**：基于 $3\sigma$ 原则（Z-score）剔除极端异常的气象数值，确保训练数据的分布稳定性。

6. **数据归一化 (Normalization)**：为了消除不同物理量纲（如气压 1000hPa 与温度 20°C）之间的差异，采用 Z-Score 标准化方法将所有特征缩放至均值为0、方差为1的标准正态分布：

$$x^{\mathrm{\prime }}=\frac{x-\mu }{\sigma }x\text{'}\; =\; \backslash frac\{x\; -\; \backslash mu\}\{\backslash sigma\}$$

其中 $\mu$ 为样本均值，$\sigma$ 为标准差。

3️⃣ 模型运行与评估

🧠 模型主要参数分类

参数类别	参数名称	值	说明
数据相关	--root_path	./dataset/China/	数据集目录
	--data_path	China.csv	数据文件路径
	--features	M	多变量预测方式（M: multivariate）
序列长度	--seq_len	96	输入时间窗口长度
	--pred_len	12 / 24 / 36 / 48	预测步长
模型结构	--d_model	256	模型隐藏维度
	--e_layers	2	编码层数
	--d_state	2	Mamba 状态维度（控制状态空间深度）
	--enc_in, --dec_in, --c_out	7/9	输入/输出通道数
训练参数	--batch_size	8	批次大小
	--learning_rate	0.00005	学习率
	--itr	1	实验轮数

📊 评估指标说明

指标	英文全称	含义
MSE	Mean Squared Error	预测值与真实值差值的平方平均，用于衡量整体拟合程度，越小越好。
MAE	Mean Absolute Error	预测值与真实值差值的绝对值平均，反映预测误差的平均水平，越小越好。

测试说明

本项目已提供完整运行脚本，只需执行以下命令即可完成 训练、测试与结果评估：

cd MambaClimateNet
bash ./scripts/multivariate_forecasting/MambaClimateNet_mu.sh

运行结果说明

模型运行完成后，系统将自动在项目根目录下生成结果文件与可视化图像，主要输出路径如下：

输出类型	保存路径	说明
日志与指标结果	./results/	包含每轮训练与测试的日志文件、最终 MSE/MAE 评估结果。
预测可视化图像	./results/vis/	以曲线形式展示预测值与真实观测值的对比，可直观观察模型拟合效果。
模型权重文件	./checkpoints/ (如启用保存)	存储训练得到的模型参数，用于后续推理或复现实验。

下图展示了 MambaClimateNet 模型在不同预测步长（pred_len = 12, 24, 36, 48）下的气象时间序列预测结果。蓝色曲线为真实观测值（Ground Truth），橙色曲线为模型预测值（Prediction）。
从图 4 中可以看到，模型在较短预测步长时能精准捕捉气象变化趋势，随着步长增大，尽管误差略有上升，但整体走势仍保持一致，体现出良好的时序建模能力与稳定性。

图 4：MambaClimateNet 预测结果可视化

实验结果

为了验证 MambaClimateNet 的有效性，我们将本模型与目前最先进（SOTA）的基线模型进行了广泛的对比评估。对比模型涵盖了基于 Mamba 的变体（S-Mamba）、基于 Transformer 的模型（iTransformer）以及经典的深度学习架构（LSTM, BiLSTM, CNN）。实验是在两个真实世界的 China（中国）和 Global（全球）气候数据集上进行的，预测长度设定为 $H \in {12, 24, 36, 48}$。

📈 性能对比

下表展示了多变量时间序列预测的详细实验数据。

说明：表格中 加粗 (Bold) 表示最优结果，下划线 (Underlined) 表示次优结果。

🇨🇳 Dataset: China (中国气象数据)

Model	MambaClimateNet	S_Mamba	iTransformer	LSTM	BiLSTM	CNN
Metric	MSE / MAE	MSE / MAE	MSE / MAE	MSE / MAE	MSE / MAE	MSE / MAE
Pred_len=12	0.686 / 0.577	0.711 / 0.585	0.799 / 0.633	0.777 / 0.603	0.803 / 0.619	0.683 / 0.578
Pred_len=24	0.731 / 0.601	0.777 / 0.620	0.884 / 0.683	0.862 / 0.646	0.889 / 0.656	0.740 / 0.614
Pred_len=36	0.784 / 0.632	0.827 / 0.645	0.884 / 0.683	0.862 / 0.646	1.010 / 0.698	0.832 / 0.666
Pred_len=48	0.866 / 0.675	0.899 / 0.670	0.960 / 0.731	0.966 / 0.706	1.122 / 0.752	0.914 / 0.712

🌍 Dataset: Global (全球气象数据)

Model	MambaClimateNet	S_Mamba	iTransformer	LSTM	BiLSTM	CNN
Metric	MSE / MAE	MSE / MAE	MSE / MAE	MSE / MAE	MSE / MAE	MSE / MAE
Pred_len=12	0.405 / 0.412	0.423 / 0.427	0.474 / 0.464	0.649 / 0.584	0.660 / 0.590	0.439 / 0.459
Pred_len=24	0.415 / 0.426	0.440 / 0.445	0.501 / 0.486	0.657 / 0.579	0.687 / 0.607	0.485 / 0.497
Pred_len=36	0.446 / 0.450	0.468 / 0.471	0.501 / 0.486	0.657 / 0.579	0.724 / 0.610	0.517 / 0.517
Pred_len=48	0.650 / 0.462	0.469 / 0.479	0.513 / 0.496	0.674 / 0.589	0.721 / 0.622	0.547 / 0.538

📝 结果分析

如上表数据所示，MambaClimateNet 在多项指标上展现了显著的优势：

• SOTA 级精度表现： 我们的模型在绝大多数实验设置下均优于所有基线模型。特别是在 Global 数据集的前中期预测（12-36步长）中，MambaClimateNet 相比最强的基线（S_Mamba）和 CNN/Transformer 类模型，MSE 误差显著降低。

• 优异的长序列鲁棒性 (Long-term Robustness)： 随着预测长度的增加（如 $T=36, 48$），传统模型（如 LSTM/CNN）受限于误差累积或感受野瓶颈，性能往往会出现大幅下滑。而 MambaClimateNet 凭借 SSM 架构对长程依赖的强大建模能力，保持了相对稳定的预测精度。即便在最具挑战性的 Global-48 任务中，我们的模型依然保持了最优的 MAE (0.462) 指标。

• 超越 Transformer 的架构优势： 与 iTransformer 相比，MambaClimateNet 不仅在准确率上全面领先（更低的 MSE/MAE），更重要的是它规避了 Transformer 二次复杂度带来的计算负担，证明了线性复杂度的 SSM 架构更适合处理高维、长周期的气候时间序列数据。

🔮 未来展望

尽管 MambaClimateNet 在多项基准测试中取得了 SOTA 性能，但针对复杂气候系统的建模仍有广阔的探索空间。我们在未来的工作中将致力于以下方向：

1. **物理约束增强 (Physics-Informed Modeling)**： 目前模型主要基于数据驱动（Data-driven）。未来我们将尝试引入物理守恒定律（如质量守恒、能量守恒）作为损失函数的正则项，以提高预测结果的物理一致性，减少违反物理常识的预测。

2. **极端天气捕捉 (Extreme Event Forecasting)**： 气候变化的痛点往往在于极端事件（如台风、寒潮）。我们将探索基于极值理论（EVT）或重尾分布的改进方案，提升模型在长尾分布数据上的表现。

3. **气候大模型预训练 (Pre-training on ERA5)**： 借鉴 NLP 领域的 Scaling Law，计划在更大规模的再分析数据集（如 ERA5）上进行无监督预训练，探索基于 Mamba 架构的气候基础模型（Climate Foundation Model），以实现跨区域、跨变量的零样本迁移（Zero-shot Transfer）。

📧 致谢

本项目在代码实现上参考了以下优秀的开源工作，特此致谢：

• Mamba-SSM: 提供了高效的状态空间模型核心实现。
• TimeXer: 提供了标准化的时间序列数据加载与评估框架。