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 :)

[junpenglao]

🤖 This reply was drafted with Claude and reviewed by me (@junpenglao).

This is a great use of the library — we'd still love to feature it in the sampling-book. One API note since this thread is a bit old: window_adaptation.run now returns (state, parameters), info (parameters, not a kernel), and you can drive the sampling loop with blackjax.util.run_inference_algorithm.

[pnkraemer]

Thanks for messaging!

Not sure I entirely understand: are you saying the current code (respectively Probdiffeq's version) breaks now, will break in the next release, or works fine but could be optimised? :)

[junpenglao]

Disregard this message - the ODE example works fine with probdiffeq in sampling-book (https://blackjax-devs.github.io/sampling-book/models/probabilistic-ode-solver-parameter-estimation/). The LLM replies was instructed by a prompt to check discussion and reply if the old posts provide information outdated due to recent API changes. I am going to ask Claude why they missed the context here.

[junpenglao]

LOL the reason was that we just read body + top-level comment only and never traversed the nested reply chain - lesson learned.

@pnkraemer how have you been? any exciting update on probdiffeq or numerical ode land in general?

[pnkraemer]

Oh, ok - glad everything works then!

That said, there will be some changes to the probabilistic numerics code in the next release (to enable some new features, like solving different kinds of ODEs). I left a reminder on Probdiffeq to revisit the sampling book example here once it's released (pnkraemer/probdiffeq#876).