conf/config.yaml

model:

    start_epoch: 0
    max_epoch: 100
    lr: 5e-4
    beta_1: 0.9
    beta_2: 0.95
    weight_decay: 1e-5
    warmup_epochs: 10
    warmup_start_lr: 1e-8
    eta_min: 1e-8

    # Stormer architecture params
    in_img_size: [128, 256]       # 1.40625° resolution (lat, lon)
    patch_size: 2
    hidden_size: 1024
    depth: 24
    num_heads: 16
    mlp_ratio: 4.0

    # Training strategy
    list_train_intervals: [6, 12, 24]  # hours, randomly chosen per batch
    steps: 1                            # number of autoregressive rollout steps
    val_lead_times: [6, 72]            # lead times for validation (hours)
    data_freq: 6                        # data frequency (hours)

    checkpoint_dir: "./data/checkpoints"
    patience: 50

    # Normalization constants directory (official Stormer .npz files)
    normalize_dir: "./data/normalize/"

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

  # dataset设定
  dataset:
    type: "hdf5"
    data_dir: './data/' # "$ONESCIENCE_DATASETS_DIR/ERA5/newh5/"

    train_time: [2000, 2001]
    val_time: [2002]
    test_time: [2003]
    img_size: [128, 256]
    verbose: true
    cache: false

    # Stormer variables: 4 surface + 5 × 13 pressure levels = 69 variables
    channels:
      # Surface variables (4)
      - "2m_temperature"
      - "10m_u_component_of_wind"
      - "10m_v_component_of_wind"
      - "mean_sea_level_pressure"

      # Geopotential at 13 pressure levels
      - "geopotential_50"
      - "geopotential_100"
      - "geopotential_150"
      - "geopotential_200"
      - "geopotential_250"
      - "geopotential_300"
      - "geopotential_400"
      - "geopotential_500"
      - "geopotential_600"
      - "geopotential_700"
      - "geopotential_850"
      - "geopotential_925"
      - "geopotential_1000"

      # U component of wind at 13 pressure levels
      - "u_component_of_wind_50"
      - "u_component_of_wind_100"
      - "u_component_of_wind_150"
      - "u_component_of_wind_200"
      - "u_component_of_wind_250"
      - "u_component_of_wind_300"
      - "u_component_of_wind_400"
      - "u_component_of_wind_500"
      - "u_component_of_wind_600"
      - "u_component_of_wind_700"
      - "u_component_of_wind_850"
      - "u_component_of_wind_925"
      - "u_component_of_wind_1000"

      # V component of wind at 13 pressure levels
      - "v_component_of_wind_50"
      - "v_component_of_wind_100"
      - "v_component_of_wind_150"
      - "v_component_of_wind_200"
      - "v_component_of_wind_250"
      - "v_component_of_wind_300"
      - "v_component_of_wind_400"
      - "v_component_of_wind_500"
      - "v_component_of_wind_600"
      - "v_component_of_wind_700"
      - "v_component_of_wind_850"
      - "v_component_of_wind_925"
      - "v_component_of_wind_1000"

      # Temperature at 13 pressure levels
      - "temperature_50"
      - "temperature_100"
      - "temperature_150"
      - "temperature_200"
      - "temperature_250"
      - "temperature_300"
      - "temperature_400"
      - "temperature_500"
      - "temperature_600"
      - "temperature_700"
      - "temperature_850"
      - "temperature_925"
      - "temperature_1000"

      # Specific humidity at 13 pressure levels
      - "specific_humidity_50"
      - "specific_humidity_100"
      - "specific_humidity_150"
      - "specific_humidity_200"
      - "specific_humidity_250"
      - "specific_humidity_300"
      - "specific_humidity_400"
      - "specific_humidity_500"
      - "specific_humidity_600"
      - "specific_humidity_700"
      - "specific_humidity_850"
      - "specific_humidity_925"
      - "specific_humidity_1000"

    # Short names for evaluation
    variables:
      - "t2m"        # 2m temperature
      - "u10"        # 10m U wind component
      - "v10"        # 10m V wind component
      - "msl"        # Mean sea level pressure
      - "z500"       # Geopotential at 500 hPa
      - "t850"       # Temperature at 850 hPa

    # 时间配置
    time_range: ["2000-01-01", "2020-12-31"]
    time_steps: 1
    time_res: 6

    # 空间配置
    spatial_resolution: [1.40625, 1.40625]

    # 采样配置
    num_samples: -1  # -1 表示使用全部数据
    shuffle: true
    random_seed: 42

    # 领域特定配置
    extra:
      levels: [500, 850, 1000]
      lat_range: [-90, 90]
      lon_range: [0, 360]

  # 数据转换配置
  transforms:
    - type: "Normalize"
      params:
        mean: [0.0, 0.0, 288.0, 101325.0, 50000.0, 270.0]
        std: [5.0, 5.0, 15.0, 1000.0, 5000.0, 10.0]
        keys: ["input", "target"]

    - type: "ToTensor"
      params:
        keys: null  # null表示转换所有numpy数组

  # DataLoader配置
  dataloader:
    mask_dtype: "float32"
    batch_size: 1
    num_workers: 1
    pin_memory: true
    drop_last: true
    shuffle: false  # 使用sampler时设为false
    prefetch_factor: 2
    persistent_workers: true

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