high-dimensional semilinear pde solver using deep bsde algorithm

src/NeuralNetDiffEq.jl

@@ -12,6 +12,9 @@ end
 nnode(chain;opt=Flux.ADAM(0.1)) = nnode(chain,opt)
 export nnode
 
+
 include("solve.jl")
+include("PdeSolve.jl")
+
 
 end # module

src/PdeSolve.jl

@@ -0,0 +1,82 @@
+function pde_solve(
+    prob,
+    grid,
+    neuralNetworkParams;
+    timeseries_errors = true,
+    save_everystep=true,
+    adaptive=false,
+    abstol = 1f-6,
+    verbose = false,
+    maxiters = 100)
+
+
+    x0 = grid[1] #initial points
+    t0 = grid[2] #initial time
+    tn = grid[3] #terminal time
+    dt = grid[4] #time step
+    d  = grid[5] # number of dimensions
+    m =  grid[6] # number of trajectories (batch size)
+
+
+    g(x) = prob[1](x)
+    f(t,x,Y,Z)  = prob[2](t,x,Y,Z)
+    μ(t,x) = prob[3](t,x)
+    σ(t,x) = prob[4](t,x)
+
+
+    data = Iterators.repeated((), maxiters)
+    ts = t0:dt:tn
+
+    #hidden layer
+    hide_layer_size = neuralNetworkParams[1]
+    opt = neuralNetworkParams[2]
+
+    U0(hide_layer_size, d) = neuralNetworkParams[3](hide_layer_size, d)
+    gradU(hide_layer_size, d) = neuralNetworkParams[4](hide_layer_size, d)
+
+    chains = [gradU(hide_layer_size, d) for i=1:length(ts)]
+    chainU = U0(hide_layer_size, d)
+    ps = Flux.params(chainU, chains...)
+
+    # brownian motion
+    dw(dt) = sqrt(dt) * randn()
+    # the Euler-Maruyama scheme
+    x_sde(x_dim,t,dwa) = [x_dim[i] + μ(t,x_dim[i])*dt + σ(t,x_dim[i])*dwa[i] for i = 1: d]
+
+    get_x_sde(x_cur,l,dwA) = [x_sde(x_cur[i], ts[l] , dwA[i]) for i = 1: m]
+    reduceN(x_dim, l, dwA) = sum([gradU*dwA[i] for (i, gradU) in enumerate(chains[l](x_dim))])
+    getN(x_cur, l, dwA) = [reduceN(x_cur[i], l, dwA[i]) for i = 1: m]
+
+    x_0 = [x0 for i = 1: m]
+
+    function sol()
+        x_cur = x_0
+        U = [chainU(x)[1] for x in x_0]
+        global x_prev
+        for l = 1 : length(ts)
+            x_prev = x_cur
+            dwA = [[dw(dt) for _=1:d] for _=1:m]
+            fa = [f(ts[l], x_cur[i], U[i], chains[l](x_cur[i])) for i= 1 : m]
+            U = U - fa*dt + getN(x_cur, l, dwA)
+            x_cur = get_x_sde(x_cur,l,dwA)
+        end
+        (U, x_prev)
+    end
+
+    function loss()
+        U0, x_cur = sol()
+        return sum(abs2, g.(x_cur) .- U0) / m
+    end
+
+
+    cb = function ()
+        l = loss()
+        verbose && println("Current loss is: $l")
+        l < abstol && Flux.stop()
+    end
+
+    Flux.train!(loss, ps, data, opt; cb = cb)
+
+    ans = chainU(x0)[1]
+    ans
+end #pde_solve

test/PdeRunTest.jl

@@ -0,0 +1,183 @@
+using Flux, Test
+using NeuralNetDiffEq
+
+# one-dimensional heat equation
+x0 = [11]  # initial points
+t0 = 0     # initial time
+T = 5      # terminal time
+dt = 0.5   # time step
+d = 1      # number of dimensions
+m = 50     # number of trajectories (batch size)
+grid = (x0, t0, T, dt, d, m)
+
+g(x) = sum(x.^2)   # terminal condition
+f(t, x, Y, Z) = 0  # function from solved equation
+μ(t,x) = 0
+σ(t,x) = 1
+prob = (g, f, μ, σ)
+
+
+hls = 10 + d #hidden layer size
+opt = Flux.ADAM(0.005)  #optimizer
+#sub-neural network approximating solutions at the desired point
+U0(hls, d) = Flux.Chain(Dense(d,hls,relu),
+                                    Dense(hls,hls,relu),
+                                    Dense(hls,1))
+# sub-neural network approximating the spatial gradients at time point
+gradU(hls, d) = Flux.Chain(Dense(d,hls,relu),
+                                       Dense(hls,hls,relu),
+                                       Dense(hls,d))
+
+# hide_layer_size
+neuralNetParam = (hls, opt, U0, gradU)
+
+
+ans = pde_solve(prob, grid, neuralNetParam, verbose = true, abstol=1e-8, maxiters = 300)
+
+u_analytical(x,t) = sum(x.^2) .+ d*t
+analytical_ans = u_analytical(x0, T)
+# error = abs(ans - analytical_ans)
+error_l2 = sqrt((ans-analytical_ans)^2/ans^2)
+
+# println("one-dimensional heat equation")
+# println("numerical = ", ans)
+# println("analytical = " ,analytical_ans)
+# println("error = ", error)
+println("error_l2 = ", error_l2, "\n")
+@test error_l2 < 0.1
+
+
+# high-dimensional heat equation
+d = 100 # number of dimensions
+x0 = fill(8,d)
+t0 = 0
+T = 2
+dt = 0.5
+m = 100 # number of trajectories (batch size)
+grid = (x0, t0, T, dt, d, m)
+
+g(x) = sum(x.^2)
+f(t,x,Y,Z) = 0
+μ(t,x) = 0
+σ(t,x) = 1
+prob = (g, f, μ, σ)
+
+
+hls = 10 + d #hidden layer size
+# hide_layer_size
+neuralNetParam = (hls, opt, U0, gradU)
+
+ans = pde_solve(prob, grid, neuralNetParam, verbose = true, abstol=1e-8, maxiters = 400)
+
+u_analytical(x,t) = sum(x.^2) .+ d*t
+analytical_ans = u_analytical(x0, T)
+# error = abs(ans - analytical_ans)
+error_l2 = sqrt((ans - analytical_ans)^2/ans^2)
+
+# println("high-dimensional heat equation")
+# println("numerical = ", ans)
+# println("analytical = " ,analytical_ans)
+# println("error = ", error)
+println("error_l2 = ", error_l2, "\n")
+@test error_l2 < 0.1
+
+
+#Black-Scholes-Barenblatt equation
+d = 100 # number of dimensions
+x0 = repeat([1, 0.5], div(d,2))
+t0 = 0
+T = 1
+dt = 0.25
+m = 100 # number of trajectories (batch size)
+grid = (x0, t0, T, dt, d, m)
+
+r = 0.05
+sigma_max = 0.4
+f(t, x, Y, Z) = r * (Y - sum(x.*Z)) # M x 1
+g(x) = sum(x.^2)  # M x D
+μ(t, x) = 0
+σ(t, x) = sigma_max*x
+prob = (g, f, μ, σ)
+
+hls  = 10 + d #hide layer size
+opt = Flux.ADAM(0.001)
+U0(hide_layer_size, d) = Flux.Chain(Dense(d,hls,relu),
+                                    Dense(hls,hls,relu),
+                                    Dense(hls,hls,relu),
+                                    Dense(hls,1))
+gradU(hide_layer_size, d) = Flux.Chain(Dense(d,hls,relu),
+                                       Dense(hls,hls,relu),
+                                       Dense(hls,hls,relu),
+                                       Dense(hls,d))
+
+neuralNetParam = (hls, opt, U0, gradU)
+
+ans = pde_solve(prob, grid, neuralNetParam, verbose = true, abstol=1e-8, maxiters = 250)
+
+u_analytical(x, t) = exp((r + sigma_max^2).*(T .- t)).*sum(x.^2)
+analytical_ans = u_analytical(x0, t0)
+# error = abs(ans - analytical_ans)
+error_l2 = sqrt((ans - analytical_ans)^2/ans^2)
+
+# println("Black Scholes Barenblatt equation")
+# println("numerical ans= ", ans)
+# println("analytical ans = " , u_analytical(x0, t0))
+# println("error = ", error)
+println("error_l2 = ", error_l2, "\n")
+@test error_l2 < 0.1
+
+#Black–Scholes Equation with Default Risk
+#...
+#
+
+# Allen-Cahn Equation
+# d = 1 # number of dimensions 20
+# x0 = fill(0,d)
+# t0 = 0
+# T = 0.3
+# dt = 0.05
+# m = 100 # number of trajectories (batch size)
+# grid = (x0, t0, T, dt, d, m)
+#
+# f(t, x, Y, Z) = - Y .+ Y.^3  # M x 1
+# g(x) = 1.0 / (2.0 + 0.4*sum(x.^2))
+# μ(t, x) = 0 #0
+# σ(t, x) = 1 #sigma_max*x
+# prob = (g, f, μ, σ)
+#
+# ans = pde_solve(prob, grid, neuralNetParam, verbose = true, abstol=1e-8, maxiters = 300)
+#
+# prob_ans = 0.30879
+# error = abs(ans - analytical_ans)
+#
+# println("Allen-Cahn equation")
+# println("numerical = ", ans)
+# println("analytical = " , prob_ans)
+# println("error = ", error)
+
+#Hamilton Jacobi Bellman Equation
+# d = 1 # number of dimensions
+# x0 = fill(0,d)
+# t0 = 0
+# T = 1
+# dt = 0.1
+# m = 100 # number of trajectories (batch size)
+# grid = (x0, t0, T, dt, d, m)
+#
+# f(t, x, Y, Z) = sum(Z.^2) # M x 1
+# g(x) = log(0.5 .+ 0.5*sum(x.^2))
+# μ(t, x) = 0 #0
+# σ(t, x) = sqrt(2) #sigma_max*x
+# prob = (g, f, μ, σ)
+#
+# ans = pde_solve(prob, grid, neuralNetParam, verbose = true, abstol=1e-8, maxiters = 300)
+
+#W =?
+# u_analytical(x, t) = log(mean(exp(-g(x + srt(2.0*abs(T-t)*W )))))
+# analytical_ans = u_analytical(x0, t0)
+# error = abs(ans - analytical_ans)
+#
+# println("Allen-Cahn equation")
+# println("numerical = ", ans)
+# println("analytical = " , prob_ans)
+# println("error = ", error)