SineBivariateVonMises generates NaN log probability for negative correlation parameter

[19MM97]

When using SineBivariateVonMises in an MCMC setting with the NUTS sampler and hard assigning the correlation parameter of the SinveBivariateVonMises to a negative value, the model can not be initialised. Sampling from a SineBivariateVonMises with a negative correlation parameter does work though.
@fehiepsi traced, that the log probability becomes NaN.

How to recreate:

data = SineBivariateVonMises(0, 0, 1, 1, -1).sample(PRNGKey(42), (1000, ))

and train this model:

@numpyro.handlers.reparam(
    config={"phi_loc": CircularReparam(), 
            "psi_loc": CircularReparam(),
           }
)
def min_example(data_2d):
    phi_loc = sample('phi_loc', VonMises(0, 10))
    psi_loc = sample('psi_loc', VonMises(0, 10))
    
    phi_conc = sample('phi_conc', Beta(1, 1))
    psi_conc = sample('psi_conc', Beta(1, 1))
    
    conc = -sample('conc', HalfNormal(1))
    
    depInd = SineBivariateVonMises(phi_loc, psi_loc, 70 * phi_conc, 70 * psi_conc, conc)
    
    obs = sample('obs', depInd, obs=data_2d)

Further information can be found in this forum post

[fehiepsi]

Could you add a full code to replicate the issue? Edit: see full code below. We have the valid parameters

{'phi_loc': DeviceArray(0.29627275, dtype=float32),
 'psi_loc': DeviceArray(-0.03310037, dtype=float32),
 'phi_concentration': DeviceArray(19.143257, dtype=float32),
 'psi_concentration': DeviceArray(44.272667, dtype=float32),
 'correlation': DeviceArray(-0.6486334, dtype=float32)}

but gets invalid log_prob with distribution

SineBivariateVonMises(phi_loc, psi_loc, phi_conc, psi_conc, conc)

and observed data_2d. I guess we don't need to whole training process or follow forum post to see the issue. Something like

from jax import random
import numpyro.distributions as dist

phi_loc = 0.29627275
psi_loc = -0.03310037
phi_conc = 19.143257
psi_conc = 44.272667
corr = -0.6486334
d = dist.SineBivariateVonMises(phi_loc, psi_loc, phi_conc, psi_conc, corr)
data = dist.SineBivariateVonMises(0, 0, 1, 1, -1).sample(random.PRNGKey(42), (1000, ))
log_prob = d.log_prob(data)
print(log_prob)

[fehiepsi]

cc @OlaRonning

[OlaRonning]

Hi @fehiepsi and @19MM97,

Thanks for the simple example. I'll spend some time debugging tomorrow morning.

[OlaRonning]

The problem is the log(corr) in the normalization. If we move the 2 factor into the log (so it becomes log(corr**2)), we will always have positive values, eliminating that source of NaNs. I'll open a PR once I've run the tests.

Edit: no reason to move the m factor; we only need to square to make it positive.

[OlaRonning]

#1515 resolves the NaNs from negative correlation. @19MM97 feel free to reopen if the problem persists.