Tweaks to #337

[patrick-kidger]

This PR has two commits:

1. The first squashes together all commits from added space-time levy area (for PR) #337. This represents all the hard work by @andyElking!
2. The second are various tidy-ups and tweaks on top -- mostly me addressing various stylistic points.

My intention is that this supersede #337, as it has a tidier commit history.

@andyElking: I'm now mostly happy with this PR. However, on the mathematics: I noticed that if I set z=0 on the levy_area=="" branch at the end of VirtualBrownianTree._evaluate_leaf, then the conditional statistics tests still seem to pass. Likewise if I chance jnp.sqrt(jnp.abs(sr * ru / su)) to just jnp.abs(sr * ru / su), then these tests again still seem to pass. Do you know what is going on here?

[andyElking]

My intention is that this supersede #337, as it has a tidier commit history.

Thanks, that's perfect!

@andyElking: I'm now mostly happy with this PR. However, on the mathematics: I noticed that if I set z=0 on the levy_area=="" branch at the end of VirtualBrownianTree._evaluate_leaf, then the conditional statistics tests still seem to pass. Likewise if I chance jnp.sqrt(jnp.abs(sr * ru / su)) to just jnp.abs(sr * ru / su), then these tests again still seem to pass. Do you know what is going on here?

If the tolerance is set to small enough, then the final interpolation should be irrelevant and can safely be skipped. I think what you're seeing is indeed that the contribution of the final interpolation is negligible. The z=0 case corresponds to having the BM piecewise constant (step function) and the jnp.abs(sr * ru / su) version makes it piecewise quadratic (as you had it initially). Maybe try all of these options with a very large tolerance and see if what you said is still the case? If you want, I can perform this experiment later when I get home.

In addition I decided to get rid of jnp.sqrt(jnp.abs, and rather use sr = jax.nn.relu(r - s) and ru = jax.nn.relu(su - sr) already when declaring them.

I also got rid of the two mentions of the "piecewise quadratic". See the comments I made in the code.
If you prefer, I can also make these edits myself.

[patrick-kidger]

On the code comments: these all look good. I'll let you make those tweaks somewhere and then merge them into this branch.

On the business of the spline: right, but I distinctly recall it being a quadratic spline being very important to get correct results before! Adding in square roots definitely gave the wrong thing. Thank you -- if you could run those experiments and figure out what's going on here then that would be great. Maybe with a test flag inside the VBT code that when enabled does the "wrong" thing, and have a test that asserts that under those circumstances, the KS test fails. (I recall getting failing KS tests with something like p<1e-50 when I had the wrong thing in my original implementation, so I don't think there should be too much tolerance-fiddling necessary to get things working reliably.)

[andyElking]

On the code comments: these all look good. I'll let you make those tweaks somewhere and then merge them into this branch.

Sounds good, I'll make the edits shortly.

On the business of the spline: right, but I distinctly recall it being a quadratic spline being very important to get correct results before! Adding in square roots definitely gave the wrong thing. Thank you -- if you could run those experiments and figure out what's going on here then that would be great. Maybe with a test flag inside the VBT code that when enabled does the "wrong" thing, and have a test that asserts that under those circumstances, the KS test fails. (I recall getting failing KS tests with something like p<1e-50 when I had the wrong thing in my original implementation, so I don't think there should be too much tolerance-fiddling necessary to get things working reliably.)

As I said, the difference between the quadratic spline and the sqrt-quadratic spline only comes out when tolerance is large. You can see a demonstration here: https://github.com/andyElking/diffrax_STLA/blob/devel/VBT_spline_test.ipynb
Maybe in your experiments you were comparing it to linear interpolation, which does give the wrong results more obviously. What I did here, however is exactly the right way mathematically speaking. If you don't trust me, ask James.

[andyElking]

I made the necessary edits and made a PR to pull those into diffrax/new_pr_branch

[patrick-kidger]

Closing as I've pushed all the updates here back to #337.