default: &DEFAULT
  num_data_workers: 1
  # model
  model: 'fno'
  depth: 5
  in_dim: 2
  out_dim: 1
  dropout: 0
  # data/domain
  Lx: !!float 1.0
  Ly: !!float 1.0
  nx: 256
  ny: 256
  # optimization
  loss_style: 'mean'
  loss_func: 'mse'
  optimizer: 'adam'
  scheduler: 'none'
  lr: !!float 1.0
  max_epochs: 500
  max_cosine_lr_epochs: 500
  batch_size: 25
  # misc
  log_to_screen: !!bool True
  save_checkpoint: !!bool False
  seed: 0
  plot_figs: !!bool False
  pack_data: !!bool False
  # Weights & Biases
  entity: 'entity_name'
  project: 'nopt'
  log_to_wandb: !!bool False
  distill: !!bool False
  subsample: 1

helmholtz: &helmholtz
  <<: *DEFAULT
  n_demos: 0
  batch_size: 128
  valid_batch_size: 128
  nx: 128
  ny: 128
  log_to_wandb: !!bool True
  save_checkpoint: !!bool True
  max_epochs: 500
  scheduler: 'cosine'
  plot_figs: !!bool True
  loss_style: 'sum'

  model: 'fno'
  layers: [64, 64, 64, 64, 64]
  modes1: [65, 65, 65, 65]
  modes2: [65, 65, 65, 65]
  fc_dim: 128

  in_dim: 2
  out_dim: 1
  mode_cut: 32
  embed_cut: 64
  fc_cut: 2

  optimizer: 'adam'

  lr: 1E-3
  pack_data: !!bool False
  

helm-64-scale-o5_15: &helm_64_o5_15
  <<: *helmholtz
  train_path:   'train_path'
  val_path:     'val_path'
  test_path:    'test_path'
  scales_path:  'scales_path'  
  train_rand_idx_path: 'train_rand_idx_path'
  batch_size: 128
  valid_batch_size: 128
  log_to_wandb: !!bool True
  in_dim: 3
  out_dim: 1
  mode_cut: 32
  embed_cut: 64
  fc_cut: 2
  lr: 1E-3
  subsample: 1
  nx: 64
  ny: 64


helm-64-pretrain-o1_20: &helm_64_o1_20_pt
  <<: *helmholtz
  train_path:   'train_path'
  val_path:     'val_path'
  test_path:    'test_path'
  scales_path:  'scales_path'  
  train_rand_idx_path: 'train_rand_idx_path'
  batch_size: 128
  valid_batch_size: 128
  log_to_wandb: !!bool True
  in_dim: 3
  out_dim: 1
  mode_cut: 32
  embed_cut: 64
  fc_cut: 2
  lr: 1E-3
  subsample: 1
  nx: 64
  ny: 64
  pt: "pretrain"
  pt_split: [46080, 8192] #[0.9, 0.1]


helm-64-finetune-o5_15: &helm_64_o5_15_ft
  <<: *helmholtz
  train_path:   'train_path'
  val_path:     'val_path'
  test_path:    'test_path'
  scales_path:  'scales_path'  
  train_rand_idx_path: 'train_rand_idx_path'
  batch_size: 128
  valid_batch_size: 128
  log_to_wandb: !!bool True
  in_dim: 3
  out_dim: 1
  mode_cut: 32
  embed_cut: 64
  fc_cut: 2
  lr: 1E-3
  subsample: 1
  nx: 64
  ny: 64
  pt: "train"
  pt_split: [46080, 8192]
  fix_backbone: False

helm-64-o5_15_ft0: &helm_64_o5_15_ft0
  <<: *helm_64_o5_15_ft
  subsample: 1

helm-64-o5_15_ft1: &helm_64_o5_15_ft1
  <<: *helm_64_o5_15_ft
  subsample: 2

helm-64-o5_15_ft2: &helm_64_o5_15_ft2
  <<: *helm_64_o5_15_ft
  subsample: 4

helm-64-o5_15_ft3_r1: &helm_64_o5_15_ft3_r1
  <<: *helm_64_o5_15_ft
  subsample: 8
  seed: 1

helm-64-o5_15_ft3_r2: &helm_64_o5_15_ft3_r2
  <<: *helm_64_o5_15_ft
  subsample: 8
  seed: 2

helm-64-o5_15_ft3_r0: &helm_64_o5_15_ft3_r0
  <<: *helm_64_o5_15_ft
  subsample: 8
  seed: 0

helm-64-o5_15_ft3_r3: &helm_64_o5_15_ft3_r3
  <<: *helm_64_o5_15_ft
  subsample: 8
  seed: 3

helm-64-o5_15_ft4_r0: &helm_64_o5_15_ft4_r0
  <<: *helm_64_o5_15_ft
  subsample: 16
  seed: 0

helm-64-o5_15_ft4_r3: &helm_64_o5_15_ft4_r3
  <<: *helm_64_o5_15_ft
  subsample: 16
  seed: 3

helm-64-o5_15_ft4_r1: &helm_64_o5_15_ft4_r1
  <<: *helm_64_o5_15_ft
  subsample: 16
  seed: 1

helm-64-o5_15_ft4_r2: &helm_64_o5_15_ft4_r2
  <<: *helm_64_o5_15_ft
  subsample: 16
  seed: 2

helm-64-o5_15_ft5_r1: &helm_64_o5_15_ft5_r1
  <<: *helm_64_o5_15_ft
  subsample: 32
  seed: 1

helm-64-o5_15_ft5_r0: &helm_64_o5_15_ft5_r0
  <<: *helm_64_o5_15_ft
  subsample: 32
  seed: 0

helm-64-o5_15_ft5_r2: &helm_64_o5_15_ft5_r2
  <<: *helm_64_o5_15_ft
  subsample: 32
  seed: 2

helm-64-o5_15_ft6_r0: &helm_64_o5_15_ft6_r0
  <<: *helm_64_o5_15_ft
  subsample: 64
  seed: 0

helm-64-o5_15_ft6_r1: &helm_64_o5_15_ft6_r1
  <<: *helm_64_o5_15_ft
  subsample: 64
  seed: 1

helm-64-o5_15_ft6_r2: &helm_64_o5_15_ft6_r2
  <<: *helm_64_o5_15_ft
  subsample: 64
  seed: 2

helm-64-o5_15_ft7_r0: &helm_64_o5_15_ft7_r0
  <<: *helm_64_o5_15_ft
  subsample: 128
  # lr: 1E-5
  batch_size: 64
  seed: 0
  valid_batch_size: 64

helm-64-o5_15_ft7_r1: &helm_64_o5_15_ft7_r1
  <<: *helm_64_o5_15_ft
  subsample: 128
  # lr: 1E-5
  batch_size: 64
  seed: 1
  valid_batch_size: 64

helm-64-o5_15_ft7_r2: &helm_64_o5_15_ft7_r2
  <<: *helm_64_o5_15_ft
  subsample: 128
  # lr: 1E-5
  batch_size: 64
  seed: 2
  valid_batch_size: 64

helm-64-o5_15_ft8_r0: &helm_64_o5_15_ft8_r0
  <<: *helm_64_o5_15_ft
  subsample: 256
  # lr: 1E-5
  batch_size: 32
  seed: 0
  valid_batch_size: 32

helm-64-o5_15_ft8_r1: &helm_64_o5_15_ft8_r1
  <<: *helm_64_o5_15_ft
  subsample: 256
  # lr: 1E-5
  batch_size: 32
  seed: 1
  valid_batch_size: 32

helm-64-o5_15_ft8_r2: &helm_64_o5_15_ft8_r2
  <<: *helm_64_o5_15_ft
  subsample: 256
  # lr: 1E-5
  batch_size: 32
  seed: 2
  valid_batch_size: 32

helm-64-o5_15_ft9_r0: &helm_64_o5_15_ft9_r0
  <<: *helm_64_o5_15_ft
  subsample: 512
  # lr: 1E-5
  batch_size: 16
  seed: 0
  valid_batch_size: 16

helm-64-o5_15_ft9_r1: &helm_64_o5_15_ft9_r1
  <<: *helm_64_o5_15_ft
  subsample: 512
  # lr: 1E-5
  batch_size: 16
  seed: 1
  valid_batch_size: 16

helm-64-o5_15_ft9_r2: &helm_64_o5_15_ft9_r2
  <<: *helm_64_o5_15_ft
  subsample: 512
  # lr: 1E-5
  batch_size: 16
  seed: 2
  valid_batch_size: 16

helm-64-o5_15_bsln: &helm_64_o5_15_baseline
  <<: *helm_64_o5_15
  pt: "train"
  pt_split: [0, 1]

helm-64-o5_15_b0: &helm-64-o1_10_ss4
  <<: *helm_64_o5_15_baseline
  subsample: 4

helm-64-o5_15_b1: &helm-64-o1_10_ss8
  <<: *helm_64_o5_15_baseline
  subsample: 8

helm-64-o5_15_b2: &helm-64-o1_10_ss16
  <<: *helm_64_o5_15_baseline
  subsample: 16

helm-64-o5_15_b3: &helm-64-o1_10_ss32
  <<: *helm_64_o5_15_baseline
  subsample: 32

helm-64-o5_15_b4: &helm-64-o1_10_ss64
  <<: *helm_64_o5_15_baseline
  subsample: 64

helm-64-o5_15_b5: &helm-64-o1_10_ss128
  <<: *helm_64_o5_15_baseline
  subsample: 128

helm-64-o5_15_b6: &helm-64-o1_10_ss256
  <<: *helm_64_o5_15_baseline
  subsample: 256

helm-64-o5_15_b7: &helm-64-o1_10_ss512
  <<: *helm_64_o5_15_baseline
  subsample: 512
  batch_size: 64
  valid_batch_size: 64

helm-64-o5_15_b8: &helm-64-o1_10_ss1024
  <<: *helm_64_o5_15_baseline
  subsample: 1024
  batch_size: 32
  valid_batch_size: 32

helm-64-o5_15_b9: &helm-64-o1_10_ss2048
  <<: *helm_64_o5_15_baseline
  subsample: 2048
  batch_size: 16
  valid_batch_size: 16
