0.7.1 release

.github/workflows/release.yml

@@ -10,14 +10,14 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Release
-        uses: patrick-kidger/action_update_python_project@v6
+        uses: patrick-kidger/action_update_python_project@v8
         with:
             python-version: "3.11"
             test-script: |
                 cp -r ${{ github.workspace }}/test ./test
                 cp ${{ github.workspace }}/pyproject.toml ./pyproject.toml
-                python -m pip install -r ./test/requirements.txt
-                python -m test
+                uv sync --extra tests --no-install-project --inexact
+                uv run --no-sync pytest
             pypi-token: ${{ secrets.pypi_token }}
             github-user: patrick-kidger
             github-token: ${{ github.token }}

.github/workflows/run_tests.yml

@@ -7,7 +7,7 @@ jobs:
   run-tests:
     strategy:
       matrix:
-        python-version: [ "3.10", "3.12" ]
+        python-version: [ "3.11", "3.13" ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}
@@ -23,13 +23,16 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -r ./test/requirements.txt
-
+          python -m pip install '.[dev,docs,tests]'
 
       - name: Checks with pre-commit
-        uses: pre-commit/action@v3.0.1
+        run: |
+          pre-commit run --all-files
 
       - name: Test with pytest
         run: |
-          python -m pip install .
           python -m test
+
+      - name: Check that documentation can be built.
+        run: |
+          mkdocs build

.gitignore

@@ -8,3 +8,4 @@ site/
 .pymon
 .idea/
 .venv/
+uv.lock

.pre-commit-config.yaml

@@ -8,15 +8,15 @@ repos:
         files: ^pyproject\.toml$
         additional_dependencies: ["toml-sort==0.23.1"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.2.2
+    rev: v0.13.0
     hooks:
       - id: ruff-format  # formatter
-        types_or: [ python, pyi, jupyter ]
+        types_or: [ python, pyi, jupyter, toml ]
       - id: ruff  # linter
-        types_or: [ python, pyi, jupyter ]
+        types_or: [ python, pyi, jupyter, toml ]
         args: [ --fix ]
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.350
+    rev: v1.1.405
     hooks:
     - id: pyright
       additional_dependencies: [equinox, jax, jaxtyping, optax, optimistix, lineax, pytest, typeguard==2.13.3, typing_extensions, wadler_lindig]

CONTRIBUTING.md

@@ -8,23 +8,15 @@ Contributions (pull requests) are very welcome! Here's how to get started.
 
 First fork the library on GitHub.
 
-Then clone and install the library in development mode:
+Then clone and install the library:
 
 ```bash
 git clone https://github.com/your-username-here/diffrax.git
 cd diffrax
-pip install -e .
+pip install -e '.[dev]'
+pre-commit install  # `pre-commit` is installed by `pip` on the previous line
 ```
 
-Then install the pre-commit hook:
-
-```bash
-pip install pre-commit
-pre-commit install
-```
-
-These hooks use ruff to lint and format the code, and pyright to type-check it.
-
 ---
 
 **If you're making changes to the code:**
@@ -34,8 +26,8 @@ Now make your changes. Make sure to include additional tests if necessary.
 Next verify the tests all pass:
 
 ```bash
-pip install -r test/requirements.txt
-pytest
+pip install -e '.[tests]'
+pytest  # `pytest` is installed by `pip` on the previous line.
 ```
 
 Then push your changes back to your fork of the repository:

benchmarks/against_scan.py

@@ -36,7 +36,7 @@ def speedtest(fn, name):
 # INTEGRATE WITH scan
 
 
-@jax.checkpoint  # pyright: ignore
+@jax.checkpoint
 def body(carry, t):
     u, v, dt = carry
     u = u + du(t, v, None) * dt

diffrax/_adjoint.py

@@ -362,8 +362,7 @@ def loop(
         if is_unsafe_sde(terms):
             kind = "lax"
             msg = (
-                "Cannot reverse-mode autodifferentiate when using "
-                "`UnsafeBrownianPath`."
+                "Cannot reverse-mode autodifferentiate when using `UnsafeBrownianPath`."
             )
         elif max_steps is None:
             kind = "lax"

diffrax/_brownian/path.py

@@ -62,9 +62,9 @@ class UnsafeBrownianPath(AbstractBrownianPath):
     """
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
-    levy_area: type[
-        BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
-    ] = eqx.field(static=True)
+    levy_area: type[BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea] = (
+        eqx.field(static=True)
+    )
     key: PRNGKeyArray
 
     def __init__(

diffrax/_brownian/tree.py

@@ -235,9 +235,9 @@ class VirtualBrownianTree(AbstractBrownianPath):
     t1: RealScalarLike
     tol: RealScalarLike
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
-    levy_area: type[
-        BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
-    ] = eqx.field(static=True)
+    levy_area: type[BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea] = (
+        eqx.field(static=True)
+    )
     key: PyTree[PRNGKeyArray]
     _spline: _Spline = eqx.field(static=True)
 

diffrax/_custom_types.py

@@ -21,7 +21,7 @@
     BoolScalarLike = bool | Array | np.ndarray
     FloatScalarLike = float | Array | np.ndarray
     IntScalarLike = int | Array | np.ndarray
-    RealScalarLike = bool | int | float | Array | np.ndarray
+    RealScalarLike = int | float | Array | np.ndarray
 else:
     BoolScalarLike = Bool[ArrayLike, ""]
     FloatScalarLike = Float[ArrayLike, ""]

diffrax/_event.py

@@ -2,6 +2,7 @@
 from collections.abc import Callable
 
 import equinox as eqx
+import jax.tree_util as jtu
 import optimistix as optx
 from jaxtyping import Array, PyTree
 
@@ -20,7 +21,33 @@ class Event(eqx.Module):
     """
 
     cond_fn: PyTree[Callable[..., BoolScalarLike | RealScalarLike]]
-    root_finder: optx.AbstractRootFinder | None = None
+    direction: PyTree[None | bool]
+    root_finder: optx.AbstractRootFinder | None
+
+    def __init__(
+        self,
+        cond_fn,
+        root_finder: optx.AbstractRootFinder | None = None,
+        direction: None | bool | PyTree[None | bool] = None,
+    ):
+        if direction in (None, False, True):
+            direction = jtu.tree_map(lambda _: direction, cond_fn, is_leaf=callable)
+
+        direction_leaves, direction_structure = jtu.tree_flatten(
+            direction, is_leaf=lambda x: x is None
+        )
+        if direction_structure != jtu.tree_structure(cond_fn, is_leaf=callable):
+            raise ValueError("Missmatch in the structure of `cond_fn` and `direction`.")
+
+        if any(x not in (None, False, True) for x in direction_leaves):
+            raise ValueError(
+                "`trig_dir` must be a `None`, `bool`, or a PyTree of `None | bool`s "
+                "with the same structure as `cond_fn`."
+            )
+
+        self.cond_fn = cond_fn
+        self.root_finder = root_finder
+        self.direction = direction
 
 
 Event.__init__.__doc__ = """**Arguments:**
@@ -39,6 +66,12 @@ class Event(eqx.Module):
     [`optimistix.Newton`](https://docs.kidger.site/optimistix/api/root_find/#optimistix.Newton)
     would be a typical choice here.
 
+- `direction`: `None` or `bool` or PyTree of `None | bool` of the same shape as
+    `cond_fn`, that decides for each `cond_fn` if it triggers an event from a
+    zero-cossing in both directions (`None`), from an upcrossing (`True`) or from a
+    downcrossing (`False`). Only needed for those `cond_fn` which return floating point
+    numbers; ignored for those `cond_fn` which return booleans.
+
 !!! Example
 
     Consider a bouncing ball dropped from some intial height $x_0$. We can model