config/Model.yaml

### base config ###
# -*- coding: utf-8 -*-
full_field: &FULL_FIELD
  loss: 'l2'
  max_epochs: 20
  batch_size: 32
  num_data_workers: 4
  dt: 1                     # how many timesteps ahead the model will predict
  n_history: 0              # how many previous timesteps to consider
  prediction_type: 'iterative'
  prediction_length: 41     # applicable only if prediction_type == 'iterative'
  n_initial_conditions: 5   # applicable only if prediction_type == 'iterative'
  ics_type: "default"       # default or datetime
  date_strings: ["2018-09-14 00:00:00"]
  save_raw_forecasts: !!bool True
  save_channel: !!bool False
  masked_acc: !!bool False
  maskpath: None
  perturb: !!bool False
  add_grid: !!bool False
  N_grid_channels: 0
  gridtype: 'sinusoidal'    # options 'sinusoidal' or 'linear'
  roll: !!bool False
  enable_nhwc: !!bool False
  optimizer_type: 'FusedAdam'

  # directory path to store training checkpoints and other output
  exp_dir: 'exp'

  train_data_path: 'data/global/train'
  valid_data_path: 'data/global/valid'  
  test_data_path:  'data/global/test'

  # land mask
  land_mask: !!bool False

  # normalization
  normalize: !!bool True
  normalization: 'zscore' # options zscore, minmax, zscore_lat

  global_means_path: 'data/global/mean.npy' 
  global_stds_path:  'data/global/std.npy' 

  # orography
  orography: !!bool False

  # noise
  add_noise: !!bool False
  noise_std: 0

  # crop
  crop_size_x: None
  crop_size_y: None

  log_to_screen: !!bool True
  log_to_wandb: !!bool True
  save_checkpoint: !!bool True
  plot_animations: !!bool False
  
  
#############################################
OneForecast: &OneForecast
  <<: *FULL_FIELD
  nettype: 'OneForecast'
  log_to_wandb: !!bool True

  # Train params
  lr: 1E-3
  batch_size: 32 
  scheduler: 'CosineAnnealingLR'

  loss_channel_wise: True 
  loss_scale: False 
  use_loss_scaler_from_metnet3: True 

  # Data params
  fusion_3d_2d: False 

  
  atmos_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68]

  in_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68]

  out_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68]
  

  out_variables: ['z0', 'z1', 'z2', 'z3', 'z4', 'z5', 'z6', 'z7', 'z8', 'z9', 'z10', 'z11', 'z12', 'q0', 'q1', 'q2', 'q3', 'q4', 'q5', 'q6', 'q7', 'q8', 'q9', 'q10', 'q11', 'q12', 't0', 't1', 't2', 't3', 't4', 't5', 't6', 't7', 't8', 't9', 't10', 't11', 't12', 'U0', 'U1', 'U2', 'U3', 'U4', 'U5', 'U6', 'U7', 'U8', 'U9', 'U10', 'U11', 'U12', 'V0', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'V7', 'V8', 'V9', 'V10', 'V11', 'V12', 'u10', 'v10', 't2m', 'msl']