Stabilize harmonic angles

[proteneer]

This PR:

1) Completely removes the usage of GROMOS harmonic cosine angles and defaults back to the harmonic angle 2) Stabilizes the harmonic angle energy and forces, using Kahan's trick to greatly improve 32-bit precision (see plot below) 3) Verifies that nitriles now behave as intended, with the average equilibrium value closer to 3 rads as opposed to 2.8 rads. 4) Modifies the stable angle to be trignometrically correct by moving the epsilon to an extra dimension of the input vectors. Before, as written, arccos and atan2 implementations differed in the result. Now they are consistent.

TODO

[ ] - Test alchemical protocol efficiency morphing a nitrile to a water.

[proteneer]

Is 2 atan2(y, x) commonly called “Kahan’s trick”

atan2(y,x) itself is not the Kahan's trick, as it's a well known primitive and documented extensively. You can use it directly:

angle = atan2(vector2.y, vector2.x) - atan2(vector1.y, vector1.x); (https://stackoverflow.com/questions/21483999/using-atan2-to-find-angle-between-two-vectors)

The real trick is in how y and x are formulated from the input vectors rji, rjk with an extra bit of normalization to improve precision:

nji = jnp.linalg.norm(rji, axis=-1, keepdims=True)
njk = jnp.linalg.norm(rjk, axis=-1, keepdims=True)
y = jnp.linalg.norm(njk * rji - nji * rjk, axis=-1)
x = jnp.linalg.norm(njk * rji + nji * rjk, axis=-1)

In previous discussions about the cos_angles part of the TM functional form, @jchodera has mentioned [Swope, Ferguson, 1992], which contains diagnosis of related numerical issues + some alternatives.

Thanks for the link. There are two issues with the harmonic angle terms in the links

1) With bond lengths that are non-zero for linear molecules, there's purely a numerical stability issue associated with angles near pi. This is analogous to the issue of computing small angles (which Kahan addresses). 2) With bond lengths that are close to zero (eg. in alchemical intermediates), there's a real singularity associated with the energies. This stems from the observation that a small perturbation, when two points are close to each other, can drastically affect the actual value of the angle (and the energy).

Kahan's trick primarily addresses 1), but also impacts 2) by changing the way epsilon is used.

FE integration — This PR, or later?

Later - I'm going to run some systematic tests (specifically with nitriles, allenes, etc.) before concluding anything on the impact of SingleTopology calcs.

Initially, the TODO in the PR description (interpolating nitrile <-> water) would be informative.

I checked this - and overlaps etc. looked fine. But wasn't sure how to write a unit test for it.