Memory allocation issue of a Euler solver implementation

General Usage Performance
, , , ,
1

Hello all,

I am implementing ODE solvers for a class assignment. Unfortunately, there is a memory allocation problem with my script, and the performance is not optimal.

Problem Description

The implementation uses the functor for the Euler algorithm with step size h, and a linear ODE problem is applied for simulation.

The memory allocation depends on h. Smaller h requires more time steps and then increases allocation.

Here is a part of script that I think the problem comes from:

function solve(f, solver, iv, tspan, h)
    ts = collect(Float64, tspan[1]:h:tspan[2])
    ys = similar(ts)
    ys[1] = iv
    for i in 1:length(ts)-1
        ys[i+1] = solver(ys[i], ts[i], f) # memory allocation
    end
    return OdeSol(ys, ts, solver)
end

Possible solution

By benchmarking the solver function, the solver function does not allocate memory, but the allocation occurs when assigning value to ys[i+1]. The for-loop creates allocations depending on the number of iterations.

Minimal reproducible code

The following is the minimum reproducible code for the Euler solver:

using BenchmarkTools

abstract type OdeSolver end

struct OdeProblem{A,B,C}
    f::A
    tspan::B
    iv::C
end

struct OdeSol{A,B,C}
    ys::A
    ts::B
    solver::C
end

# Solver
struct Euler <: OdeSolver
    h :: Float64 # To avoid memory allocation for functor
end

function (sv::Euler)(y:: Float64, t::Float64, f) 
    y_nxt = y + sv.h * f(y, t)
    return y_nxt
end

# Solver utils
function solve(f, solver, iv, tspan, h)
    ts = collect(Float64, tspan[1]:h:tspan[2])
    ys = similar(ts)
    ys[1] = iv
    for i in 1:length(ts)-1
        ys[i+1] = solver(ys[i], ts[i], f)
    end
    return OdeSol(ys, ts, solver)
end

solve(prob::OdeProblem, solver::OdeSolver) = solve(prob.f, solver, prob.iv, prob.tspan, solver.h)


# ODE Problem
function f(y, t)
    dy = t - y + 2.
    return dy
end

prob = OdeProblem(f, (0., 1.), 2.);
h = 0.001



# Benchmarking

@btime Euler(h)(3., 2., f)
"""
3.042 ns (0 allocations: 0 bytes)
3.001
"""

@btime solve(prob, Euler(h))
"""
  77.375 μs (4494 allocations: 86.20 KiB)
OdeSol{Vector{Float64}, Vector{Float64}, Euler}([2.0, 2.0, 2.000001, 2.0000029990000003, 2.000005996001, 2.0000099900049992, 2.0000149800149942, 2.000020965034979, 2.000027944069944, 2.0000359161258743  …  2.3620212907402465, 2.362650269449506, 2.3632796191800565, 2.3639093395608763, 2.3645394302213156, 2.3651698907910945, 2.3658007209003036, 2.3664319201794033, 2.367063488259224, 2.367695424770965], [0.0, 0.001, 0.002, 0.003, 0.004, 0.005, 0.006, 0.007, 0.008, 0.009  …  0.991, 0.992, 0.993, 0.994, 0.995, 0.996, 0.997, 0.998, 0.999, 1.0], Euler(0.001))
"""

Goals

  1. Eeliminate the memory allocations that depend on step size h.
  2. Keep generic typing as much as possible

Thank you!

2

Try changing your function signature to function solve(f::F, solver, iv, tspan, h) where F. The issue is likely related to cases where Julia avoids specializing on argument types. It’s possible you might need to sprinkle this pattern a couple more places in inner functions, as well.

2 Likes
3

You’re right! The issue is solved:

@btime solve(prob, Euler(h))

5.306 μs (4 allocations: 16.05 KiB)

Revised script

The following is the revised script:

using BenchmarkTools

abstract type OdeSolver end

struct OdeProblem{A,B,C}
    f::A
    tspan::B
    iv::C
end

struct OdeSol{A,B,C}
    ys::A
    ts::B
    solver::C
end

# Solver
struct Euler <: OdeSolver
    h :: Float64 # To avoid memory allocation for functor
end

function (sv::Euler)(y:: Float64, t::Float64, f) 
    y_nxt = y + sv.h * f(y, t)
    return y_nxt
end

# Solver utils
function solve(f::F, solver, iv, tspan, h) where F
    ts = collect(Float64, tspan[1]:h:tspan[2])
    ys = similar(ts)
    ys[1] = iv
    for i in 1:length(ts)-1
        ys[i+1] = solver(ys[i], ts[i], f)
    end
    return OdeSol(ys, ts, solver)
end

solve(prob::OdeProblem, solver::OdeSolver) = solve(prob.f, solver, prob.iv, prob.tspan, solver.h)


# ODE Problem
function f(y, t)
    dy = t - y + 2.
    return dy
end

prob = OdeProblem(f, (0., 1.), 2.);
h = 0.001



# Benchmarking

@btime Euler(h)(3., 2., f)
"""
3.042 ns (0 allocations: 0 bytes)
3.001
"""

@btime solve(prob, Euler(h))
"""
5.306 μs (4 allocations: 16.05 KiB)
OdeSol{Vector{Float64}, Vector{Float64}, Euler}([2.0, 2.0, 2.000001, 2.0000029990000003, 2.000005996001, 2.0000099900049992, 2.0000149800149942, 2.000020965034979, 2.000027944069944, 2.0000359161258743  …  2.3620212907402465, 2.362650269449506, 2.3632796191800565, 2.3639093395608763, 2.3645394302213156, 2.3651698907910945, 2.3658007209003036, 2.3664319201794033, 2.367063488259224, 2.367695424770965], [0.0, 0.001, 0.002, 0.003, 0.004, 0.005, 0.006, 0.007, 0.008, 0.009  …  0.991, 0.992, 0.993, 0.994, 0.995, 0.996, 0.997, 0.998, 0.999, 1.0], Euler(0.001))
"""