Consider stabilizing periodic_torsion for nearly collinear inputs

[maxentile]

periodic_torsion becomes undefined for collinear inputs, and the associated force magnitude becomes large for nearly collinear inputs.

Such configurations are extremely rare in typical simulations, due to large force constants of the harmonic bond and angle terms keeping distances > 0 and angles % pi != 0.

However, configurations where periodic_torsion is unstable can be sampled in practice during valence parameter interpolation, for example, if a periodic_torsion term is active for atoms (a, b, c, d) but one or more of the spanned harmonic bonds [(a, b), (b, c), (c, d)] or angles [(a, b, c), (b, c, d)] have force constants close to 0. Avoiding instability during valence parameter interpolation requires workarounds such as staging, as in #820 .

Three small examples, where atoms are (1) collinear (with angle = 0), (2) collinear (with angle = pi), or (3) simply collocated:

import jax
jax.config.update("jax_enable_x64", True)

from jax import grad, jit, numpy as jnp
from timemachine.potentials.bonded import periodic_torsion
import numpy as np

np.random.seed(0)


def u(x):
    idxs = np.array([[0,1,2,3]])
    # ks, phases, periods = params.T
    params = np.array([[1.0, 0.0, 1]])
    return periodic_torsion(x, params, None, 0.0, idxs)


# example 1: collinear with angle(b, c, d) == 0
print("exactly collinear, with angle(b, c, d) == 0")
x_0 = np.array([
    [0,1,0],
    [1,0,0],
    [3,0,0],
    [2,0,0],
], dtype=float)
print(f"energy = {u(x_0)},\nforce = {grad(u)(x_0)}") # (2, nans)

# example 2: collinear with angle(b, c, d) == pi
print("\nexactly collinear, with angle(b, c, d) == pi")
x_pi = np.array([
    [0,1,0],
    [1,0,0],
    [2,0,0],
    [3,0,0],
], dtype=float)
print(f"energy = {u(x_pi)},\nforce = {grad(u)(x_pi)}") # (2, nans)

# example 3: collocated, a == b == c == d
print("\nexactly collocated, with a == b == c == d")
x_collocated = np.array([
    [0,0,0],
    [0,0,0],
    [0,0,0],
    [0,0,0],
], dtype=float)
print(f"energy = {u(x_collocated)},\nforce = {grad(u)(x_collocated)}") # (nan, nans)

# nearly collinear -> very large forces
eps = 1e-30
print(f"\n\nperturbations by N(0, epsilon) noise, with epsilon = {eps}")

# example 1 + noise
x_eps = x_0 + np.random.randn(*x_0.shape) * eps
energy = float(u(x_eps))
force_norm = np.linalg.norm(grad(u)(x_eps))
print(f"energy = {energy},\t|force| = {force_norm}") # (~2, ~1e30)

# example 2 + noise
x_eps = x_pi + np.random.randn(*x_pi.shape) * eps
energy = float(u(x_eps))
force_norm = np.linalg.norm(grad(u)(x_eps))
print(f"energy = {energy},\t|force| = {force_norm}") # (~2, ~1e30)

# example 3 + noise
x_eps = x_collocated + np.random.randn(*x_collocated.shape) * eps
energy = float(u(x_eps))
force_norm = np.linalg.norm(grad(u)(x_eps))
print(f"energy = {energy},\t|force| = {force_norm}") # (~2, ~1e30)