torchmd / torchmd-net

Training neural network potentials
MIT License
330 stars 74 forks source link

Periodic boundary conditions #92

Closed peastman closed 9 months ago

peastman commented 2 years ago

I want to try applying a model to a larger system that uses periodic boundary conditions. Currently that isn't supported. It looks to me like the only class that would need to be modified is Distance? But it implements the calculation with torch_cluster.radius_graph(), which doesn't support periodic boundary conditions. So either we would need to add that feature to torch_cluster, or else we would need to rewrite that class to work differently.

Any suggestions on the best way to approach this?

giadefa commented 2 years ago

Given that raimondas is rewriting that function for the faster version, we could directly implement it there maybe.

On Mon, May 16, 2022, 22:37 Peter Eastman @.***> wrote:

I want to try applying a model to a larger system that uses periodic boundary conditions. Currently that isn't supported. It looks to me like the only class that would need to be modified is Distance? But it implements the calculation with torch_cluster.radius_graph(), which doesn't support periodic boundary conditions. So either we would need to add that feature to torch_cluster, or else we would need to rewrite that class to work differently.

Any suggestions on the best way to approach this?

— Reply to this email directly, view it on GitHub https://github.com/torchmd/torchmd-net/issues/92, or unsubscribe https://github.com/notifications/unsubscribe-auth/AB3KUOUPCB52FTAO2DWKIXLVKKW2HANCNFSM5WCWC3OA . You are receiving this because you are subscribed to this thread.Message ID: @.***>

raimis commented 2 years ago

The proof-of-concept kernel is here (https://github.com/torchmd/torchmd-net/pull/61). It was optimized for small molecules, but at the moment neither batching, nor the periodic boundary condition are not supported.

peastman commented 2 years ago

If we can add periodic boundary conditions to it, that would be great. For now I think it's fine to only support rectangular boxes, though we should keep open the possibility of triclinic boxes in the future. Or we could just include that from the start, since it only adds minimal complexity.

giadefa commented 2 years ago

Could you make a PR on this? We are able to work on this right now.

On Tue, May 17, 2022, 16:20 Peter Eastman @.***> wrote:

If we can add periodic boundary conditions to it, that would be great. For now I think it's fine to only support rectangular boxes, though we should keep open the possibility of triclinic boxes in the future. Or we could just include that from the start, since it only adds minimal complexity.

— Reply to this email directly, view it on GitHub https://github.com/torchmd/torchmd-net/issues/92#issuecomment-1128933155, or unsubscribe https://github.com/notifications/unsubscribe-auth/AB3KUOXA2DOCF3EMXHXN3K3VKOTMFANCNFSM5WCWC3OA . You are receiving this because you commented.Message ID: @.***>

peastman commented 2 years ago

Since #61 is still under development, I think it would make more sense to include the changes there. The code is very simple. After computing the delta between two positions, you just add a few extra lines to adjust it. Here's the code from OpenMM for a rectangular box.

delta.x -= floor(delta.x*invPeriodicBoxSize.x+0.5f)*periodicBoxSize.x;
delta.y -= floor(delta.y*invPeriodicBoxSize.y+0.5f)*periodicBoxSize.y;
delta.z -= floor(delta.z*invPeriodicBoxSize.z+0.5f)*periodicBoxSize.z;

For a triclinic box it's just a little bit more complicated.

real scale3 = floor(delta.z*invPeriodicBoxSize.z+0.5f);
delta.x -= scale3*periodicBoxVecZ.x;
delta.y -= scale3*periodicBoxVecZ.y;
delta.z -= scale3*periodicBoxVecZ.z;
real scale2 = floor(delta.y*invPeriodicBoxSize.y+0.5f);
delta.x -= scale2*periodicBoxVecY.x;
delta.y -= scale2*periodicBoxVecY.y;
real scale1 = floor(delta.x*invPeriodicBoxSize.x+0.5f);
delta.x -= scale1*periodicBoxVecX.x;
giadefa commented 2 years ago

I guess that the more involved part is the batching

On Tue, May 17, 2022 at 5:27 PM Peter Eastman @.***> wrote:

Since #61 https://github.com/torchmd/torchmd-net/pull/61 is still under development, I think it would make more sense to include the changes there. The code is very simple. After computing the delta between two positions, you just add a few extra lines to adjust it. Here's the code from OpenMM for a rectangular box.

delta.x -= floor(delta.xinvPeriodicBoxSize.x+0.5f)periodicBoxSize.x; delta.y -= floor(delta.yinvPeriodicBoxSize.y+0.5f)periodicBoxSize.y; delta.z -= floor(delta.zinvPeriodicBoxSize.z+0.5f)periodicBoxSize.z;

For a triclinic box it's just a little bit more complicated.

real scale3 = floor(delta.zinvPeriodicBoxSize.z+0.5f); delta.x -= scale3periodicBoxVecZ.x; delta.y -= scale3periodicBoxVecZ.y; delta.z -= scale3periodicBoxVecZ.z; real scale2 = floor(delta.yinvPeriodicBoxSize.y+0.5f); delta.x -= scale2periodicBoxVecY.x; delta.y -= scale2periodicBoxVecY.y; real scale1 = floor(delta.xinvPeriodicBoxSize.x+0.5f); delta.x -= scale1*periodicBoxVecX.x;

— Reply to this email directly, view it on GitHub https://github.com/torchmd/torchmd-net/issues/92#issuecomment-1129011069, or unsubscribe https://github.com/notifications/unsubscribe-auth/AB3KUOVV2FJZB43ZYFD6DWLVKO3FBANCNFSM5WCWC3OA . You are receiving this because you commented.Message ID: @.***>

raimis commented 2 years ago

Thanks @peastman! I'll try to move the code to NNPOps and make something working tomorrow.

peastman commented 1 year ago

What's the status of this? We're working on a project that needs periodic boundary conditions.

FranklinHu1 commented 1 year ago

I have a basic implementation of periodic boundary conditions using square boxes that is working for liquid systems (e.g. water and imidazole), but it could definitely be optimized for better integration with the torchmd-net architecture. It would be great to collaborate with someone on the torchmd-net team about getting this feature optimized and added!

RaulPPelaez commented 1 year ago

I will add triclinic PBC from NNPops to #169. The new distance module there is a drop in replacement for Distance.

FranklinHu1 commented 1 year ago

Thanks! Roughly what would the timeline be for PBCs to be usable within torchmd-net?

RaulPPelaez commented 1 year ago

I am actively working on it. In that PR I have an N^2 and a cell list backend. The N^2 I have ready, with the cell list I have cubic boxes working but triclinic is giving me a a bit of a headache. I will clean up a bit and upload with cubic box support ASAP.

peastman commented 1 year ago

When you say cubic, I assume you mean rectangular? What's the issue with triclinic?

RaulPPelaez commented 1 year ago

I meant rectangular, yes. The cell list requires to bin the domain and identify the bin for each particle. In a rectangular cell this is straight forward, say:

int cell_x = floor((pos_x + 0.5*L_x)/cutoff)

I am struggling to achieve this in a triclinic box in a sane way.

giadefa commented 1 year ago

We don't care about triclinic

On Fri, May 5, 2023, 20:44 Raul @.***> wrote:

I meant rectangular, yes. The cell list requires to bin the domain and identify the bin for each particle. In a rectangular cell this is straight forward, say:

int cell_x = floor((pos_x + 0.5*L_x)/cutoff)

I am struggling to achieve this in a triclinic box in a sane way.

— Reply to this email directly, view it on GitHub https://github.com/torchmd/torchmd-net/issues/92#issuecomment-1536637385, or unsubscribe https://github.com/notifications/unsubscribe-auth/AB3KUOXC57YVRU3HY2EANK3XEVDADANCNFSM5WCWC3OA . You are receiving this because you commented.Message ID: @.***>

FranklinHu1 commented 1 year ago

Is it possible to just release rectangular box PBCs first?

peastman commented 1 year ago

You can bin it in exactly the same way. A triclinic box is identical to a rectangular box with offsets between copies. The only part you need to handle differently is when identifying which bins in other periodic copies you need to search. See https://doi.org/10.1002/(SICI)1096-987X(19971130)18:15%3C1930::AID-JCC8%3E3.0.CO;2-P for the theory of how it works. And see https://github.com/openmm/openmm/blob/master/platforms/reference/src/SimTKReference/ReferenceNeighborList.cpp for a reference implementation in OpenMM.

We don't care about triclinic

But we do. :)

RaulPPelaez commented 1 year ago

Thanks for the refs @peastman!. Since rectangular seems more time pressing I will work to finish that ASAP and then I will finish triclinic.

RaulPPelaez commented 1 year ago

Lets move the discussion to #169