zero emb for zero edges

ACEsuit · Mar 9, 2023 · 6cd52fd · 6cd52fd
1 parent d78dc6c
commit 6cd52fd
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/mace_jax/modules/blocks.py b/mace_jax/modules/blocks.py
@@ -98,7 +98,9 @@ def func(lengths):
                 )
                 factor = jnp.mean(func(samples) ** 2).item() ** -0.5
 
-        embedding = factor * func(edge_lengths)  # [n_edges, num_basis]
+        embedding = factor * jnp.where(
+            edge_lengths == 0.0, 0.0, func(edge_lengths)
+        )  # [n_edges, num_basis]
         return e3nn.IrrepsArray(f"{embedding.shape[-1]}x0e", embedding)
 
 

diff --git a/mace_jax/modules/models.py b/mace_jax/modules/models.py
@@ -129,8 +129,8 @@ def __call__(
                 self.radial_embedding(lengths),
                 e3nn.spherical_harmonics(
                     self.sh_irreps,
-                    vectors / lengths[..., None],
-                    normalize=False,
+                    vectors,
+                    normalize=True,
                     normalization="component",
                 ),
             ]

diff --git a/mace_jax/tools/utils.py b/mace_jax/tools/utils.py
@@ -134,7 +134,7 @@ def unflatten_dict(xs, sep=None):
 def safe_norm(x: jnp.ndarray, axis: int = None, keepdims=False) -> jnp.ndarray:
     """nan-safe norm."""
     x2 = jnp.sum(x**2, axis=axis, keepdims=keepdims)
-    return jnp.where(x2 == 0, 1, x2) ** 0.5
+    return jnp.where(x2 == 0.0, 0.0, jnp.where(x2 == 0, 1.0, x2) ** 0.5)
 
 
 def compute_mean_std_atomic_inter_energy(