Commit c3a52a0d authored by Maximilian LECHNER's avatar Maximilian LECHNER
Browse files

Removed run_minization_scan_neighbor_list as it is not needed anymore due to a change in jax-md.

parent c6528be8
......@@ -52,32 +52,6 @@ def run_minimization_while(
return state.position, get_maxgrad(state), num_iterations
# This version is fully differentiable and internally jitted
def run_minimization_scan_neighbor_list(
energy_fn, neighbor_fn, nbrs, R_init, shift, num_steps=5000, **kwargs
):
init, apply = minimize.fire_descent(jit(energy_fn), shift, **kwargs)
apply = jit(apply)
@jit
def get_maxgrad(state):
return jnp.amax(jnp.abs(state.force))
@jit
def body_fn(state_nbrs, t):
state, nbrs = state_nbrs
nbrs = neighbor_fn(state.position, nbrs)
state = apply(state, neighbor=nbrs)
return (state, nbrs), 0
state = init(R_init, neighbor=nbrs)
state_nbrs, _ = lax.scan(body_fn, (state, nbrs), jnp.arange(num_steps))
state, nbrs = state_nbrs
return state.position, get_maxgrad(state), nbrs
# This version is internally jitted, differentiable and uses neighbor lists.
# The benefit is that it terminates when properly minimized.
def run_minimization_while_neighbor_list(
......
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