Update _compute_pairwise_periodic to allow double backward of the energies

[JMorado]

Description

This PR is related to #149. It recomputes the distances in a way that preserves gradients, thus allowing the computation of derivatives of the energy with respect to the FF parameters.

P.S.: @lilyminium, I now understand what you meant by having to retain the distances to determine are_interacting. I was indeed missing that. Thanks!

Status

• Ready to go

[lilyminium]

(blocking) Thanks @JMorado! The changes to the code look good to me by eye. Could you please add a test specifically that resolves the issue in #149? e.g.

def test_compute_pairwise_periodic_double_backward():
    conformer = torch.tensor(..., requires_grad=True)
    box_vectors = torch.tensor(...)
    cutoff = torch.tensor(...)

    pairwise = _compute_pairwise_periodic(conformer, box_vectors, cutoff)
    # first backward (forces)
    grad = torch.autograd.grad(pairwise.distances.sum(), conformer, create_graph=True)[0]
    # second backward
    grad.sum().backward()

I've approved so after you've added the test please feel free to merge -- thanks!

[JMorado]

Thanks @lilyminium! The tricky bit is that this issue only occurs on CUDA. Would it still be worth adding it even though, as far as I can tell, the CI won’t test it?

[lilyminium]

Ah... yes, that's a good point. I think it would be still worth adding even though the issue can't be reproduced without CUDA, because with this update the test will still run through the code path.

[JMorado]

I've added a test. All tests are passing -- I'll merge this PR. Thanks!