ODE non autonomous interpolation needed

Marta_Pardo_Araujo · January 24, 2025, 5:02pm

I am trying to optimize my code where I integrate a system of non-autonomous ODEs. The problem is the interpolations that I am doing in order to have the value of a parameter at each time step needed by the integrator.
This is the code:

# Code that integrates the Hanski model to obtain the solutions 
# Load pkgs and remove data -----------------------------------------------------
import Pkg
using DifferentialEquations
using DataFrames
using CSV
using Plots
using Shapefile
import GeoDataFrames as GDF
using LinearAlgebra
using ODE
using Interpolations
using DiffEqParamEstim
using SparseArrays
#using RecursiveArrayTools
using Optimization
using Statistics
using Dates


# Constants
N = 5
m_c = 0.001 # probability of mosquito in a car
alp = 1/200 # average natural dispersal distance

# Create fake data
end_ind = 5

# Create fake mobility matrix
eta =  [0.0       7.74415   5.24014     8.57707     5.21319;
        6.9175    0.0       7.23771    45.2212     16.8025;
        6.17319   7.426     0.0        17.7236      5.49988;
        8.40241  44.21     20.1597      0.0      5309.46;
        5.8309   17.5283    5.36841  5226.36        0.0]

#  IC
pop_init = zeros(end_ind)
pop_init[1] = 1

# Fake time series
num_years = 1
num_days = 365*num_years
days = 1:num_days
days_per_year = 365
dates = collect(1:num_days)

#  Fake seasonal
function fake_timeseries!(interpolated_functions, days ,days_per_year,num_days,dates)
    R_Ms = Matrix{Float64}(undef, num_days,end_ind)
    R_Ms[1:end,1] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days))./10
    R_Ms[1:end,2] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days))./30
    R_Ms[1:end,3] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days))./40
    R_Ms[1:end,4] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days))./10
    R_Ms[1:end,5] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days))./100
    
    # Perform interpolation for each location
    for i in 1:end_ind
        # Extract temperature values for the current location
        R_Ms_val = R_Ms[:, i]
        dates_num = 1:size(dates, 1)

        # Perform linear interpolation
        itp = LinearInterpolation(dates_num, R_Ms_val,extrapolation_bc=Flat())
        # Store the interpolated function
        push!(interpolated_functions, itp)

    end
end

# Create an array to store interpolated functions
interpolated_functions = []
fake_timeseries!(interpolated_functions, days ,days_per_year,num_days,dates)
interpolated_functions_mint = []
fake_timeseries!(interpolated_functions_mint, days ,days_per_year,num_days,dates)

# Parameters
R_M_vec = zeros(N)
min_temp = zeros(N) 
p = [0.001,0.000000001,5, -0.5,-15] 

#  Non autonomous sigmoidal model ---------------------------------------------------------
# Withouth interpolation
function fun_na1!(du, u, p, t)
  R_M_vec .= 1#[interpolated_functions[i](t) for i in 1:N] # Time series RM
  min_temp .= 1#[interpolated_functions_mint[i](t) for i in 1:N] # Time series min temperature
  du.= R_M_vec.*(p[1]./(1.0 .+ exp.(-p[2]*m_c.*eta.+p[3]))*u).* (1 .- u) - exp.(p[4].+(p[5].*min_temp)).* u
end

# With just one interpolation
function fun_na2!(du, u, p, t)
  R_M_vec .= [interpolated_functions[i](t) for i in 1:N] # Time series RM
  min_temp .= 1#[interpolated_functions_mint[i](t) for i in 1:N] # Time series min temperature
  du.= R_M_vec.*(p[1]./(1.0 .+ exp.(-p[2]*m_c.*eta.+p[3]))*u).* (1 .- u) - exp.(p[4].+(p[5].*min_temp)).* u
end
# With both interpolations
function fun_na3!(du, u, p, t)
  R_M_vec .= [interpolated_functions[i](t) for i in 1:N] # Time series RM
  min_temp .= [interpolated_functions_mint[i](t) for i in 1:N] # Time series min temperature
  du.= R_M_vec.*(p[1]./(1.0 .+ exp.(-p[2]*m_c.*eta.+p[3]))*u).* (1 .- u) - exp.(p[4].+(p[5].*min_temp)).* u
end

p1 = [0.001,0.000000001,5, -0.5,-15, R_M_vec, min_temp] 

#  Non autonomous sigmoidal model ---------------------------------------------------------
function fun_na1!(du, u, p1, t)
  p1[6] .= [interpolated_functions[i](t) for i in 1:N] # Time series RM
  p1[7] .= [interpolated_functions_mint[i](t) for i in 1:N] # Time series min temperature
  du.= R_M_vec.*(p1[1]./(1.0 .+ exp.(-p1[2]*m_c.*eta.+p1[3]))*u).* (1 .- u) - exp.(p1[4].+(p1[5].*min_temp)).* u
end

# Set parameters
t0=0.0f0
tf=365.0f0
tspan = (t0, tf)
t_vect=1:tf
u0 = pop_init

# Create the ode model to solve it ---------------------------------------------------------
# Pkg.add("BenchmarkTools")
using BenchmarkTools
prob1 = DifferentialEquations.ODEProblem(fun_na1!, u0, tspan, p)
prob2 = DifferentialEquations.ODEProblem(fun_na2!, u0, tspan, p)
prob3 = DifferentialEquations.ODEProblem(fun_na3!, u0, tspan, p)
prob4 = DifferentialEquations.ODEProblem(fun_na3!, u0, tspan, p)

@benchmark sol = solve(prob1,alg_hints=[:stiff])
@benchmark sol = solve(prob2,alg_hints=[:stiff])
@benchmark sol = solve(prob3,alg_hints=[:stiff])
@benchmark sol = solve(prob4,alg_hints=[:stiff])
# plot(sol)

This is the output:

The memory allocated is multiply by an order of magnitude aprox. when introducing the interpolation inside the ode function. Is there any other way to do this which is more efficient? I have tried to used all the recommendations given in the differentialequations.jl youtube tutorial

stevengj · January 25, 2025, 3:14am

Note that you are allocating a new array every time this line executes, and then assigning it to p1[6]. You could just write a loop to get rid of the allocation, or use e.g. map!(i -> interpolated_functions[i](t), p1[6], 1:N).

Similarly for your other functions. This is one of the reaons why you have so many allocations.

(Or just get rid of the arrays of interpolated values completely, and instead write into your du array directly. This is one of those cases where you are probably making life hard for yourself by trying so hard to write “vectorized” operations. Just write a loop. Loops are fast in Julia.)

You’re also killing performance by (a) using untyped global variables, and (b) using an abstractly typed container [] (the same as Any[]). See the Julia performance tips. This will also cause a lot of allocations because of the dynamic dispatch it forces.

Using piecewise-linear interpolation is probably going to kill the performance of the ODE algorithm, because it’s not smooth — fast ODE solvers try to use high-order polynomials to converge quickly, but whenever your piecewise-linear function has a kink this will backfire.

If you can use a smooth interpolant or fit it should converge a lot faster. Otherwise you might have to mess around with callbacks to let the ODE solver know about the kinks.

ChrisRackauckas · January 25, 2025, 10:20am

All of what Steven said, and then the next big performance thing would be to switch to DataInterpolations (preferably something with a continuous derivative like a cubic spline). DataInterpolations has some caching features so that repeated “nearby” interpolations should be about an order of magnitude faster, which is a special case that is common when used with ODE integration that the library is prepared for.

stevengj · January 26, 2025, 5:50pm

What’s the best way to tell the ODE solver where the knots of the spline are (so that it doesn’t get confused by the discontinuities in the third derivative there)?

ChrisRackauckas · January 26, 2025, 5:52pm

d_discontinuities or tstops. Though it depends on how frequent they are: if they are so frequent then stopping the solver at every discontinuity could acutally hurt, and instead using a lower order integrator that doesn’t assume the existence of higher derivatives may make sense.

Marta_Pardo_Araujo · January 27, 2025, 10:53am

Thanks for the help.
I have followed your advice and I have done the following:
. Use DataInterpolation.jl with cubicSpline
. Define global variables
. Do a loop instead of the vectorized form

The code looks now:

# Code that integrates the Hanski model to obtain the solutions 
# With input: (from code input_Hanski.R)
#     . Flow matrix
#     . R_M time series
#     . t_min time series
# Load pkgs  -----------------------------------------------------
import Pkg
using DifferentialEquations
using DataFrames
using CSV
using Plots
using Shapefile
import GeoDataFrames as GDF
using LinearAlgebra
using ODE
using DataInterpolations  # Ensure the package is imported
using DiffEqParamEstim
using SparseArrays
#using RecursiveArrayTools
using Optimization
using Statistics
using Dates

# Constants
const N = 5
const m_c = 0.001 # probability of mosquito in a car

# Create fake data
const end_ind = 5

# Create fake mobility matrix
const eta =  [0.0       7.74415   5.24014     8.57707     5.21319;
        6.9175    0.0       7.23771    45.2212     16.8025;
        6.17319   7.426     0.0        17.7236      5.49988;
        8.40241  44.21     20.1597      0.0      5309.46;
        5.8309   17.5283    5.36841  5226.36        0.0]

# Choose number of patches and IC
pop_init = zeros(end_ind)
pop_init[1] = 1

# Fake time series
num_years = 1
num_days = 365*num_years
days = 1:num_days
days_per_year = 365
dates = collect(1:num_days)

function fake_timeseries!(interpolated_functions, days, days_per_year, num_days, dates, end_ind)
    # Create a matrix to store fake seasonal data
    R_Ms = Matrix{Float64}(undef, num_days, end_ind)

    # Populate the seasonal data with different amplitudes
    R_Ms[:, 1] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days)) ./ 10
    R_Ms[:, 2] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days)) ./ 30
    R_Ms[:, 3] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days)) ./ 40
    R_Ms[:, 4] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days)) ./ 10
    R_Ms[:, 5] .= (30 .+ 10 * sin.(2 * π * days / days_per_year) + 2. * randn(num_days)) ./ 100

    # Create interpolations for each location
    for i in 1:end_ind
        # Extract seasonal data for the current location
        R_Ms_val = R_Ms[:, i]
        dates_num = collect(1:num_days)  # Convert dates to numeric indices

        # Perform smooth spline interpolation
        # itp = interpolate((dates_num,), R_Ms_val, Spline(Quadratic(Line())))  # Quadratic spline with linear extrapolation
        itp = DataInterpolations.CubicSpline(R_Ms_val, dates_num, extrapolate = true)
        # Store the interpolated function
         interpolated_functions[i] = itp
    end
    # # Test
    # new_x = collect(1:num_days)
    # p = plot(new_x, interpolated_functions[4].(new_x), label = "Cubic Spline")
    # plot!(p,new_x, R_Ms[:,4], label = "Data")
end

# Interpolate RM and min temp. First create a global variable
InterpType = DataInterpolations.CubicSpline{Vector{Float64}, Vector{Int64}, Vector{Int64}, Vector{Float64}, Float64}
const interpolated_functions = Vector{InterpType}(undef, N)
fake_timeseries!(interpolated_functions, days, days_per_year, num_days, dates, 5)
# plot(interpolated_functions)
const interpolated_functions_mint = Vector{InterpType}(undef, N)
fake_timeseries!(interpolated_functions_mint, days, days_per_year, num_days, dates, 5)

# Parameters
const R_M_vec = Vector{Float64}(undef, N)
const min_temp = Vector{Float64}(undef, N)
p = [0.001,0.000000001,5, -0.5,-15] 

function fun_na1!(du, u, p, t)
  R_M_vec .= [interpolated_functions[i](t) for i in 1:N] # Time series RM
  min_temp .= [interpolated_functions_mint[i](t) for i in 1:N] # Time series min temperature
  du.= R_M_vec.*(p[1]./(1.0 .+ exp.(-p[2]*m_c.*eta.+p[3]))*u).* (1 .- u) -
   exp.(p[4].+(p[5].*min_temp)).* u
end

#  Non autonomous sigmoidal model ---------------------------------------------------------
function fun_na2!(du, u, p, t)
  du.= [interpolated_functions[i](t) for i in 1:N].*(p[1]./(1.0 .+
   exp.(-p[2]*m_c.*eta.+p[3]))*u).* (1 .- u) -
   exp.(p[4].+(p[5].*[interpolated_functions_mint[i](t) for i in 1:N])).* u
end

# const du = Vector{Float64}(undef, N)
function fun_na3!(du, u, p, t)
  for i in 1:N
      du[i] = interpolated_functions[i](t) * sum(p[1]./ (1.0 .+
       exp.(-p[2].*m_c.* eta[i,:] .+ p[3])) .* u) * (1 - u[i]) -
              exp(p[4] + p[5] * interpolated_functions_mint[i](t)) * u[i]
  end
end

# Set parameters
t0=0.0f0
tf=365.0f0
tspan = (t0, tf)
t_vect=1:tf

# Create the ode model to solve it ---------------------------------------------------------
# Pkg.add("BenchmarkTools")
prob1 = DifferentialEquations.ODEProblem(fun_na1!, pop_init, tspan, p)
prob2 = DifferentialEquations.ODEProblem(fun_na2!, pop_init, tspan, p)
prob3 = DifferentialEquations.ODEProblem(fun_na3!, pop_init, tspan, p)

using BenchmarkTools
@benchmark sol = solve(prob1,alg_hints=[:stiff])
@benchmark sol = solve(prob2,alg_hints=[:stiff])
@benchmark sol = solve(prob3,alg_hints=[:stiff])

# plot(sol)

Now it is much faster and it allocates almost half of the memory than before. How ever the model with the non vectorized model, model 3, with a loop it is the faster but it allocates more memory. Maybe I did not understand that part properly and it is possible to optimize the code. Also should I choose the faster or the faster? This is a test the system or ODEs is N =300. So I am very concern about both. This is the benchmark output:

ChrisRackauckas · January 27, 2025, 12:10pm

If you make a temporary variable the third can easily have 0 allocations.

Marta_Pardo_Araujo · January 27, 2025, 12:27pm

Thanks for the answer. Do you mean du[i].= ?

stevengj · January 27, 2025, 12:29pm

Marta_Pardo_Araujo:

function fun_na3!(du, u, p, t)
  for i in 1:N
      du[i] = interpolated_functions[i](t) * sum(p[1]./ (1.0 .+
       exp.(-p[2].*m_c.* eta[i,:] .+ p[3])) .* u) * (1 - u[i]) -
              exp(p[4] + p[5] * interpolated_functions_mint[i](t)) * u[i]
  end
end

On every loop iteration you are allocating a temporary array to pass to sum. You can replace the sum call with:

sum(j -> p[1] / (1 + exp(-p[2] * m_c * eta[i,j] + p[3])) * u[j], 1:N)

to get rid of this.

(You might also want to transpose your eta matrix, since accessing it as eta[i,j] is in the wrong order for locality.)

Marta_Pardo_Araujo · January 27, 2025, 1:04pm

Thanks a lot. Now with your advice, it got better. The best one is the sum(j->) than transposing:

function fun_na3!(du, u, p, t)
  for i in 1:N
      du[i] = interpolated_functions[i](t) * sum(p[1]./ (1.0 .+
       exp.(-p[2].*m_c.* eta1[:,i] .+ p[3])) .* u) * (1 - u[i]) -
              exp(p[4] + p[5] * interpolated_functions_mint[i](t)) * u[i]
  end
end

function fun_na4!(du, u, p, t)
  for i in 1:N
      du[i]= interpolated_functions[i](t) * sum(j -> p[1] / (1 + exp(-p[2] * m_c * eta1[j,i] + p[3])) * u[j], 1:N) * (1 - u[i]) -
              exp(p[4] + p[5] * interpolated_functions_mint[i](t)) * u[i]
  end
end

The allocation memmory and the time is much better:

The part of the temporary variable that will make 0 allocations I did not get it

stevengj · January 27, 2025, 1:14pm

I would benchmark just your rhs function, not the solve. Otherwise you can’t distinguish allocations coming from your code (your rhs) rather than allocations coming from insider the ODE solver itself.

Marta_Pardo_Araujo · January 27, 2025, 2:01pm

Like this you mean?

function fun_na3!(du, u, p, t)
  for i in 1:N
      du[i] = interpolated_functions[i](t) * sum(p[1]./ (1.0 .+
       exp.(-p[2].*m_c.* eta1[:,i] .+ p[3])) .* u) * (1 - u[i]) -
              exp(p[4] + p[5] * interpolated_functions_mint[i](t)) * u[i]
  end
end

function fun_na4!(du, u, p, t)
  for i in 1:N
      du[i]= interpolated_functions[i](t) * sum(j -> p[1] / (1 + exp(-p[2] * m_c * eta1[j,i] + p[3])) * u[j], 1:N) * (1 - u[i]) -
              exp(p[4] + p[5] * interpolated_functions_mint[i](t)) * u[i]
  end
end


# Set parameters
t0=0.0f0
tf=365.0f0
tspan = (t0, tf)
t_vect=1:tf
u0 = pop_init

# Create the ode model to solve it ---------------------------------------------------------
# Pkg.add("BenchmarkTools")
prob3 = DifferentialEquations.ODEProblem(fun_na3!, u0, tspan, p)
prob4 = DifferentialEquations.ODEProblem(fun_na4!, u0, tspan, p)

@benchmark fun_na3!(Vector{Float64}(undef,N), u0, p, 1)
@benchmark fun_na4!(Vector{Float64}(undef,N), u0, p, 1)

The second its much faster now. But still allocates memory.

stevengj · January 27, 2025, 2:33pm

You need to interpolate the array construction (otherwise you are timing it too) and interpolate globals (to avoid dynamic dispatch):

@btime fun_na4!($(Vector{Float64}(undef,N)), $u0, $p, 1);

(Similarly for @benchmark, but I typically use @btime rather than @benchmark; the latter provides more information than is normally needed.)

Marta_Pardo_Araujo · January 27, 2025, 2:42pm

Ok, now it is allocating I have used @btime. As you mentioned and it gives me this:
645.024 ns (0 allocations: 0 bytes)

Perfect! Thanks a lot. However I am curious about what Chris R. meant when he said producing a temporary array…

stevengj · January 27, 2025, 6:19pm

I think he meant eliminating the temporary array?

Marta_Pardo_Araujo · January 27, 2025, 6:36pm

Ok, that make sense.

Topic		Replies	Views
Problem not finding initial step in Bayesian Differential Equations Turing General Usage turing , differentialequation	6	155	July 9, 2024
Problem in solve() DifferentialEquations. None algorithm define Modelling & Simulations question , ode , differentialequation	11	197	April 25, 2024
Interpolation problem Modelling & Simulations diffeq , interpolationsjl	2	263	March 23, 2023
Solving ODEs in a loop with DifferentialEquations.jl Performance question , diffeq	10	2809	April 21, 2021
Integrating ODE system over gridded data Modelling & Simulations	1	321	November 9, 2020

ODE non autonomous interpolation needed

Related topics