# FGN (Functional Generative Networks) 训练配置示例
# 论文配置：第二阶 Markov（输入两个先验状态帧 x_{t-2},x_{t-1}）预测下一状态 x_t；
# 网格 0.25°（1440×721），潜空间为 6 次细分球面 icosahedral 网格，潜维度 768、
# 处理器 24 层、6 attention heads（每模型种子约 180M 参数）；训练 4 个模型种子，
# 评估时每个种子生成 14 个集合成员（共 56）；目标为公平 CRPS（Eq.4），训练时 N=2。
# 时间步长 6h，15 天预报共 60 帧。
# 当前为连通性验证小配置：虚拟数据 32×32 网格、潜网格 8×8、latent_dim=64。
model:
  start_epoch: 0
  max_epoch: 100
  lr: 1E-3                 # 论文 peak lr=8e-4，cosine 退火（warmup 1000 步）
  patience: 50
  checkpoint_dir: "./data/checkpoints"

  # FGN 结构参数（论文值见注释）
  in_channels: 6           # 通道数（论文 6 个大气变量×13 层 + 6 个表面变量 = 84 通道）
  out_channels: 6
  input_steps: 2           # 第二阶 Markov，输入先验两帧（论文 input_steps=2）
  output_steps: 2          # 自回归预报帧数（论文 6h 步长 15 天 = 60 帧）
  grid_shape: [32, 32]     # 观测网格尺寸（论文 0.25° 1440×721）
  mesh_shape: [8, 8]       # 潜网格尺寸（论文 6 次细分 icosahedral 网格约 40k 节点）
  latent_dim: 64           # 潜空间维度（论文 768）
  num_encoder_layers: 2    # 编码器 GNN 消息传递层数
  num_decoder_layers: 2    # 解码器 GNN 消息传递层数
  num_processor_blocks: 2  # 处理器 graph-transformer 块数（论文 24 层）
  n_heads: 4               # 处理器 attention 头数（论文 6）
  hidden_dim: 64           # 处理器 FFN / MLP 隐藏维度
  noise_dim: 32            # 全局噪声向量维度（论文 32，经条件 LayerNorm 注入）
  num_ensemble_models: 1   # 深度集成的模型种子数（论文 4）
  num_members: 2           # 推理时每个模型种子生成的集合成员数（论文每种子 14，共 56）
  channel_weights: [1, 1, 1, 1, 1, 1]  # CRPS 逐通道权重（论文取 GraphCast/GenCast 权重）

# 整个数据读取流程
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

    # 气象变量（论文 t/z/q/u/v/w 大气变量 + 2t/10u/10v/msl/sst/tp 表面变量；
    # 此处使用 6 个表面变量占位）
    channels: ['2m_temperature', '10m_u_component_of_wind', '10m_v_component_of_wind',
               'mean_sea_level_pressure', 'sea_surface_temperature', 'total_precipitation']

  # DataLoader 配置
  dataloader:
    mask_dtype: "float32"
    batch_size: 2
    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