dataset_name: pretrain
dim_coords: 3
root_path: ${oc.env:DATA_DIR}/pretrain_stable
prop: stability
num_targets: 1
# prop: scaled_lattice
# num_targets: 6
niggli: true
primitive: False
graph_method: crystalnn
lattice_scale_method: scale_length
preprocess_workers: 30
readout: mean
max_atoms: 52
otf_graph: false
eval_model_name: mp20
tolerance: 0.1

use_space_group: false
use_pos_index: false
train_max_epochs: 1000
early_stopping_patience: 100000
teacher_forcing_max_epoch: 300

# diff from mpts
md_dataset_name: mpt
root_path_md: ${oc.env:DATA_DIR}/${data.md_dataset_name}
require_order: False
prop_md: 
  - energy
  - forces
energy_only: False
t_mode: flow

datamodule:
  _target_: UniPath.datamodule_uni.CrystDataModule # diff from mpts 

  task_mode: uni # options: uni/csp/md
  datasets:
    train:
      _target_: diffcsp.pl_data.dataset.CrystDataset
      name: Formation energy train
      path: ${data.root_path}/train.csv
      save_path: ${data.root_path}/train_ori.pt
      prop: ${data.prop}
      niggli: ${data.niggli}
      primitive: ${data.primitive}
      graph_method: ${data.graph_method}
      tolerance: ${data.tolerance}
      use_space_group: ${data.use_space_group}
      use_pos_index: ${data.use_pos_index}
      lattice_scale_method: ${data.lattice_scale_method}
      preprocess_workers: ${data.preprocess_workers}

    val:
      - _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy val
        path: ${data.root_path}/val.csv
        save_path: ${data.root_path}/val_ori.pt
        prop: ${data.prop}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}

    test:
      - _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy test
        path: ${data.root_path}/test.csv
        save_path: ${data.root_path}/test_ori.pt
        prop: ${data.prop}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}

    train_md: # options: uni/csp/md
      _target_: diffcsp.pl_data.dataset.CrystDataset
      name: Formation energy train
      path: ${data.root_path_md}/train.csv
      save_path: ${data.root_path_md}/train_ori.pt
      require_order: ${data.require_order}
      prop: ${data.prop_md}
      niggli: ${data.niggli}
      primitive: ${data.primitive}
      graph_method: ${data.graph_method}
      tolerance: ${data.tolerance}
      use_space_group: ${data.use_space_group}
      use_pos_index: ${data.use_pos_index}
      lattice_scale_method: ${data.lattice_scale_method}
      preprocess_workers: ${data.preprocess_workers}

    val_md: # options: uni/csp/md
      - _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy val
        path: ${data.root_path_md}/val.csv
        save_path: ${data.root_path_md}/val_ori.pt
        require_order: ${data.require_order}
        prop: ${data.prop_md}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}

  num_workers:
    train: 40
    val: 40
    test: 40

  batch_size:
    train: 64
    val: 512
    test: 256
    train_md: 256 # diff from mp20 
    val_md: 1024 # diff from mp20