# yamllint disable rule:line-length
# Smoke-test config for NequIP on OneScience.
# This config uses a tiny synthetic Cu dataset and a small model to verify the
# training pipeline (data loading, model build, forward, backward, checkpoint).

run: [train, test]

cutoff_radius: 4.0

num_layers: 2
l_max: 1
num_features: 8

model_type_names: [Cu]
chemical_species: ${model_type_names}

monitored_metric: val0_epoch/weighted_sum

# ============
#     DATA
# ============
data:
  _target_: onescience.datapipes.materials.nequip.datamodule.ASEDataModule
  seed: 456

  split_dataset:
    file_path: ${demo_dir:reference_data/smoke.xyz}
    train: 0.75
    val: 0.125
    test: 0.125

  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: 2
    num_workers: 0
    shuffle: true
  val_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 2
    num_workers: 0
  test_dataloader: ${data.val_dataloader}

  stats_manager:
    _target_: onescience.datapipes.materials.nequip.CommonDataStatisticsManager
    dataloader_kwargs:
      batch_size: 2
    type_names: ${model_type_names}

# =============
#    TRAINER
# =============
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: true

  callbacks:
    - _target_: lightning.pytorch.callbacks.ModelCheckpoint
      monitor: ${monitored_metric}
      dirpath: ${hydra:runtime.output_dir}/checkpoints
      filename: best
      save_last: true

# =====================
#    TRAINING MODULE
# =====================
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
    seed: 456
    model_dtype: float32
    type_names: ${model_type_names}
    r_max: ${cutoff_radius}
    num_bessels: 4
    bessel_trainable: false
    polynomial_cutoff_p: 6
    num_layers: ${num_layers}
    l_max: ${l_max}
    parity: false
    num_features: ${num_features}
    radial_mlp_depth: 1
    radial_mlp_width: 16
    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

# Slurm / launch metadata used by demo/run.sh
name: nequip_smoke
launch:
  mode: local
  num_nodes: 1
  num_gpus: 1
slurm:
  partition: hx1hdnormal01
  nodelist: a01r1n02
  time: "00:10:00"
  cpus_per_task: 8
env:
  OMP_NUM_THREADS: 1