Gradients are nan for reasonable parameters in quasisep.Carma (pull #90)

[nevencaplar]

@ywx649999311 and others I hope you can have a look. I have been testing the pull #90. As one of the tests, I generated a simple lightcurve with a damped random walk process in EzTao and tried to fit it with tinygp. I am unable to find a gradient for the model for any reasonable parameters, which leads to wrong solutions when trying to find optimal parameters. I also tried a bit more complex models (CARMA (2,0)), but I seem to have similar problems. Same problem if not applying jit. Please find below the most minimal example that I have managed to create.

I hesitate to open it as a bug, as this is still only a pull, and it could be that I am doing something wrong. I would greatly appreciate your help.

import numpy as np

import jax
# import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

import tinygp
from tinygp import GaussianProcess
from tinygp.kernels import quasisep
print(tinygp.__version__)

import eztao
from eztao.carma import DRW_term
from eztao.ts import gpSimRand
print(eztao.__version__)
---------
0.2.3.dev57+g60dde94
0.4.0

# simulate stochastic component
amp = 0.2
tau = 100
DRW_kernel = DRW_term(np.log(amp), np.log(tau))
t, y, yerr = gpSimRand(DRW_kernel, 10, 365*10, 200)

alpha_1, beta_0 = DRW_kernel.get_carma_parameter()
print('alpha_1, beta_0: ' + str(alpha_1) +','+ str(beta_0))
---------
alpha_1, beta_0: 0.009999999999999995,0.02828427124746189

# multiplying with 0.9 so that I am sure that I am not hitting correct values
theta_drw_init = {   
     "alpha": alpha_1*0.9,
    "beta": beta_0*0.9
}

# 1. create Gaussian process
def build_drw_gp(theta, t, y, yerr):
   
    carma_kernel = quasisep.CARMA.init(
        alpha=theta["alpha"], beta=theta["beta"]
    )
        
    return GaussianProcess(carma_kernel, t, diag=yerr, mean=np.mean(y))

# 2. create loss function
@jax.jit
def neg_log_likelihood_drw(theta, t, y, yerr):
    gp_drw = build_drw_gp(theta, t, y, yerr)
    return -gp_drw.log_probability(y)

# 3. evaluate in situ
obj_drw = jax.jit(jax.value_and_grad(neg_log_likelihood_drw))


print(f"Initial negative log likelihood: {obj_drw(theta_drw_init, t, y, yerr)[0]}")
print(
    f"Gradient of the negative log likelihood, \
    wrt the parameters:\n{obj_drw(theta_drw_init, t, y, yerr)[1]}"
---------
Initial negative log likelihood: -143.14509220819014    
Gradient of the negative log likelihood, wrt the parameters:
{'alpha': DeviceArray(nan, dtype=float64), 'beta': DeviceArray(nan, dtype=float64)}

[ywx649999311]

@nevencaplar I can confirm the problem comes from the CARMA kernel. I just tried running your code after replacing the CARMA kernel with an equivalent Exp kernel. It returns proper gradients. I will have to have a deeper look.

[nevencaplar]

I have not had enough time to solve it but, I am pretty sure that problems are coming from h2 = jnp.sqrt(h2_2) and from square root in line h1 = (c * h2 - jnp.sqrt(a * d2 - s2 * h2_2)) / (d + eta * real_mask).

I presume that roots become numerically negative? when values are small (they are all 0. in the examples I tried), this produces nan and breaks everything.

I will continue looking into it tomorrow.

P.S. For the longest time I thought it was connected to this issue https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where, because real_mask has a definition that includes where. The problem persists if one removes where from all CARMA definitions.