# Prithvi WxC 训练配置示例
# 论文配置为 0.5°×0.625°（360/361×576）网格、160 动态变量、embed_dim=2560、
# 编码器 13 本地+12 全局块、解码器 3 本地+2 全局块（约 23 亿参数，需大规模显存）。
# 当前为连通性验证小配置（约百万级参数）。
model:
  start_epoch: 0
  max_epoch: 100
  lr: 1E-3
  patience: 50
  checkpoint_dir: "./data/checkpoints"

  # Prithvi WxC 结构参数
  in_channels: 6                 # 论文为 160（20 单层 + 10 变量 ×14 层）
  input_size_time: 2             # 两个输入时刻
  in_channels_static: 4          # 静态通道数（论文 MEMRA-2 为 4）
  n_lats_px: 32                  # 数据纬度方向像元数（论文约 360）
  n_lons_px: 64                  # 数据经度方向像元数（论文 576）
  patch_size_px: [2, 2]          # token 尺寸（论文 2×2 像素）
  mask_unit_size_px: [8, 8]      # 掩码单元尺寸
  mask_ratio_inputs: 0.0         # 预训练为 0.5；预报滚动微调为 0.0
  embed_dim: 32                  # 隐藏维度（论文 2560）
  n_blocks_encoder: 1            # 编码器本地-全局对数量（论文 13 对 -> 25 块）
  n_blocks_decoder: 1            # 解码器本地-全局对数量（论文 5 块）
  mlp_multiplier: 4.0
  n_heads: 4                     # 注意力头数（论文 16）
  dropout: 0.0
  drop_path: 0.0
  parameter_dropout: 0.0
  residual: "none"               # none / temporal / climate
  masking_mode: "global"
  positional_encoding: "absolute"
  encoder_shifting: false
  decoder_shifting: false

# 整个数据读取流程
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, 64]
    verbose: true
    cache: false

    # 气象变量（论文 160 通道配置的子集，用于连通性验证）
    channels: ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature',
               'mean_sea_level_pressure', 'geopotential_500', 'temperature_850']

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