mini-ode 0.1.4

A minimalistic ODE solvers library built on top of PyTorch

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

Solver Class	Method	Suitable For	Implicit	Adaptive Step
EulerMethodSolver	Euler	Simple, fast, and educational use.	❌	❌
RK4MethodSolver	Runge-Kutta 4th Order (RK4)	General-purpose with fixed step size.	❌	❌
ImplicitEulerMethodSolver	Implicit Euler	Stiff or ill-conditioned problems.	✅	❌
GLRK4MethodSolver	Gauss-Legendre RK (Order 4)	High-accuracy, stiff problems.	✅	❌
RKF45MethodSolver	Runge-Kutta-Fehlberg 4(5)	Adaptive step size control.	❌	✅
ROW1MethodSolver	Rosenbrock-Wanner (Order 1)	Fast semi-implicit method for stiff systems.	semi	❌

📦 Building the Library

Rust

To build the core Rust library:

cd mini-ode
cargo build --release

Python

To build and install the Python package (in a virtual environment or Conda environment):

cd mini-ode-python
LIBTORCH_USE_PYTORCH=1 maturin develop

This builds the Python bindings using maturin and installs the package locally.

🐍 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
📧 [e-mail]
🔗 https://github.com/antoniprzybylik