Upload weights and configs for quartz

quartz/config_cueq.yaml

@@ -0,0 +1,97 @@
+model:
+  precision: 64
+  num_chem_elements: 2
+  edge_cutoff: 4.5
+  num_edge_rbf: 8
+  num_edge_poly_cutoff: 6
+  num_vel_rbf: 8
+  vel_max: 0.035
+  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/quartz/data
+  name: quartz_300K_N5000_dt30.0_rc4.5_train
+  cutoff_radius: 4.5
+  files:
+  - quartz_300K_rep1_timestep300_train1667.extxyz
+  - quartz_300K_rep2_timestep300_train1667.extxyz
+  - quartz_300K_rep3_timestep300_train1666.extxyz
+  rename: true
+  atom_type_mapper:
+    8: 0
+    14: 1
+training:
+  seed: 506
+  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/quartz/data
+        name: quartz_300K_N5000_dt30.0_rc4.5_val
+        files:
+        - quartz_300K_rep4_timestep300_val1250.extxyz
+        cutoff_radius: 4.5
+        rename: true
+        atom_type_mapper:
+          8: 0
+          14: 1

quartz/config_e3nn.yaml

@@ -0,0 +1,35 @@
+model:
+  precision: 64
+  num_chem_elements: 2
+  edge_cutoff: 4.5
+  num_edge_rbf: 8
+  num_edge_poly_cutoff: 6
+  num_vel_rbf: 8
+  vel_max: 0.035
+  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

quartz/state_dict_cueq.pt

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

quartz/state_dict_e3nn.pt

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