JSoC 2019: Using neural SDE for solving high-dimensional PDE

Introduction

Basically, it is improved method that was described in the previous post. This post will present a brief description of the idea, either the case of solve of Allen-Cahn equation and plots comparing of the deep BSDE method and the new approach.

Brief description of the idea

The method also usage training neural networks in algorithm scheme but instead bsde and sde, are represented as a form of a single neural SDE. New method should provide an efficient and general approach for solving high-dimensional PDEs

Result

Firstly, to need to upload all the packages that are used and check version of Julia.

"$VERSION"

0.5s

Julia

"1.2.0"

using Pkg
Pkg.add("Flux")
Pkg.add("StochasticDiffEq")
Pkg.add("LinearAlgebra")
Pkg.add("Statistics")
Pkg.add("NeuralNetDiffEq")

21.4s

Julia

using Flux
using StochasticDiffEq
using LinearAlgebra
using Statistics
using NeuralNetDiffEq

0.3s

Julia

Ok , done it.

Now let consider Allen-Cahn equation:

\frac{\partial u}{\partial t}(t,x) = \Delta u(t,x) + u(t,x) - u(t,x)^{3}

with the initial condition: for .

Looking for a solution for u(t = 0.3, x = (0, ..., 0)).

Below is the code of solves 20-dimensional Allen-Cahn Equation that use implementation of new algorithm :

288.5s

Julia

# Allen-Cahn Equation
d = 20 # number of dimensions
x0 = fill(0.0f0,d)
tspan = (0.3f0,0.6f0)
dt = 0.015 # time step
m = 100 # number of trajectories (batch size)

g(X) = 1.0 / (2.0 + 0.4*sum(X.^2))
f(X,u,σᵀ∇u,p,t) = u .- u.^3
μ(X,p,t) = zero(X) #Vector d x 1
σ(X,p,t) = Diagonal(ones(Float32,d)) #Matrix d x d
prob = TerminalPDEProblem(g, f, μ, σ, x0, tspan)

hls = 10 + d #hidden layer size
opt = Flux.ADAM(5^-4)  #optimizer
#sub-neural network approximating solutions at the desired point
u0 = Flux.Chain(Dense(d,hls,relu),
                Dense(hls,hls,relu),
                Dense(hls,1))

# sub-neural network approximating the spatial gradients at time point
σᵀ∇u = Flux.Chain(Dense(d+1,hls,relu),
                  Dense(hls,hls,relu),
                  Dense(hls,d))
pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)

@time ans = solve(prob, pdealg, verbose=true, maxiters=200, trajectories=m,
                            sde_algorithm=EM(), dt=dt)

prob_ans = 0.30879
error_l2 = sqrt((ans - prob_ans)^2/ans^2)

println("Allen-Cahn equation")
println("numerical = ", ans)
println("prob_ans = " , prob_ans)
println("error_l2 = ", error_l2, "\n")

288.5s

Julia

Plots of approximation of u(t=0.3, x=(0, . . . , 0)) against the number of iteration steps in the case of the 20-dimensional Allen-Cahn equation for deep bsde algorithm(left pic.1) and for new algorithm with neural sde (right pic.1). The shaded area depicts the mean ± the standard deviation of approximation of u(t=0.3, x=(0, . . . , 0)) for 5 independent runs.

Left(pic.2): Relative error of the deep BSDE method for u(t=0.3, x=(0, . . . , 0)) against the number of iteration steps in the case of the 20-dimensional Allen-Cahn equation with 20 equidistant time steps (dt=0.015) and learning rate 0.0005. The shaded area depicts the mean ± the standard deviation of the relative error for 5 different runs. Right(pic.2): the same for the new method using neural sde.

Notice that new algorithm (right pic.2) has better convergence compared to deep bsde (left pic.2) with the same parameters for the solution.

Conclusion

This is the last week, and this was the last task in the framework of JSoC, but I don't want to stop there and would like to continue down this route. There is still a lot of work left, and I hope that much more will be done. Also