Selectively fix atom positions about jax-md HOT 2 CLOSED

Have you tried just using a customer shift function?

from jax-md.

Hey @Noahyt,

Sorry about the delay in response, we were busy getting ready for ICML. As @cpgoodri, it should be possible to do this without really modifying jax md. The basic premise is that we should use a new shift function that only moves a subset of the particles. By passing an explicit mask to the particles, it should be possible to adjust the frozen particles dynamically. One caveat is that this won't work with simulation environments that need to use the number of degrees of freedom as a parameter (for example, NVT simulations); if this use-case is important for you please let me know and I'll see what I can do!

In any case, here is schematically how I would write such a shift function,

displacement, shift = space.periodic(box_size)
def masked_shift(R, dR, is_mobile=None):
  if is_mobile is None:
    return shift(R, dR)
  return np.where(is_mobile[:, None], shift(R, dR), R)

Then this masked_shift function can be passed to minimizers. The is_mobile array can be passed as a keyword argument to the minimizer and it should just work.

fire_init, fire_apply = minimize.fire_descent(energy_fn, masked_shift)
fire_apply = jit(fire_apply)
fire_state = fire_init(R)

# Freeze all but the first 10 particles.
is_mobile = np.arange(N) > 10 
for i in range(steps):
  fire_state = fire_apply(fire_state, is_mobile=is_mobile)

Here is an example colab notebook that puts all of this together. Please let me know if this isn't what you had in mind or if you get stuck!

from jax-md.