[Models] Vectorized edges length and force computation to avoid nan

src/jax_cem/equilibrium/models.py

@@ -404,7 +404,7 @@ def edges_length(
         """
         vectors = edges_vector(xyz, structure.connectivity)
 
-        return vector_length(vectors)
+        return vmap(vector_length)(vectors)
 
     # ------------------------------------------------------------------------------
     # Edge forces
@@ -440,12 +440,13 @@ def trails_force(
         The force in the trail edges of a structure.
         """
         residuals = jnp.concatenate(residuals)
-        forces = trail_force(residuals)
+        forces = vmap(trail_force)(residuals)
 
         lengths = jnp.reshape(jnp.concatenate(lengths), (-1, 1))
 
         return jnp.copysign(forces, lengths)
 
+
 # ------------------------------------------------------------------------------
 # Helpers
 # ------------------------------------------------------------------------------

src/jax_cem/geometry.py

@@ -6,7 +6,13 @@ def vector_length(v: jax.Array, keepdims: bool = True) -> jax.Array:
     """
     Calculate the length of a vector over its last dimension.
     """
-    return jnp.linalg.norm(v, axis=-1, keepdims=keepdims)
+    v = jnp.nan_to_num(v)
+    is_zero_vector = jnp.allclose(v, 0.0)
+    d = jnp.where(is_zero_vector, jnp.ones_like(v), v)  # replace d with ones if is_zero
+
+    length = jnp.where(is_zero_vector, 0.0, jnp.linalg.norm(d, axis=-1, keepdims=keepdims))
+
+    return length
 
 
 def vector_normalized(u: jax.Array) -> jax.Array: