Specifying a PhysicsInformedNN with variables in different dimensions

[killah-t-cell]

Maybe this issue should live in GalacticOptim (if so, let me know)

Say I have a system with variables with different dimensions (e.g. a system that is partially in configuration space and partially in phase space – u1(t, x, vx) and u2(t, x)). When I try to build this system with NeuralPDE, GalacticOptim throws a DimensionMismatch error. Is there a way to get around the fact that one of the variables has a higher dimension? How does one go about doing that?

Concretely, this is an example of a Vlasov-Poisson solver I wrote:

using NeuralPDE, ModelingToolkit, Symbolics, DomainSets, LinearAlgebra, Flux, GalacticOptim, Optim, DiffEqFlux
using Quadrature,Cubature
@parameters t, x, y, z, vx, vy, vz
@variables f(..), V(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dzz = Differential(z)^2
Dtt = Differential(t)^2
Dx = Differential(x)
Dy = Differential(y)
Dz = Differential(z)
Dt = Differential(t)
# Physics
μ_0 = 1.25663706212e-6 # N A⁻²
ε_0 = 8.8541878128e-12 # F ms⁻¹
q   = 1.602176634e-19 # Coulombs
m   = 3.34358377241e-27 # Kg
# Integrals
Ix = Integral(x, ClosedInterval(0, 1)) 
Iy = Integral(y, ClosedInterval(0, 1))
Iz = Integral(z, ClosedInterval(0, 1)) 
# Helpers
divergence(a)  = Dx(a[1]) + Dy(a[2]) + Dz(a[3])
gradV = [Dx(V(t, x, y, z)), Dy(V(t, x, y, z)), Dz(V(t, x, y, z))]
LaplacianV = Dxx(V(t, x, y, z)) + Dyy(V(t, x, y, z)) + Dzz(V(t, x, y, z))
# Equations
v = [vx, vy, vz]
E = - gradV
ρ  = q * Ix(f(t, x, y, z, vx, vy, vz)) * Iy(f(t, x, y, z, vx, vy, vz)) * Iz(f(t, x, y, z, vx, vy, vz))
vp1 = Dt(f(t, x, y, z, vx, vy, vz)) ~ - divergence(v .* f(t, x, y, z, vx, vy, vz)) - divergence(q .* E .* f(t, x, y, z, vx, vy, vz))
vp2 = LaplacianV ~ - ρ
eqs = [vp1, vp2]
bcs = [f(0, x, y, z, vx, vy, vz) ~ sin(pi*x),
       f(t, 0, y, z, vx, vy, vz) ~ exp(-t),
       f(t, 1, y, z, vx, vy, vz) ~ -exp(-t),
       f(t, x, 0, z, vx, vy, vz) ~ exp(-t),
       f(t, x, 1, z, vx, vy, vz) ~ -exp(-t),
       f(t, x, y, 0, vx, vy, vz) ~ exp(-t),
       f(t, x, y, 1, vx, vy, vz) ~ -exp(-t),
       f(t, x, y, z, 0, vy, vz) ~ exp(-t),
       f(t, x, y, z, 1, vy, vz) ~ -exp(-t),
       f(t, x, y, z, vx, 0, vz) ~ exp(-t),
       f(t, x, y, z, vx, 1, vz) ~ -exp(-t),
       f(t, x, y, z, vx, vy, 0) ~ exp(-t),
       f(t, x, y, z, vx, vy, 1) ~ -exp(-t),
       V(0, x, y, z) ~ sin(pi*x),
       V(t, 0, y, z) ~ exp(-t),
       V(t, 1, y, z) ~ -exp(-t),
       V(t, x, 0, z) ~ exp(-t),
       V(t, x, 1, z) ~ -exp(-t),
       V(t, x, y, 0) ~ exp(-t),
       V(t, x, y, 1) ~ -exp(-t)]
# Space and time domains
domains = [t ∈ Interval(0.0,1.0),
           x ∈ Interval(0.0,1.0),
           y ∈ Interval(0.0,1.0),
           z ∈ Interval(0.0,1.0),
           vx ∈ Interval(0.0,1.0),
           vy ∈ Interval(0.0,1.0),
           vz ∈ Interval(0.0,1.0)]
# Neural network
input_ = length(domains)
n = 15
chain =[FastChain(FastDense(input_,n,Flux.σ),FastDense(n,n,Flux.σ),FastDense(n,1)) for _ in 1:2]
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
_strategy = QuadratureTraining()
discretization = PhysicsInformedNN(chain, _strategy, init_params= initθ)
pde_system = PDESystem(eqs,bcs,domains,[t,x, y, z, vx, vy, vz],[f,V])
prob = discretize(pde_system,discretization)
sym_prob = symbolic_discretize(pde_system,discretization)
pde_inner_loss_functions = prob.f.f.loss_function.pde_loss_function.pde_loss_functions.contents
bcs_inner_loss_functions = prob.f.f.loss_function.bcs_loss_function.bc_loss_functions.contents
cb = function (p,l)
    println("loss: ", l )
    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
    return false
end
res = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=5000)
phi = discretization.phi

The error in this example occurs in res = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=5000). The error is a dimension mismatch from 7 to 4 dimensions.

[ChrisRackauckas]

@KirillZubov this should generally work because we're doing it with the P2D right?

[KirillZubov]

It's relate to WIP that @zoemcc do https://github.com/SciML/NeuralPDE.jl/pull/298

[killah-t-cell]

Merged https://github.com/SciML/NeuralPDE.jl/pull/374