mini-ode

A minimalistic, multi-language library for solving Ordinary Differential Equations (ODEs). mini-ode is designed with a shared Rust core and a consistent interface for both Rust and Python users. It supports explicit, implicit, fixed step and adaptive step algorithms.

✨ Features

• Dual interface: call the same solvers from Rust or Python
• PyTorch-compatible: define the derivative function using PyTorch
• Multiple solver methods: includes explicit, implicit, and adaptive-step solvers
• Modular optimizers: implicit solvers allow flexible optimizer configuration

🧠 Supported Solvers

📦 Building the Library

mini-ode can be built against two different libtorch installations:

• a system-installed libtorch, or
• the libtorch bundled with a PyTorch installation.

It is important that the headers used at compile time and the shared libraries loaded at runtime come from the same installation.

Rust (system libtorch)

To build the Rust library against your system-installed libtorch:

cd mini-ode
cargo build --release

This is the recommended approach if you are developing a pure Rust application.

Python

To build the Python bindings using the PyTorch installation from your active Python environment:

cd mini-ode-python
LIBTORCH_USE_PYTORCH=1 maturin develop

This builds the Python bindings using maturin and installs the package locally. LIBTORCH_USE_PYTORCH=1 tells torch-sys to obtain the include directories and libraries from the currently active Python installation.

Building Rust code against a PyTorch installation

If you want to build the Rust version of mini-ode against the libtorch that ships with your Conda or virtual environment PyTorch installation, you must configure both the compiler and the runtime linker.

Using only

LIBTORCH_USE_PYTORCH=1 cargo build

is not enough.

While this causes torch-sys to compile against the headers from your Python installation, system will typically still load the system libtorch (for example /usr/lib/libtorch.so) at runtime. Simply activating a Conda environment does not guarantee that the Conda libtorch will be used. Mixing the compile-time headers from one PyTorch installation with the runtime libraries from another would lead to runtime failures.

Correct way to build and run:

export LIBTORCH_USE_PYTORCH=1
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python*/site-packages/torch/lib:$LD_LIBRARY_PATH"

cargo build
cargo run

This ensures that both the compiler and the runtime use the same libtorch installation.

Tip

If your Rust code is loaded as a Python extension, importing torch before importing your extension ensures that Python has already loaded the correct libtorch libraries.

🐍 Python Usage Overview

To use mini-ode from Python:

1. Define the derivative function f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
   • x is a scalar tensor (rank 0, shape ()).
   • y is a 1D tensor (rank 1, shape (n,) where n is the dimension of the state vector).
   • The function must return a 1D tensor of the same shape as y (i.e., (n,)).
2. Trace the function using torch.jit.trace to convert it to TorchScript. Provide example inputs matching the shapes (e.g., torch.tensor(0.) for x and torch.tensor([0.] * n) for y).
3. Create a solver instance, configuring parameters like step size or optimizer as needed.
4. Call solver.solve(traced_f, x_span, y0):
   • x_span: A tuple (start, end) for the integration interval.
   • y0: Initial state as a 1D tensor (shape (n,)).
   • Returns: (xs, ys) where xs is a 1D tensor of x-values (shape (num_points,)), and ys is a 2D tensor of y-values (shape (num_points, n)).
     • For fixed-step solvers, xs is evenly spaced (e.g., via torch.linspace).
     • For adaptive-step solvers, xs has a variable number of points based on error control.

Example usage flow (not full code):

import torch
import mini_ode

# 1. Define derivative function using PyTorch
def f(x: torch.Tensor, y: torch.Tensor):
    return y.flip(0) - torch.tensor([0, 1]) * (y.flip(0) ** 3)

# 2. Trace the function to TorchScript
traced_f = torch.jit.trace(f, (torch.tensor(0.), torch.tensor([0., 0.])))

# 3. Create a solver instance
solver = mini_ode.RK4MethodSolver(step=0.01)

# 4. Solve the ODE
xs, ys = solver.solve(traced_f, (0., 5.), torch.tensor([1.0, 0.0]))

🔧 Using Optimizers (Implicit Solvers Only)

Some solvers like GLRK4MethodSolver or ImplicitEulerMethodSolver require an optimizer for nonlinear system solving:

optimizer = mini_ode.optimizers.CG(
    max_steps=5,
    gtol=1e-8,
)

solver = mini_ode.GLRK4MethodSolver(step=0.2, optimizer=optimizer)

🦀 Rust Usage Overview

In Rust, solvers use the same logic as in Python - but you pass in a tch::CModule representing the TorchScripted derivative function.

Example 1: Load a TorchScript model from file

This approach uses a model traced in Python (e.g., with torch.jit.trace) and saved to disk.

use mini_ode::Solver;
use tch::{Tensor, CModule};

fn main() -> anyhow::Result<()> {
    let solver = Solver::Euler { step: 0.01 };
    let model = CModule::load("my_traced_function.pt")?;
    let x_span = (0.0, 2.0);
    let y0 = Tensor::from_slice(&[1.0f64, 0.0]);

    let (xs, ys) = solver.solve(model, x_span, y0)?;
    println!("{:?}", xs);
    Ok(())
}

Example 2: Trace the derivative function directly in Rust

You can also define and trace the derivative function in Rust using CModule::create_by_tracing.

use mini_ode::Solver;
use tch::{Tensor, CModule};

fn main() -> anyhow::Result<()> {
    // Initial value for tracing
    let y0 = Tensor::from_slice(&[1.0f64, 0.0]);

    // Define the derivative function closure
    let mut closure = |inputs: &[Tensor]| {
        let x = &inputs[0];
        let y = &inputs[1];
        let flipped = y.flip(0);
        let dy = &flipped - &(&flipped.pow_tensor_scalar(3.0) * Tensor::from_slice(&[0.0, 1.0]));
        vec![dy]
    };

    // Trace the model directly in Rust
    let model = CModule::create_by_tracing(
        "ode_fn",
        "forward",
        &[Tensor::from(0.0), y0.shallow_clone()],
        &mut closure,
    )?;

    // Use an adaptive solver, for example
    let solver = Solver::RKF45 {
        rtol: 0.00001,
        atol: 0.00001,
        min_step: 1e-9,
        safety_factor: 0.9
    };
    let x_span = Tensor::from_slice(&[0.0f64, 5.0]);

    let (xs, ys) = solver.solve(model, x_span, y0)?;
    println!("Final state: {:?}", ys);
    Ok(())
}

📁 Project Structure

mini-ode/           # Core Rust implementation of solvers
mini-ode-python/    # Python bindings using PyO3 + maturin
example.ipynb       # Jupyter notebook demonstrating usage

📄 License

This project is licensed under the GPL-2.0 License.

👤 Author

Antoni Michał Przybylik
📧 antoni@taon.io
🔗 https://github.com/antoniprzybylik