Python

The PyO3 and maturin crates allow you to easily create Python bindings for Rust libraries, and Diffsol is no exception. For a full example of how to use Diffsol from Python, see the Diffsol Python example. Below is a brief overview of some of the key features from this example.

Getting started

You can install maturin using pip:

pip install maturin

Then, you can create a template maturin project using pyo3:

maturin new -b pyo3 python-diffsol

cd python-diffsol
cargo add diffsol

Wrapping a Diffsol problem

First, lets define some types that we'll use, including a matrix M, vector V and context type C for the linear algebra operations, as well as a linear solver type LS. For the JIT compilation backend, we'll swap between the LLVM and Cranelift backends, depending on whether the diffsol-llvm feature is enabled.

use std::sync::{Arc, Mutex};

use diffsol::{
    error::DiffsolError, NonLinearOp, OdeBuilder, OdeEquations, OdeSolverMethod, OdeSolverProblem,
    Op, Vector, VectorHost,
};
use numpy::{
    ndarray::{s, Array2},
    IntoPyArray, PyArray2, PyReadonlyArray1,
};
use pyo3::{exceptions::PyValueError, prelude::*};

type M = diffsol::NalgebraMat<f64>;
type V = diffsol::NalgebraVec<f64>;
type C = diffsol::NalgebraContext;
type LS = diffsol::NalgebraLU<f64>;
#[cfg(not(feature = "diffsol-llvm"))]
type CG = diffsol::CraneliftJitModule;
#[cfg(feature = "diffsol-llvm")]
type CG = diffsol::LlvmModule;
type Eqn = diffsol::DiffSl<M, CG>;

Now lets create a simple struct that we're going to wrap to use from Python. This will just store a Diffsol problem that we can solve later. Since we're using the DiffSL equations type (see Eqn above), this isn't threadsafe, so we'll use an Arc<Mutex<_>> to allow us to share the problem between threads safely.

#[pyclass]
struct PyDiffsol {
    problem: Arc<Mutex<OdeSolverProblem<Eqn>>>,
}

Solving the problem

For our implementation for PyDiffsol, we'll create two methods: one for creating the problem from a DiffSL string (new), and another for solving the problem (solve). The solve method will take a set of parameters and a set of times to solve the problem at. It then creates an Array2 to store the solution (we are using the numpy crate to allow us to return a NumPy array), and then iterates over the times, stepping the solver and interpolating the solution at each time. If the problem has an output function, it will call that function to get the output, otherwise it will just return the state vector.

#[pymethods]
impl PyDiffsol {
    #[new]
    fn new(code: &str) -> Result<Self, PyDiffsolError> {
        let problem = OdeBuilder::<M>::new().build_from_diffsl(code)?;
        Ok(Self {
            problem: Arc::new(Mutex::new(problem)),
        })
    }

    #[pyo3(signature = (params, times))]
    fn solve<'py>(
        &mut self,
        py: Python<'py>,
        params: PyReadonlyArray1<'py, f64>,
        times: PyReadonlyArray1<'py, f64>,
    ) -> Result<Bound<'py, PyArray2<f64>>, PyDiffsolError> {
        let mut problem = self
            .problem
            .lock()
            .map_err(|e| PyDiffsolError(DiffsolError::Other(e.to_string())))?;
        let times = times.as_array();
        let params = V::from_slice(params.as_array().as_slice().unwrap(), C::default());
        problem.eqn.set_params(&params);
        let mut solver = problem.bdf::<LS>()?;
        let nout = if let Some(_out) = problem.eqn.out() {
            problem.eqn.nout()
        } else {
            problem.eqn.nstates()
        };
        let mut sol = Array2::zeros((nout, times.len()));
        for (i, &t) in times.iter().enumerate() {
            while solver.state().t < t {
                solver.step()?;
            }
            let y = solver.interpolate(t)?;
            let out = if let Some(out) = problem.eqn.out() {
                out.call(&y, t)
            } else {
                y
            };
            sol.slice_mut(s![.., i])
                .iter_mut()
                .zip(out.as_slice().iter())
                .for_each(|(a, b)| *a = *b);
        }
        Ok(sol.into_pyarray(py))
    }
}

Error handling

In our implementation, we need to handle errors that may occur when working with the Diffsol library. We'll create a custom error type PyDiffsolError that wraps the DiffsolError type, and implement the From trait to convert between the two types and the Python PyErr type. This will allow us to easily convert DiffsolError errors into Python exceptions.

struct PyDiffsolError(DiffsolError);

impl From<PyDiffsolError> for PyErr {
    fn from(error: PyDiffsolError) -> Self {
        PyValueError::new_err(error.0.to_string())
    }
}

impl From<DiffsolError> for PyDiffsolError {
    fn from(other: DiffsolError) -> Self {
        Self(other)
    }
}

The Python module

Finally, we need to create the Python module that will be imported by Python.

#[pymodule]
fn python_diffsol(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyDiffsol>()?;
    Ok(())
}

Testing from Python

To build the Python module, we can use the maturin tool. This will compile the Rust code and create a Python wheel.

maturin develop

Then we can write and run a simple Python test to ensure everything is working correctly:

from python_diffsol import PyDiffsol
import numpy as np
import unittest


class TestStringMethods(unittest.TestCase):
    def test_solve(self):
        model = PyDiffsol(
            """
            in = [r, k]
            r { 1 } k { 1 }
            u_i { y = 0.1 }
            F_i { (r * y) * (1 - (y / k)) }
            """
        )
        times = np.linspace(0.0, 1.0, 100)
        k = 1.0
        r = 1.0
        y0 = 0.1
        y = model.solve(np.array([r, k]), times)
        soln = k / (1.0 + (k - y0) * np.exp(-r * times) / y0)
        np.testing.assert_allclose(y[0], soln, rtol=1e-5)
        
if __name__ == '__main__':
    unittest.main()