dataset_name: ??
elem_refs: ??
normalizer_rmsd: ??

train_dataset:
  splits:
    train:
      src: ??
  format: ase_db
  transforms:
    common_transform:
      dataset_name: ${data.dataset_name}
    stress_reshape_transform:
      dataset_name: ${data.dataset_name}

val_dataset:
  splits:
    val:
      src: ??
  format: ase_db
  transforms:
    common_transform:
      dataset_name: ${data.dataset_name}
    stress_reshape_transform:
      dataset_name: ${data.dataset_name}

regress_stress: False
pass_through_head_outputs: False

heads:
  energy:
    module: onescience.modules.head.uma_head.MLP_Energy_Head

tasks_list:
- _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
  name: energy
  level: system
  property: energy
  loss_fn:
    _target_: onescience.modules.loss.uma_loss.DDPMTLoss
    loss_fn:
      _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
    coefficient: 20
  out_spec:
    dim: [1]
    dtype: float32
  normalizer:
    _target_: onescience.utils.uma.normalization.normalizer.Normalizer
    mean: 0.0
    rmsd: ${data.normalizer_rmsd}
  element_references:
    _target_: onescience.utils.uma.normalization.element_references.ElementReferences
    element_references:
      _target_: torch.DoubleTensor
      _args_:
        - ${data.elem_refs}
  datasets:
    - ${data.dataset_name}
  metrics:
    - mae
    - per_atom_mae