Commit 94470c06 authored by Carl Goodrich's avatar Carl Goodrich
Browse files

fix norms

parent 090ddfd5
......@@ -160,7 +160,11 @@ def setup_DNEB_force(true_energy_fn, spring_energy_fn):
gtilde_perp = gtilde - gtilde_parallel
#calculate gtilde_star via eq. (13). this is the "second nudge"
g_perp_hat = g_perp / jnp.linalg.norm(g_perp, axis=(1, 2), keepdims=True)
g_perp_norm = jnp.linalg.norm(g_perp, axis=(1, 2), keepdims=True)
g_perp_norm = jnp.where(g_perp_norm < 1e-14, 1e-14, g_perp_norm)
#g_perp_norm += 1e-14
g_perp_hat = g_perp / g_perp_norm
#g_perp_hat = g_perp / jnp.linalg.norm(g_perp, axis=(1, 2), keepdims=True)
gtilde_star = gtilde_perp - vmap_projection(gtilde_perp, g_perp_hat)
#get the total gradient via eq. (12) with some parts (g_parallel and gtilde_perp) projected out
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment