# Eight-DCU DDP training plan for the official NequIP fcu.xyz tutorial.
# DataLoader batch size is per rank: batch_size=5 gives a global batch of 40.
# The learning rate remains conservative at the official 0.01; tune it for
# production based on validation behavior rather than scaling it blindly.
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: 8
  num_nodes: 1
  strategy:
    _target_: lightning.pytorch.strategies.DDPStrategy
  enable_checkpointing: true
  max_epochs: 1000
  max_time: "03:00:00:00"
  log_every_n_steps: 20
  logger:
    _target_: lightning.pytorch.loggers.CSVLogger
    save_dir: ${hydra:runtime.output_dir}
    name: metrics
    version: 0
  enable_progress_bar: true
  callbacks:
    - _target_: onescience.utils.nequip.train.callbacks.PlainTextMetricsLogger
    - _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: 1.0e-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_8dcu
launch:
  mode: auto
  num_nodes: 1
  num_gpus: 8
slurm:
  partition: hx1hdnormal01
  nodelist: ""
  time: "3-00:00:00"
  cpus_per_task: 8
env:
  OMP_NUM_THREADS: 8
  NCCL_DEBUG: "WARN"