FireDescent should use velocity and not momentum when calculating P

[MarcBerneman]

This note was left in the code (jax_md.minimize.fire_descent) for fire descent and I think that it's valid.
NOTE: In the original FIRE algorithm, the quantity that determines when to reset the momenta is F.V rather than F.P. However, all of the JAX MD simulations are in momentum space for easier agreement with prior work /rigid body physics. We only use the sign of F.P here, which shouldn't differ from F.V, however if there are regressions then we should reconsider this choice.
Working with F.P leads to non-convergence when working with infinite masses (jnp.inf). In this case, the velocity of the particles with infinite mass are 0, but their momentum is not. Their momentum are quite large in force balance and can thus have a big influence on the sign of F.P. When I changed the code to use F.V, where V is velocity, this solved my problem.