Issue tracker for PyMC implementation of INLA

[theorashid]

trafficstars

This is for https://github.com/pymc-devs/pymc/issues/3242 and https://github.com/pymc-devs/pymc/issues/6992. cc: @ricardoV94 @zaxtax

The first hackathon for this will be on Friday 31st May. But we plan to continue development on this beyond that day.

Approximate the marginal posterior distribution of some subset of the parameters, referred to as the marginal Laplace approximation. Then, integrate out the remaining parameters using another method.

This is great for latent Gaussian models.

Reading list for those who are interested

• jax implementation (500 lines of code ish, and then some sparse stuff)
• pymc3 effort
• Stan paper
• Dan Simpson's blog
• my rough notes

1. Laplace approximation (and misc)

• [x] Implement https://github.com/pymc-devs/pymc-experimental/issues/341 #345
• [x] Update https://github.com/pymc-devs/pymc-experimental/issues/342
• [x] PyMC MvNormal with efficient calculation for logp (no inverse) https://github.com/pymc-devs/pymc/pull/7345
  • _precision_mv_normal_logp should be obtained by a call to a PrecisionMVNormal class instance https://github.com/pymc-devs/pymc/pull/7895

2. Marginal Laplace approximation

• [ ] Implement https://github.com/pymc-devs/pymc-experimental/issues/344 using step sampler interface or pmx.MarginalModel(). https://github.com/pymc-devs/pymc-extras/pull/533
• [x] Use https://github.com/pymc-devs/pymc-experimental/issues/343 https://github.com/pymc-devs/pymc-extras/pull/513. https://github.com/pymc-devs/pytensor/pull/1182 https://github.com/pymc-devs/pytensor/pull/978
• [ ] Optimise logp runtime https://github.com/pymc-devs/pymc-extras/issues/568
• [ ] MinimiseOp fails with varying parameter shapes https://github.com/pymc-devs/pytensor/issues/1550 https://github.com/pymc-devs/pytensor/pull/1582
• [ ] Un-marginalise the latent field https://github.com/pymc-devs/pymc-extras/issues/569
• [ ] Improvement: obtain dimension of the latents more robustly https://github.com/pymc-devs/pymc-extras/issues/570

3. API

• [ ] Glue all the above together and get it working in pm.Model. https://github.com/pymc-devs/pymc-extras/pull/533
• [ ] Create an interface with bambi that can be accessed like R-INLA

4. Sparse matrix operations

INLA can work without it, but this is what will make it very quick and scalable and get it nearer to R-INLA performance. This would lie in https://github.com/pymc-devs/pytensor/tree/main/pytensor/sparse. There is a jax implementation of all the parts we need.

• [ ] Implement sparse matrix in pytensor with ability to:
  • add to the diagonal
  • multiply by a vector
  • solve a linear system
  • compute a log determinant
• [ ] Implement a MVN distribution with a sparse precision matrix

5. Documentation and examples

• [ ] generalised linear mixed model example https://github.com/pymc-devs/pymc-extras/pull/533
• [ ] Spatial stats example (maybe @elizavetasemenova)
• [ ] pymc ICAR example but rewritten using INLA
• [ ] Time series example setting up AR model with ICAR (see Dan footnote 41) and more... https://github.com/pymc-devs/pymc-extras/pull/533

Note, I will update and link to the issues/PRs once they are made. If you want to tackle one of these issues, comment below and I will update the list with your name.

If you have any more things to add, please comment and I will add them to the list and create issues.

[elizavetasemenova]

I'm interested in working on #341 and spatial stats

[ColtAllen]

@theorashid @zaxtax do you feel the recent merge of PR 513 is enough to check off the Newton Solver box under Section 1. Laplace Approximation?

My thoughts are that the minimum requirements are met, but in the future we'll want to expand the API for better Hessian support and perhaps pytensor.tensor.optimize.root. The analytical expression for the "full" Laplace approximate requires a Hessian, but for now we're good to proceed with a simplified approximation (i.e., a third order Taylor series expansion).

[asifzubair]

Thank you, @theorashid , for sending me here! I was wondering if I could pick up the sparse matrix work. I've been wanting to contribute to pytensor and this seems quite meaningful. Concurrently, I've also been working on a PR ( pymc-devs/pymc#7879 ) on CAR/ICAR distribution in pymc, which I believe was borne out of a discussion with you. I'd appreciate it if you could provide some review comments on the PR. Thanks, again 🙏

[asifzubair]

Hi @theorashid , just a gentle reminder to ask if I can pick up the sparse matrix work ? Thank you! 🙏

[theorashid]

Oh yeah, if you want to do sparse matrix stuff go for it. I won't stop you

[ColtAllen]

Thank you, @theorashid , for sending me here! I was wondering if I could pick up the sparse matrix work. I've been wanting to contribute to pytensor and this seems quite meaningful. Concurrently, I've also been working on a PR ( pymc-devs/pymc#7879 ) on CAR/ICAR distribution in pymc, which I believe was borne out of a discussion with you. I'd appreciate it if you could provide some review comments on the PR. Thanks, again 🙏

Awesome! I left a comment on your PR as well as some notes in the pertitent pytensor issue for you 🚀