Predator-prey model example

[oarriaga]

Hi I would like to send a PR to blackjax with this example from numpyro without using numpyro's initialize_model model function. However, it seems that with some seeds the window adaption outputs kernels with close to zero step_size and inverse_mass_matrix values. I don't know exactly what I am doing wrong. The original example doesn't seem to do any reparametrization either. Any input would be appreciated :)

import jax
jax.config.update('jax_platform_name', 'cpu')
from collections import namedtuple
import jax.numpy as jnp
from jax.experimental.ode import odeint
from tensorflow_probability.substrates import jax as tfp
tfd = tfp.distributions

params = namedtuple('params', ['uinit', 'vinit', 'alpha', 'betta',
                               'gamma', 'delta', 'u_sigma', 'v_sigma'])

uinit_prior = tfd.LogNormal(jnp.log(10.0), 1.0)
vinit_prior = tfd.LogNormal(jnp.log(10.0), 1.0)
alpha_prior = tfd.TruncatedNormal(1.00, 0.50, 0.0, jnp.inf)
betta_prior = tfd.TruncatedNormal(0.05, 0.05, 0.0, jnp.inf)
gamma_prior = tfd.TruncatedNormal(1.00, 0.50, 0.0, jnp.inf)
delta_prior = tfd.TruncatedNormal(0.05, 0.05, 0.0, jnp.inf)
u_sigma_prior = tfd.LogNormal(-1, 1)
v_sigma_prior = tfd.LogNormal(-1, 1)


def lotka_volterra_ODE(z, t, alpha, betta, gamma, delta):
    u, v = z[0], z[1]
    du_dt = (+alpha - betta * v) * u
    dv_dt = (-gamma + delta * u) * v
    return jnp.stack([du_dt, dv_dt])


def Model(ODE, rtol=1e-6, atol=1e-5, max_steps=1000):
    def apply(z_init, year_args, alpha, betta, gamma, delta):
        z = odeint(ODE, z_init, year_args, alpha, betta, gamma, delta,
                   rtol=rtol, atol=atol, mxstep=max_steps)
        return z
    return apply


def target_log_prob(model, data):
    year_args = jnp.arange(len(data), dtype=jnp.float32)

    def apply(params):
        uinit_log_prob = uinit_prior.log_prob(params.uinit)
        vinit_log_prob = vinit_prior.log_prob(params.vinit)
        alpha_log_prob = alpha_prior.log_prob(params.alpha)
        betta_log_prob = betta_prior.log_prob(params.betta)
        gamma_log_prob = gamma_prior.log_prob(params.gamma)
        delta_log_prob = delta_prior.log_prob(params.delta)
        u_sigma_log_prob = u_sigma_prior.log_prob(params.u_sigma)
        v_sigma_log_prob = v_sigma_prior.log_prob(params.v_sigma)
        zinit = jnp.array([params.uinit, params.vinit])
        args = (params.alpha, params.betta, params.gamma, params.delta)
        z = model(zinit, year_args, *args)
        sigmas = jnp.array([params.u_sigma, params.v_sigma])
        log_likelihood = tfd.LogNormal(jnp.log(z), sigmas).log_prob(data)
        return (uinit_log_prob + vinit_log_prob + alpha_log_prob +
                betta_log_prob + gamma_log_prob + delta_log_prob +
                u_sigma_log_prob + v_sigma_log_prob + log_likelihood.sum())
    return apply


def sample(key):
    keys = jax.random.split(key, 8)
    uinit = uinit_prior.sample(seed=keys[0])
    vinit = vinit_prior.sample(seed=keys[1])
    alpha = alpha_prior.sample(seed=keys[2])
    betta = betta_prior.sample(seed=keys[3])
    gamma = gamma_prior.sample(seed=keys[4])
    delta = delta_prior.sample(seed=keys[5])
    u_sigma = u_sigma_prior.sample(seed=keys[6])
    v_sigma = v_sigma_prior.sample(seed=keys[7])
    return params(uinit, vinit, alpha, betta, gamma, delta, u_sigma, v_sigma)


def inference_loop(rng_key, kernel, initial_state, num_samples):

    @jax.jit
    def one_step(state, rng_key):
        state, info = kernel(rng_key, state)
        return state, (state, info)

    keys = jax.random.split(rng_key, num_samples)
    _, (states, infos) = jax.lax.scan(one_step, initial_state, keys)
    return states, infos


if __name__ == "__main__":
    from numpyro.examples.datasets import LYNXHARE, load_dataset
    import blackjax

    fetch = load_dataset(LYNXHARE, shuffle=False)[1]
    year, data = fetch()

    key = jax.random.PRNGKey(1)
    key_init, key_warm, key_loop, key_pred = jax.random.split(key, 4)
    warmup_steps = 1000
    num_samples = 1000

    model = Model(lotka_volterra_ODE)
    log_prob = target_log_prob(model, data)
    position = sample(key_init)
    adapt = blackjax.window_adaptation(
        algorithm=blackjax.nuts,
        logprob_fn=log_prob,
        is_mass_matrix_diagonal=False,
        initial_step_size=1.0,
        progress_bar=True)

    state, kernel, kernel_params = adapt.run(key_warm, position, warmup_steps)
    print('Kernel params', kernel_params)
    states, infos = inference_loop(key_loop, kernel, state, num_samples)

[rlouf]

Thank you for your interest in contributing. Before trying to run the example locally I have a few questions:

• Are there divergences during warmup? Where do they occur in the sample space (close to the border of the truncated normal)?
• Have you tried running Numpyro’s example with different seeds to see if the same thing happens?
• What do you mean by close to zero? Have you tried sampling with these values, what happens?

It would be interesting to understand what’s happening, but then we probably should reparametrize the model anyway.

[oarriaga]

Thank you for the fast response! I found an error that I fixed while posting this. The error was a RNG_KEY being used twice. I thought that this would not be too problematic, but now I don't seem to get any close to zero step_size or inverse_mass_matrix. The posted code here works for me now when I change to keys. But I would still like to answer you :)

1. I don't know :( How can I check this?
2. Yes. I tried multiple seeds and it always obtained similar parameters and posterior samples.
3. I am currently trying to sample from the posterior to replicate the numpyro image. I am building the posterior with tensorflow's Empirical distribution. Would you recommend any other way to do it?

[rlouf]

Thank you for the fast response! I found an error that I fixed while posting this. The error was a RNG_KEY being used twice.

Ah, yes that could explain it. Which one? (Re-read the code and can’t find it).

1. I don't know :( How can I check this?

Yeah it seems we can’t do it with the high-level interface. That’s a problem and we should open an issue for that.

3. I am currently trying to sample from the posterior to replicate the numpyro image. I am building the posterior with tensorflow's Empirical distribution. Would you recommend any other way to do it?

Posterior predictive? Yeah you’d need to implement the forward sampling version of your model and vmap this sampling function over the samples from the posterior distribution given by Blackjax.

[oarriaga]

It was this line v_sigma = v_sigma_prior.sample(seed=keys[7]). I used keys[6] repeating it for u_sigma. But I fixed it while posting this. Yes the posterior predictive, I will update this issue in case I encounter another problem! Thank you.

[rlouf]

Np. Looking forward to the PR :)

[johannahaffner]

Hi!

I found this example quite useful and was able to adapt it to construct a loglikelihood for a related problem (ODE fitting with generalised least squares). Thanks for showing how to use closures to obtain a logdensity function!

If I'm not mistaken, this example requires a change in inference loop since #501, after which adapt.run no longer returns a kernel.
Building the kernel and initialising the kernel would now be implemented as described in the "Getting Started" example, right?

[oarriaga]

Hi @johannahaffner I am happy to hear it helped you! This was some time ago. Most likely is not working with the new API. I didn't continue to send the PR since I was unable to obtain the same posterior results as in numpyro.

If you are interested please feel free to continue the work started here.

[johannahaffner]

Hi @oarriaga,

good to know! Do you think it would make sense to make a smaller example that focuses on the construction of a loglikelihood, samples from that and returns a covariance matrix? In that case, I could fit a curve to synthetic data, fit it with least squares and then sample around the estimate to quantify the uncertainty of the estimate.
I have working code that does that.

Do you think that would be useful to blackjax users?
Also, do you have a good guess as to why you did not get the same posterior results? I have no experience with numpyro, and no time to get into it right now.