Implement generic O(2) population density support

[maedoc]

From

• https://pubmed.ncbi.nlm.nih.gov/19013532
• https://www.sciencedirect.com/science/article/pii/S0895717700000364

it seems straightforward to use the auto-diff in Jax to implement 2nd order ODEs for some arbitrary Langevin w/ dX_i = f_i(X)dt + σ_i dW_i, D_i = σ_i²,

• dμ_i/dt = f_i(μ) + (1/2)∑_j,k (∂²f_i/∂x_j∂x_k)|_μ C_jk
• dC_ij/dt = ∑_k [(∂f_i/∂x_k)|_μ C_kj + (∂f_j/∂x_k)|_μ C_ik] + D_iδ_ij

with a quick guess at initial implementation,

import jax
import jax.numpy as jnp
from functools import partial

def make_moment_odes(f, D):  
    def moment_odes(state):

        # TODO handle (u,C) directly as pytree instead of flattened vector
        n = len(D)
        μ = state[:n]
        C = state[n:].reshape(n, n)
        
        # Compute Jacobian and Hessian at current means
        J = jax.jacfwd(f)(μ)
        H = jax.hessian(f)(μ)
        
        # Compute mean derivatives
        # dμ_i/dt = f_i(μ) + (1/2)∑_j,k (∂²f_i/∂x_j∂x_k)|_μ C_jk
        dμ = f(μ) + 0.5 * jnp.einsum('ijk,jk->i', H, C)
        
        # Compute covariance derivatives
        # dC_ij/dt = ∑_k [(∂f_i/∂x_k)|_μ C_kj + (∂f_j/∂x_k)|_μ C_ik] + D_iδ_ij
        dC = jnp.einsum('ik,kj->ij', J, C) + jnp.einsum('jk,ki->ij', J, C) + jnp.diag(D)
        
        # Return concatenated derivatives
        return jnp.concatenate([dμ, dC.flatten()])
    
    return moment_odes

NB both math & code were drafted by Claude and require manual & numerical verification