dataset_params:
  batch_size: 1
  database_type: db
  max_num_nbr: 32
  radius: 6.0
  split_file: null
  test_ratio: 0.1
  train_ratio: 0.8
  val_ratio: 0.1
  graph_data_path: /public/home/zhongyang/EPC_test/Universal_Hamltonian/Tests_new/raw_hetro_MoS2_WSe2_cp_calculated/ # graph_data.npz所在的目录

losses_metrics:
  losses:
  - loss_weight: 1.0
    metric: mae
    prediction: hamiltonian
    target: hamiltonian
  # - loss_weight: 0.01
  #   metric: mae
  #   prediction: band_gap
  #   target: band_gap
  # - loss_weight: 0.01
  #   metric: mae
  #   prediction: band_energy
  #   target: band_energy
  #- loss_weight: 1.0
  #  metric: mae
  #  prediction: overlap
  #  target: overlap
  #- loss_weight: 1.0
  #  metric: mae
  #  prediction: peak
  #  target: peak
  #- loss_weight: 0.0
  #  metric: mae
  #  prediction: hamiltonian_imag
  #  target: hamiltonian_imag
  #- loss_weight: 0.0001
  #  metric: abs_mae
  #  prediction: wavefunction
  #  target: wavefunction
  metrics:
  - metric: mae
    prediction: hamiltonian
    target: hamiltonian
  # - metric: mae
  #   prediction: band_gap
  #   target: band_gap
  #- metric: mae
  #  prediction: peak
  #  target: peak
  #- metric: mae
  #  prediction: overlap
  #  target: overlap
  #- metric: mae
  #  prediction: hamiltonian_imag
  #  target: hamiltonian_imag
  #- metric: mae
  #  prediction: hamiltonian_imag
  #  target: hamiltonian_imag
  # - metric: mae
  #   prediction: band_energy
  #   target: band_energy
  #- metric: abs_mae
  #  prediction: wavefunction
  #  target: wavefunction

# 一般optim_params模块只需要设置初始学习率lr即可
optim_params:
  lr: 0.0001 #初始学习率
  lr_decay: 0.5
  lr_patience: 5
  gradient_clip_val: 0.0
  max_epochs: 3000
  min_epochs: 10
  stop_patience: 30

output_nets:
  output_module: HamGNN_out
  HamGNN_out:
    create_graph: false # 预留参数，不用管，暂时无用
    return_forces: false # 预留参数，不用管，暂时无用
    include_triplet: false # 预留参数，不用管，暂时无用
    ham_only: true # true: 只计算哈密顿量H； 'false': 拟合H和S
    ham_type: openmx # openmx: 拟合openmx的哈密顿量, abacus: 拟合abacus的哈密顿量
    nao_max: 26 # 根据数据集中原子最大轨道数进行修改，可选14, 19或27
    add_H0: True # 一般时候为true，预测H_scf加上H_nonscf
    symmetrize: true # 对哈密顿量施加厄密性约束，保证输出哈密顿量满足厄密性
    calculate_band_energy: False # 是否计算能带以对能带进行训练
    num_k: 5 # 当计算能带时，所使用的k点数
    band_num_control: 6 # dict:控制每种原子在能带计算时考虑多少个基底； int: [VBM-num, VBM+num]； null: all bands
    k_path: null # 'auto': 自动确定k点路径. null: 随机k点. list: 用户自己确定的k点路径列表
    soc_switch: false # true：拟合SOC哈密顿量
    nonlinearity_type: norm # norm or gate
    export_reciprocal_values: false # 计算电声耦合矩阵epc时才有用，平时保持false

profiler_params:
  progress_bar_refresh_rat: 1
  train_dir: /public/home/zhongyang/EPC_test/Universal_Hamltonian/Tests_new/raw_hetro_MoS2_WSe2_cp_calculated/  #预测结果以及训练好的网络参数最后存放的目录

representation_nets:
  # 网络参数通常无需改动
  HamGNN_pre:
    cutoff: 26.0
    resnet: True
    cutoff_func: cos
    edge_sh_normalization: component
    edge_sh_normalize: true
    ######## Irreps set 1 (crystal): ################
    feature_irreps_hidden: 32x0o+128x0e+128x1o+64x1e+128x2e+32x2o+64x3o+32x3e+32x4o+32x4e+16x5o+8x5e+8x6e
    irreps_edge_output: 32x0o+128x0e+128x1o+64x1e+128x2e+32x2o+64x3o+32x3e+32x4o+32x4e+16x5o+8x5e+8x6e
    irreps_edge_sh: 0e + 1o + 2e + 3o + 4e + 5o + 6e
    irreps_node_features: 32x0o+128x0e+128x1o+64x1e+128x2e+32x2o+64x3o+32x3e+32x4o+32x4e+16x5o+8x5e+8x6e
    irreps_node_output: 32x0o+128x0e+128x1o+64x1e+128x2e+32x2o+64x3o+32x3e+32x4o+32x4e+16x5o+8x5e+8x6e
    irreps_triplet_output: 32x0o+128x0e+128x1o+64x1e+128x2e+32x2o+64x3o+32x3e+32x4o+32x4e+16x5o+8x5e+8x6e
    invariant_layers: 3
    invariant_neurons: 128
    num_interaction_layers: 5
    num_radial: 64
    num_spherical: 32
    export_triplet: false
    rbf_func: bessel
    set_features: true
    add_edge_tp: false
    num_types: 120
    irreps_node_prev: 16x0o+16x0e+8x1o+8x1e+8x2e+8x2o+8x3o+8x3e+8x4o+8x4e
    num_node_attr_feas: 64

setup:
  GNN_Net: HamGNN_pre # HamGNN_pre, HamGNN_pre2
  accelerator: null
  ignore_warnings: true
  checkpoint_path: /public/home/zhongyang/EPC_test/Universal_Hamltonian/Tests_new/raw_hetro_MoS2_WSe2_cp_calculated/version_0/checkpoints/epoch=11-val_loss=0.000000.ckpt
  # checkpoint_path: /public/home/zhongyang/sxliu/software/Universal_HamGNN/universal_Hamiltonian.ckpt
  # checkpoint_path: /public/home/zhongyang/EPC_test/Universal_Hamltonian/mp_all_splits_calculated/QMOF_graph/version_3/checkpoints/epoch=71BF-val_loss=0.000000.ckpt
  l_minus_mean: false
  load_from_checkpoint: True
  resume: False
  num_gpus: 1
  precision: 32
  property: hamiltonian
  stage: test