# DGMR 训练配置示例
# 论文配置：4 帧雷达观测上下文（num_context=4）预测 18 帧（forecast_steps=18），
# 空间 256×256 单通道，latent_channels=768 / context_channels=384（生成器约 13.4M 参数）。
# 当前默认配置为连通性验证小配置：空间 128×128，预测 6 帧；fake_data.py
# 根据 num_context、forecast_steps 与 batch_size 自动推导每年所需时间长度 T。
model:
  start_epoch: 0
  max_epoch: 100
  lr: 1E-4                 # 论文使用 Adam lr=1e-4 分别训练生成器与判别器
  lr_disc: 1E-4
  patience: 50
  checkpoint_dir: "./data/checkpoints"

  # DGMR 结构参数（论文值见注释）
  forecast_steps: 6        # 预测帧数，论文为 18；小配置取 6
  num_context: 4           # 上下文（观测）帧数，论文为 4
  input_channels: 1        # 每帧通道数，论文为 1（雷达回波）
  output_shape: 128        # 空间尺寸（需 32 整除），论文为 256
  conv_type: "standard"    # standard / coord / 3d
  latent_channels: 384     # 论文为 768
  context_channels: 192    # 论文为 384
  generation_steps: 6      # 网格单元正则器 Monte-Carlo 采样次数，论文为 6
  grid_lambda: 20.0        # 网格单元正则器权重，论文为 20
  precip_weight_cap: 24.0  # 网格单元权重函数上限，论文为 24

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

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

    # 气象变量（DGMR 面向单通道雷达回波场，此处为占位变量名）
    channels: ['rain_radar']

  # DataLoader 配置
  # 注意：DGMR 判别器内部含 BatchNorm1d，训练 batch 必须 >= 2。
  # fake_data.py 会保证每个年份至少产生 batch_size 个完整时间窗口。
  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