Skip to content

Commit

Permalink
docs: updates to book and readme (#117)
Browse files Browse the repository at this point in the history
* docs: misc book fixes

* continue on failure in beta and nightly

* update README.md

* update readme

* add links to readme

* readme

* readme
  • Loading branch information
martinjrobins authored Dec 17, 2024
1 parent 091a8af commit 13c15a1
Show file tree
Hide file tree
Showing 19 changed files with 654 additions and 86 deletions.
5 changes: 5 additions & 0 deletions .github/workflows/rust.yml
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,11 @@ jobs:
unit-tests:
name: Tests - ${{ matrix.os }} - ${{ matrix.toolchain }}
runs-on: ${{ matrix.os }}
continue-on-error: ${{ matrix.experimental }}
strategy:
matrix:
experimental:
- false
toolchain:
- stable
os:
Expand All @@ -37,8 +40,10 @@ jobs:
include:
- toolchain: beta
os: ubuntu-latest
experimental: true
- toolchain: nightly
os: ubuntu-latest
experimental: true

steps:
- uses: actions/checkout@v4
Expand Down
45 changes: 19 additions & 26 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,25 @@

# DiffSol

Diffsol is a library for solving ordinary differential equations (ODEs) or
semi-explicit differential algebraic equations (DAEs) in Rust. You can use it
out-of-the-box with vectors and matrices from the
[nalgebra](https://nalgebra.org) or [faer](https://github.com/sarah-ek/faer-rs) crates, or you can implement your own types by
implementing the various vector and matrix traits in diffsol.
Diffsol is a library for solving ordinary differential equations (ODEs) or semi-explicit differential algebraic equations (DAEs) in Rust. It can solve equations in the following form:

## Features
```math
M \frac{dy}{dt} = f(t, y, p)
```

where $M$ is a (possibly singular and optional) mass matrix, $y$ is the state vector, $t$ is the time and $p$ is a vector of parameters.

The equations can be given by either rust closures or the [DiffSL](https://martinjrobins.github.io/diffsl/) Domain Specific Language (DSL). The DSL uses automatic differentiation using [Enzyme](https://enzyme.mit.edu/) to calculate the necessary jacobians, and JIT compilation (using either [LLVM](https://llvm.org/) or [Cranelift](https://cranelift.dev/)) to generate efficient native code at runtime. The DSL is ideal for using DiffSol from a higher-level language like Python or R while still maintaining similar performance to pure rust.

You can use DiffSol out-of-the-box with vectors, matrices and linear solvers from the [nalgebra](https://nalgebra.org) or [faer](https://github.com/sarah-ek/faer-rs) crates, or you can implement your own types or solvers by implementing the required traits.

## Installation and Usage

See installation instructions on the [crates.io page](https://crates.io/crates/diffsol).

The [DiffSol book](https://martinjrobins.github.io/diffsol/) describes how to use DiffSol using examples taken from several application areas (e.g. population dynamics, electrical circuits and pharmacological modelling), as well as more detailed information on the various APIs used to specify the ODE equations. For a more complete description of the API, please see the [docs.rs API documentation](https://docs.rs/diffsol).

## Solvers

DiffSol implements the following solvers:
- A variable order Backwards Difference Formulae (BDF) solver, suitable for stiff problems and singular mass matrices.
Expand All @@ -30,23 +42,4 @@ All solvers feature:
- forward sensitivity analysis,
- backwards or adjoint sensitivity analysis,

For comparison, the BDF solvers are similar to MATLAB's `ode15s` solver, the `bdf` solver in SciPy's `solve_ivp` function, or the BDF solver in SUNDIALS.
The ESDIRK solver using the provided `tr_bdf2` tableau is similar to MATLAB's `ode23t` solver.

Users can specify the equations to solve in the following ODE form, either using closures or the [DiffSL](https://martinjrobins.github.io/diffsl/) Domain Specific Language (DSL):

```math
M \frac{dy}{dt} = f(t, y, p)
```

where $M$ is a (possibly singular) mass matrix, $y$ is the state vector, $t$ is the time, $p$ is a
vector of parameters, and $f$ is the right-hand side function. The mass matrix
$M$ is optional (assumed to be the identity matrix if not provided).

## Installation

See instructions on the [crates.io page](https://crates.io/crates/diffsol).

## Usage

For more documentation and examples, see the [API documentation](https://docs.rs/diffsol/latest/diffsol/).
For comparison, the BDF solvers are similar to MATLAB's `ode15s` solver, the `bdf` solver in SciPy's `solve_ivp` function, or the BDF solver in SUNDIALS. The ESDIRK solver using the provided `tr_bdf2` tableau is similar to MATLAB's `ode23t` solver.
26 changes: 13 additions & 13 deletions book/src/benchmarks/python.md
Original file line number Diff line number Diff line change
@@ -1,38 +1,38 @@
# Python (Diffrax & Casadi)

[Diffrax](https://docs.kidger.site/diffrax/) is a Python library for solving ODEs and SDEs implemented using JAX. [Casadi](https://web.casadi.org/) is a C++ library with Python and MATLAB bindings, for solving ODEs and DAEs, nonlinear optimisation and algorithmic differentiation. In this benchmark we compare the performance of the DiffSol implementation with the Diffrax and Casadi libraries.
[Diffrax](https://docs.kidger.site/diffrax/) is a Python library for solving ODEs and SDEs implemented using JAX. [Casadi](https://web.casadi.org/) is a C++ library with Python and MATLAB bindings for solving ODE and nonlinear optimisation problems. In this benchmark we compare the performance of the DiffSol implementation with the Diffrax and Casadi libraries.

As well as demonstrating the performance of the DiffSol library, this benchmark also serves as an example of how to wrap and use DiffSol in other languages. The code for this benchmark can be found [here](https://github.com/martinjrobins/diffsol_python_benchmark). The [maturin](https://www.maturin.rs/) library was used to generate a template for the Python bindings and the CI/CD pipeline neccessary to build the bindings, run pytest tests and build the wheels ready for distribution on PyPI. The [pyo3](https://github.com/PyO3/pyo3) library was used to wrap the DiffSol library in Python.
As well as demonstrating the performance of the DiffSol library, this benchmark also serves as an example of how to wrap and use DiffSol in other languages. The code for this benchmark can be found [here](https://github.com/martinjrobins/diffsol_python_benchmark). The [maturin](https://www.maturin.rs/) library was used to generate a template for the Python bindings and the CI/CD pipeline neccessary to build the wheels ready for distribution on PyPI. The [pyo3](https://github.com/PyO3/pyo3) library was used to wrap the DiffSol library in Python.

## Problem setup

We will use the `robertson_ode` problem for this benchmark. This is a stiff ODE system with 3 equations and 3 unknowns, and is a common benchmark problem for ODE solvers. To illustrate the performance over a range of problem sizes we duplicated the equations by a factor of `ngroups`, so the number of equations is `3 * ngroups`.
We will use the `robertson_ode` problem for this benchmark. This is a stiff ODE system with 3 equations and 3 unknowns, and is a common benchmark problem. To illustrate the performance over a range of problem sizes, the Robertson equations were duplicated a factor of `ngroups`, so the total number of equations solved is `3 * ngroups`.

For the Diffrax implementation we based this on the [example](https://docs.kidger.site/diffrax/examples/stiff_ode/) from the Diffrax documentation, extending this to include the `ngroups` parameter. As with the example, we used the `Kvaerno5` method for the solver. You can see the final implementation of the model [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/diffsol_python_benchmark/diffrax_models.py).
The Diffrax implementation was based this on the [example](https://docs.kidger.site/diffrax/examples/stiff_ode/) in the Diffrax documentation, which was further extending to include the `ngroups` parameter. As is already used in the example, the `Kvaerno5` method was used for the solver. You can see the final implementation of the model [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/diffsol_python_benchmark/diffrax_models.py).

For the Casadi implementation we wrote this from scratch using the libraries Python API. You can see the final implementation of the model [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/diffsol_python_benchmark/casadi_models.py).
The Casadi implementation was written from scratch using Casadi's python API. You can see the final implementation of the model [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/diffsol_python_benchmark/casadi_models.py).

The DiffSol implementation of the model was done using the DiffSL language, and you can see the final implementation of the model [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/diffsol_python_benchmark/diffsol_models.py).
The DiffSol implementation of the model written using the DiffSL language, you can see the final implementation of the model [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/diffsol_python_benchmark/diffsol_models.py).

The final implementation of the benchmark using these models is done [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/bench/bench.py). The DiffSol benchmark is done using the `bdf` solver. For `ngroup` < 20 it uses the `nalgebra` dense matrix and LU solver, and for `ngroups` >= 20 the `faer` sparse matrix and LU solver.
The full implementation of the benchmark presented below can be seen [here](https://github.com/martinjrobins/diffsol_python_benchmark/blob/main/bench/bench.py). The DiffSol benchmark is performed using the `bdf` solver. For `ngroup` < 20 it uses the `nalgebra` dense matrix and LU solver, and for `ngroups` >= 20 the `faer` sparse matrix and LU solver are used.

## Differences between implementations

There are a number of differences between the Diffrax, Casadi and DiffSol implementations that may affect the performance of the solvers. The main differences are:
There are a few key differences between the Diffrax, Casadi and DiffSol implementations that may affect the performance of the solvers. The main differences are:
- The Casadi implementation uses sparse matrices, whereas the DiffSol implementation uses dense matrices for `ngroups` < 20, and sparse matrices for `ngroups` >= 20. This will provide an advantage for DiffSol for smaller problems.
- I'm unsure if the Diffrax implementation uses sparse or dense matrices, but it is most likely dense, as JAX only has experimental support for sparse matrices. This will provide an advantage for DiffSol for larger problems.
- The Diffrax implementation uses the `Kvaerno5` method, which is a 5th order implicit Runge-Kutta method. This is different from the BDF method used by both the Casadi and DiffSol implementations.
- Each library was allowed to use multiple threads according to their default settings. The only part of the DiffSol implementation that takes advantage of multiple threads is the `faer` sparse LU solver and matrix. Both the `nalgebra` LU solver / matrix, and the DiffSL generated code are single-threaded only. Diffrax uses JAX, which takes advantage of multiple threads (CPU only, no GPUs were used in these benchmarks). The Casadi implementation also uses multiple threads, but I'm unsure of the details.
- I'm unsure if the Diffrax implementation uses sparse or dense matrices, but it is most likely dense as JAX only has experimental support for sparse matrices. Treating the Jacobian as dense will be a disadvantage for Diffrax for larger problems as the Jacobian is very sparse.
- The Diffrax implementation uses the `Kvaerno5` method (a 5th order implicit Runge-Kutta method). This is different from the BDF method used by both the Casadi and DiffSol implementations.
- Each library was allowed to use multiple threads according to their default settings. The only part of the DiffSol implementation that takes advantage of multiple threads is the `faer` sparse LU solver and matrix. Both the `nalgebra` LU solver, matrix, and the DiffSL generated code are all single-threaded. Diffrax uses JAX, which takes advantage of multiple threads (CPU only, no GPUs were used in these benchmarks). The Casadi implementation also uses multiple threads.


## Results

The benchmarks were run on a Dell PowerEdge R7525 2U rack server, with dual AMD EPYC 7343 3.2Ghz 16C CPU and 128GB Memory. Each benchmark was run using both a low (1e-8) and high (1e-4) tolerances for both `rtol` and `atol`, and with `ngroup` ranging between 1 - 10,000. The results are presented in the following graphs, where the x-axis is the size of the problem `ngroup` and the y-axis is the time taken to solve the problem relative to the time taken by the DiffSol implementation (so `10^0` is the same time as DiffSol, `10^1` is 10 times slower etc.)
The benchmarks were run on a Dell PowerEdge R7525 2U rack server, with dual AMD EPYC 7343 3.2Ghz 16C CPU and 128GB Memory. Each benchmark was run using both a low (1e-8) and high (1e-4) tolerances for both `rtol` and `atol`, and with `ngroup` ranging between 1 - 10,000. The results are presented in the following graph, where the x-axis is the size of the problem `ngroup` and the y-axis is the time taken to solve the problem relative to the time taken by the DiffSol implementation (so `10^0` is the same time as DiffSol, `10^1` is 10 times slower etc.).

![Python](./images/python_plot.svg)

DiffSol is faster than both the Casadi and Diffrax implementations over the range of problem sizes and tolerances tested, although the Casadi and DiffSol implementations converge to be similar for larger problems (`ngroups` > 100).

The region that DiffSol really outperforms the other implementations is for smaller problems (`ngroups` < 5), at `ngroups` = 1, Casadi and Diffrax are between 3 - 40 times slower than DiffSol. This small size region are where the dense matrix and solver used is more appropriate for the problem, and the overhead of the other libraries is more significant. The Casadi library needs to traverse a graph of operations to calculate each rhs or jacobian evaluation, whereas the DiffSL JIT compiler will compile to native code using the LLVM backend, along with low-level optimisations that are not available to Casadi. Diffrax as well is significantly slower than DiffSol for smaller problems, and this might be due to (a) Diffrax being a ML library and not optimised for solving stiff ODEs, and (b) double precision is used, which again is not a common use case for ML libraries.
The DiffSol implementation outperforms the other implementations significantly for small problem sizes (`ngroups` < 5). E.g. at `ngroups` = 1, Casadi and Diffrax are between 3 - 40 times slower than DiffSol. At these small problem sizes, the dense matrix and solver used by DiffSol provide an advantage over the sparse solver used by Casadi. Casadi also has additional overhead to evaluate each function evaluation, as it needs to traverse a graph of operations to calculate each rhs or jacobian evaluation, whereas the DiffSL JIT compiler will compile to native code using the LLVM backend, along with low-level optimisations that are not available to Casadi. Diffrax is also significantly slower than DiffSol for smaller problems, this might be due to (a) Diffrax being a ML library and not optimised for solving stiff ODEs, or (b) double precision is used, which again is not a common use case for ML libraries, or (c) perhaps the different solver methods (Kvaerno5 vs BDF) are causing the difference.

As the problem sizes get larger, the performance of Diffrax and Casadi improve rapidly relative to DiffSol, but after `ngroups` > 10 the performance of Diffrax drops off again, probably due to JAX not taking advantage of the sparse structure of the problem. The performance of Casadi continues to improve, and for `ngroups` > 100 it is comparable to DiffSol. By the time `ngroups` = 10,000, the performance of Casadi is identical to DiffSol.
2 changes: 1 addition & 1 deletion book/src/choosing_a_solver.md
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ Each solver's state struct implements the [`OdeSolverState`](https://docs.rs/dif

For example, say that you wish to bypass the initialisation of the state as you already have the algebraic constraints and so don't need to solve for them. You can use the `new_without_initialise` method on the `OdeSolverState` trait to create a new state without initialising it. You can then use the `as_mut` method to get a mutable reference to the state and set the values manually.

Note that each state struct has a [`as_ref`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/trait.OdeSolverState.html#tymethod.as_ref) and [`as_mut`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/trait.OdeSolverState.html#tymethod.as_mut) methods that return a [`StateRef`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/struct.StateRef.html) or ['StateRefMut`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/struct.StateRefMut.html) struct respectively. These structs provide a solver-independent way to access the state values so you can use the same code with different solvers.
Note that each state struct has a [`as_ref`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/trait.OdeSolverState.html#tymethod.as_ref) and [`as_mut`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/trait.OdeSolverState.html#tymethod.as_mut) methods that return a [`StateRef`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/struct.StateRef.html) or [`StateRefMut`](https://docs.rs/diffsol/latest/diffsol/ode_solver/state/struct.StateRefMut.html) struct respectively. These structs provide a solver-independent way to access the state values so you can use the same code with different solvers.

```rust
# use diffsol::OdeBuilder;
Expand Down
4 changes: 2 additions & 2 deletions book/src/primer/compartmental_models_of_drug_delivery.md
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ These are often referred to as ADME, and taken together describe the drug concen
The body itself is modelled as one or more *compartments*, each of which is defined as a kinetically homogeneous unit (these compartments do not relate to specific organs in the body, unlike Physiologically based pharmacokinetic, PBPK, modeling). There is typically a main *central* compartment into which the drug is administered and from which the drug is excreted from the body, combined with zero or more *peripheral* compartments to which the drug can be distributed to/from the central compartment (See Fig 2). Each
peripheral compartment is only connected to the central compartment.

![Fig 2](https://sabs-r3.github.io/software-engineering-projects/fig/pk2.svg)
![Fig 2](images/pk2.svg)

The following example PK model describes the two-compartment model shown diagrammatically in the figure above. The time-dependent variables to be solved are the drug quantity in the central and peripheral compartments, $q_c$ and $q_{p1}$ (units: [ng]) respectively.

Expand Down Expand Up @@ -74,7 +74,7 @@ For the dose function, we will specify a dose of 1000 ng at regular intervals of
V_c = 1000 \text{ mL}, \quad V_{p1} = 1000 \text{ mL}, \quad CL = 100 \text{ mL/h}, \quad Q_{p1} = 50 \text{ mL/h}
\\]

Let's now solve this system of ODEs using DiffSol.
Let's now solve this system of ODEs using DiffSol. To implement the discrete dose events, we set a stop time for the simulation at each dose event using the [OdeSolverMethod::set_stop_time](https://docs.rs/diffsol/latest/diffsol/ode_solver/method/trait.OdeSolverMethod.html#tymethod.set_stop_time) method. During timestepping we can check the return value of the [OdeSolverMethod::step](https://docs.rs/diffsol/latest/diffsol/ode_solver/method/trait.OdeSolverMethod.html#tymethod.step) method to see if the solver has reached the stop time. If it has, we can apply the dose and continue the simulation.

```rust
# fn main() {
Expand Down
Loading

0 comments on commit 13c15a1

Please sign in to comment.