_target_: onescience.datapipes.materials.mattergen.datamodule.CrystDataModule
_recursive_: true
properties: []
  # Supported properties:
  # - dft_bulk_modulus
  # - dft_band_gap
  # - dft_mag_density

transforms:
- _target_: onescience.datapipes.materials.mattergen.transform.symmetrize_lattice
  _partial_: true
- _target_: onescience.datapipes.materials.mattergen.transform.set_chemical_system_string
  _partial_: true

dataset_transforms:
  - _target_: onescience.datapipes.materials.mattergen.dataset_transform.filter_sparse_properties
    _partial_: true

average_density: 0.05771451654022283 # atoms/Angstrom**3 : this is used in models/scripts/run.py to set lattice_limit_density
root_dir: ${oc.env:ONESCIENCE_DATASETS_DIR}/matchem/mattergen/cache/mp_20

train_dataset:
  _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
  cache_path: ${data_module.root_dir}/train
  properties: ${data_module.properties}
  transforms: ${data_module.transforms}
  dataset_transforms: ${data_module.dataset_transforms}

val_dataset:
  _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
  cache_path: ${data_module.root_dir}/val
  properties: ${data_module.properties}
  transforms: ${data_module.transforms}
  dataset_transforms: ${data_module.dataset_transforms}

test_dataset:
  _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
  cache_path: ${data_module.root_dir}/test
  properties: ${data_module.properties}
  transforms: ${data_module.transforms}
  dataset_transforms: ${data_module.dataset_transforms}

num_workers:
  train: 0
  val: 0
  test: 0

batch_size:
  # total batch size of 512, adjust for number of devices, nodes, and gradient accumulation
  train:  ${eval:'(512 // ${trainer.accumulate_grad_batches}) // (${trainer.devices} * ${trainer.num_nodes})'}
  val: ${eval:'(64 // (${trainer.devices} * ${trainer.num_nodes}))'}
  test: ${eval:'(64 // (${trainer.devices} * ${trainer.num_nodes}))'}

max_epochs: 900