Debug extremely large inertial restraint force

[proteneer]

There's an instability in the inertial restraints, causing excess shear on the ligand. This happens once every several hundred thousand steps. Doing a force decomposition in the timemachine reveals a single frame causing extremely large shear forces

<timemachine.lib.potentials.HarmonicBond object at 0x7f8ac019cdd8> max norm of force 5447.236777093574 nrg 13841.26315214916
<timemachine.lib.potentials.HarmonicAngle object at 0x7f8ac0b4f048> max norm of force 3127.8959150129162 nrg 8331.309307040987
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816080> max norm of force 1512.9616097552519 nrg 5006.375628526592
<timemachine.lib.potentials.HarmonicAngle object at 0x7f8b3f816668> max norm of force 1478.5844982713984 nrg 462.306427993346
<timemachine.lib.potentials.HarmonicBond object at 0x7f8b3f816780> max norm of force 2907.5026530646105 nrg 73.67688241275027
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816828> max norm of force 94.8941503518715 nrg 4.278436809283448
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816d30> max norm of force 486.3823185046202 nrg 65.98731434749789
<timemachine.lib.potentials.HarmonicAngle object at 0x7f8b3f816be0> max norm of force 1790.6036490985516 nrg 418.7167803249322
<timemachine.lib.potentials.HarmonicBond object at 0x7f8b3f816da0> max norm of force 2896.7413993855725 nrg 55.711845652840566
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816e48> max norm of force 38.95069521368167 nrg 0.41211473328075954
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816a20> max norm of force 579.3667071740321 nrg 59.750894849279575
<timemachine.lib.potentials.Nonbonded object at 0x7f8b3f812438> max norm of force 3611.793539461736 nrg -116490.19066778762
<timemachine.lib.potentials.CentroidRestraint object at 0x7f8b3f812240> max norm of force 10.112349050081152 nrg 4.735674220051356
<timemachine.lib.potentials.InertialRestraint object at 0x7f8b3f8125c0> max norm of force 316.8126342054399 nrg 1.279563750300979
step 21254
<timemachine.lib.potentials.HarmonicBond object at 0x7f8ac019cdd8> max norm of force 5326.834655507455 nrg 13721.756436221687
<timemachine.lib.potentials.HarmonicAngle object at 0x7f8ac0b4f048> max norm of force 2761.0284539040445 nrg 8214.45094851031
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816080> max norm of force 1544.8239567279231 nrg 5018.621901277427
<timemachine.lib.potentials.HarmonicAngle object at 0x7f8b3f816668> max norm of force 1441.3473967255659 nrg 460.2551430631429
<timemachine.lib.potentials.HarmonicBond object at 0x7f8b3f816780> max norm of force 2699.670794791638 nrg 66.04856574231962
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816828> max norm of force 94.8615849790374 nrg 4.335874354102998
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816d30> max norm of force 416.76786537643306 nrg 64.43818149171966
<timemachine.lib.potentials.HarmonicAngle object at 0x7f8b3f816be0> max norm of force 1683.1994686068492 nrg 415.01639457786223
<timemachine.lib.potentials.HarmonicBond object at 0x7f8b3f816da0> max norm of force 2959.2834977738667 nrg 46.94698089189478
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816e48> max norm of force 28.89665320297211 nrg 0.27348836621206374
<timemachine.lib.potentials.PeriodicTorsion object at 0x7f8b3f816a20> max norm of force 435.7834046742932 nrg 58.89304468404225
<timemachine.lib.potentials.Nonbonded object at 0x7f8b3f812438> max norm of force 3657.233676879651 nrg -116196.67184198985
<timemachine.lib.potentials.CentroidRestraint object at 0x7f8b3f812240> max norm of force 10.118233676196171 nrg 4.741187435721552
<timemachine.lib.potentials.InertialRestraint object at 0x7f8b3f8125c0> max norm of force 207489.7612308369 nrg 28.336251455720802

Going to debug a little deeper into the eigenvalues/eigenvectors.

[proteneer]

This is due to degeneracy of eigenvalues. See the stripped frame and repro here:

https://github.com/google/jax/discussions

And one possible solution by increasing the "degeneracy" cutoff check:

jax-ml/jax#1665

The jvp solution will require about 9 times more effort to compute compared to the vjp solution.

[proteneer]

It looks like we also need to modify the loss function to smoothen it out. Consider the problem of aligning a single object to the standard basis when the inertial tensor is near symmetry. Let's plot the potential as a function of small perturbations for each element in the 3x3 inertia tensor:

The default version not only has a sharp jump but also extremely large gradients near the singularity:

We can use a cutoff check for degeneracy instead to make this piece wise near the high curvature regions:

An even better approach however, is to rescale using the normalized eigenvalues differences as a scale factor:

from jax.config import config; config.update("jax_enable_x64", True)
import numpy as onp
import jax.numpy as np
import jax
import functools

def loss(tensor):
    w, v = np.linalg.eigh(tensor)

    is_degen = np.abs(np.expand_dims(w, axis=-1) - np.expand_dims(w, axis=0))
    is_degen += np.eye(3)

    # pick one of these three
    # scale = np.ones(3) # unnormalized
    # scale = np.amin(is_degen, axis=0) > 1e-3 # cutoff
    scale = np.amin(is_degen, axis=0) / np.linalg.norm(w) # rescaled

    return np.sum(scale*(1-np.abs(np.diag(v)))**2)

import matplotlib.pyplot as plt

def test():

    for i in range(3):
        for j in range(3):

            xs = []
            ys = []

            for delta in np.linspace(-0.01, 0.01, 200):

                tensor = onp.array([
                    [ 0.06683436, -0.05386403,  0.0127573 ],
                    [-0.05386403,  0.10799386,  0.00877524],
                    [ 0.0127573 ,  0.00877524,  0.14300045]]
                )

                tensor[i][j] += delta

                xs.append(tensor[i][j])
                ys.append(loss(tensor))

            plt.plot(xs, ys, label="delta ["+str(i)+"]["+str(j)+"]")
    
    plt.xlabel("element[i][j]")
    plt.ylabel("loss")
    plt.show()

test()

[proteneer]

This will be closed in favor of RMSD-based restraints.