# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


## Training options
# Reynolds number parameter
reynolds_number: 100

load_ckpt: False
output_dir: './outputs/Re${reynolds_number}/figures/'

###################
## Model options
model_params:
  layers: 8
  modes: 8
  num_fno_layers: 4
  fc_dim: 128
  decoder_layers: 1
  in_dim: 6 # 3 + in_fields
  out_dim: 3
  dimension: 3
  activation: 'gelu'
  pad_x: 5
  pad_y: 0
  pad_z: 0
  input_norm: [1.0, 1.0, 1.0, 1.0, 1.0, 0.00025]
  output_norm: [1.0, 1.0, 0.00025]

  #TensorLy arguments
  rank: 0.5
  factorization: 'cp'
  fixed_rank_modes: null
  decomposition_kwargs: {}

###################
## Dataset options
dataset_params:
  data_dir: '/data/mhd_data/simulation_outputs_Re${reynolds_number}'
  field_names: ['velocity', 'vector potential']
  output_names: 'output-????'
  dataset_type: 'mhd'
  name: 'MHDVecPot_TFNO_Re${reynolds_number}'
  num: -1  # -1 means use full dataset, otherwise specify total number
  num_train: 0.8  # percentage of dataset for training
  num_test: 0.2   # percentage of dataset for testing
  sub_x: 1
  sub_t: 1
  ind_x: null
  ind_t: null
  nin: 3
  nout: 3
  fields: ['u', 'v', 'A']

###################
## Dataloader options
train_loader_params:
  batch_size: 4
  shuffle: True
  num_workers: 8
  pin_memory: True

val_loader_params:
  batch_size: 4
  shuffle: False
  num_workers: 8
  pin_memory: True

test_loader_params:
  batch_size: 4
  shuffle: False
  num_workers: 8
  pin_memory: True

###################
## Loss options
loss_params:
  nu: 0.004
  eta: 0.004
  rho0: 1.0

  data_weight: 5.0
  ic_weight: 1.0
  pde_weight: 1.0
  constraint_weight: 10.0

  use_data_loss: True
  use_ic_loss: True
  use_pde_loss: True
  use_constraint_loss: True

  u_weight: 1.0
  v_weight: 1.0
  A_weight: 1.0

  Du_weight: 1.0
  Dv_weight: 1.0
  DA_weight: 1_000_000

  div_B_weight: 1.0
  div_vel_weight: 1.0

  Lx: 1.0
  Ly: 1.0
  tend: 1.0

  use_weighted_mean: False

###################
## Optimizer options
optimizer_params:
  betas: [0.9, 0.999]
  lr: 5.0e-4
  milestones: [20, 40, 60, 80, 100]
  gamma: 0.5


###################
## Train params
train_params:
  epochs: 100
  ckpt_freq: 10
  ckpt_path: './outputs/checkpoints/Re${reynolds_number}/'

###################
## log params
log_params:
  log_dir: 'logs'
  log_project: 'MHD_PINO'
  log_group: 'MHDVecPot_TFNO_Re${reynolds_number}'
  log_num_plots: 1
  log_plot_freq: 5
  log_plot_types: ['ic', 'pred', 'true', 'error']

test:
  batchsize: 1
  ckpt_path: './outputs/checkpoints/Re${reynolds_number}/'