From 5ae43d69bad598999ec394e6a7879b16c765c653 Mon Sep 17 00:00:00 2001
From: Simon Blackburn <simon.blackburn@mila.quebec>
Date: Tue, 25 Jun 2024 10:57:23 -0400
Subject: [PATCH] config files update

---
 .../config_files/diffusion/config_diffusion_mace.yaml     | 6 +++++-
 .../diffusion/config_diffusion_mace_orion.yaml            | 8 ++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/examples/config_files/diffusion/config_diffusion_mace.yaml b/examples/config_files/diffusion/config_diffusion_mace.yaml
index ab81edb1..92e3f784 100644
--- a/examples/config_files/diffusion/config_diffusion_mace.yaml
+++ b/examples/config_files/diffusion/config_diffusion_mace.yaml
@@ -27,6 +27,8 @@ model:
     r_max: 5.0
     num_bessel: 8
     num_polynomial_cutoff: 5
+    num_edge_hidden_layers: 0
+    edge_hidden_irreps: 8x0e
     max_ell: 2
     interaction_cls: RealAgnosticResidualInteractionBlock
     interaction_cls_first: RealAgnosticInteractionBlock
@@ -38,10 +40,12 @@ model:
     correlation: 3
     gate: silu
     radial_MLP: [8, 8, 8]
-    radial_type: bessel
+    radial_type: gaussian
     conditional_prob: 0.0
     conditional_gamma: 2
     condition_embedding_size: 64
+    use_batchnorm: False
+    tanh_after_interaction: True
   noise:
     total_time_steps: 100
     sigma_min: 0.001  # default value
diff --git a/examples/config_files/diffusion/config_diffusion_mace_orion.yaml b/examples/config_files/diffusion/config_diffusion_mace_orion.yaml
index c5b0ca6e..a1ec43c0 100644
--- a/examples/config_files/diffusion/config_diffusion_mace_orion.yaml
+++ b/examples/config_files/diffusion/config_diffusion_mace_orion.yaml
@@ -25,8 +25,10 @@ model:
     architecture: diffusion_mace
     number_of_atoms: 8
     r_max: 5.0
-    num_bessel: 8
+    num_bessel: 'orion~choices([128, 256, 512])'
     num_polynomial_cutoff: 5
+    num_edge_hidden_layers: 0
+    edge_hidden_irreps: 8x0e
     max_ell: 2
     interaction_cls: RealAgnosticResidualInteractionBlock
     interaction_cls_first: RealAgnosticInteractionBlock
@@ -38,10 +40,12 @@ model:
     correlation: 3
     gate: silu
     radial_MLP: 'orion~choices([[8, 8, 8], [32, 32, 32], [64, 64]])'
-    radial_type: bessel
+    radial_type: 'orion~choices(["bessel", "gaussian"])'
     conditional_prob: 'orion~choices([0.0, 0.25, 0.5, 0.75])'
     conditional_gamma: 2
     condition_embedding_size: 'orion~choices([32, 64])'
+    use_batchnorm: False
+    tanh_after_interaction: True
   noise:
     total_time_steps: 100
     sigma_min: 0.001  # default value