Performance when using DifferentialEquations.solve repeatedly inside a loop

I am using the split step method to solve the following equation:

The first 3 terms are straightforward when using the split step method. For the fourth term, I basically need to solve this equation:

after evolving the equation for the first three terms. i.e. I need to repeatedly use DifferentialEquations.solve inside an outer time loop to evolve the solution a single step. This is very slow, unfortunately.

This is a simple example of how this works (I can provide a full MWE if needed but it’s on the longer side)

function solve!(...)
    # Prepare some stuff
    # ...

    # Differential Operators
    α = 0.1
    x_order = 2
    D = CenteredDifference(1, x_order, dx, Nx) 
    Q = PeriodicBC(Float64)

    function M!(du, u,p,t)
        du = 6*α*D*Q*u.*abs2.(u)
    end

    for i = 1:Nt-1
        @views ψ[:, i+1] .= T(ψ[:, i], dt, M!, ...)
    end

    return nothing
end #solve

function T(ψ, W, dt, F, F̃, M!)
    # Perform some operations in-place on ψ

    # Calling this every loop is slow and allocates a lot
    prob = ODEProblem(M!, ψ, (0.0, dt))
    sol = solve(prob, BS3(), save_everystep = false, dt = dt) 
    ψ = sol.u[end]

    return ψ
end

So as you can see, the function T evolves time by a single step, so it performs some operations on ψ then has to solve the ODE. This happens every single time step, so re-initializing the problem and then solving it makes the whole process very time consuming. I made the function M! in place and it helped a bit, but I am stuck at this point.

3 posts - 2 participants

Read full topic