# Inverse problem with NeuralPDE and GPU support

**URL:** <https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051>\
**Category:** Machine Learning\
**Tags:** gpu, gpuarrays, pde, sciml, inverse\
**Created:** [January 20, 2024, 5:53pm UTC](https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051 "2024-01-20T17:53:26Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![Rene\_Schenkendorf](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rene_schenkendorf/32/206381_2.png) [@Rene\_Schenkendorf](https://discourse.julialang.org/u/Rene_Schenkendorf)\
**Post date:** [January 20, 2024, 5:53pm UTC](https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051/1 "2024-01-20T17:53:26Z")

</div>

Hello everyone. I am new to Julia (at least that’s how it feels from time to time). Using the GPU according to tutorial [https://docs.sciml.ai/NeuralPDE/stable/tutorials/gpu/](https://docs.sciml.ai/NeuralPDE/stable/tutorials/gpu/) works; the inverse problem (first-principles model’s parameter identification) as described here [https://docs.sciml.ai/NeuralPDE/stable/tutorials/param\_estim/](https://docs.sciml.ai/NeuralPDE/stable/tutorials/param_estim/) also works without any problems. However, as soon as I want to solve the inverse problem with GPU support, a _Scalar indexing_ error occurs (having `CUDA.allowscalar(false)`) when executing `prob = NeuralPDE.discretize(pde_system, discretization)` :

```julia
import Pkg
using NeuralPDE, ModelingToolkit, Optimization, OptimizationOptimJL, OrdinaryDiffEq,
      Plots, Lux, LuxCUDA, ComponentArrays, LuxDeviceUtils, Random, QuasiMonteCarlo
import ModelingToolkit: Interval, infimum, supremum
CUDA.allowscalar(false)
@parameters t, σ_, β, ρ
@variables x(..), y(..), z(..)
Dt = Differential(t)
eqs = [Dt(x(t)) ~ σ_ * (y(t) - x(t)),
    Dt(y(t)) ~ x(t) * (ρ - z(t)) - y(t),
    Dt(z(t)) ~ x(t) * y(t) - β * z(t)]

bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0]
domains = [t ∈ Interval(0.0, 1.0)]
dt = 0.01

input_ = length(domains)
n = 8
chain1 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                   Dense(n, 1))
chain2 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                   Dense(n, 1))
chain3 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                   Dense(n, 1))

chain = [chain1, chain2, chain3]
names = :x, :y, :z

init_params = Lux.initialparameters.(Random.default_rng(), chain)
init_params = NamedTuple{names}(init_params)
init_params = ComponentArray(init_params)

strategy = NeuralPDE.GridTraining(dt)

ps1 = Lux.setup(Random.default_rng(), chain1)[1]
ps1 = ps1 |> ComponentArray |> gpu_device() .|> Float64

ps2 = Lux.setup(Random.default_rng(), chain2)[1]
ps2 = ps2 |> ComponentArray |> gpu_device() .|> Float64

ps3 = Lux.setup(Random.default_rng(), chain3)[1]
ps3 = ps3 |> ComponentArray |> gpu_device() .|> Float64             

function lorenz!(du, u, p, t)
    du[1] = 10.0 * (u[2] - u[1])
    du[2] = u[1] * (28.0 - u[3]) - u[2]
    du[3] = u[1] * u[2] - (8 / 3) * u[3]
end

u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 1.0)
prob = ODEProblem(lorenz!, u0, tspan)
sol = solve(prob, Tsit5(), dt = 0.1)
ts = [infimum(d.domain):dt:supremum(d.domain) for d in domains][1]
function getData(sol)
    data = []
    us = hcat(sol(ts).u...)
    ts_ = hcat(sol(ts).t...)
    return [us, ts_]
end
data = getData(sol)

(u_, t_) = data
len = length(data[2])                   

depvars = [:x, :y, :z]
function additional_loss(phi, θ, p)
    return sum(sum(abs2, phi[i](t_, θ[depvars[i]]) .- u_[[i], :]) / len for i in 1:1:3)
end

discretization = NeuralPDE.PhysicsInformedNN(chain, strategy, init_params=[ps1, ps2, ps3],param_estim = true,
additional_loss = additional_loss) 

@named pde_system = PDESystem(eqs, bcs, domains, [t], [x(t), y(t), z(t)], [σ_, ρ, β],
                              defaults = Dict([p .=> 1.0 for p in [σ_, ρ, β]]))
prob = NeuralPDE.discretize(pde_system, discretization)
callback = function (p, l)
    println("Current loss is: $l")
    return false
end
res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 2500)
p_ = res.u[(end - 2):end] # p_ = [9.93, 28.002, 2.667]

```

I am certainly making a typical beginner’s mistake here 😬 However, I don’t see how I can define the model parameters differently.

v1.9.4

```julia
[34f89e08] LuxDeviceUtils v0.1.11
[315f7962] NeuralPDE v5.9.0

```

---

<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:** [January 22, 2024, 9:50am UTC](https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051/2 "2024-01-22T09:50:53Z")

</div>

We just released a new version, NeuralPDE.jl v5.10, which should fix this issue.

---

<div class="post-metadata">

**Author:** ![Rene\_Schenkendorf](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rene_schenkendorf/32/206381_2.png) [@Rene\_Schenkendorf](https://discourse.julialang.org/u/Rene_Schenkendorf)\
**Post date:** [January 22, 2024, 11:19am UTC](https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051/3 "2024-01-22T11:19:49Z")

</div>

Great that you looked at it; much appreciated! But even with the update v5.10 I get the error message with the scalar indexing when discretizing:

```julia
prob = NeuralPDE.discretize(pde_system, discretization)

```

Do I pass the defaults correctly?

```julia
defaults = Dict([p .=> 1.0 for p in [σ_, ρ, β]])

```

---

<div class="post-metadata">

**Author:** ![gabo-di](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gabo-di/32/206427_2.png) [@gabo-di](https://discourse.julialang.org/u/gabo-di)\
**Post date:** [January 22, 2024, 11:13pm UTC](https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051/4 "2024-01-22T23:13:01Z")

</div>

Hello, after spending the whole afternoon and trying some hacks on NeuralPDE and ComponentArrays ( XD ). I found a clean solution.

```import
using NeuralPDE, ModelingToolkit, Optimization, OptimizationOptimJL, OrdinaryDiffEq,
      Plots, Lux, LuxCUDA, ComponentArrays, LuxDeviceUtils, Random, QuasiMonteCarlo
import ModelingToolkit: Interval, infimum, supremum
CUDA.allowscalar(false)
@parameters t, σ_, β, ρ
@variables x(..), y(..), z(..)
Dt = Differential(t)
eqs = [Dt(x(t)) ~ σ_ * (y(t) - x(t)),
    Dt(y(t)) ~ x(t) * (ρ - z(t)) - y(t),
    Dt(z(t)) ~ x(t) * y(t) - β * z(t)]

bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0]
domains = [t ∈ Interval(0.0, 1.0)]
dt = 0.01

input_ = length(domains)
n = 8
chain1 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                   Dense(n, 1))
chain2 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                   Dense(n, 1))
chain3 = Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, n, Lux.σ),
                   Dense(n, 1))

chain = [chain1, chain2, chain3]
names = :x, :y, :z

init_params = Lux.initialparameters.(Random.default_rng(), chain)
init_params = NamedTuple{names}(init_params)
init_params = ComponentArray(init_params)

strategy = NeuralPDE.GridTraining(dt)

ps1 = Lux.setup(Random.default_rng(), chain1)[1]
#ps1 = ps1 |> ComponentArray |> gpu_device() .|> Float64

ps2 = Lux.setup(Random.default_rng(), chain2)[1]
#ps2 = ps2 |> ComponentArray |> gpu_device() .|> Float64

ps3 = Lux.setup(Random.default_rng(), chain3)[1]
#ps3 = ps3 |> ComponentArray |> gpu_device() .|> Float64             

function lorenz!(du, u, p, t)
    du[1] = 10.0 * (u[2] - u[1])
    du[2] = u[1] * (28.0 - u[3]) - u[2]
    du[3] = u[1] * u[2] - (8 / 3) * u[3]
end

u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 1.0)
prob = ODEProblem(lorenz!, u0, tspan)
sol = solve(prob, Tsit5(), dt = 0.1)
ts = [infimum(d.domain):dt:supremum(d.domain) for d in domains][1]
function getData(sol)
    data = []
    us = hcat(sol(ts).u...)
    ts_ = hcat(sol(ts).t...)
    return [us, ts_]
end
data = getData(sol)

(u_, t_) = data
len = length(data[2])                   
u_ = u_ |> gpu_device()
depvars = [:x, :y, :z]
function additional_loss(phi, θ, p)
    return sum(sum(abs2, phi[i](t_, θ[depvars[i]]) .- u_[[i], :]) / len for i in 1:1:3)
end

discretization = NeuralPDE.PhysicsInformedNN(chain, strategy, init_params=[ps1, ps2, ps3],param_estim = true,
additional_loss = additional_loss) 

@named pde_system = PDESystem(eqs, bcs, domains, [t], [x(t), y(t), z(t)], [σ_, ρ, β],
                              defaults = Dict([p .=> 1.0 for p in [σ_, ρ, β]]))
prob = NeuralPDE.discretize(pde_system, discretization)
prob = remake(prob; u0 = prob.u0 |> gpu_device() )

callback = function (p, l)
    println("Current loss is: $l")
    return false
end
res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 2500)
p_ = res.u[(end - 2):end] # p_ = [9.93, 28.002, 2.667]

```

Basically you pass to gpu `prob.u0` after creating the NeuralPDE problem. Also `u_` so `additional_loss` does not have problems

---

<div class="post-metadata">

**Author:** ![Rene\_Schenkendorf](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rene_schenkendorf/32/206381_2.png) [@Rene\_Schenkendorf](https://discourse.julialang.org/u/Rene_Schenkendorf)\
**Post date:** [January 23, 2024, 8:11am UTC](https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051/5 "2024-01-23T08:11:16Z")

</div>

> [@gabo-di](#):
>
> `prob = remake(prob; u0 = prob.u0 |> gpu_device() )`

Okay, that’s interesting. So technically it works now! 👍 (At least that was the aim of the implementation; whether GPU brings an advantage here is another story.) @gabo-di & @ChrisRackauckas Thanks so much for looking into this; great Julia community experience 🤝

---

<div class="post-metadata">

**Author:** ![Rene\_Schenkendorf](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rene_schenkendorf/32/206381_2.png) [@Rene\_Schenkendorf](https://discourse.julialang.org/u/Rene_Schenkendorf)\
**Post date:** [January 23, 2024, 8:29am UTC](https://discourse.julialang.org/t/inverse-problem-with-neuralpde-and-gpu-support/109051/6 "2024-01-23T08:29:00Z")

</div>

Perhaps for the sake of completeness, it should be briefly noted that the original problem with the “scalar indexing error” message during discretization still occurs, of course, if you follow the example of the GPU implementation [https://docs.sciml.ai/NeuralPDE/stable/tutorials/gpu/](https://docs.sciml.ai/NeuralPDE/stable/tutorials/gpu/) but extended with the inverse problem. This is avoided by discretizing first and then passing/assigning the initial parameters to the GPU in “post-processing” via remake.