Using OrdinaryDiffEq.jl as a native ODE solver library from Rust with JuliaC --trim

General Usage
1

I didn’t find an exemple for this so I’m posting this here just in case someone needs it. People kept saying OrdinaryDiffEq is compatible with JuliaC but I wanted to see how it looks like and if someone ask its here now. Also, didn’t want to code an adaptative alg nor learn rust libs for ode solvers.

julia :

using OrdinaryDiffEqSSPRK
using SciMLBase
using SciMLLogging

Base.errormonitor(t::Task) = t

Base.@ccallable function solve_ode(
    pout::Ptr{Float64},
    capacity::UInt64,
    pf::Ptr{Cvoid},
    pu0::Ptr{Float64},
    u0len::UInt64,
    t0::Float64,
    T::Float64,
    pp::Ptr{Cvoid},
    nsave::UInt64
)::Int64
    if u0len == 0 || nsave < 2 ||
       u0len > typemax(Int) || nsave > typemax(Int) ||
       !isfinite(t0) || !isfinite(T) || T <= t0
        return -1
    end
    rhs(du, u, p, t) = ccall(
        pf, Cvoid,
        (Ptr{Float64}, Ptr{Float64}, UInt64, Ptr{Cvoid}, Float64),
        pointer(du), pointer(u), u0len, p, t
    )
    u0 = unsafe_wrap(Vector{Float64}, pu0, (Int(u0len),))
    ode = ODEProblem{true, SciMLBase.FullSpecialize}(
        rhs, u0, (t0, T), pp
    )
    sol = solve(
        ode, SSPRK432(),
        saveat=range(t0, T; length=Int(nsave)),
        verbose=SciMLLogging.None(),
        abstol=1e-12,
        reltol=1e-12
    )
    N = length(sol.u)
    if UInt64(N) > capacity ÷ u0len
        return -2
    end
    for i in 1:N
        for j in 1:Int(u0len)
            unsafe_store!(
                pout,
                sol.u[i][j],
                (i - 1) * Int(u0len) + j
            )
        end
    end
    return Int64(N)
end

wrapper rust :

use std::ffi::c_void;
use std::sync::atomic::{AtomicBool, Ordering};
use std::panic::{catch_unwind, AssertUnwindSafe};
#[link(name = "ssplib", kind = "dylib")]
unsafe extern "C" {
    fn solve_ode(
        pout: *mut f64,
        capacity: u64,
        pf: *const c_void,
        pu0: *mut f64,
        u0len: u64,
        t0: f64,
        t1: f64,
        pp: *const c_void,
        nsave: u64,
    ) -> i64;
}
struct Context<'a, P, F> {
    params: &'a P,
    rhs: F,
    panicked: AtomicBool,
}
extern "C" fn trampoline<P, F>(
    du: *mut f64,
    u: *const f64,
    len: u64,
    pp: *mut c_void,
    t: f64,
)
where
    P: Sync,
    F: Fn(&mut [f64], &[f64], &P, f64) + Sync,
{
    let ctx = unsafe { &*(pp as *const Context<'_, P, F>) };
    let n = len as usize;
    if ctx.panicked.load(Ordering::Relaxed) {
        unsafe { std::ptr::write_bytes(du, 0, n) };
        return;
    }
    let result = catch_unwind(AssertUnwindSafe(|| {
        let input = unsafe {
            std::slice::from_raw_parts(u, n)
        }.to_vec();
        let output = unsafe {
            std::slice::from_raw_parts_mut(du, n)
        };
        (ctx.rhs)(output, &input, ctx.params, t);
    }));
    if result.is_err() {
        ctx.panicked.store(true, Ordering::Relaxed);
        unsafe { std::ptr::write_bytes(du, 0, n) };
    }
}

pub fn solve<P, F>(
    u0: &[f64],
    params: &P,
    tspan: (f64, f64),
    nsave: usize,
    rhs: F,
) -> Result<Vec<Vec<f64>>, &'static str>
where
    P: Sync,
    F: Fn(&mut [f64], &[f64], &P, f64) + Sync,
{
    let capacity = u0.len()
        .checked_mul(nsave)
        .ok_or("Dimension error")?;
    let mut initial = u0.to_vec();
    let mut output = vec![0.0_f64; capacity];
    let ctx = Context {
        params,
        rhs,
        panicked: AtomicBool::new(false),
    };
    let n = unsafe {
        solve_ode(
            output.as_mut_ptr(),
            capacity as u64,
            trampoline::<P, F> as *const () as *const c_void,
            initial.as_mut_ptr(),
            initial.len() as u64,
            tspan.0,
            tspan.1,
            &ctx as *const _ as *const c_void,
            nsave as u64,
        )
    };
    Ok(output
        .chunks_exact(u0.len())
        .map(|row| row.to_vec())
        .collect())
}

main rust

mod solver;
use solver::solve;
struct Params {
    alpha: f64,
    beta: f64,
    gamma: i32,
}
fn rhs(du: &mut [f64], u: &[f64], p: &Params, _t: f64) {
    for i in 0..u.len() {
        du[i] = -p.alpha * u[i] + p.beta;
    }
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
    let u0 = vec![1.0, 2.0, 3.0];
    let params = Params {
        alpha: 2.0,
        beta: 0.5,
        gamma: 3,
    };
    let sol = solve(&u0, &params, (0.0, 1.0), 101, rhs)?;
    println!("Solution finale : {:?}", sol.last().unwrap());
    let exact: Vec<f64> = u0.iter()
        .map(|x| {
            params.beta / params.alpha
                + (x - params.beta / params.alpha)
                * (-params.alpha).exp()
        })
        .collect();

    println!("Solution exacte : {:?}", exact);
    Ok(())
}

compile :

juliac --output-lib ssplib --project="." ssprk.jl --trim

ln -sf ssplib.dylib libssplib.dylib

RUSTFLAGS="-L native=$(pwd) -C link-arg=-Wl,-rpath,$(pwd)" cargo run --release

and you get an adaptive ode solver from julia in rust with very similar api than julia and any kind of parameter (c representable).

One thing a little sad : Base.errormonitor(t::Task) = t hope this gets avoided at some point. Btw I didn’t put the tols in but its straighforwar and left as an exercise for the readers lol.
edit : sorry for the symbolic link just change the name rust wants libsomething.dylib.

2 Likes