Commit b40f0a81 authored by Carl Goodrich's avatar Carl Goodrich
Browse files

update to new neighbor_fn api

parent 3c17295a
......@@ -67,13 +67,13 @@ def run_brownian_neighbor_list(
measure_fn=lambda R, nbrs: R,
**static_kwargs
):
nbrs = neighbor_fn(R_init)
nbrs = neighbor_fn.allocate(R_init)
init, apply = simulate.brownian(energy_fn, shift, dt=dt, kT=kT, gamma=gamma)
def body_fn(state_nbrs, t):
state, nbrs = state_nbrs
nbrs = neighbor_fn(state.position, nbrs)
nbrs = neighbor_fn.update(state.position, nbrs)
state = apply(state, neighbor=nbrs, **static_kwargs)
return (state, nbrs), 0
......@@ -93,7 +93,7 @@ def run_brownian_neighbor_list(
# the simulation.
if nbrs.did_buffer_overflow:
print("Buffer overflow.")
nbrs = neighbor_fn(state.position)
nbrs = neighbor_fn.allocate(state.position)
else:
state = new_state
step += step_inc
......
......@@ -65,7 +65,7 @@ def run_minimization_while_neighbor_list(
verbose=False,
**kwargs
):
nbrs = neighbor_fn(R_init)
nbrs = neighbor_fn.allocate(R_init)
init, apply = minimize.fire_descent(jit(energy_fn), shift, **kwargs)
apply = jit(apply)
......@@ -77,7 +77,7 @@ def run_minimization_while_neighbor_list(
@jit
def body_fn(state_nbrs, t):
state, nbrs = state_nbrs
nbrs = neighbor_fn(state.position, nbrs)
nbrs = neighbor_fn.update(state.position, nbrs)
state = apply(state, neighbor=nbrs)
return (state, nbrs), 0
......@@ -93,7 +93,7 @@ def run_minimization_while_neighbor_list(
# the simulation.
if nbrs.did_buffer_overflow:
print("Buffer overflow.")
nbrs = neighbor_fn(state.position)
nbrs = neighbor_fn.allocate(state.position)
else:
state = new_state
step += step_inc
......
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