Imposing Constraints on Physics-Informed Neural Network (PINN) Solutions

Let's consider the Fokker-Planck equation:

\[- \frac{∂}{∂x} \left [ \left( \alpha x - \beta x^3\right) p(x)\right ] + \frac{\sigma^2}{2} \frac{∂^2}{∂x^2} p(x) = 0 \, ,\]

which must satisfy the normalization condition:

\[\Delta t \, p(x) = 1\]

with the boundary conditions:

\[p(-2.2) = p(2.2) = 0\]

with Physics-Informed Neural Networks.

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
using Integrals, IntegralsCubature
import ModelingToolkit: Interval, infimum, supremum
# the example is taken from this article https://arxiv.org/abs/1910.10503
@parameters x
@variables p(..)
Dx = Differential(x)
Dxx = Differential(x)^2

α = 0.3
β = 0.5
_σ = 0.5
x_0 = -2.2
x_end = 2.2
# Discretization
dx = 0.01

eq  = Dx((α*x - β*x^3)*p(x)) ~ (_σ^2/2)*Dxx(p(x))

# Initial and boundary conditions
bcs = [p(x_0) ~ 0. ,p(x_end) ~ 0.]

# Space and time domains
domains = [x ∈ Interval(x_0,x_end)]

# Neural network
inn = 18
chain = Lux.Chain(Dense(1,inn,Lux.σ),
                  Dense(inn,inn,Lux.σ),
                  Dense(inn,inn,Lux.σ),
                  Dense(inn,1))

lb = [x_0]
ub = [x_end]
function norm_loss_function(phi,θ,p)
    function inner_f(x,θ)
         dx*phi(x, θ) .- 1
    end
    prob = IntegralProblem(inner_f, lb, ub, θ)
    norm2 = solve(prob, HCubatureJL(), reltol = 1e-8, abstol = 1e-8, maxiters =10);
    abs(norm2[1])
end

discretization = PhysicsInformedNN(chain,
                                   GridTraining(dx),
                                   additional_loss=norm_loss_function)

@named pdesystem = PDESystem(eq,bcs,domains,[x],[p(x)])
prob = discretize(pdesystem,discretization)
phi = discretization.phi

sym_prob = NeuralPDE.symbolic_discretize(pdesystem, discretization)

pde_inner_loss_functions = sym_prob.loss_functions.pde_loss_functions
bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions
aprox_derivative_loss_functions = sym_prob.loss_functions.bc_loss_functions

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))
    println("additional_loss: ", norm_loss_function(phi,p,nothing))
    return false
end

res = Optimization.solve(prob,LBFGS(),callback = cb_,maxiters=400)
prob = remake(prob,u0=res.u)
res = Optimization.solve(prob,BFGS(),callback = cb_,maxiters=2000)
u: ComponentVector{Float64}(layer_1 = (weight = [0.3311585138458213; 4.038695358008086; … ; 5.539832659024947; 3.7345939556788164;;], bias = [8.152219776150925; 6.674946973123868; … ; -0.5852231574423145; 2.5664457883276395;;]), layer_2 = (weight = [2.906012483285053 -0.11538325790376724 … 0.8190790459144631 0.31628253056222844; -0.1243539493136044 -0.3820878086231485 … -1.3521824918001366 -0.6971203140283717; … ; -1.1591667277900737 -2.3897214187956717 … -0.32448295309596054 0.8533284846415512; 1.1885869862433887 -0.9877402307036307 … -0.7633484787120767 0.8452606315793193], bias = [-0.8065137549924094; -1.113860492312126; … ; -0.9983160861837724; 1.3349450144198551;;]), layer_3 = (weight = [-3.882741286076276 -0.06239097795466473 … -2.9679631701295173 -1.3391652024199157; 0.8057592559517719 1.1020824066553372 … 0.6555808891069365 0.6942119898969178; … ; 3.862245116012606 0.2712703608405679 … -0.6682952727293601 4.2740849601191595; -6.170798921720984 -0.3808879015609576 … 0.04363365411371203 -4.276965464815126], bias = [-1.6908242378254397; 1.6186548718126161; … ; 3.5855665029502974; -4.400540993226269;;]), layer_4 = (weight = [-10.92405116672532 6.896852012287165 … 11.16693077798775 -11.657582126758168], bias = [10.49590362030567;;]))

And some analysis:

using Plots
C = 142.88418699042 #fitting param
analytic_sol_func(x) = C*exp((1/(2*_σ^2))*(2*α*x^2 - β*x^4))

xs = [infimum(d.domain):dx:supremum(d.domain) for d in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.u)) for x in xs]

plot(xs ,u_real, label = "analytic")
plot!(xs ,u_predict, label = "predict")

fp