Confusion with the jacobian

General
1

Dear all,

I have been working on identifying the correlation between two distributions using Frank copula, and in my code I identify the parameter theta_frank that can be both negative and positive. Then I put the jacobian on logit of the correlation coefficient Kendall tau via the following:

real Debye_1st_integral_ = ode_rk45(Debye_1st_function, [1e-6]', 1e-6, {1.0}, theta_frank)[1,1];
tau_Kendall = 1. - 4 / theta_frank + 4 * Debye_1st_integral_;
logit_tau_Kendall_raw = logit((tau_Kendall + 1) / 2);
jacobian += log(4.0) - 2 * log(abs(theta_frank)) + log(abs(1 + theta_frank / expm1(theta_frank) - 2 * theta_frank * Debye_1st_integral_)) - log1m(tau_Kendall) - log1p(tau_Kendall) + log(2);

Then I impose prior on

logit_tau_Kendall_raw ~ normal(0, 2);

Fitting the data, gave me a reasonable theta_frank around -13 and Kendall’s tau around -0.72.

However, initially I made a mistake in my code and I declared the variable log_theta_frank, then exponentiate it to theta_frank, making theta_frank being only positive. In this case I wrote:

theta_frank = exp(log_theta_frank);
jacobian += log(4.0 / theta_frank) + log(abs(1 + theta_frank / expm1(theta_frank) - 2 * theta_frank * Debye_1st_integral_)) - log1m(tau_Kendall) - log1p(tau_Kendall) + log(2);

Fitting the data still gave me Kendall’s tau around -0.72. It was surprising since Kendall’s tau should always be positive. But the posterior distribution for theta_frank was at 0.00001 (95% CrI: 0 to 0.00001). So the latter was kind of nonsense.

However, I wonder if imposing that prior on logit of transformed Kendall’s tau somehow made the contraints on theta_frank > 0 to be ignored. Am I right with this thought?

Another idea of mine was that the Debye_1st_integral was unstable around theta_frank zero, which made some erroneous calculations. However, the same value for the posterior mean of Kendall’s tau in both cases is suspicious.

Thank you for your help and time in advance!

Andrei

    vector Debye_1st_function(real s, vector y, real theta) {  
        return [s / expm1(theta * s)]';
    }
2

How did you derive the log absolute value of the Jacobian here?

Using gemini, it gives

\ln|J_z| = \frac{1}{2}\ln(1 + x^4) - \ln | x - 2 + 2 x y| - \ln |1 - xy|

where x = \tt{theta\_frank}, y = \tt{Debye\_1st\_integral\_}

3

I derived it first by hand and then verified my formula with Claude AI

image
image1982×908 58.4 KB

4

I suspect this is all due to the integrator and tough geometry around where you are sampling. Do you have any other diagnostic issues? Check some pairs plots and try to narrow down where in the region the issue might be arising. I suggest placing informative priors around theta to see if you can get reasonable results in a small region and expanding out from there.

On the jacobian piece, I missed the function definition for the integrator and see that this is a 1d change of variables. Whenever I do these now I make sure to numerically confirm, even if that’s written by an llm we can make sure it’s all written down correctly and that the math checks out. It turns out to be the same as what you have which is good.

import jax
import jax.numpy as jnp
from jax.scipy.special import logit
from jax.experimental.ode import odeint

# 1. Define the ODE system (The Debye Integrand)
def debye_system(y, s, theta):
    # s is the integration variable (0 to 1)
    # y is the accumulated integral
    # We use expm1 for numerical stability near s=0
    return s / jnp.expm1(theta * s)

# 2. The full transformation using the ODE solver
def transform_with_ode(log_theta_frank):
    theta = jnp.exp(log_theta_frank)
    
    # Solve ODE: odeint(func, y0, t_points, params)
    # This matches Stan's ode_rk45(system, initial_state, t0, t_final, params)
    t_points = jnp.array([1e-6, 1.0]) 
    d_array = odeint(debye_system, 0.0, t_points, theta)
    d = d_array[1] # The value at s=1.0
    
    tau_kendall = 1.0 - 4.0 / theta + 4.0 * d
    return logit((tau_kendall + 1.0) / 2.0)

# 3. Analytical Formula (for comparison)
def analytical_log_det(log_theta):
    theta = jnp.exp(log_theta)
    # We still need the value of d to compute the derivative
    d = odeint(debye_system, 0.0, jnp.array([1e-6, 1.0]), theta)[1]
    tau = 1.0 - 4.0 / theta + 4.0 * d
    
    # dz/dx = [8 / (1 - tau^2)] * [1/theta + 1/(exp(theta)-1) - 2d]
    term = (1.0 / theta) + (1.0 / jnp.expm1(theta)) - 2.0 * d
    jacobian = (8.0 / (1.0 - jnp.square(tau))) * term
    return jnp.log(jnp.abs(jacobian))

# Testing
log_theta_val = 1.5
jax_log_det = jnp.log(jnp.abs(jax.grad(transform_with_ode)(log_theta_val)))
formula_log_det = analytical_log_det(log_theta_val)

print(f"JAX (Auto-diff through ODE): {jax_log_det:.10f}")
print(f"Analytical Formula:         {formula_log_det:.10f}")
print(f"Difference:                 {jnp.abs(jax_log_det - formula_log_det):.2e}")

I run that and I get

>>> %run -- "/Users/sean/logdet_jac.py"
JAX (Auto-diff through ODE): -0.2878684700
Analytical Formula:         -0.2878643274
Difference:                 4.14e-06
5

Can you share the Stan code? Where was the constraint declared? If it was a parameter constraint declared in the parameters block in Stan, it won’t be ignored. If you’re doing it yourself, then it depends how you coded it.

As is, your code requires tau_Kendall to be in (-1, 1) for the logit to be well defined after adding 1 and dividing by 2. I don’t know what it can actually be given theta_frank and Debye_1st_integral.