# Inverse problem with NeuralPDE and GPU support

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/ works; the inverse problem (first-principles model’s parameter identification) as described here 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)` :

``````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]
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,

@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

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

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

1 Like

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:

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

Do I pass the defaults correctly?

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

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

``````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]
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,

@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

1 Like

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

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/ 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.

1 Like