protocol: non_official_protocol
seed: 42

paths:
  data_root: data
  static_root: data/static
  checkpoint: data/checkpoint/model_bak.pth
  training_metrics: output/training/metrics.json
  inference_output: output/inference/forecast.nc
  visualization_output: figures/fuxi21_t2m.png

data:
  input_steps: 2
  output_steps: 1
  time_step_hours: 6
  grid_size: [721, 1440]
  crop_size: null
  num_workers: 0
  splits:
    train:
      years: [2021,2022]
      time_steps: 10
    val:
      years: [2023]
      time_steps: 10
    inference:
      years: [2024]
      time_steps: 10
  hdf5:
    chunks: [1, 1, 64, 64]
    compression: gzip
    compression_level: 1

variables:
  pressure_levels: [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
  pressure: [z, t, u, v, q]
  surface: [msl, t2m, d2m, sst, ws10m, ws100m, u10m, v10m, u100m, v100m, lcc, mcc, hcc, tcc, ssr, ssrd, fdir, ttr, tcw, tp]
  diagnostic: [ssr, ssrd, fdir, ttr, tp]
  mapping:
    pressure_order: variable_major
    pressure_name: "{variable}{level}"
    surface_name: "{variable}"
    source_dataset: fields
    units: source_native
    transform: identity

model:
  profile: full
  static_fields_file: data/static/static_fields.npy
  channel_mask_file: data/static/channel_mask.npy
  activation_checkpointing: true
  profiles:
    full:
      grid_size: [721, 1440]
      patch_size: 6
      embed_dim: 1536
      depth: 30
      num_heads: 24
      mlp_dim: 4096
      window_size: 20
    smoke:
      grid_size: [13, 12]
      patch_size: 6
      embed_dim: 32
      depth: 2
      num_heads: 4
      mlp_dim: 64
      window_size: 2

training:
  device: auto
  epochs: 5
  batch_size: 1
  precision: fp32
  distributed_strategy: ddp
  fsdp_sharding: full_shard
  gradient_accumulation_steps: 1
  optimizer: AdamW
  learning_rate: 0.0001
  min_learning_rate: 0.000001
  weight_decay: 0.01
  scheduler: CosineAnnealingLR
  gradient_clip_norm: 1.0
  channel_weights: null
  checkpoint_mode: scratch
  load_checkpoint: null
  save_checkpoint: data/checkpoint/model_bak.pth

inference:
  device: auto
  checkpoint: data/checkpoint/model_bak.pth
  split: inference
  steps: 1
  zero_diagnostic_feedback: false
  output_file: output/inference/forecast.nc

visualization:
  input_file: output/inference/forecast.nc
  output_file: figures/fuxi21_t2m.png
  channel: t2m
  cmap: coolwarm