water/config_cueq.yaml

model:
  precision: 64
  num_chem_elements: 2
  edge_cutoff: 6
  num_edge_rbf: 8
  num_edge_poly_cutoff: 6
  num_vel_rbf: 8
  vel_max: 0.135
  max_rotation_order: 2
  num_hidden_channels: 64
  num_mp_layers: 4
  edge_mlp_kwargs:
    n_neurons:
    - 64
    - 64
    - 64
    activation: silu
  vel_mlp_kwargs:
    n_neurons:
    - 64
    - 64
    - 64
    activation: silu
  nl_gate_kwargs:
    activation_scalars:
      o: tanh
      e: silu
    activation_gates:
      e: silu
  conserve_ang_mom: false
  o3_backend: cueq
  net_lin_mom:
  - 0.0
  - 0.0
  - 0.0
data:
  root: /dccstor/chemistry_ai/trajcast/paper/water/data
  name: water_300K_N5000_dt5.0_rc6_train
  cutoff_radius: 6
  files:
  - water_300K_rep1_timestep50_train1667.extxyz
  - water_300K_rep2_timestep50_train1667.extxyz
  - water_300K_rep3_timestep50_train1666.extxyz
  rename: true
  atom_type_mapper:
    1: 0
    8: 1
training:
  seed: 2303
  model_type: EfficientTrajCastModel
  device: cuda
  restart_latest: true
  target_field: target
  reference_fields:
  - displacements
  - update_velocities
  batch_size: 2
  max_grad_norm: 0.5
  num_epochs: 1500
  criterion:
    loss_type:
      main_loss: mse
    learnable_weights: false
  optimizer: adam
  optimizer_settings:
    lr: 0.01
    amsgrad: true
  scheduler:
  - ReduceLROnPlateau
  scheduler_settings:
    ReduceLROnPlateau:
      factor: 0.8
      patience: 25
      min_lr: 0.0001
  chained_scheduler_hp:
    milestones:
    - 100000000
    per_epoch: true
    monitor_lr_scheduler: false
  tensorboard_settings:
    loss: true
    lr: true
    weights:
      every: 5
    gradients:
      every: 5
    loss_validation:
      data:
        root: /dccstor/chemistry_ai/trajcast/paper/water/data
        name: water_300K_N5000_dt5.0_rc6_val
        files:
        - water_300K_rep4_timestep50_val1250.extxyz
        cutoff_radius: 6
        rename: true
        atom_type_mapper:
          1: 0
          8: 1