Inverse problem with NeuralPDE and GPU support

Rene_Schenkendorf · January 20, 2024, 5:53pm

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

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

ChrisRackauckas · January 22, 2024, 9:50am

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

Rene_Schenkendorf · January 22, 2024, 11:19am

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 [σ_, ρ, β]])

gabo-di · January 22, 2024, 11:13pm

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

Rene_Schenkendorf · January 23, 2024, 8:11am

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

Rene_Schenkendorf · January 23, 2024, 8:29am

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.