Improved Levy area documentation

patrick-kidger · May 19, 2024 · 80d1f34 · 80d1f34
1 parent 852dbed
commit 80d1f34
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 5 deletions.
diff --git a/diffrax/_brownian/path.py b/diffrax/_brownian/path.py
@@ -45,7 +45,19 @@ class UnsafeBrownianPath(AbstractBrownianPath):
     motion. Hence the restrictions above. (They describe the general case for which the
     correlation structure isn't needed.)
 
-    Depending on the `levy_area` argument, this can also be used to generate Levy area.
+    !!! info "Levy Area"
+
+        Can be initialised with `levy_area` set to `diffrax.BrownianIncrement`, or
+        `diffrax.SpaceTimeLevyArea`. If `levy_area=diffrax.SpaceTimeLevyArea`, then it
+        also computes space-time Lévy area `H`. This is an additional source of
+        randomness required for certain stochastic Runge--Kutta solvers; see
+        [`diffrax.AbstractSRK`][] for more information.
+
+        An error will be thrown during tracing if Lévy area is required but is not
+        available.
+
+        The choice here will impact the Brownian path, so even with the same key, the
+        trajectory will be different depending on the value of `levy_area`.
     """
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)

diff --git a/diffrax/_brownian/tree.py b/diffrax/_brownian/tree.py
@@ -137,10 +137,19 @@ def _split_interval(
 class VirtualBrownianTree(AbstractBrownianPath):
     """Brownian simulation that discretises the interval `[t0, t1]` to tolerance `tol`.
 
-    Can be initialised with `levy_area` set to `""`, or `"space-time"`.
-    If `levy_area="space_time"`, then it also computes space-time Lévy area `H`.
-    This will impact the Brownian path, so even with the same key, the trajectory will
-    be different depending on the value of `levy_area`.
+    !!! info "Levy Area"
+
+        Can be initialised with `levy_area` set to `diffrax.BrownianIncrement`, or
+        `diffrax.SpaceTimeLevyArea`. If `levy_area=diffrax.SpaceTimeLevyArea`, then it
+        also computes space-time Lévy area `H`. This is an additional source of
+        randomness required for certain stochastic Runge--Kutta solvers; see
+        [`diffrax.AbstractSRK`][] for more information.
+
+        An error will be thrown during tracing if Lévy area is required but is not
+        available.
+
+        The choice here will impact the Brownian path, so even with the same key, the
+        trajectory will be different depending on the value of `levy_area`.
 
     ??? cite "Reference"