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

[SiLiKhon]

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:

m3gnet/m3gnet/utils/_math.py

Line 249 in 87173c1

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.