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 ModelingToolkit, NeuralPDE, SciMLBase, Lux, Optimization, OptimizationOptimJL, LineSearches
using Optim: BFGS
using DomainSets: Interval
using IntervalSets: leftendpoint, rightendpoint
# 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

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

# Initial and boundary conditions
bcs = [p(x_0) ~ 0.0, p(x_end) ~ 0.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
# Use a simple trapezoidal rule for the normalization constraint.
# This avoids AD issues with Integrals.jl's C-based quadrature solvers.
norm_xs = collect(range(lb, ub, length = 200))
norm_dx = Float64(norm_xs[2] - norm_xs[1])
function norm_loss_function(phi, θ, p)
    # Evaluate phi at quadrature points (each point as a 1-element vector)
    s = sum(1:length(norm_xs)) do i
        first(phi([norm_xs[i]], θ))
    end
    norm_val = 0.01 * s * norm_dx
    abs(norm_val - 1)
end

discretization = PhysicsInformedNN(chain,
    QuadratureTraining(),
    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
approx_derivative_loss_functions = sym_prob.loss_functions.bc_loss_functions

cb_ = function (p, l)
    println("loss: ", l)
    println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
    println("additional_loss: ", norm_loss_function(phi, p.u, nothing))
    return false
end

res = Optimization.solve(
    prob, BFGS(linesearch = BackTracking()), callback = cb_, maxiters = 600)

retcode: MaxIters
u: ComponentVector{Float64}(layer_1 = (weight = [-1.7637658321942284; 5.1884369705057125; … ; 0.3243319626438653; -1.6364256522154312;;], bias = [5.634551206900745, 4.05085374505729, -4.568657989718057, -3.2570544668399157, 10.90469117219218, -6.285048781626668, 3.0567869054627392, 4.118207030413935, -5.154477732256597, 1.010991393103439, -3.36150818692179, -13.764736064855548, -1.0446344344600411, -4.233621257553112, 3.3089692724013884, 4.2086357167965, -3.629145451204005, 0.7298565231827783]), layer_2 = (weight = [-3.543601686696463 -2.9951457951878573 … -0.2585434175139751 -7.056235410881007; 3.5950223638878533 -2.9529194558443947 … 0.33164825995927494 -12.288720953495137; … ; 5.379755520050217 -3.1079603423645334 … 0.40918319389895463 -13.628713457741142; -1.7263326503037464 -0.4889152479619072 … -0.5692105116746256 1.5905729161659048], bias = [-2.2962052891896145, 0.4053118271523102, -1.1274605684546644, 1.8568055157433725, -0.006123088375005551, -0.6206452139264886, -4.003153164423823, 0.2798092947247811, -1.3829874625427143, -0.1457701865176874, 0.898415762767931, 0.5096321764029209, -1.5445494794304708, 1.0853916572575517, -1.5615536926649565, -0.6764277485611006, 0.36346897008382045, -0.6966420934696665]), layer_3 = (weight = [8.111818774462455 1.0616334562713379 … 2.9778399130374673 1.1223217730552477; 2.3925129638532123 0.5405577462607243 … 0.8836494983350897 0.48730548090114373; … ; -1.6913977361701409 0.9297163429281696 … -1.9829101660044361 -0.4786578710925447; 4.339799521340788 0.5446336800025802 … 1.6098399343826773 0.7917806865700528], bias = [0.16429899682894117, 0.1597362194722964, 0.022568120313005233, -0.5304007607890495, 2.850303372312381, -0.6263795320077379, -3.1102543654764223, -0.0564766117796214, 1.023029444272841, -1.466274631717512, 0.1859562367410734, -1.1287989686658753, -0.5141007064713896, 1.014273808658336, 0.5159387761358747, 0.0011091088404393287, 2.8034132077622127, -0.014511562860940713]), layer_4 = (weight = [5.45519582611596 -0.12157537861371066 … 37.2834575247768 -8.885341422866981], bias = [-4.229091399199208]))

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 = [leftendpoint(d.domain):0.01:rightendpoint(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")