default: &DEFAULT
  num_data_workers: 4
  # 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 True
  distill: !!bool False
  subsample: 1

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

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

  in_dim: 4
  out_dim: 1
  mode_cut: 16
  embed_cut: 64
  fc_cut: 2

  optimizer: 'adam'

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


poisson-64-scale-e5_15: &poisson_64_e5_15
  <<: *poisson
  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
  lr: 1E-3

  mode_cut: 32
  embed_cut: 64
  fc_cut: 2 
  subsample: 1
  nx: 64
  ny: 64
  
  pt: "train"
  pt_split: [46080, 8192]


pois-64-pretrain-e1_20: &pois_64_e1_20_pt
  <<: *poisson
  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
  mode_cut: 32
  embed_cut: 64
  fc_cut: 2 
  subsample: 1
  nx: 64
  ny: 64
  lr: 1E-3
  pt: "pretrain"
  pt_split: [46080, 8192]
  blur: [0, 1]


pois-64-finetune-e5_15: &pois_64_e5_15_ft
  <<: *poisson
  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
  mode_cut: 32
  embed_cut: 64
  fc_cut: 2 
  subsample: 1
  nx: 64
  ny: 64
  lr: 1E-3
  pt: "train"
  pt_split: [46080, 8192]
  fix_backbone: False

pois-64-e5_15_ft0: &pois_64_e5_15_ft0
  <<: *pois_64_e5_15_ft
  subsample: 1

pois-64-e5_15_ft1: &pois_64_e5_15_ft1
  <<: *pois_64_e5_15_ft
  subsample: 2

pois-64-e5_15_ft2: &pois_64_e5_15_ft2
  <<: *pois_64_e5_15_ft
  subsample: 4

pois-64-e5_15_ft3: &pois_64_e5_15_ft3
  <<: *pois_64_e5_15_ft
  subsample: 8

pois-64-e5_15_ft4: &pois_64_e5_15_ft4
  <<: *pois_64_e5_15_ft
  subsample: 16

pois-64-e5_15_ft5: &pois_64_e5_15_ft5
  <<: *pois_64_e5_15_ft
  subsample: 32

pois-64-e5_15_ft6: &pois_64_e5_15_ft6
  <<: *pois_64_e5_15_ft
  subsample: 64

pois-64-e5_15_ft7: &pois_64_e5_15_ft7
  <<: *pois_64_e5_15_ft
  subsample: 128
  # lr: 1E-5
  batch_size: 64
  valid_batch_size: 64

pois-64-e5_15_ft8: &pois_64_e5_15_ft8
  <<: *pois_64_e5_15_ft
  subsample: 256
  # lr: 1E-5
  batch_size: 32
  valid_batch_size: 32

pois-64-e5_15_ft9: &pois_64_e5_15_ft9
  <<: *pois_64_e5_15_ft
  subsample: 512
  # lr: 1E-5
  batch_size: 16
  valid_batch_size: 16

pois-64-pretrain-e1_20_m0: &pois-64-e1_20_pt_m0
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.

pois-64-pretrain-e1_20_m1: &pois-64-e1_20_pt_m1
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.1

pois-64-pretrain-e1_20_m2: &pois-64-e1_20_pt_m2
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.2

pois-64-pretrain-e1_20_m3: &pois-64-e1_20_pt_m3
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.3

pois-64-pretrain-e1_20_m4: &pois-64-e1_20_pt_m4
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.4

pois-64-pretrain-e1_20_m5: &pois-64-e1_20_pt_m5
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.5

pois-64-pretrain-e1_20_m6: &pois-64-e1_20_pt_m6
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.6

pois-64-pretrain-e1_20_m7: &pois-64-e1_20_pt_m7
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.7

pois-64-pretrain-e1_20_m8: &pois-64-e1_20_pt_m8
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.8

pois-64-pretrain-e1_20_m9: &pois-64-e1_20_pt_m9
  <<: *pois_64_e1_20_pt
  mask_ratio: 0.9



poisson-64-e5_15_bsln: &pois_64_e5_15_baseline
  <<: *poisson_64_e5_15

# 8192
poisson-64-e5_15_b0: &pois_64_e5_15_ss4
  <<: *pois_64_e5_15_baseline
  subsample: 1

poisson-64-e5_15_b1: &pois_64_e5_15_ss8
  <<: *pois_64_e5_15_baseline
  subsample: 2

poisson-64-e5_15_b2: &pois_64_e5_15_ss16
  <<: *pois_64_e5_15_baseline
  subsample: 4

poisson-64-e5_15_b3: &pois_64_e5_15_ss32
  <<: *pois_64_e5_15_baseline
  subsample: 8

poisson-64-e5_15_b4: &pois_64_e5_15_ss64
  <<: *pois_64_e5_15_baseline
  subsample: 16

poisson-64-e5_15_b5: &pois_64_e5_15_ss128
  <<: *pois_64_e5_15_baseline
  subsample: 32

poisson-64-e5_15_b6: &pois_64_e5_15_ss256
  <<: *pois_64_e5_15_baseline
  subsample: 64

poisson-64-e5_15_b7: &pois_64_e5_15_ss512
  <<: *pois_64_e5_15_baseline
  subsample: 128
  batch_size: 64
  valid_batch_size: 64

poisson-64-e5_15_b8: &pois_64_e5_15_ss1024
  <<: *pois_64_e5_15_baseline
  subsample: 256
  lr: 5E-3
  batch_size: 32
  valid_batch_size: 32

poisson-64-e5_15_b9: &pois_64_e5_15_ss2048
  <<: *pois_64_e5_15_baseline
  subsample: 512
  lr: 5E-3
  batch_size: 16
  valid_batch_size: 16