Added spline test for VBT, fixed a bug in levy_tree_transpose, and added references to VBT #342

andyElking · 2023-12-20T15:48:53Z

@patrick-kidger This includes the updates you suggested. For the timeing tests, however, I made a branch identical to this one, but with an additional Jupyter notebook:
Timing tests

I'm not really sure where you wanted me to put the timeing code, so I put it in a separate branch, but I can also add it here.

…ded references in VBT

patrick-kidger · 2023-12-21T22:17:50Z

diffrax/_brownian/tree.py

+         ??? cite "Reference"
+        Based on section 6.1 of
+        ```bibtex
+        @phdthesis{foster2020a,
+          publisher = {University of Oxford},
+          school = {University of Oxford},
+          title = {Numerical approximations for stochastic differential equations},
+          author = {Foster, James M.},
+          year = {2020}
+        }
+
+        In particular see Theorem 6.1.6.
+        ```
+
+


This should be formatted as:

??? cite "References" ```bibtex ... ```

i.e. with a new line, with an indent, and with only a single blank line between the reference section and the arguments section.

Oops, not sure what happened with the indent there. Will fix

patrick-kidger · 2023-12-21T22:19:08Z

diffrax/_custom_types.py

@@ -1,9 +1,9 @@
 import typing
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, Literal, Optional, TYPE_CHECKING, TypeAlias, Union


TypeAlias should be imported from typing_extensions for now, as we don't yet require Python 3.10.

Aha, I see. Didn't know it was new in 3.10.

patrick-kidger · 2023-12-21T22:19:19Z

diffrax/_custom_types.py

@@ -52,6 +52,8 @@ class Real(AbstractDtype):
 BufferDenseInfos = dict[str, PyTree[eqxi.MaybeBuffer[Shaped[Array, "times ..."]]]]
 sentinel: Any = eqxi.doc_repr(object(), "sentinel")

+_LA: TypeAlias = Literal["", "space-time"]


This shouldn't start with an underscore if it's imported into another module. I'd suggest calling it just LevyArea.

Good point, yes

patrick-kidger · 2023-12-21T22:21:31Z

test/test_brownian.py

-@pytest.mark.parametrize("use_levy", (False, True))
-def test_conditional_statistics(levy_area, use_levy):
+def conditional_statistics(
+    levy_area: _LA, use_levy: bool, tol=2**-6, spacing=2**-6, spline: _Spline = "sqrt"


Let's remove the default arguments here as I don't think they're ever used.

This is a general principle -- prefer not to use default arguments where possible, in particular in internal APIs, as they're a common source of surprising behaviour.

I know, thanks for the reminder. I guess old habits die hard.

patrick-kidger · 2023-12-21T22:21:39Z

test/test_brownian.py

    # Get >80 randomly selected points; not too close to avoid discretisation error.
    t0 = 0.3
    t1 = 8.7
    ts = jr.uniform(sample_key, shape=(100,), minval=t0, maxval=t1)
+    # ts = jnp.array([1.0, 3.0, 6.0, 7.0])


Oops, I tried to catch all the random debugging comments, seems this one escaped.

patrick-kidger · 2023-12-21T22:23:09Z

test/test_brownian.py

+    assert jnp.all(pvals_w1 > 0.1 / pvals_w1.shape[0])
+    if levy_area == "space-time" and use_levy:
+        assert jnp.all(pvals_w2 > 0.1 / pvals_w2.shape[0])
+        assert jnp.all(pvals_h > 0.1 / pvals_h.shape[0])


else: assert len(pvals_w2) == 0 assert len(pvals_h) == 0

?

Yeah, might as well.

patrick-kidger · 2023-12-21T22:23:47Z

test/test_brownian.py

            continue
        prev_ti = ti
        ts.append(ti)
    ts = jnp.stack(ts)
-    assert len(ts) > 80
+    assert len(ts) > min(0.2 * (8.0 / spacing), 75)  # for spacing = 2**-5, this is 51


maybe just provide a lower bound as an explicit argument to this function, rather than using a heuristic like this?

patrick-kidger · 2023-12-21T22:24:45Z

test/test_brownian.py

+
+    for spline in splines:
+        pvals_w1, pvals_w2, pvals_h = conditional_statistics(
+            levy_area, use_levy=True, tol=2**-5, spacing=2**-6, spline=spline


Can we parameterise this test by use_levy=False/True?

I think it might be a bit pointless, but sure, it's not that expensive.

patrick-kidger · 2023-12-21T22:25:14Z

test/test_brownian.py

+        else:
+            # make sure that for incorrect splines at least one p-value is
+            # below 0.01 (subject to multiple-testing correction) and the
+            # average p-value is below 0.2.


I think the average p-value should probably be a lot smaller than this?

This is a bit of a subtle issue. So normally yes, it is a lot smaller, exept in one case. When levy_area="spacetime" then pvals_w1 actually are quite good even with spline="zero", due to the randomness the w_r receives from h_su.

What I'm saying is, that the variance that pval_w1 sees is the variance of the Brownian parabola, but without accounting for the conditioning on H, i.e.
$$
\E[ ( \E[W_{s,r} | W_{s,u}, H_{s,u} ] )^2 - ( \E[W_{s,r} | W_{s,u} ] )^2 ]
$$
Which is actually a very good approximation of the actual variance. In pval_w2 the influence of H is accounted for in the mean, which is subtracted, and hence the variance can only come from x1 (which is zero when spline="zero"). I hope this makes sense.

So what I will do, is to just keep this more permissive upper bound to pvals_w1 when levy_area="spacetime", and add a stricter upper bound to all the other cases.

patrick-kidger · 2023-12-21T22:25:45Z

For the timing code, I'd suggest adding it as a standalone .py file inside benchmarks/.

andyElking · 2023-12-22T08:11:12Z

For the timing code, I'd suggest adding it as a standalone .py file inside benchmarks/.

Honestly, I wasn't even aware the benchmarks/ folder even existed until now. I added it, and these are the results (on my laptop):
New Shallow BM: 2.406
New Shallow STLA: 4.481
New Deep BM: 6.489
New Deep STLA: 12.068
Old Shallow: 3.123
Old Deep: 7.904

patrick-kidger

Minor comments (really teaching points) on the new benchmark/test, but otherwise LGTM. :)

patrick-kidger · 2023-12-23T18:54:21Z

benchmarks/brownian_tree_times.py

+import jax.random as jr
+import jax.tree_util as jtu
+from diffrax import VirtualBrownianTree
+from diffrax._brownian.base import AbstractBrownianPath


This is accessible as diffrax.AbstractBrownianpath.

I think that was done automatically by the IDE, and I haven't noticed.

patrick-kidger · 2023-12-23T18:56:26Z

benchmarks/brownian_tree_times.py

+from diffrax import VirtualBrownianTree
+from diffrax._brownian.base import AbstractBrownianPath
+from diffrax._custom_types import RealScalarLike
+from diffrax._misc import default_floating_dtype, is_tuple_of_ints, split_by_tree


I think we can skip the is_tuple_of_ints and split_by_tree imports given what we initialise OldVBT(..., shape=(100,)).

default_floating_dtype is now available from lineax.internal with the latest Lineax release.

(I'm just trying to minimise private imports here.)

Yes, that makes sense

patrick-kidger · 2023-12-23T19:00:05Z

benchmarks/brownian_tree_times.py

+    tree = tree_cls(t0=t0, t1=t1, tol=tol, shape=(100,), key=key, levy_area=levy_area)
+
+    def f():
+        return jax.block_until_ready(vec_eval(tree, ts))


I think there should be a jax.jit here as vec_eval involves a jax.vmap, which is a JAX operation happening outside of vmap.

Also, for timing benchmarks like this it is better to do many repeats and take a minimum. (Noise can only increase things from the best possible time, so you want a minimum, not a mean.)

Thus I think you want something like:

@jax.jit def run(ts): return jax.vmap(lambda _t: tree.evaluate(_t, use_levy=True))(_ts) return min(timeit.repeat(lambda: jax.block_until_ready(run(ts)), number=1, repeat=100))

Oh I see, so any JAX operation should be wrapped inside a jit?

Yup! At least, if you want them to fast.

So for typical neural network training, I often won't bother putting a JIT when __init__ialising a model (which only happens once), but will definitely put one around the entirety of the forward pass.

patrick-kidger · 2023-12-23T19:07:28Z

test/test_brownian.py

+@pytest.mark.parametrize("spline", ("quad", "sqrt", "zero"))
+def test_spline(levy_area: LevyArea, use_levy, spline):
+    if levy_area == "space-time" and spline == "quad":
+        pytest.skip("Quad spline is not implemented for space-time Levy area")


FYI, a skipped test is usually used to refer to something that is currently broken, but for which we just don't to worry about right now. In particular it still shows up in the output test log each time.
This means that skips should usually only be used in the short-term, not in the long-term.

Anyway, here you probably want this:

def _levy_area_spline(): for levy_area in ("", "space-time"): for spline in ("quad", "sqrt", "zero"): if levy_area == "space-time" and spline == "quad": continue yield levy_area, spline @pytest.mark.parametrize("levy_area,spline", _levy_area_spline())

which won't even generate the test to be skipped.

Thanks, that's good to know.

andyElking · 2023-12-23T19:44:10Z

Minor comments (really teaching points) on the new benchmark/test, but otherwise LGTM. :)

Thanks! These all make sense. I won't be able to make these edits today or tomorrow, however, so is it okay if you fix these in your new_pr_branch and merge if you're already at it? Otherwise I can do it on Monday 😊.

patrick-kidger · 2023-12-23T20:52:48Z

Haha, no worries! I'll let you make the changes after Christmas.

Merry Christmas!

andyElking · 2023-12-23T21:58:35Z

Merry Christmas to you too!

andyElking · 2023-12-24T11:17:29Z

I managed to do it today after all 😊.

Should I squash them all together, or do you intend to reorganise the commits yourself anyway?

patrick-kidger · 2023-12-24T18:43:52Z

Alright, LGTM! I've squashed these together and made a few tweaks, and you can see the result back in #337. Let's continue the discussion over there :)

Added spline test for VBT, fixed a bug in levy_tree_transpose, and ad…

35a4907

…ded references in VBT

patrick-kidger reviewed Dec 21, 2023

View reviewed changes

andyElking added 2 commits December 22, 2023 09:51

Nits

200d712

Nit

34334c6

patrick-kidger reviewed Dec 23, 2023

View reviewed changes

Nits

2dd6e74

patrick-kidger merged commit f84e731 into patrick-kidger:new_pr_branch Dec 24, 2023
0 of 4 checks passed

Added spline test for VBT, fixed a bug in levy_tree_transpose, and added references to VBT #342

Added spline test for VBT, fixed a bug in levy_tree_transpose, and added references to VBT #342

Conversation

andyElking commented Dec 20, 2023

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

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

patrick-kidger commented Dec 21, 2023

andyElking commented Dec 22, 2023 • edited Loading

patrick-kidger left a comment • 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

Choose a reason for hiding this comment

andyElking commented Dec 23, 2023

patrick-kidger commented Dec 23, 2023

andyElking commented Dec 23, 2023 via email • edited Loading

andyElking commented Dec 24, 2023 • edited Loading

patrick-kidger commented Dec 24, 2023

andyElking commented Dec 22, 2023 •

edited

Loading

patrick-kidger left a comment •

edited

Loading

andyElking commented Dec 23, 2023 via email •

edited

Loading

andyElking commented Dec 24, 2023 •

edited

Loading