Inconsistent typing making adding arrays and adding a vector to the end of an array difficult

We can get another 10x performance increase by using StaticArrays as recommended by @mikkoku and @Sukera. The edited code below will run in ~14 ms (@btime). Note that for this problem I did not observe any time improvement by defining in struct, t or vars as ::Array{Float64,N} versus ::Array{Float64} as recommended by @DNF, but do not understand the implications of this in general case.

using LinearAlgebra, Plots, StaticArrays
using BenchmarkTools

struct solObj
    t::Array{Float64}
    vars::Array{Float64}
end

struct paramObj
    sigma::Float64
    beta::Float64
    rho::Float64
end

function f(params, t, vars)
    sigma, beta, rho = params.sigma, params.beta, params.rho;
    x, y, z = vars[1], vars[2], vars[3];
    dx    = sigma*(y-x);
    dy    = x*(rho-z) - y;
    dz    = x*y - beta*z;
    return  @SVector [dx, dy, dz];
end

function rkf45(f::Function, params, t0::Float64, tf::Float64, conds, epsilon::Float64, dtInitial::Float64)
    dt = dtInitial;
    t = Float64[t0]
    vars = [conds]
    i = 1
    ti = t0
    while ( ti < tf )
        varsi =  vars[i]
        dt = minimum((dt, tf-ti))
        K1 = dt*f(params, ti, varsi)
        K2 = dt*f(params, ti + dt/4, varsi + K1/4)
        K3 = dt*f(params, ti + 3*dt/8, varsi + 3*K1/32 + 9*K2/32)
        K4 = dt*f(params, ti + 12*dt/13, varsi + 1932*K1/2197 - 7200*K2/2197 + 7296*K3/2197)
        K5 = dt*f(params, ti + dt, varsi + 439*K1/216 - 8*K2 + 3680*K3/513 - 845*K4/4104)
        K6 = dt*f(params, ti + dt/2, varsi - 8*K1/27 + 2*K2 - 3544*K3/2565 + 1859*K4/4104 - 11*K5/40)
        vars1 = varsi + 25*K1/216 + 1408*K3/2565 + 2197*K4/4104 - K5/5
        vars2 = varsi + 16*K1/135 + 6656*K3/12825 + 28561*K4/56430 - 9*K5/50 + 2*K6/55
        R = maximum(abs.(vars2 - vars1))/dt
        s = (epsilon/(2*R))^0.25
        if (R <= epsilon)
            push!(t, ti+dt)
            push!(vars, vars1)
            i += 1
            ti = t[i]
        end
        dt *= s
    end
    vars = transpose(reduce(hcat, vars))   # Optional transpose
    return solObj(t, vars);
end

# Problem parameters
sigma = 10;
beta = 8/3;
rho = 28;
params = paramObj(sigma, beta, rho);
t0, tf = 0.0, 60.0;
x0, y0, z0 = 1.0, 1.0, 1.0;
conds = @SVector [x0,y0,z0];
epsilon = 1e-9;
dtInitial = 0.1;

sol = rkf45(f, params, t0, tf, conds, epsilon, dtInitial)
# @btime  rkf45(f, params, t0, tf, conds, epsilon, dtInitial)

Plots.plot(sol.vars[:,1],sol.vars[:,2])
1 Like