Inverse problem with NeuralPDE and GPU support

Rene_Schenkendorf
Sat, 20 Jan 2024 17:53:26 +0000
Machine Learning

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

6 posts - 3 participants