schmidtjonathan / pnmol-experiments Goto Github PK

PNMOL with independent prior dimensions and a polynomial kernel should equal the output of a traditional ODE filter.

It seems like it does not.

Why?

Linting and tests

It would be useful to have a method PDEProblem.to_tornadox_ivp() which removes the "PN" from "PNMOL". We can use it to benchmark pnmol against tornadox.

One difficulty will be the boundary conditions. We cannot just transfer f, df, t0, etc., but have to do the padding-with-zeros-or-the-last-row as in tornadox for the "new" IVP problem.

We need PDE problems to test PNMOL solves on:

1. linear: 1D / 2D heat equation
2. non-linear: 2D inviscid Burgers' equation

There are two hacks left, one in each solver, that overfit the implementation to a system of PDEs with 3 states (the SIR model)

There are still a few numpy imports. I suggest we use only jax.numpy, because this way we can go all in with jit and vmap (the latter will be quite important I guess)

I think we need a few more PDE problem classes. As far as I can judge at the moment, we should have

• DiscretisedLinearPDE (d_t u = d_x u * const)
• DiscretisedSemiLinearPDE (d_t u = d_x u * const + g(u))
• DiscretisedQuasiLinearPDE (d_t u = d_x u * g(u))
• DiscretisedNonlinearPDE (d_t u = g(u, d_x u))

And for each of those, we can implement EK0 and EK1 cleverly.

This is too little explored territory for PNMOL thus far. Let's only initialise y, \dot y, and \ddot y (while conditioning y on the boundary conditions?) and hope for the best.

Currently, RK init is implemented in a way that exploits the Kronecker structure in the system matrices a little too much -- it assumes cov = cov \otimes I_d, which in the PNMOL case seems wrong to me.

While it is not clear what the perfect input scale should be (see the figure in the paper), I think it is a decent estimate to use the maximum likelihood estimate based on the initial condition? This is easy to implement (via a grid search, for example; this is used to generate figure 2)