project:
  name: FourCastNet_v2
  seed: 42
  output_dir: ./result
  checkpoint_dir: ./data/checkpoint

data:
  dataset_dir: ./data/era5_fake
  train_years: [2014,2015]
  val_years: [2016]
  test_years: [2018]
  official_splits:
    train_years: [2014, 2015]
    val_years: [2016, 2017]
    test_years: [2018]
  input_steps: 1
  output_steps: 1
  time_step_hours: 6
  grid_shape: [721, 1440]
  normalize: true
  variables:
    - u10m
    - v10m
    - u100m
    - v100m
    - t2m
    - sp
    - msl
    - tcwv
    - u50
    - u100
    - u150
    - u200
    - u250
    - u300
    - u400
    - u500
    - u600
    - u700
    - u850
    - u925
    - u1000
    - v50
    - v100
    - v150
    - v200
    - v250
    - v300
    - v400
    - v500
    - v600
    - v700
    - v850
    - v925
    - v1000
    - z50
    - z100
    - z150
    - z200
    - z250
    - z300
    - z400
    - z500
    - z600
    - z700
    - z850
    - z925
    - z1000
    - t50
    - t100
    - t150
    - t200
    - t250
    - t300
    - t400
    - t500
    - t600
    - t700
    - t850
    - t925
    - t1000
    - r50
    - r100
    - r150
    - r200
    - r250
    - r300
    - r400
    - r500
    - r600
    - r700
    - r850
    - r925
    - r1000

fake_data:
  time_steps_per_year: 20
  chunk_time_steps: 1
  fill_value: 0.0
  materialize_pattern: true

model:
  profile: smoke
  profiles:
    full_resolution:
      img_size: [721, 1440]
      in_channels: 73
      out_channels: 73
      spectral_transform: sht
      filter_type: non-linear
      scale_factor: 6
      embed_dim: 256
      num_layers: 12
      num_blocks: 8
      normalization_layer: instance_norm
      mlp_mode: distributed
      spectral_layers: 3
      complex_activation: real
      hard_thresholding_fraction: 1.0
      big_skip: true
    smoke:
      img_size: [16, 32]
      in_channels: 73
      out_channels: 73
      spectral_transform: sht
      filter_type: linear
      scale_factor: 4
      embed_dim: 8
      num_layers: 2
      num_blocks: 1
      normalization_layer: instance_norm
      mlp_mode: serial
      spectral_layers: 1
      complex_activation: real
      hard_thresholding_fraction: 0.5
      big_skip: true
      # The smoke profile keeps the same equations but reduces the grid/model.

checkpoint:
  initialize_from: scratch
  prefix: model_bak
  finetune_from: ./data/checkpoint/one_step/model_bak.pt
  strict: true

training:
  stage: one_step
  epochs: 3
  batch_size: 1
  num_workers: 0
  learning_rate: 0.0006
  weight_decay: 0.0
  optimizer_betas: [0.9, 0.95]
  max_grad_norm: 32.0
  scheduler: cosine
  amp: false
  max_train_batches: null
  max_val_batches: null
  finetune:
    autoregressive_steps: 2
    epochs: 3
    learning_rate: 0.0001

distributed:
  backend: nccl
  master_addr: 127.0.0.1
  master_port: 29500

inference:
  checkpoint_path: ./data/checkpoint/finetune/model_bak.pt
  rollout_steps: 1
  max_samples: 1
  save_normalized: false
  output_dir: ./result/output

visualization:
  variable: t2m
  sample_index: 0
  output_dir: ./result/figures
  cmap: coolwarm

slurm:
  job_name: fcnv2_train
  nodes: 1
  gpus_per_node: 8
  cpus_per_task: 8
  time: "24:00:00"
  partition: null
  conda_env: fourcastnetv2_develop