BlackJAX & Probabilistic Numerical ODE Solvers

[pnkraemer]

Hey guys,

Big fan of BlackJAX here! I recently published some probabilistic numerics JAX code. I made a tutorial notebook about combining BlackJAX with probabilistic ODE solvers. The sampling bit worked like a charm!

Wanted to share this in case probabilistic numerics or differential equations are of interest to anyone here :)

Nico

[albcab]

This is amazing! And a unique use of the library. It would be great to have your notebook featured on the Sampling book if possible. Notice window_adaptation returns parameters of the kernel and not the kernel itself on the latest version of the library 😄

[rlouf]

Awesome, this kind of application is exactly what I had in mind when designing Blackjax. Happy to hear it's helpful :)

[pnkraemer]

Awesome, glad you like it! :) (And thanks for the window_adaptation pointer.)

It would be great to have your notebook featured on the Sampling book if possible.

Sure. What would we need to do?

[albcab]

Normally you would open an issue on blackjax-devs/sampling-book but since you already have a working markdown file, skip this and open a PR directly.

Place the .md file in the sampling-book/book/models directory, maybe change its title to highlight its use for probabilistic numerical ODE solvers. It would be great if you could include some of the details in part 1 of your implementation for better context, like the bibliography (follow the example of Periodic Orbital MCMC). Call the file from the TOC. Then build the docs on you local machine to make sure it all runs correctly:

pip install -r requirements.txt
jupyter-book build --all book/

(you might need to mkdir jupyter_cache if there isn't one), and open a PR. Thanks!

[pnkraemer]

PR has been opened. Thanks for the explanations, @albcab, that accelerated the entire process massively :)