You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
When using a NeuralODE model within Turing.jl, specifically with parameters and data in Float64, the auto-promotion of parameters to Dual type for AD seems not to work correctly. This results in a type conversion error (MethodError related to converting ForwardDiff.Dual to Float64) during the sampling process.
Expected behavior
The expectation is that the parameters of the NeuralODE would automatically promote to the appropriate Dual type for AD, as they do in the usual ODEProblem. This should happen without manual intervention, ensuring compatibility with Turing.jl's sampling process.
Minimal Reproducible Example 👇
You can download trained NeuralODE parameters from here: nn_parameters.zip
using DifferentialEquations, LinearAlgebra, JLD2, Turing
using Lux, Statistics, Plots, ForwardDiff
using Random, DiffEqFlux, ComponentArrays, Optimization, OptimizationOptimisers
C_0 =0.0014
rel_start =0
rel_end =10
N =10
tspan = (rel_start,rel_end)
datasize =50
tsteps =range(tspan[1], tspan[2]; length = datasize)
y_0 =fill(C_0, N)
functiondiffusion!(dOx, Ox, p, t)
C_eq, h, N, D_z = p
h0 = h / N
N =Int(N)
dOx[1] = (D_z * (Ox[2] - Ox[1]) - D_z * (Ox[1] - C_eq)) / h0^2
dOx[2:N-1] = (D_z * (Ox[3:N] - Ox[2:N-1]) - D_z * (Ox[2:N-1] - Ox[1:N-2])) / h0^2
dOx[N] =- D_z * (Ox[N] - Ox[N-1]) / h0^2end
p_ode = [40.6205, 0.0015, 10, 1.9340e-08]
J_prototype =Tridiagonal(fill(1.0, N-1), fill(1.0, N), fill(1.0, N-1))
fun =ODEFunction(diffusion!, jac_prototype=J_prototype)
prob =ODEProblem(fun, y_0, tspan, p_ode)
ode_data =Array(solve(prob,Rodas5();saveat=tsteps))
true_ode_data =mean(ode_data, dims=1)
struct DiffusionLayer <:Lux.AbstractExplicitLayer
C_eq::Float64
h::Float64
D_z::Float64endfunctionDiffusionLayer(C_eq, h, D_z)
returnDiffusionLayer(C_eq, h, D_z)
end# No trainable parameters in this layer
Lux.initialparameters(::AbstractRNG, ::DiffusionLayer) =NamedTuple()
# States are the ODE parameters C_eq, h, D_z, and N
Lux.initialstates(::AbstractRNG, layer::DiffusionLayer) = (C_eq=layer.C_eq, h=layer.h, D_z=layer.D_z)
function (layer::DiffusionLayer)(x, ps, st)
N =length(x)
h0 = st.h / N
D = st.D_z / h0^2
main_diag = [i ==1?-3* D : i == N ?-D :-2* D for i in1:N]
off_diag =fill(D, N -1)
A =Tridiagonal(off_diag, main_diag, off_diag)
modified_x =vcat(x[1] - st.C_eq, x[2:end])
dOx = A * modified_x
return dOx, st
end
nn = Lux.Chain(
DiffusionLayer(40.6205, 0.0015, 1.9340e-08),
Lux.Dense(N, 50, tanh),
Lux.Dense(50, N)
)
rng =Xoshiro(31)
p_nn , st = Lux.setup(rng, nn)
neural_ode =NeuralODE(nn, tspan, Tsit5(), saveat=tsteps)
neural_ode(y_0, p_nn, st)
functionpredict_neuralode(y_0,p_nn,st)
solution =first(neural_ode(y_0, p_nn, st))
returnreduce(hcat,solution.u), solution.retcode
end
nn_param =load_object("nn_parameters.jld2")
@modelfunctioninverse_model(data)
global st
# Prior distribution
D_z ~Normal(1.9340e-08 , 1e-7)
ct = (layer_1 = (C_eq =40.6205, h =0.0015, D_z = D_z), layer_2 =NamedTuple(),
layer_3 =NamedTuple())
predicted, retcode =predict_neuralode(y_0, nn_param.u, ct)
avg_concentration =mean(predicted, dims=1)[:]
if retcode !== ReturnCode.Success
Turing.@addlogprob!-Infreturnnothingend
data[:] ~MvNormal(avg_concentration,1)
returnnothingend
model =inverse_model(true_ode_data)
chain =sample(model, NUTS(), MCMCSerial(), 100, 1; progress=true)
Julia Version 1.9.3
Commit bed2cd540a (2023-08-2414:43 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Windows (x86_64-w64-mingw32)
CPU:8×Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz
WORD_SIZE:64
LIBM: libopenlibm
LLVM: libLLVM-14.0.6 (ORCJIT, skylake)
Threads:8 on 8 virtual cores
Environment:
JULIA_EDITOR = code
JULIA_NUM_THREADS =8
Additional context
In a discussion with @ChrisRackauckas, it was suggested that the issue might be related to auto-promotion not working as expected in the NeuralODE context with Turing.jl. The workaround eltype(D_z).(nn_param.u) was suggested and worked, implying that the auto-promotion should have handled this case.
The text was updated successfully, but these errors were encountered:
Describe the bug 🐞
When using a
NeuralODE
model within Turing.jl, specifically with parameters and data inFloat64
, the auto-promotion of parameters to Dual type for AD seems not to work correctly. This results in a type conversion error (MethodError related to converting ForwardDiff.Dual to Float64) during the sampling process.Expected behavior
The expectation is that the parameters of the
NeuralODE
would automatically promote to the appropriate Dual type for AD, as they do in the usualODEProblem
. This should happen without manual intervention, ensuring compatibility with Turing.jl's sampling process.Minimal Reproducible Example 👇
You can download trained
NeuralODE
parameters from here:nn_parameters.zip
Error & Stacktrace⚠️
Environment (please complete the following information):
using Pkg; Pkg.status()
using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)
versioninfo()
Additional context
In a discussion with @ChrisRackauckas, it was suggested that the issue might be related to auto-promotion not working as expected in the NeuralODE context with Turing.jl. The workaround
eltype(D_z).(nn_param.u)
was suggested and worked, implying that the auto-promotion should have handled this case.The text was updated successfully, but these errors were encountered: