ChrisRackauckas

DiffEqFlux.jl (Julia) vs Jax on an Epidemic Model

The Jax developers optimized a differential equation benchmark in this issue which used DiffEqFlux.jl as a performance baseline. The Julia code from there was updated to include some standard performance tricks and is the benchmark code here. Thus both codes have been optimized by the library developers.

Results

Forward Pass

ChrisRackauckas

Julia @ccallable C Header Generator

This collection of Julia scripts automatically generates C header files from Julia functions marked with the @ccallable macro. This is useful when creating Julia libraries that need to be called from C/C++ code.

Files Description

1. ccallable_header_generator.jl

The main generator script that provides a simple and robust approach to generating C headers. It uses a manual definition approach where you explicitly define function signatures, avoiding complex AST parsing.

Features:

ChrisRackauckas

SciPy+Numba odeint vs Julia DifferentialEquations.jl vs NumbaLSODA Summary

All are solved at reltol=1e-3, abstol=1e-6 using the fastest ODE solver of the respective package for the given problem.

Absolute Performance Numbers:

• SciPy LSODA through odeint takes ~489μs
• SciPy LSODA through odeint with Numba takes ~257μs
• NumbaLSODA takes ~25μs
• DifferentialEquations.jl Rosenbrock23 takes ~9.2μs

ChrisRackauckas

torchdiffeq vs Julia DiffEqFlux Neural ODE Training Benchmark

The spiral neural ODE was used as the training benchmark for both torchdiffeq (Python) and DiffEqFlux (Julia) which utilized the same architecture and 500 steps of ADAM. Both achived similar objective values at the end. Results:

• DiffEqFlux defaults: 7.4 seconds
• DiffEqFlux optimized: 2.7 seconds
• torchdiffeq: 288.965871299999 seconds

ChrisRackauckas

Torchdiffeq vs DifferentialEquations.jl (/ DiffEqFlux.jl) Neural ODE Compatible Solver Benchmarks

Only non-stiff ODE solvers are tested since torchdiffeq does not have methods for stiff ODEs. The ODEs are chosen to be representative of models seen in physics and model-informed drug development (MIDD) studies (quantiative systems pharmacology) in order to capture the performance on realistic scenarios.

Summary

Below are the timings relative to the fastest method (lower is better). For approximately 1 million ODEs and less, torchdiffeq was more than an order of magnitude slower than DifferentialEquations.jl

ChrisRackauckas

using OrdinaryDiffEq
using Plots
using Flux, DiffEqFlux, Optim

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end

ChrisRackauckas

using Cassette, DiffRules
using Core: CodeInfo, SlotNumber, SSAValue, ReturnNode, GotoIfNot

const printbranch = true

Cassette.@context HasBranchingCtx

function Cassette.overdub(ctx::HasBranchingCtx, f, args...)
    if Cassette.canrecurse(ctx, f, args...)
        return Cassette.recurse(ctx, f, args...)

ChrisRackauckas

torchsde vs DifferentialEquations.jl / DiffEqFlux.jl (Julia)

This example is a 4-dimensional geometric brownian motion. The code for the torchsde version is pulled directly from the torchsde README so that it would be a fair comparison against the author's own code. The only change to that example is the addition of a dt choice so that the simulation method and time step matches between the two different programs.

The SDE is solved 100 times. The summary of the results is as follows:

ChrisRackauckas

using OrdinaryDiffEq, RecursiveArrayTools, LinearAlgebra, Test, SparseArrays, SparseDiffTools, Sundials

# Define the constants for the PDE
const α₂ = 1.0
const α₃ = 1.0
const β₁ = 1.0
const β₂ = 1.0
const β₃ = 1.0
const r₁ = 1.0
const r₂ = 1.0

ChrisRackauckas

ERROR: MethodError: Cannot `convert` an object of type Float32 to an object of type Vector{Float32}
Closest candidates are:
  convert(::Type{Array{T, N}}, ::SizedArray{S, T, N, N, Array{T, N}}) where {S, T, N} at C:\Users\accou\.julia\packages\StaticArrays\0T5rI\src\SizedArray.jl:121
  convert(::Type{Array{T, N}}, ::SizedArray{S, T, N, M, TData} where {M, TData<:AbstractArray{T, M}}) where {T, S, N} at C:\Users\accou\.julia\packages\StaticArrays\0T5rI\src\SizedArray.jl:115
  convert(::Type{<:Array}, ::LabelledArrays.LArray) at C:\Users\accou\.julia\packages\LabelledArrays\lfn1b\src\larray.jl:133
  ...
Stacktrace:
  [1] setproperty!(x::OrdinaryDiffEq.ODEIntegrator{Tsit5, false, Vector{Float32}, Float32}, f::Symbol, v::Float32)
    @ Base .\Base.jl:43
  [2] initialize!(integrator::OrdinaryDiffEq.ODEIntegrator{Tsit5, false, Vector{Float32}, Float32})