Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Derivatives of forces are NaN when there's a right angle in a structure #58

Open
SiLiKhon opened this issue Feb 16, 2023 · 0 comments
Open

Comments

@SiLiKhon
Copy link

The problem occurs in the spherical harmonics calculation:

import tensorflow as tf
from m3gnet.utils import SphericalHarmonicsFunction

sph = SphericalHarmonicsFunction(3, use_phi=False)

def Ylm_2nd_der(x):
    with tf.GradientTape() as t0:
        with tf.GradientTape() as t:
            t.watch(x)
            t0.watch(x)
            y = sph([x], [0.0])

        dydx = t.gradient(y, x)
    d2ydx2 = t0.gradient(dydx, x)

    return d2ydx2

for costheta in [-1.0, -0.5, 0.0, 0.5, 1.0]:
    print(Ylm_2nd_der(tf.convert_to_tensor(costheta, dtype="float32")))

Output:

tf.Tensor(1.8923494, shape=(), dtype=float32)
tf.Tensor(1.8923494, shape=(), dtype=float32)
tf.Tensor(nan, shape=(), dtype=float32)
tf.Tensor(1.8923494, shape=(), dtype=float32)
tf.Tensor(1.8923494, shape=(), dtype=float32)

It happens due to casting to complex here:

costheta = tf.cast(costheta, dtype=tf.dtypes.complex64)

Then, the actual harmonics functions involve a pow in (DT_COMPLEX64, DT_COMPLEX64) -> DT_COMPLEX64 signature, which probably utilizes polar form, hence undefined derivative in 0.

I also notice that in some cases tf.function seems to optimize it out and returns correct value of the derivative, but it's very unpredictable as to when this optimization takes place (e.g., it may happen in t.gradient, but not in t.jacobian).

A possible workaround would be to avoid casting to complex when self.use_phi is False.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant