Metadata-Version: 2.4
Name: geo_correction
Version: 0.1.0
Summary: Meteorological grid correction training and inference toolkit
Author: geo_correction maintainers
License-Expression: MIT
Requires-Python: >=3.8
Requires-Dist: numpy>=1.21
Requires-Dist: pyyaml>=6.0
Requires-Dist: torch>=1.11
Provides-Extra: test
Requires-Dist: pytest>=7; extra == 'test'
Description-Content-Type: text/markdown

# 气象 AI 订正训练工程（`train/`）说明

本文件是基于当前 `train/` 目录代码结构整理的新版 README 草案，目标是帮助你快速理解工程组织、主流程、配置方式，以及当前仓库里需要注意的兼容性问题。

## 1. 项目定位

该工程用于气象格点 AI 订正，覆盖以下典型场景：

- 常规地面要素订正（温度、湿度、风、气压、能见度等）
- 降水订正（TP）
- 高空层（Upar）订正
- 海洋/海浪/海流相关订正
- 在线业务推理（从模式数据生成订正产品）

工程同时包含训练、评估、推理、数据集制作、结果可视化、物理量转换等模块。

## 2. 工程结构总览

建议优先理解以下目录：

```text
train/
├─ main.py / main_ddp.py                # 常规训练入口（单卡 / DDP）
├─ eval.py / eval_tp.py                 # 离线评估入口（常规/降水）
├─ inference.py                         # 在线推理入口（实时业务链路）
├─ module/
│  ├─ pretrain.py                       # 训练器（Trainer/LRESOTrainer/Ocean...）
│  └─ forecast.py                       # 推理器（WeatherForecast/TPForecast/...）
├─ data/
│  ├─ dataset.py                        # Dataset 与归一化/反归一化逻辑
│  ├─ makeData.py                       # 数据集制作（离线）
│  └─ makeOnlineData.py                 # 在线模式数据读取/拼装
├─ model/                               # 网络结构（AFNO/UNet/HAN/Swin/GraphCast/MOS/...）
├─ loss/                                # 损失函数与指标
├─ tools/
│  ├─ config.py                         # 配置读取 + 环境路径 + 区域定义
│  └─ utils.py                          # 通用工具（索引、归一化读取、插值等）
├─ yaml/                                # 训练/测试 YAML 配置（085/095/...）
├─ json_defs/                           # 数据制作任务定义（JSON）
├─ saved_models/                        # 模型权重（示例）
├─ run/                                 # TensorBoard 日志目录
└─ images/                              # 验证图、误差图等输出（运行后生成）
```

## 3. 核心流程（训练 / 评估 / 推理）

### 3.1 训练流程（常规入口）

主入口是 `main.py`（常规）和 `main_ddp.py`（DDP / 多卡）。

训练流程大致为：

1. 读取 YAML 配置
2. 根据 `Input_ELE / Output_ELE` 从 `mean_std_feature.json` 计算特征索引
3. 构建 Trainer（陆地或海洋）
4. Trainer 内部创建 Dataset / DataLoader、模型、优化器、损失函数
5. 训练 + 验证 + TensorBoard + 保存最优 checkpoint

代码关系（关键路径）：

- `main.py` -> `module.pretrain.LRESOTrainer / OceanHighWindTrainer`
- `module/pretrain.py` -> `data/dataset.py` (`ECMWFData` / `OceanData`)
- `module/pretrain.py` -> `model/*`（AFNO/HAN/Swin/UNet/TransUNet/OceanTransformer/MOS 等）

### 3.2 评估流程（离线）

- `eval.py`: 常规、T799、海面大风评估入口
- `eval_tp.py`: 降水评估入口

两者都会进入 `module/forecast.py` 的推理器类，逻辑通常包括：

1. 自动查找 `saved_models/...` 下最佳 checkpoint
2. 加载权重
3. 执行推理
4. 反归一化 / 后处理 / 重塑范围
5. 保存 NetCDF 产品和专题图（默认开启）

### 3.3 在线推理流程（业务化）

入口为 `inference.py`，整体流程：

1. 根据业务时间参数（起报时间、时效）读取在线模式数据
2. 通过 `data.makeOnlineData` 的 `RUNNER_REGISTRY` 调用对应数据拼装器（`ECMWF` / `T799` / `T799Ocean` / `T1279`）
3. 选择 YAML 指定的特征通道
4. 推理预处理（含降水单位处理、归一化）
5. 调用 `module.forecast` 对应 Forecast 类完成推理与产品落盘

## 4. 配置体系（YAML）

`yaml/` 下有大量按任务拆分的配置文件，例如：

- `yaml/085/config_085_HAN_Conver_Train.yaml`
- `yaml/085/config_085_SWIN_TP_Test.yaml`
- `yaml/095/config_095_AFNO_Conver_Test.yaml`
- `yaml/095/config_095_OceanTransUNet_..._Test.yaml`

### 4.1 常见字段

`model`

- `name`: 模型名称（如 `AFNONet` / `HAN` / `SwinTransformerV2Cr` / `TransUNet` / `OceanTransformer` / `MOS-Global`）

`data`

- `root_dir`: 数据集根目录（通常包含 `train/valid/test` 子目录）
- `statistic_dir`: 统计量目录（`x_mean.npy/x_std.npy/y_mean.npy/y_std.npy` 所在根目录）
- `mean_std_feature`: 特征与目标索引映射 JSON（`Feature/Target`）
- `AreaType`: 区域（如 `CHN`、`095_Area`、海洋区域字符串）
- `DataType`: 任务数据类型（如 `Conver` / `TP` / `Upar` / `Ocean`）
- `TaskType`: `regression` 或 `classification`
- `DZ_Befor_ELE`: 订正前要素名列表
- `Input_ELE`: 模型输入要素名列表
- `Output_ELE`: 模型输出要素名列表
- `grid_resolution`: 网格分辨率
- `batch_size`, `num_workers`, `shuffle`
- `isnormal`: 是否归一化
- `scale`: 是否降尺度（`1` 表示不缩放）
- `skip_normal_tp`: 降水任务可选，跳过 TP 归一化

`train_config`

- `gpu`, `gpuid`, `seed`
- `s_epoch`, `epoch`, `lr`
- `lat_h/lon_w`: 输入尺寸
- `lat_label_h/lon_label_w`: 标签尺寸
- `loss_name`
- 任务特定字段（如 `is_reshape_range`, `pos_weight`, `is_logistic` 等）

## 5. 数据集约定（HDF5 / 归一化）

### 5.1 `data/dataset.py` 的主要 Dataset

- `ECMWFData`: 常规任务
- `ECMWFDataTP`: 降水任务（支持 `skip_normal_tp`）
- `OceanData`: 海洋任务（会对 NaN 做 `nan_to_num`）

### 5.2 HDF5 常用键（按代码读取逻辑）

输入：

- `meteo_ecmwf`
- `ftime`
- `dtime`

标签（会按顺序自动匹配其一）：

- `meteo_era5`
- `meteo_cldas`
- `meteo_era5_hycom`
- `meteo_nbs`
- `meteo_kd`

### 5.3 归一化与索引

- 训练/评估大量逻辑依赖 `mean_std_feature.json` 中的 `Feature` / `Target` 字段
- `tools.utils.get_indices_data()` 会根据 YAML 中要素名列表映射到通道索引
- 同时依赖 `x_mean.npy/x_std.npy/y_mean.npy/y_std.npy`

## 6. 模型与任务分支（当前代码支持）

`model/__init__.py` 当前导出的主要模型包括：

- `HAN`
- `SmaAt_UNet`
- `NestedUNet`
- `FFNet5`
- `GraphCast`
- `AFNONet`
- `SwinTransformerV2Cr`
- `TransUNet`
- `OceanTransformer`

训练器与推理器的任务分支主要在：

- `module/pretrain.py`
- `module/forecast.py`

其中：

- 常规场：`LRESOTrainer` / `WeatherForecast`
- 降水：`TPForecast`（推理）
- 095 项目：`Weather095Forecast`
- 海洋大风：`OceanHighWindTrainer` / `OceanHighWindForcast`
- Ocean 其他：`OceanForcast`

## 7. 输出目录与产物

### 7.1 训练输出

- `saved_models/<数据源>/<区域>/<模型>/<DataType...>/`
  - 权重命名示例：`AFNONet_0.533_20.pth`
- `run/<数据源>/<区域>/<模型>/<DataType...>/`
  - TensorBoard 日志

### 7.2 验证可视化输出

- `images/vailding/...`
  - 单轮可视化图
  - 误差对比图
  - 误差箱线图

### 7.3 推理产品输出

`module/forecast.py` 会将推理结果保存为：

- NetCDF (`.nc`)
- 专题图 (`.png`)
- 部分场景会追加站点插值 CSV（取决于具体 Forecast 子类）

## 8. 快速上手（建议版）

### 8.1 环境依赖（现状）

仓库根部有一个 `requirment.sh`，目前只写了旧版 PyTorch 安装示例（`torch==1.11.0 + cu102`），不足以覆盖当前工程依赖。

建议你后续补一个完整 `requirements.txt` 或 `environment.yml`。

### 8.2 训练（常规）

示例（085/095 类配置）：

```bash
python main.py --train_config config_095_AFNO_Conver_Train --mode train --area land
```

海洋训练：

```bash
python main.py --train_config config_095_OceanTransUNet_Ocean_hjocean_-15.0N_30.0E-30.0N_100.0E_Train --mode train --area ocean
```

DDP（海洋大风类训练器）：

```bash
torchrun --nproc_per_node=2 main_ddp.py --train_config config_xxx --mode train --area ocean
```

### 8.3 评估

常规场：

```bash
python eval.py --test_config config_095_AFNO_Conver_Test --mode ecmwf
```

T799/海洋大风：

```bash
python eval.py --test_config config_xxx --mode t799
python eval.py --test_config config_xxx --mode oceanwins
```

降水：

```bash
python eval_tp.py --test_config config_085_SWIN_TP_Test
```

### 8.4 在线推理

```bash
python inference.py -F 20240101 -H 0 -D 6 -C config_095_AFNO_Conver_Test -R
```

批量时效模式（代码中使用 `-D -1` 触发多时效）：

```bash
python inference.py -F 20240101 -H 0 -D -1 -C config_xxx -R
```

## 9. 已知注意事项（基于当前代码扫描）

以下是我根据当前工程文件确认的事项，建议在正式使用前统一整理：

### 9.1 启动目录 / 导入路径不一致

部分脚本使用：

- `from train.tools.config import load_config`

另一些脚本使用：

- `from tools.config import load_config`

这会导致“从 `train/` 目录运行”与“从上级目录把 `train` 当包运行”之间存在不一致。建议后续统一导入风格（推荐统一为包内相对导入或统一运行入口）。

### 9.2 `tools/config.py` 有导入时副作用，强依赖外部环境配置

`tools/config.py` 在 import 时会立即读取：

- `.../env/<env>/application.yml`

如果该文件不存在，很多脚本会在 import 阶段直接失败。建议：

- 将环境路径配置与 `load_config()` 解耦
- 或提供本地默认 `application.yml` 模板

### 9.3 部分旧入口脚本与当前 `module/pretrain.py` 可能已不一致

例如：

- `main_tp.py` 依赖 `TPTrainer`, `get_criterion`
- `main_mos.py` 依赖 `MOSTrainer`

但在当前 `module/pretrain.py` 中未检索到对应定义（看起来像历史重构后的遗留入口）。

同时，`main_tp.py` / `main_mos.py` 默认将 YAML 路径拼接为 `./yaml/<config>.yaml`，而对应配置多数实际位于 `yaml/085/` 或 `yaml/095/` 子目录。

### 9.4 `inference.py` 存在 095 分支类名未导入风险

`inference.py` 中使用了 `Weather095Forecast` 分支，但顶部导入列表里未包含该类，需补充导入后再用。

### 9.5 `data/makeOnlineData.py` 的 `T799Ocean` 分支需再检查

`run_t799Ocean_online_data()` 中出现了对 `np/xr` 和 `ds_ocean_mask` 的使用，但该文件顶部未看到 `numpy/xarray` 导入，且 `ds_ocean_mask` 初始化行被注释掉，存在直接运行报错风险。

### 9.6 `module/pretrain.py` 的 `get_optimizer()` 需确认 `optim` 导入

`get_optimizer()` 中使用 `optim.Adam / optim.AdamW / optim.SGD`，建议确认文件顶部是否显式导入 `torch.optim as optim`（当前扫描结果看起来存在缺失风险）。

### 9.7 测试脚本可能滞后于当前接口

例如 `tests/test_utils.py` 中对 `get_indices_data()` 的调用未传 `mean_std_feature`，而当前函数实现会检查该路径是否存在。

## 10. 推荐的后续整理方向（可选）

如果你准备把这个工程对外或给新同事使用，建议优先做这 5 件事：

1. 统一所有入口脚本的导入路径（`tools.config` vs `train.tools.config`）
2. 为 `tools/config.py` 提供本地开发默认配置或延迟加载
3. 清理/修复失效入口（`main_tp.py`、`main_mos.py` 等）并统一 YAML 路径策略
4. 补一个可复现环境文件（`environment.yml`）
5. 增加最小可运行 smoke test（至少覆盖 `load_config`、Dataset、一个模型前向）

