Alternative inference machinery

[afmagee42]

In exploring means to improve MCMC, I tested out Laplace approximations to the posterior for use as a mass matrix. As a mass matrix it was a failure, but at some point it might be nice to enable some form(s) of approximate inference. Development cycles could be notably faster with the right one, bypassing MCMC entirely.

This, put where MCMC is currently run in retrospective_forecasting/main.py, works to get samples from a Laplace approximation. (Note that they are apparently unnamed, unlike those returned from MCMC, so some effort to determine the order may be necessary.)

    model_class = linmod.models.__dict__[model_name]
    model = model_class(model_data)

    guide = numpyro.infer.autoguide.AutoLaplaceApproximation(model.numpyro_model)
    svi = numpyro.infer.SVI(
        model.numpyro_model,
        guide,
        numpyro.optim.Adam(0.3),
        numpyro.infer.Trace_ELBO(),
    )
    svi_res = svi.run(jax.random.key(0), 1000)
    laplace_approx = guide.get_posterior(svi_res.params)
    samples = laplace_approx.sample(jax.random.key(0), sample_shape=(500,))