Add Better Discretizers

[DanWaxman]

Supersedes #226. Includes better discretization schemes to be used, and a notebook showing a simple comparison.

I've gone through & iterated on the resulting code structure & included methods quite closely, which looks good to me. What remains is a closer look at the mathematical correctness + some interrogation of the config files, which I think got a bit bloated with overly technical details.

[mattlevine22]

should be "...250 evenly spaced...". I caught that w mine own eyes too!

[mattlevine22]

and on interval [0, 10]

[DanWaxman]

fixed (though now 100 evenly spaced on [0, 4]; finer discretization took a long time otherwise).

[mattlevine22]

Overall I think this is great! Can you please (a) update the notebook a bit per below and (b) take a look at codex's complaint below and decide if it is worth chasing down?

Some simple suggestions for improving the notebook:

1. make the notation for all the different discretization the same and canonical:

p(x_k+1 | x_k, u_k, t_k, t_k+1) \approx ...

2. Run diffrax-Euler with dt=1e-3 and show that it does better

3. Give some recommendations at the end:

• for high-dimensional systems, ...
• for slow/fast sample-rate systems, ...
• for nearly-linear systems ...
• for stiff systems, ...

Also flagging a Codex complaint:

[P1] The Van Loan covariance calculation is unstable in default float32. In exact_affine.py (line 71), the block exponential contains -F.T. For a stable scalar system F=-100, L=1, and h=1, the exact covariance is approximately 0.005, but the implementation overflows and raises “non-finite transition covariance.” This also affects local linearization and is especially problematic because automatic routing now sends every eligible affine model through this path. It needs a numerically stable covariance algorithm and a stiff-system regression test.

[mattlevine22]

@DanWaxman sorry to add 1 more thing, but can we add to the existing SDE L63 tutorial that uses a new Discretizer w Cuthbert EnKF and compare the posterior and the compute time to cd-dynamax EnKF?

[DanWaxman]

Thanks! Made the updates! Also introduced a better scaled version of the van Loan solver. Solves over a smaller interval if the matrix norm is large; e.g., solve over h/2 instead of h, then apply A_{h/2} A_{h/2} = A_h.

[DanWaxman]

@DanWaxman sorry to add 1 more thing, but can we add to the existing SDE L63 tutorial that uses a new Discretizer w Cuthbert EnKF and compare the posterior and the compute time to cd-dynamax EnKF?

Sure, done!