From 80d1f34f7608629246600831f9397eb6fc5973a3 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sun, 19 May 2024 04:17:06 +0200
Subject: [PATCH] Improved Levy area documentation

---
 diffrax/_brownian/path.py | 14 +++++++++++++-
 diffrax/_brownian/tree.py | 17 +++++++++++++----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/diffrax/_brownian/path.py b/diffrax/_brownian/path.py
index 3fa481aa..1e18f0e5 100644
--- 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
index 3d40e186..6488242b 100644
--- 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"