# GraphDOP 训练配置示例
# 论文配置：输入一个 12 小时观测窗口（O96 reduced Gaussian 网格约 1°、40320 潜节点），
# 经 GNN 编码器映射到潜网格、Transformer 处理器推进时间、GNN 解码器预测下一窗口观测；
# 潜空间通道 1024，WMSE 目标，18 年数据（2004-2021）训练，64×H100 70k 步。
# 当前为连通性验证小配置：虚拟数据 32×32 网格、潜网格 8×8、latent_dim=64。
model:
  start_epoch: 0
  max_epoch: 100
  lr: 1E-3                 # 论文起始 lr=1e-3，cosine 退火到 3e-7（warmup 1000 步）
  patience: 50
  checkpoint_dir: "./data/checkpoints"

  # GraphDOP 结构参数（论文值见注释）
  in_channels: 6           # 观测通道数（论文为多仪器通道，如 ATMS/AMSU-A/IASI 等）
  out_channels: 6          # 预报通道数
  input_steps: 2           # 输入窗口帧数（论文为单个 12h 窗口；time_step=6h 故取 2 帧）
  output_steps: 2          # 输出窗口帧数（论文为下一个 12h 窗口）
  grid_shape: [32, 32]     # 观测网格尺寸（论文 O96 约 1°）
  mesh_shape: [8, 8]       # 潜网格尺寸（论文 O96 reduced Gaussian 40320 节点）
  latent_dim: 64           # 潜空间通道数（论文 1024）
  num_encoder_layers: 2    # 编码器 GNN 消息传递层数
  num_decoder_layers: 2    # 解码器 GNN 消息传递层数
  num_processor_blocks: 1  # 处理器 Transformer 块数
  n_heads: 4
  hidden_dim: 64
  channel_weights: [1, 1, 1, 1, 1, 1]  # WMSE 逐通道权重（论文 w_{c,i} 经验值）

# 整个数据读取流程
datapipe:
  name: "ERA5"
  task: "weather_forecasting"

  dataset:
    type: "hdf5"
    data_dir: './data/'
    train_time: [1951, 1952]
    val_time: [1953]
    test_time: [1954]
    img_size: [32, 32]
    verbose: true
    cache: false

    # 气象变量（论文观测类型占位：ATMS 亮温/掩星弯角/散射计后向散射/雷达高度计有效波高/常规观测）
    channels: ['atms_brightness_temperature', 'gpsro_bending_angle', 'ascat_sigma0',
               'significant_wave_height', '2m_temperature', '10m_wind_speed']

  # DataLoader 配置
  dataloader:
    mask_dtype: "float32"
    batch_size: 4
    num_workers: 1
    pin_memory: true
    drop_last: true
    shuffle: false
    prefetch_factor: 2
    persistent_workers: true

  # 分布式配置
  distributed:
    enabled: true
    sampler: "DistributedSampler"
    rank: 0
    world_size: 2
    shuffle: true
    seed: 42
    drop_last: true