jax-md / jax-md

Differentiable, Hardware Accelerated, Molecular Dynamics
Apache License 2.0
1.18k stars 197 forks source link

Forces which aren't the gradient of a pair potential #130

Open saridout opened 3 years ago

saridout commented 3 years ago

Is there a canonical procedure in mind for adding forces which are not the gradient of a pair potential, e.g. a uniform gravitational field or friction? (...I guess the latter may also require handling rotation/torques, so maybe it's more of an effort)

cpgoodri commented 3 years ago

You should be able to just define a custom force function and pass it to any of the dynamics functions instead of an energy function. JAXMD uses quantity.canonicalize_force to automatically determine if its a force or energy function, so behavior will be as expected.

saridout commented 3 years ago

Ok, I understand. An arbitrary function of the coordinates of the particles will work. On the other hand, it doesn't look like the force function can use the velocities, to produce some kind of drag? (Of course, as a kludge one can manually modify the simulation loops to also pass the velocities to the force function.) Would it be useful to add a flag to all the simulation loops to allow others to use this kind of behaviour?

On Thu, Mar 25, 2021 at 11:33 AM Carl Goodrich @.***> wrote:

You should be able to just define a custom force function and pass it to any of the dynamics functions instead of an energy function. JAXMD uses quantity.canonicalize_force to automatically determine if its a force or energy function, so behavior will be as expected.

— You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub https://github.com/google/jax-md/issues/130#issuecomment-806958467, or unsubscribe https://github.com/notifications/unsubscribe-auth/AG6OKO46Z3IAG36W7552NLLTFNJUZANCNFSM4ZZQGCGQ .

sschoenholz commented 3 years ago

Just to reiterate, all of the simulations accept arbitrary functions of positions that produce either an energy (in which case the output is a scalar) or a force (in which case the output is an (N, spatial_dimension) array. Also, Carl wrote a really nice colab notebook.

You're right, however, that we don't really have much support for rotation / torques or for drag-like forces that are functions of the velocity (aside from Langevin dynamics which naturally has such a term). I think it might be worth separating the issues of rotation / torques from drag-like forces since my understanding is that they involve different modifications of the integrators. I am currently working on a refactor of the simulation code that will make it easier to describe custom simulations. Do you have a specific use-case that you'd be interested in enabling?

saridout commented 3 years ago

I didnt realize there was a problem with the integrators; i had naively thought we could just replace positions with angles, moments with angular momenta, forces with torques, etc [edit: I misunderstood you I think]

Basically someone in the group is looking for a way to do granular simulations - so they need some kind of dry friction (of which I think there are models which are a bit easier to handle in a simulation than the high-school physics law), uniform gravity, and closed or open rather than periodic boundaries. The latter two seem easy to handle with what you and Carl told me.

I realize that maybe isn't enough to be quite helpful. We can do some research into the possible friction models and get back to you.

On Thu, Mar 25, 2021, 1:26 PM sschoenholz @.***> wrote:

Just to reiterate, all of the simulations accept arbitrary functions of positions that produce either an energy (in which case the output is a scalar) or a force (in which case the output is an (N, spatial_dimension) array. Also, Carl wrote a really nice colab notebook https://colab.sandbox.google.com/github/google/jax-md/blob/master/notebooks/customizing_potentials_cookbook.ipynb .

You're right, however, that we don't really have much support for rotation / torques or for drag-like forces that are functions of the velocity (aside from Langevin dynamics which naturally has such a term). I think it might be worth separating the issues of rotation / torques from drag-like forces since my understanding is that they involve different modifications of the integrators. I am currently working on a refactor of the simulation code that will make it easier to describe custom simulations. Do you have a specific use-case that you'd be interested in enabling?

— You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub https://github.com/google/jax-md/issues/130#issuecomment-807139455, or unsubscribe https://github.com/notifications/unsubscribe-auth/AG6OKO72MA5G5JS2QFYHMJTTFNW3XANCNFSM4ZZQGCGQ .

sdillavou commented 3 years ago

The simulation we're trying to make is of 2D circular grains flowing through an opening, which will occasionally clog. So we can probably get away without worrying about angular momentum/rotation, but a friction force is definitely necessary.

sschoenholz commented 3 years ago

Sorry for the delay, I was a bit busy with the ICML rebuttal deadline. I'd love to help get this working for you all, and I do think that it would be a nice simulation that would work well in JAX MD. If you can say a bit more about the friction model that you are trying to implement (perhaps with a link to a particular paper) then I'd be more than happy to help you get it working.

sdillavou commented 3 years ago

Not a problem! I think actually it's best if I do a bit more background work on this so I don't waste your time. I haven't done a simulation like this (and not much of any other kinds) before, and the idea is still a bit vague. The eventual goal is to have a simulation/ANN feedback loop that lets the ANN learn the system's dynamics.