Upload water models

water/config_cueq.yaml

@@ -0,0 +1,97 @@
+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

water/config_e3nn.yaml

@@ -0,0 +1,35 @@
+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: e3nn
+  net_lin_mom:
+  - 0.0
+  - 0.0
+  - 0.0

water/state_dict_cueq.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a83d0eb52654225321f37cff199ece49546dad894205124ed9b81c3329d36e2a
+size 18592910

water/state_dict_e3nn.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:769d00ab4b8836e19cffe00accb84e13ad7400c8c1065e0d57f42631bd7f59bd
+size 18717216