# Minimum Working Example (MWE) showing error in Universal Differential Equation (UDE) implementation

**URL:** <https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366>\
**Category:** Modelling & Simulations\
**Tags:** sciml, autodiff, differentialequation, lux\
**Created:** [January 29, 2025, 11:20pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366 "2025-01-29T23:20:45Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![Ashima\_Kalathingal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ashima_kalathingal/32/212264_2.png) [@Ashima\_Kalathingal](https://discourse.julialang.org/u/Ashima_Kalathingal)\
**Post date:** [January 29, 2025, 11:20pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/1 "2025-01-29T23:20:45Z")

</div>

The following code gives a Minimum Working Example for UDE which I wrote. But unfortunately it is showing error. When I run the code in VS Code the terminal crashes.

```julia
using OrdinaryDiffEq , SciMLSensitivity ,Optimization, OptimizationOptimisers,OptimizationOptimJL, LineSearches
using Statistics
using StableRNGs, Lux, Zygote , Plots , ComponentArrays

rng = StableRNG(11)

# Generating training data
function actualODE!(du,u,p,t,T∞,I)
    
    Cbat = 5*3600 
    du[1] = -I/Cbat

    C₁ = -0.00153 # Unit is s-1
    C₂ = 0.020306 # Unit is K/J

    R0 = 0.03 # Resistance set a 30mohm

    Qgen =(I^2)*R0

    du[2] = (C₁*(u[2]-T∞)) + (C₂*Qgen)

end

t1 = collect(0:1:3400)
T∞1,I1 = 298.15,5

actualODE1!(du,u,p,t) = actualODE!(du,u,p,t,T∞1,I1)

prob = ODEProblem(actualODE1!,[1.0,T∞1],(t1[1],t1[end]))
solution = solve(prob,Tsit5(),saveat = t1)
X = Array(solution)
T1 = X[2,:]
# Plotting the results
plot(solution[2,:],color = :red,label = ["True Data" nothing])

# Defining the neural network
const U = Lux.Chain(Lux.Dense(3,20,tanh),Lux.Dense(20,20,tanh),Lux.Dense(20,1))
_para,st = Lux.setup(rng,U)
const _st = st

function NODE_model!(du,u,p,t,T∞,I)

    Cbat = 5*3600
    du[1] = -I/Cbat

    C₁ = -0.00153
    C₂ = 0.020306

    G = I*(U([u[1],u[2],I],p,_st)[1][1])

    du[2] = (C₁*(u[2]-T∞)) + (C₂*G)

end

NODE_model1!(du,u,p,t) = NODE_model!(du,u,p,t,T∞1,I1)
prob1 = ODEProblem(NODE_model1!,[1.0,T∞1],(t1[1],t1[end]),_para)

function loss(θ)
    _prob1 = remake(prob1,p=θ)
    _sol = Array(solve(_prob1,Tsit5(),saveat = t1))
    loss1 = mean(abs2,T1.-_sol[2,:])
    return loss1
end

losses = Float64[]

callback = function(state,l)
    push!(losses,l)
    println("RMSE Loss at iteration $(length(losses)) is $sqrt(l)")
    
    return false

end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p) -> loss(x),adtype)
optprob = Optimization.OptimizationProblem(optf,ComponentVector{Float64}(_para))

res1 = Optimization.solve(optprob, OptimizationOptimisers.Adam(),callback = callback,maxiters = 500)

```

Before crashing a warning about EnzymeVJP is shown there after a lot of messages come rapidly and terminal crashes. Due to the crashing, I couldn’t copy the messages. But I took some screenshots which I am attaching.

 ![error 1](https://global.discourse-cdn.com/julialang/original/3X/1/b/1b2c0a159cb1f12fd44d66157bd4161731ea89ad.jpeg)  
 ![error 2](https://global.discourse-cdn.com/julialang/original/3X/0/c/0c9a9faecd85f722192f4640bffe1f4503ece9d3.png)

Does anybody know why this happens? Is the same issue occuring in your system?

---

<div class="post-metadata">

**Author:** ![Ashima\_Kalathingal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ashima_kalathingal/32/212264_2.png) [@Ashima\_Kalathingal](https://discourse.julialang.org/u/Ashima_Kalathingal)\
**Post date:** [January 30, 2025, 3:16pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/2 "2025-01-30T15:16:40Z")

</div>

Update on the code. I was able to run it when I changed the solver settings. The loss function was modified to the following

```julia
function loss(θ)
    _prob1 = remake(prob1,p=θ)
    _sol = Array(solve(_prob1,Tsit5(),saveat = t1,abstol = 1e-6, reltol = 1e-6,sensealg = QuadratureAdjoint(autojacvec = ReverseDiffVJP(true))))
    loss1 = mean(abs2,T1.-_sol[2,:])
    return loss1
end

```

I tried using the BFGS algorithm after ADAM using the following lines of code

```julia
optprob2 = Optimization.OptimizationProblem(optf,res1.u)
res2 = Optimization.solve(optprob2,BFGS(),callback=callback,maxiters=50)

```

The loss function is reducing and the parameters are updating but it is showing `retcode:Failure` . Does anyone know why this happens?

Please let me know if you have any idea about the issues I posted. Any help would be much appreciated

---

<div class="post-metadata">

**Author:** ![apo383](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/apo383/32/11272_2.png) [@apo383](https://discourse.julialang.org/u/apo383)\
**Post date:** [January 30, 2025, 8:03pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/3 "2025-01-30T20:03:16Z")

</div>

I have no idea what is happening, but a few comments. First, explain what you’re trying to do. I’m guessing you want a neural network to generate the forcing function for an ODE, but I’m too lazy to step through everything and figure it out.

Second, your MWE isn’t all that minimal. It would be easier for others if you simplify further, e.g. get rid of numerical parameters, use first-order ODE, one where you know the optimum. Then maybe you can produce by hand the NODE that works, to use as a reference for the optimization. Also before you use NODE, why not just try to optimize a constant forcing function, so you’re testing Enzyme without an ANN.

Third, I personally wouldn’t use a callback to log data. The ODE solvers already provide ways to log the info you need, which lets them manage allocations. Here the `push!` allocates, and it will store every query the solver attempts. Not sure if it’s a problem, but I wouldn’t recommend it.

When the terminal crashes, I suspect running out of memory and other system-like problems. Failed retcode suggests the optimization didn’t converge, which is why it would help to make the MWE more minimal.

Sorry I can’t say more, just commenting as a naive non-expert.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [February 7, 2025, 2:18pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/4 "2025-02-07T14:18:28Z")

</div>

Are you using a standard Julia implementation from [Download Julia](http://julialang.org/downloads) or juliaup?

---

<div class="post-metadata">

**Author:** ![Ashima\_Kalathingal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ashima_kalathingal/32/212264_2.png) [@Ashima\_Kalathingal](https://discourse.julialang.org/u/Ashima_Kalathingal)\
**Post date:** [February 8, 2025, 12:53pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/5 "2025-02-08T12:53:41Z")

</div>

Yes I am. I dowloaded Julia using the following command

```julia
winget install julia -s msstore

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [February 8, 2025, 3:21pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/6 "2025-02-08T15:21:38Z")

</div>

Show ]st?

---

<div class="post-metadata">

**Author:** ![Ashima\_Kalathingal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ashima_kalathingal/32/212264_2.png) [@Ashima\_Kalathingal](https://discourse.julialang.org/u/Ashima_Kalathingal)\
**Post date:** [February 8, 2025, 8:01pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/7 "2025-02-08T20:01:02Z")

</div>

```julia
(@v1.11) pkg> st
Status `C:\Users\Kalath_A\.julia\environments\v1.11\Project.toml` (empty project)

```

This is the output. I created another environment in another folder to run my code.This is the status of the package in the directory I am working.

```julia
julia> cd("E:/PhD Ashima/Neural ODE/Julia/env_NODE")

julia> using Pkg

julia> Pkg.activate(".")
  Activating project at `E:\PhD Ashima\Neural ODE\Julia\env_NODE`

(env_NODE) pkg> st
Status `E:\PhD Ashima\Neural ODE\Julia\env_NODE\Project.toml`
⌃ [b0b7db55] ComponentArrays v0.15.22
  [0c46a032] DifferentialEquations v7.15.0
  [033835bb] JLD2 v0.5.11
  [d3d80556] LineSearches v7.3.0
  [b2108857] Lux v1.6.0
  [7f7a1694] Optimization v4.1.0
  [36348300] OptimizationOptimJL v0.4.1
  [42dfb2eb] OptimizationOptimisers v0.3.7
  [1dea7af3] OrdinaryDiffEq v6.90.1
  [91a5bcdd] Plots v1.40.9
  [1ed8b502] SciMLSensitivity v7.72.0
  [860ef19b] StableRNGs v1.0.2
  [10745b16] Statistics v1.11.1
⌅ [e88e6eb3] Zygote v0.6.75
  [2a0f44e3] Base64 v1.11.0
  [9e88b42a] Serialization v1.11.0
Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated`

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [February 9, 2025, 12:32am UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/8 "2025-02-09T00:32:38Z")

</div>

If you try to up componentarrays what happens?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [February 11, 2025, 1:24pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/9 "2025-02-11T13:24:21Z")

</div>

Okay I see that this is a segfault introduced in the latest Julia release:

> <https://github.com/SciML/FunctionProperties.jl/issues/18>
>
> Segfaults on v1.11:
> 
> \`\`\`julia
> using OrdinaryDiffEq , SciMLSensitivity ,Optimizat…ion, OptimizationOptimisers,OptimizationOptimJL, LineSearches
> using Statistics
> using StableRNGs, Lux, Zygote , Plots , ComponentArrays
> 
> rng = StableRNG(11)
> 
> \# Generating training data
> function actualODE!(du,u,p,t,T∞,I)
>     
> Cbat = 5\*3600 
> du\[1\] = -I/Cbat
> 
> C₁ = -0.00153 # Unit is s-1
> C₂ = 0.020306 # Unit is K/J
> 
> R0 = 0.03 # Resistance set a 30mohm
> 
> Qgen =(I^2)\*R0
> 
> du\[2\] = (C₁\*(u\[2\]-T∞)) + (C₂\*Qgen)
> 
> end
> 
> t1 = collect(0:1:3400)
> T∞1,I1 = 298.15,5
> 
> actualODE1!(du,u,p,t) = actualODE!(du,u,p,t,T∞1,I1)
> 
> prob = ODEProblem(actualODE1!,\[1.0,T∞1\],(t1\[1\],t1\[end\]))
> solution = solve(prob,Tsit5(),saveat = t1)
> X = Array(solution)
> T1 = X\[2,:\]
> \# Plotting the results
> plot(solution\[2,:\],color = :red,label = \["True Data" nothing\])
> 
> 
> \# Defining the neural network
> const U = Lux.Chain(Lux.Dense(3,20,tanh),Lux.Dense(20,20,tanh),Lux.Dense(20,1))
> \_para,st = Lux.setup(rng,U)
> const \_st = st
> 
> function NODE\_model!(du,u,p,t,T∞,I)
> 
> Cbat = 5\*3600
> du\[1\] = -I/Cbat
> 
> C₁ = -0.00153
> C₂ = 0.020306
> 
> G = I\*(U(\[u\[1\],u\[2\],I\],p,\_st)\[1\]\[1\])
> 
> du\[2\] = (C₁\*(u\[2\]-T∞)) + (C₂\*G)
> 
> end
> 
> NODE\_model1!(du,u,p,t) = NODE\_model!(du,u,p,t,T∞1,I1)
> prob1 = ODEProblem(NODE\_model1!,\[1.0,T∞1\],(t1\[1\],t1\[end\]),\_para)
> 
> function loss(θ)
> \_prob1 = remake(prob1,p=θ)
> \_sol = Array(solve(\_prob1,Tsit5(),saveat = t1))
> loss1 = mean(abs2,T1.-\_sol\[2,:\])
> return loss1
> end
> 
> losses = Float64\[\]
> 
> callback = function(state,l)
> push!(losses,l)
> println("RMSE Loss at iteration $(length(losses)) is $sqrt(l)")
>     
> return false
> 
> end
> 
> adtype = Optimization.AutoZygote()
> optf = Optimization.OptimizationFunction((x,p) -\> loss(x),adtype)
> optprob = Optimization.OptimizationProblem(optf,ComponentVector{Float64}(\_para))
> 
> res1 = Optimization.solve(optprob, OptimizationOptimisers.Adam(),callback = callback,maxiters = 500)
> \`\`\`

I’m going to have to find a way around this.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [February 11, 2025, 2:18pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/10 "2025-02-11T14:18:45Z")

</div>

I workaround for now is to just choose the sensealg directly:

```julia
function loss(θ)
    _prob1 = remake(prob1,p=θ)
    _sol = Array(solve(_prob1,Tsit5(),saveat = t1, sensealg = QuadratureAdjoint(autojacvec=ReverseDiffVJP(true))))
    loss1 = mean(abs2,T1.-_sol[2,:])
    return loss1
end

```

```julia
u: ComponentVector{Float64}(layer_1 = (weight = [0.5868697166442871 -0.15597978234291077 -0.9511871933937073; -0.9119327664375305 -0.6244392991065979 0.033625759184360504; … ; -0.43237248063087463 1.2632182836532593 -0.4303165078163147; -0.6320686936378479 0.43183207511901855 -1.2375504970550537], bias = [0.048540983349084854, 0.06472337990999222, 0.1911112517118454, 0.1919097602367401, -0.4226461946964264, 0.19537760317325592, 0.4472426474094391, -0.26705870032310486, -0.04794495552778244, 0.561735987663269, 0.5390963554382324, -0.5670568346977234, -0.12485334277153015, -0.351481169462204, -0.42042452096939087, 0.49070412098247174, 0.2708824872970581, 0.3420274257659912, 0.03861868381500244, 0.29712456464767456]), layer_2 = (weight = [0.1544748395664798 -0.2996096909062759 … -0.4314056932909592 -0.3035692572577103; -0.14554913234359354 0.1017788882886734 … 0.6032390136683623 -0.05910374704354658; … ; -0.12327671254753542 -0.3033059781300973 … -0.2575799578917075 -0.0741627792012744; 0.235487428029535 -0.5690564992823629 … -0.6127333161435099 -0.5991582390866251], bias = [-0.011010955086182364, -0.17882417250030905, 0.14381766406152968, -0.1367216425038503, 0.17397684976890115, 0.018784809632670055, -0.10860474003211544, -0.03374763143127334, 0.07022429099183829, 0.03744440423428891, -0.03669568702174315, 0.09550134675459442, 0.10763176988630602, -0.14078365081581634, -0.018990078863569185, -0.14813947903183441, 0.12806158644128623, -0.04716030758742096, -0.10837891494155455, 0.20525834290170888]), layer_3 = (weight = [-0.052892020295968754 0.06296751098405273 … 3.545035572751222e-5 0.00634621667079882], bias = [-0.09864665768500594]))

```

In the meantime I’ll get the real thing fixed.

---

<div class="post-metadata">

**Author:** ![Ashima\_Kalathingal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ashima_kalathingal/32/212264_2.png) [@Ashima\_Kalathingal](https://discourse.julialang.org/u/Ashima_Kalathingal)\
**Post date:** [February 11, 2025, 3:04pm UTC](https://discourse.julialang.org/t/minimum-working-example-mwe-showing-error-in-universal-differential-equation-ude-implementation/125366/11 "2025-02-11T15:04:56Z")

</div>

Okay. Thank you for the help and suggestions. 😄