Non-constant types in out-of-place ODE solve with multshooting?

New to Julia
1

Hi, I’m trying to use multiple shooting for my analysis. I’ve been trying to reproduce the simple script shown of the Julia SciML documentation and I keep running into an error:
ERROR: Detected non-constant types in an out-of-place ODE solve, i.e. for
du = f(u,p,t) we see typeof(du) !== typeof(u/t). This is not
supported by OrdinaryDiffEq.jl’s solvers. Please either make f
type-constant (i.e. typeof(du) === typeof(u/t)) or use the mutating
in-place form f(du,u,p,t) (which is type-constant by construction).

Note that one common case for this is when computing with GPUs, using
Float32 for u0 and Float64 for tspan. To correct this, ensure
that the element type of tspan matches the preferred compute type,
for example ODEProblem(f,0f0,(0f0,1f0)) for Float32-based time.

typeof(du) = Vector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{false, SciMLBase.AutoSpecialize, var"#53#54", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float32}, Float32, 12}}I am running Julia version 1.10.8 on a 2024 MacBook with Apple Silicon. My script modifies the code from the documentation to ensure all floating point values are Float64.

I keep think this must a bug in either… 1.10.8 or on the Apple Silicon hardware… or my packages are misaligned. Any help would be appreciated.

using ComponentArrays
using Lux
using DiffEqFlux
using Optimization
using OptimizationPolyalgorithms
using OrdinaryDiffEq
using Plots
using DiffEqFlux: group_ranges
using Random

rng = Random.default_rng() #Xoshiro(123)

datasize = 30
u0 = [2.0, 0.0]
tmin = 0.0
tmax = 5.0
tspan = (0.0, 5.0)
tsteps = range(tspan[1], tspan[2]; length = datasize)

function trueODEfunc!(du, u, p, t)
true_A = [-0.1 2.0 ; -2.0 -0.1]
#du .= ((u .^3)’ * true_A)’
du .= ((u.^3)‘true_A)’
#du .= true_A * (u.^3)
end

function trueODEfunc2!(du, u, p, t)
du[1] = -0.1 * u[1]^3 - 2.0 * u[2]^3
du[2] = 2.0 * u[1]^3 - 0.1 * u[2]^3
end
prob_trueode = ODEProblem(trueODEfunc!, u0, tspan)
ode_data = Array(solve(prob_trueode, Tsit5(); saveat = tsteps))

#define the neural network
nn = Lux.Chain(x → x .^3, Lux.Dense(2, 16, tanh), Lux.Dense(16, 2))
p_init, st = Lux.setup(rng, nn)

neuralode = NeuralODE(nn, tspan, Tsit5(); saveat = tsteps)
prob_node = ODEProblem((u, p, t) → nn(u, p, st)[1], u0, tspan, ComponentArray(p_init))

function plot_multiple_shoot(plt, preds, group_size)
step = group_size - 1
ranges = group_ranges(datasize, group_size)

for (i, ig) in enumerate(range)
    plot!(plt, tsteps[rg], preds[i][1,:]; markershape = :circle, label = "Group $(i)")
end

end

anim = Plots.Animation()
iter = 0
callback = function (p, l, preds; doplot = true)
display(l)
global iter
iter += 1
if doplot && iter % 1 == 0
plot the original data
plt = scatter(tsteps, ode_data[1,:]; label = “Data”)
plot the different predictions for individual shoot
plot_multiple_shoot(plt, preds, group_size)
frame(anim)
display(plot(plt))
end
return false
end

#define parameters for multiple multiple shooting
group_size = 3
continuity_term = 200

function loss_function(data, pred)
return sum(abs2, data - pred)
end

ps = ComponentArray(p_init)
pd, pax = getdata(ps), getaxes(ps)

function loss_multiple_shooting(p)
ps = ComponentArray(p, pax)
return multiple_shoot(
ps,
ode_data,
tsteps,
prob_node,
loss_function,
Tsit5(),
group_size; continuity_term
)
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p) → loss_multiple_shooting(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pd)
res_ms = Optimization.solve(optprob, PolyOpt(); callback = callback)
gif(anim, “multiple_shooting.gif”; fps = 15)

2

Make this ComponentArray{Float64}(p_init)? You may be mixing Float32 and Float64

3

No change. Could something be mismatched with my package versions?

image
image670×356 42.8 KB

I’m working with Raj Dandekar through his SciML bootcamp. I understand you were his PhD advisor. He told me he’s reproduced the error on his (x86) MacBook Air running Julia 1.6.7. He’s puzzled as well

4

This shot shows the code as posted online.

image
image950×642 83.2 KB

It’s mixing Float32 values in u0 and tspan but I would think the coefficients in true_A would be read as Float64. I tried to correct for that by making everything Float64 in my code but… problem still occurs

5

Show the code you ran, and use triple backticks so it’s easier to read.

6

Hi I think see what’s triggering the “non-constant types out-of-place error”.
It happens if I define u0 and tspan as Float64 by default (as in u0 = [2.0, 0.0]…)

If I define u0 and tspan as Float32
u0 = Float32[2.0, 0.0]
and tspan = (0.0f0, 5.0f0)

the out-of-place ODE solve error isn’t thrown.

using ComponentArrays
using Lux
using DiffEqFlux
using Optimization
using OptimizationPolyalgorithms
using OrdinaryDiffEq
using Plots
using DiffEqFlux: group_ranges
using Random

rng = Random.default_rng() #Xoshiro(123)

datasize = 30
u0 = [2.0, 0.0] #defining u0  and tspan as Float64 by default
tspan = (0.0, 5.0)

#u0 = Float32[2.0, 0.0] #defining u0 and tspan as Float32
#tspan = (0.0f0, 5.0f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)

function trueODEfunc!(du, u, p, t)
    true_A = [-0.1 2.0 ; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end

prob_trueode = ODEProblem(trueODEfunc!, u0, tspan)
ode_data = Array(solve(prob_trueode, Tsit5(); saveat = tsteps))

#define the neural network
nn = Lux.Chain(x -> x .^3, Lux.Dense(2, 16, tanh), Lux.Dense(16, 2))
p_init, st = Lux.setup(rng, nn)

neuralode = NeuralODE(nn, tspan, Tsit5(); saveat = tsteps)
prob_node = ODEProblem((u, p, t) -> nn(u, p, st)[1], u0, tspan, ComponentArray{Float64}(p_init))

function plot_multiple_shoot(plt, preds, group_size)
    step = group_size - 1
    ranges = group_ranges(datasize, group_size)

    for (i, ig) in enumerate(range)
        plot!(plt, tsteps[rg], preds[i][1,:]; markershape = :circle, label = "Group $(i)")
    end
end

anim = Plots.Animation()
iter = 0
callback = function (p, l, preds; doplot = true)
    display(l)
    global iter 
    iter += 1
    if doplot && iter % 1 == 0
        #plot the original data
        plt = scatter(tsteps, ode_data[1,:]; label = "Data")
        #plot the different predictions for individual shoot
        plot_multiple_shoot(plt, preds, group_size)
        frame(anim)
        display(plot(plt))
    end
    return false
end

#define parameters for multiple multiple shooting 
group_size = 3
continuity_term = 200

function loss_function(data, pred)
    return sum(abs2, data - pred)
end

ps = ComponentArray(p_init)
pd, pax = getdata(ps), getaxes(ps)

function loss_multiple_shooting(p)
    ps = ComponentArray(p, pax)
    return multiple_shoot(
                     ps,
                     ode_data,
                      tsteps,
                      prob_node,
                      loss_function,
                      Tsit5(),
                      group_size; continuity_term
    )
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p) -> loss_multiple_shooting(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pd)
res_ms = Optimization.solve(optprob, PolyOpt(); callback = callback)
gif(anim, "multiple_shooting.gif"; fps = 15)

However, that means the script runs up to the second to last line (89) where
res_ms = Optimization.solve(optprob, PolyOpt); callback = callback)
Triggers:

image
image1084×550 71.8 KB

I had though I could solve this problem by keeping all floating point values as float64.

So…
(a) should declaring all floating point variables Float64 have triggered the non-constant types error?
(b) Any advice for resolving the issue with PolyOpt()?

I appreciate any assistance

7

That’s because you’re mixing Float64 and Float32 values, which will upconvert to Float64. Make everything Float64 or Float32 if you want to make things easier. Neural network parameters are Float32 by default unless you say otherwise.

8

I understand that mixing Float32 and Float64 leads to an error with the optimizer. The code on the SciML site explicitly showed u0 and tspan declared as Float32 with the coefficients in the trueODEfunc defaulting to Float64. l I attempted to make all floating point values Float64. This resulted in the “non-constant types” error.

On Raj’s advice I also tried making all floating point values Float32:

using ComponentArrays
using Lux
using DiffEqFlux
using Optimization
using OptimizationPolyalgorithms
using OrdinaryDiffEq
using Plots
using DiffEqFlux: group_ranges
using Random

rng = Random.default_rng() #Xoshiro(123)

#Define initial conditions and time steps
datasize = 30
u0 = Float32[2.0,0.0]
tspan = (0.0f0, 5.0f0)
tsteps = range(tmin, tmax; length = datasize) #will be Float32

#Get the data
function trueODEfunc(du, u, p, t)
    #first = Float32(-0.1)
    #second = Float32(2.0)
    #third = Float32(-2.0)
    #fourth = Float32(-0.1)
    #true_A = [first second ; third fourth]
    true_A = Float32[-0.1 2.0; -2.0 -0.1]
    du .= ((u .^3)'true_A)'
    
end

prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
ode_data = Array(solve(prob_trueode, Tsit5(); saveat = tsteps))

#define the neural network
nn = Lux.Chain(x -> x .^3, Lux.Dense(2 => 16, tanh), Lux.Dense(16 => 2))
#p_init, st = Lux.setup(rng, nn)
p64, st = Lux.setup(rng, nn)
p_init = Float32.(p64) #broadcast convert... THIS TRIGGERS AN ERROR

neuralode = NeuralODE(nn, tspan, Tsit5(); saveat = tsteps)
prob_node = ODEProblem((u, p, t) -> nn(u, p, st)[1], u0, tspan, ComponentArray(p_init))

function plot_multiple_shoot(plt, preds, group_size)
    step = group_size - 1
    ranges = group_ranges(datasize, group_size)

    for (i, ig) in enumerate(range)
        plot!(plt, tsteps[rg], preds[i][1,:]; markershape = :circle, label = "Group $(i)")
    end
end

anim = Plots.Animation()
iter = 0
callback = function (p, l, preds; doplot = true)
    display(l)
    global iter 
    iter += 1
    if doplot && iter % 1 == 0
        #plot the original data
        plt = scatter(tsteps, ode_data[1,:]; label = "Data")
        #plot the different predictions for individual shoot
        plot_multiple_shoot(plt, preds, group_size)
        frame(anim)
        display(plot(plt))
    end
    return false
end

#define parameters for multiple multiple shooting 
group_size = 3
continuity_term = 200

function loss_function(data, pred)
    return sum(abs2, data - pred)
end

ps = ComponentArray(p_init)
pd, pax = getdata(ps), getaxes(ps)

function loss_multiple_shooting(p)
    ps = ComponentArray(p, pax)
    return multiple_shoot(
                     ps,
                     ode_data,
                      tsteps,
                      prob_node,
                      loss_function,
                      Tsit5(),
                      group_size; continuity_term
    )
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p) -> loss_multiple_shooting(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pd)
res_ms = Optimization.solve(optprob, PolyOpt(); callback = callback)
gif(anim, "multiple_shooting.gif"; fps = 15)

This does NOT trigger the non-constant types error but it does throw an error when I try to cast p_64 to Float32 as in:

image
image1780×492 70.4 KB

Are these errors expected? I understand that all floating point values need to be consistently Float64 or Float32

9

Well the error message should be rather clear, is it not? You cannot convert the parameters to Float32 by doing Float32.(p64) because you cannot broadcast over a NamedTuple.

However, I think it might just be unnecessary as Chris said:

So try removing this conversion.

10

Since it’s nested named tuples, that’s the reason for the ComponentArray. p_init = ComponentArray{Float32}(p64) etc.