STTLA PR #414

andyElking · 2024-05-08T21:21:56Z

This PR adds space-time-time Levy area to VirtualBrownianTree and UnsafeBrownianPath. The changes to the overall logic are minor, I mostly just added the math for the space-time-time case.

The additions to test_brownian.py are more extensive:

In test_conditional_statistics we test the pvals of the marginals for W, H, and K separately. This confirms that they are all Gaussian, so their joint distribution is exactly determined by their means and covariances, which we also test for. All of these are conditioned on the sample just before and the one just after the current sample (as it was before). I moved the code for computing the statistics into the conditional_statistics function.
Since the code of test_conditional_statistics itself is based on thm 3.7 from the Single-seed paper, I chose to test its correctness from a different angle in test_whk_interpolation. This test tries three things: a) just final interpolation with no halving (brownian bridge) steps, b) only one halving step followed by a final interpolation, c) 15 halving steps (i.e. tol = 2**-15) and no final interpolation (aka. "zero" spline). Since they all have the same distribution as predicted by thm 3.7, I conclude that indeed everything is as it should be. This still relies on the assumption that the halving steps are correct (i.e. thm 3.5), but that was computed by James in his thesis, so I suppose that's a fair assumption :).

In the third commit I added shape annotations to arrays in AbstractBrownianIncrement and its subclasses. This is optional, so if you prefer we can remove this commit.

In the fourth commit I just removed two unnecessary # pyright: ignore comments, so pyright doesn't complain. These have nothing to do with the rest of this PR.

I hope this PR is less challenging to review. No hurry though :)

patrick-kidger · 2024-05-25T21:00:19Z

Okay, happy to start thinking about reviewing this PR now! If you can rebase on top of main then I'll take a look :)

andyElking · 2024-05-26T20:33:07Z

That's great news! I rebased it on top of main and changed the base branch :)

patrick-kidger

Awesome stuff! I've just quickly looked over this. I'll let you address my first round of comments, but my initial impression is that this is a fairly small change, so it should be pretty easy to land.

patrick-kidger · 2024-05-29T20:02:10Z

diffrax/_autocitation.py

+        return r"""
+% You are simulating Brownian motion using Levy area, the formulae for which
+% are due to:
+@misc{jelinčič2024singleseed,
+  title={Single-seed generation of Brownian paths and integrals
+  for adaptive and high order SDE solvers},
+  author={Andraž Jelinčič and James Foster and Patrick Kidger},
+  year={2024},
+  eprint={2405.06464},
+  archivePrefix={arXiv},
+  primaryClass={math.NA}
+}
+
+% and Theorem 6.1.6 of
+@phdthesis{foster2020a,
+  publisher = {University of Oxford},
+  school = {University of Oxford},
+  title = {Numerical approximations for stochastic differential equations},
+  author = {Foster, James M.},
+  year = {2020}
+}


I think these citations deserve to also go VirtualBrownianTree.__doc__! :) You did a lot of work making this happen!

Do you mean in the autocitation for the VBT? It is already in the doc of VBT, but I did add it to autocitation as well.

What I mean is that I think these citations should appear in the docstring, and be parsed out with _parse_reference_multi.

patrick-kidger · 2024-05-29T20:03:22Z

test/test_brownian.py

@@ -13,6 +13,11 @@
 import scipy.stats as stats


+levy_areas = (


Nit: should start with an underscore, it's private to this file.

patrick-kidger · 2024-05-29T20:05:06Z

test/test_brownian.py

+    ):
+        # VBT with STTLA does not support float16 or complex dtypes
+        # because it uses jax.random.multivariate_normal
+        shapes_dtypes = shapes_dtypes[:6]


This seems a little error prone. Can this be switched into two groups shapes_dtypes1 and shapes_dtypes2, which are then combined or not appropriately?

Yes, good point.

patrick-kidger · 2024-05-29T20:08:14Z

diffrax/_brownian/tree.py

+# in addition to the non-rescaled ones, for the purposes of
+# taking the difference between two Levy areas.
+class _AbstractLevyVal(eqx.Module):
+    dt: eqx.AbstractVar[Inexact[Array, ""]]
+    W: eqx.AbstractVar[Array]
+
+
+class _BMLevyVal(_AbstractLevyVal):
+    dt: Inexact[Array, ""]
+    W: Array
+
+
+class _AbstractSpaceTimeLevyVal(_AbstractLevyVal):
+    H: eqx.AbstractVar[Array]
+    bar_H: eqx.AbstractVar[Array]
+
+
+class _SpaceTimeLevyVal(_AbstractSpaceTimeLevyVal):
+    dt: Inexact[Array, ""]
+    W: Array
+    H: Array
+    bar_H: Array
+
+
+class _SpaceTimeTimeLevyVal(_AbstractSpaceTimeLevyVal):
+    dt: Inexact[Array, ""]
+    W: Array
+    H: Array
+    bar_H: Array
+    K: Array
+    bar_K: Array


Possibly slightly simpler might be something like:

class _LevyVal(eqx.Module): brownian_increment: AbstractBrownianIncrement bar_H: Optional[Array] bar_K: Optional[Array]

?
Possibly adding some Generic[...] typevars for each of the three arguments if required.

Fair enough, I suppose I was trying to do it very properly, but since it's never exposed to the user, you're right, I should go for the simpler option.

patrick-kidger · 2024-05-29T20:08:29Z

diffrax/_brownian/tree.py

 class _State(eqx.Module):
    level: IntScalarLike  # level of the tree
-    s: RealScalarLike  # starting time of the interval
+    s: Inexact[Array, ""]  # starting time of the interval


Times should always be real, not complex.

I was trying to avoid a myriad of with jax.numpy_dtype_promotion("standard"):, which spaghettify the code a bit, and I thought the times which are only used internally could safely be complex. But I guess it could be a source of bugs, so fair enough.

andyElking · 2024-05-30T12:10:57Z

Thanks for the quick review! I made the changes you suggested.

I also added a little subsection on Levy areas in docs/api/brownian/md. Mainly to clear up how solver.minimal_levy_area interacts with VBT, etc.

patrick-kidger · 2024-06-04T05:11:51Z

diffrax/_autocitation.py

+        return r"""
+% You are simulating Brownian motion using Levy area, the formulae for which
+% are due to:
+@misc{jelinčič2024singleseed,
+  title={Single-seed generation of Brownian paths and integrals
+  for adaptive and high order SDE solvers},
+  author={Andraž Jelinčič and James Foster and Patrick Kidger},
+  year={2024},
+  eprint={2405.06464},
+  archivePrefix={arXiv},
+  primaryClass={math.NA}
+}
+
+% and Theorem 6.1.6 of
+@phdthesis{foster2020a,
+  publisher = {University of Oxford},
+  school = {University of Oxford},
+  title = {Numerical approximations for stochastic differential equations},
+  author = {Foster, James M.},
+  year = {2020}
+}


What I mean is that I think these citations should appear in the docstring, and be parsed out with _parse_reference_multi.

patrick-kidger · 2024-06-04T05:12:42Z

diffrax/_brownian/path.py

        use_levy: bool,
    ):
        w_std = jnp.sqrt(t1 - t0).astype(shape.dtype)
+        key_w, key_hh, key_kk = jr.split(key, 3)


Can we do the split inside the if blocks, to avoid the overhead in the common case of no Levy area?

I wanted to keep the generated path the same regardless of the Levy area setting, but fair, I guess this isn't really relevant for UBP.

And I think I was using the wrong key to generate the bm anyway, oops.

patrick-kidger · 2024-06-04T05:13:58Z

diffrax/_brownian/tree.py

+    K: Optional[Array]
+    bar_K: Optional[Array]
+
+    def __post_init__(self):


Nit: prefer __check_init__ where possible. This is an Equinox-specific extension that (a) runs even when you inherit, and (b) isn't ignored if you define a custom __init__ method, and (c) doesn't allow you to still mutate the class whilst it is running.

Oh, I understand, will do.

patrick-kidger · 2024-06-04T05:14:24Z

diffrax/_brownian/tree.py

-    - `x0`: `LevyVal` at time `s`.
-    - `x1`: `LevyVal` at time `u`.
+    - `x0`: `_AbstractLevyVal` at time `s`.
+    - `x1`: `_AbstractLevyVal` at time `u`.


Remove the Abstract.

Oops, will do.

patrick-kidger · 2024-06-04T05:16:52Z

diffrax/_brownian/tree.py

    """
+    dtype = jnp.dtype(x0.W)


I'd prefer using jnp.result_type as the argument is an array. I forget exactly what goes wrong but I recall getting bit by jnp.dtype at some point before.

patrick-kidger · 2024-06-04T05:31:56Z

diffrax/_brownian/tree.py

+            )
+
+            if self._spline == "sqrt":
+                # NOTE: not compatible with jnp.float16


I think we should still aim to handle this. Probably this can be done by casting the dtype?

dtype_atleast32 = jnp.result_type(dtype, jnp.float32) hat_y = jr.multivariate_normal(..., dtype=dtype_atleast32) hat_y = hat_y.astype(dtype)

FWIW I think generating a multivariate normal shouldn't be too tricky if we wanted to just do it manually.

I also don't love that this is apparently calling out to SVD. Admittedly, though, I have just done a quick google for an explicit form for the square root of a 3x3 posdef matrix and it looks like there isn't one...

I did the Cholesky decomp manually (with the help of some obscure symbolic algebra package), but it's very complicated and in my very basic experiments it took slightly longer to compute it with that formula (so computing the decomposition directly from s, r, u), as opposed to first computing the cov matrix, and then doing SVD on that. The issue with Cholesky is that it is very imprecise when the matrix is close to singular (which happens when r is close to s or to u). SVD is the only good option for near-singular matrices. Alternatively we'd have to do a separate case for when r-s or u-r is small, which is probably not ideal.

Yes, the type-casting idea is very good.

patrick-kidger · 2024-06-04T05:36:46Z

diffrax/_brownian/tree.py

+            hat_w_sr, hat_hh_sr, hat_kk_sr = [
+                x.squeeze(axis=-1) for x in jnp.split(hat_y, 3, axis=-1)
+            ]


I thnk just hat_w_sr, hat_hh_sr, hat_kk_sr = hat_y will work?

No, I don't think it will, because we're splitting by the last dimension, whereas tuple unpacking splits along the first.

patrick-kidger · 2024-06-04T05:39:00Z

diffrax/_solver/dopri5.py

-    ] = _Dopri5Interpolation
+    interpolation_cls: ClassVar[Callable[..., _Dopri5Interpolation]] = (
+        _Dopri5Interpolation
+    )


Can we skip the spurious reformatting?

(Let's put that in a separate PR if you like? Perhaps this is coming from bumping ruff.)

Yes, I'll drop it, sorry.

patrick-kidger · 2024-06-04T05:40:09Z

docs/api/brownian.md

+
+## Levy areas
+
+Brownian controls can return certain types of Levy areas. These are iterated integrals


Levy should have a diacritic: Lévy.

I don't always include this when writing informally (including in your name, I know!) but I try to get it right in persisent documentation!

Hmm yes, I think there's many places where I should fix this.

patrick-kidger · 2024-06-04T05:41:03Z

docs/api/brownian.md

+For example if `solver.minimal_levy_area` returns an `AbstractSpaceTimeLevyArea`, then
+the Brownian motion (which is either an `UnsafeBrownianPath` or 
+a `VirtualBrownianTree`) should be initialized with `levy_area=SpaceTimeLevyArea` or 
+`levy_area=SpaceTimeTimeLevyArea`. Note that for the BM, a concrete class must be


Nit: I try to avoid the "BM" abbreviation in documentation.

andyElking · 2024-06-12T18:35:50Z

Hi Patrick, sorry this round of corrections took so long. I made the fixes you mentioned, including adding diacritics on Lévy, which changed some files outside this PR, but we're still only on 12 files changed, so I hope it's fine 😅. I also rebased on top of the current main.

patrick-kidger · 2024-06-15T11:32:20Z

No worries -- and the great news is that I think this is ready to be merged :D
Can you rebase on top of the latest main, and we can get this in!

andyElking · 2024-06-18T09:40:41Z

That's very good to hear, I rebased it and squashed some of the commits. I can squash all of them if needed. Hopefully the tests pass.

…itic on Levy

andyElking · 2024-06-23T19:06:27Z

Hi Patrick,
Sorry for yet another delay. Seems like the tests passed, so I hope this means this PR is ready for a merge. Let me know if there's any more fixes I should make.

patrick-kidger · 2024-06-24T21:21:19Z

Aaaaaand merged! :D
This is really great stuff, it's awesome to see how much you're adding to Diffrax.

I've just merged this into a new dev branch (currently equal to main + your PR) in preparation for the next release! I'm aiming to merge in #387, and possibly #436, and then do a new release. What do your own future plans look like here?

andyElking · 2024-06-25T14:19:58Z

That's great news, thanks!

I will probably open a new PR with Langevin solvers sometime over the next few days. That one is a bit larger than this one, but in my opinion still more lightweight than the SRK PR. But feel free to finish the next release first and turn to Langevin PR afterwards if you prefer.

And after that I'll give you some peace as far as Diffrax is concerned. Upcoming projects are then Langevin MCMC, some signature stuff and finally diffusion models.

andyElking force-pushed the sttla_pr branch 2 times, most recently from 7b9b1e7 to a1f2e2a Compare May 9, 2024 14:07

andyElking force-pushed the sttla_pr branch from 2e28bea to 6368de8 Compare May 22, 2024 12:15

andyElking changed the base branch from dev to main May 26, 2024 20:32

patrick-kidger reviewed May 29, 2024

View reviewed changes

patrick-kidger reviewed Jun 4, 2024

View reviewed changes

andyElking force-pushed the sttla_pr branch 2 times, most recently from 9f33807 to fd946c5 Compare June 12, 2024 18:33

andyElking added 3 commits June 18, 2024 10:24

Added space-time-time Levy area to the VirtualBrownianTree

8418610

added SpaceTimeTimeLevyArea to UnsafeBrownianPath and added tests for it

1edbad5

added shape checking to arrays in Brownian return types

a84cbfa

andyElking force-pushed the sttla_pr branch 2 times, most recently from d4a7e4c to fa83ab4 Compare June 18, 2024 09:39

Added single-seed paper to docs and autocite, some minor fixes, diacr…

4ad40ca

…itic on Levy

andyElking force-pushed the sttla_pr branch from fa83ab4 to 4ad40ca Compare June 23, 2024 15:07

patrick-kidger changed the base branch from main to dev June 24, 2024 21:18

patrick-kidger merged commit da5031d into patrick-kidger:dev Jun 24, 2024
2 checks passed

andyElking deleted the sttla_pr branch June 26, 2024 09:09


		## Levy areas

		Brownian controls can return certain types of Levy areas. These are iterated integrals

STTLA PR #414

STTLA PR #414

Conversation

andyElking commented May 8, 2024 • edited Loading

patrick-kidger commented May 25, 2024

andyElking commented May 26, 2024

patrick-kidger left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andyElking May 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andyElking commented May 30, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andyElking Jun 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andyElking commented Jun 12, 2024

patrick-kidger commented Jun 15, 2024

andyElking commented Jun 18, 2024 • edited Loading

andyElking commented Jun 23, 2024

patrick-kidger commented Jun 24, 2024 • edited Loading

andyElking commented Jun 25, 2024

andyElking commented May 8, 2024 •

edited

Loading

andyElking May 30, 2024 •

edited

Loading

andyElking Jun 12, 2024 •

edited

Loading

andyElking commented Jun 18, 2024 •

edited

Loading

patrick-kidger commented Jun 24, 2024 •

edited

Loading