# NequIP 0.19 tutorial reproduction using the official fcu.xyz dataset.
# The validated smoke schedule uses two complete epochs. Set max_epochs to 1000
# to match the upstream tutorial's full training schedule.
run: [train, test]

cutoff_radius: 5.0
num_layers: 4
l_max: 1
num_features: 32
model_type_names: [C, H, O, Cu]
chemical_species: ${model_type_names}
monitored_metric: val0_epoch/weighted_sum

data:
  _target_: onescience.datapipes.materials.nequip.datamodule.ASEDataModule
  seed: 456
  split_dataset:
    file_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/matchem/NequIP/fcu.xyz
    train: 0.8
    val: 0.1
    test: 0.1
  transforms:
    - _target_: onescience.datapipes.materials.nequip.transforms.ChemicalSpeciesToAtomTypeMapper
      model_type_names: ${model_type_names}
    - _target_: onescience.datapipes.materials.nequip.transforms.NeighborListTransform
      r_max: ${cutoff_radius}
  train_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 5
    num_workers: 0
    shuffle: true
  val_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 10
    num_workers: 0
  test_dataloader: ${data.val_dataloader}
  stats_manager:
    _target_: onescience.datapipes.materials.nequip.CommonDataStatisticsManager
    dataloader_kwargs:
      batch_size: 10
    type_names: ${model_type_names}

trainer:
  _target_: lightning.Trainer
  accelerator: gpu
  devices: 1
  num_nodes: 1
  enable_checkpointing: true
  max_epochs: 2
  log_every_n_steps: 1
  logger: false
  enable_progress_bar: false
  callbacks:
    - _target_: lightning.pytorch.callbacks.EarlyStopping
      monitor: ${monitored_metric}
      min_delta: 1e-3
      patience: 20
    - _target_: lightning.pytorch.callbacks.ModelCheckpoint
      monitor: ${monitored_metric}
      dirpath: ${hydra:runtime.output_dir}/checkpoints
      filename: best
      save_last: true

training_module:
  _target_: onescience.utils.nequip.train.EMALightningModule
  ema_decay: 0.999
  loss:
    _target_: onescience.utils.nequip.train.EnergyForceLoss
    per_atom_energy: true
    coeffs:
      total_energy: 1.0
      forces: 1.0
  val_metrics:
    _target_: onescience.utils.nequip.train.EnergyForceMetrics
    coeffs:
      total_energy_mae: 1.0
      forces_mae: 1.0
  train_metrics: ${training_module.val_metrics}
  test_metrics: ${training_module.val_metrics}
  optimizer:
    _target_: torch.optim.Adam
    lr: 0.01
  lr_scheduler:
    scheduler:
      _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
      factor: 0.6
      patience: 5
      threshold: 0.2
      min_lr: 1e-6
    monitor: ${monitored_metric}
    interval: epoch
    frequency: 1
  model:
    _target_: onescience.models.nequip.model.NequIPGNNModel
    compile_mode: eager
    seed: 456
    model_dtype: float32
    type_names: ${model_type_names}
    r_max: ${cutoff_radius}
    num_bessels: 8
    bessel_trainable: false
    polynomial_cutoff_p: 6
    num_layers: ${num_layers}
    l_max: ${l_max}
    parity: true
    num_features: ${num_features}
    radial_mlp_depth: 2
    radial_mlp_width: 64
    avg_num_neighbors: ${training_data_stats:num_neighbors_mean}
    per_type_energy_scales: ${training_data_stats:per_type_forces_rms}
    per_type_energy_shifts: ${training_data_stats:per_atom_energy_mean}
    per_type_energy_scales_trainable: false
    per_type_energy_shifts_trainable: false
    pair_potential:
      _target_: onescience.models.nequip.nn.pair_potential.ZBL
      units: metal
      chemical_species: ${chemical_species}

name: nequip_fcu_tutorial
launch:
  mode: local
  num_nodes: 1
  num_gpus: 1
slurm:
  partition: hx1hdnormal01
  nodelist: a01r1n02
  time: "00:30:00"
  cpus_per_task: 8
env:
  OMP_NUM_THREADS: 8