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