API PINO PDE

[KirillZubov]

This issue is for more of an API discussion before I dig in to implement PINO PDE. Here I provide examples of supposed API for Physics Informed Neural operator (PINO) problem.

articles
https://arxiv.org/abs/2103.10974
https://arxiv.org/abs/2111.03794

Relate to #806, #575

@ChrisRackauckas , @sathvikbhagavan

using NeuralPDE, Lux, ModelingToolkit, Optimization, NeuralOperators

##example ODE
@parameters t
@variables u(..)
@parameters p [bounds = (0.1f0, pi)]
Dt = Differential(t)
eq = Dt(u(t)) ~ cos(p * t)
bc = u(0) ~ 1.0f0

domain = t ∈ Interval(0.0, 1.0)
neural_operator = SomeNeuralOperator(some_args)
pino = PhysicsInformedNO(neural_operator, sometraining)

@named pde_system = PDESystem(eq, bc, domain, [t], [u(t)], [p])
hasbounds(pde_system.ps[1])
getbounds(pde_system.ps[1])

prob = discretize(pde_system, pino)
res = Optimization.solve(prob, ADAM(0.1);maxiters=4000)
phi = discretization.phi

##example Poisson equation
@parameters x y
@variables u(..)
@parameters p [bounds = (-1, 1)] a [bounds = (-pi, pi)]
hasbounds(a)
Dxx = Differential(x)^2
Dyy = Differential(y)^2

eq = Dxx(u(x, y)) + Dyy(u(x, y)) + p  ~ -sin(pi * x) * sin(pi * y)

bcs = [u(0, y) ~ a, u(1, y) ~ 0, u(x, 0) ~ 0.0, u(x, 1) ~ 0]

domains = [x ∈ Interval(0.0, 1.0),
           y ∈ Interval(0.0, 1.0)]

neural_operator = SomeNeuralOperator(some_args)

pino = PhysicsInformedNO(neural_operator, SomeTraining)

@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)], [p ,a])
hasbounds(pde_system.ps[1])
getbounds(pde_system.ps[1])

prob = discretize(pde_system, pino)

res = Optimization.solve(prob, ADAM(0.1); maxiters=4000)
phi = discretization.phi

##example Additional loss learn with data
##Burgers’ Equation
@parameters t, x
@variables u(..)
Dt = Differential(t)
Dx = Differential(x)
Dxx = Differential(x)^2

ν = 0.1
eq = Dt(u(t, x)) + u(t, x) * Dx(u(t, x)) - ν * Dxx(u(t, x)) ~ 0
bcs = []

domains = [t ∈ Interval(0.0, 1.0), x ∈ Interval(0.0, 1.0)]

neural_operator = SomeNeuralOperator(some_args)

pino = PhysicsInformedNO(neural_operator, SomeTraining)

@named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])

data =DataLoader("Burger_data")
function data_loss(phi, θ)
    a, u0_data = data
    u0 = phi(a, θ)
    loss(u0, u0_data)
end
pino = PhysicsInformedNN(neural_operator, SomeTraining; additional_loss=data_loss)
res = Optimization.solve(prob, ADAM(0.1); maxiters=4000)
phi = discretization.phi

[finmod]

@KirillZubov Can I suggest a PINO-based spatiotemporal MFG example and test https://www.mdpi.com/2227-7390/12/6/803

It is an example of system of coupled forward-backward PDEs consisting of an HJB and a FKP equation where PINO solves the bc of an initial distribution and a final value of the MFGs.
The implementation of this PINO for ST-MFG algorithm is available in python here: https://github.com/LovelyBuggies/MFG-PINO

[ChrisRackauckas]

Burger's is missing an initial condition?

[ChrisRackauckas]

I think the thing missing here is showing how to use a parameterized initial condition in order to sample the space of initial conditions as well.

[KirillZubov]

Burger's is missing an initial condition?

yes, it is example where use data

I think the thing missing here is showing how to use a parameterized initial condition in order to sample the space of initial conditions as well.

it isn't missing, it is in first example bcs = u(0, y) ~ a

[KirillZubov]

##example parameterized initial/boundary condition
@parameters t, x
@variables u(..)
Dt = Differential(t)
Dx = Differential(x)
Dxx = Differential(x)^2
@parameters a [bounds = (-1, 1)] b [bounds = (-1, 1)]
@parameters c [bounds=(2, cos(2))] d [bounds = (4, 5)]

ν = 0.1
eq = Dt(u(t, x)) + u(t, x) * Dx(u(t, x)) - ν * Dxx(u(t, x)) ~ 0
init_cond1 = u(0, x) ~ a
init_cond2 = u(t, x) ~ cos(a)*c + b
init_cond3 = u(t, 0) + d ~  some_func(t,x,c)
bcs = [init_cond1, init_cond2, init_cond3, init_cond4]

domains = [t ∈ Interval(0.0, 1.0), x ∈ Interval(0.0, 1.0)]

neural_operator = SomeNeuralOperator(some_args)

pino = PhysicsInformedNO(neural_operator, SomeTraining)

@named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t, x)], [a,b,c,d])

pino = PhysicsInformedNN(neural_operator, SomeTraining)
res = Optimization.solve(prob, ADAM(0.1); maxiters=4000)
phi = discretization.phi