docs: fix docs

docs/Project.toml

@@ -2,6 +2,8 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 HighDimPDE = "57c578d5-59d4-4db8-a490-a9fc372d19d2"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [compat]
 Documenter = "1"

docs/make.jl

@@ -6,7 +6,7 @@ cp("./docs/Project.toml", "./docs/src/assets/Project.toml", force = true)
 include("pages.jl")
 
 makedocs(sitename = "HighDimPDE.jl",
-    authors = "Victor Boussange",
+    authors = "#",
     pages = pages,
     clean = true, doctest = false, linkcheck = true,
     format = Documenter.HTML(assets = ["assets/favicon.ico"],

docs/src/NNKolmogorov.md

@@ -1,4 +1,4 @@
-# [The `NNKolmogorov` algorithm](@id nn_komogorov)
+# [The `NNKolmogorov` algorithm](@id nn_kolmogorov)
 
 ### Problems Supported:
 1. [`ParabolicPDEProblem`](@ref)

docs/src/NNParamKolmogorov.md

@@ -1,4 +1,4 @@
-# [The `NNParamKolmogorov` algorithm](@id nn_komogorov)
+# [The `NNParamKolmogorov` algorithm](@id nn_paramkolmogorov)
 
 ### Problems Supported:
 1. [`ParabolicPDEProblem`](@ref)

docs/src/assets/Project.toml

@@ -0,0 +1,11 @@
+[deps]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+HighDimPDE = "57c578d5-59d4-4db8-a490-a9fc372d19d2"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
+
+[compat]
+Documenter = "1"
+Flux = "0.13, 0.14"
+HighDimPDE = "2"

docs/src/getting_started.md

@@ -32,8 +32,8 @@ x0 = fill(0., d)  # initial point
 g(x) = exp(-sum(x.^2)) # initial condition
 μ(x, p, t) = 0.0 # advection coefficients
 σ(x, p, t) = 0.1 # diffusion coefficients
-f(x, y, v_x, v_y, ∇v_x, ∇v_y, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_x)) # nonlocal nonlinear part of the
-prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining the problem
+f(x, v_x, ∇v_x, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_x)) # nonlocal nonlinear part of the
+prob = ParabolicPDEProblem(μ, σ, x0, tspan; g, f) # defining the problem
 
 ## Definition of the algorithm
 alg = MLP() # defining the algorithm. We use the Multi Level Picard algorithm
@@ -117,7 +117,7 @@ sol = solve(prob,
 
 [`DeepSplitting`](@ref deepsplitting) can run on the GPU for (much) improved performance. To do so, just set `use_cuda = true`.
 
-```@example DeepSplitting_gpu
+```@example DeepSplitting_non_local_PDE
 sol = solve(prob, 
             alg, 
             0.1, 

docs/src/index.md

@@ -3,25 +3,25 @@
 
 **HighDimPDE.jl** is a Julia package to **solve Highly Dimensional non-linear, non-local PDEs** of the forms: 
 
-1. Partial Integro Differential Equations: 
+### 1. Partial Integro Differential Equations: 
 ```math
 \begin{aligned}
     (\partial_t u)(t,x) &= \int_{\Omega} f\big(t,x,{\bf x}, u(t,x),u(t,{\bf x}), ( \nabla_x u )(t,x ),( \nabla_x u )(t,{\bf x} ) \big) \, d{\bf x} \\
     & \quad + \big\langle \mu(t,x), ( \nabla_x u )( t,x ) \big\rangle + \tfrac{1}{2} \text{Trace} \big(\sigma(t,x) [ \sigma(t,x) ]^* ( \text{Hess}_x u)(t, x ) \big).
 \end{aligned}
 ```
 
-where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large.
+where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large. These equations are defined using the [`PIDEProblem`](@ref)
 
-2. Parabolic Partial Differential Equations: 
+### 2. Parabolic Partial Differential Equations: 
 ```math
 \begin{aligned}
     (\partial_t u)(t,x) &=  f\big(t,x, u(t,x), ( \nabla_x u )(t,x )\big) 
     + \big\langle \mu(t,x), ( \nabla_x u )( t,x ) \big\rangle + \tfrac{1}{2} \text{Trace} \big(\sigma(t,x) [ \sigma(t,x) ]^* ( \text{Hess}_x u)(t, x ) \big).
 \end{aligned}
 ```
 
-where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large.
+where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large. These equations are defined using the [`ParabolicPDEProblem`](@ref)
 
 !!! note
     The difference between the two problems is that in Partial Integro Differential Equations, the integrand is integrated over **x**, while in Parabolic Integro Differential Equations, the function `f` is just evaluated for `x`.
@@ -32,8 +32,11 @@ where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject
 
 * the [Multi-Level Picard iterations scheme](@ref mlp)
 
-* the Deep BSDE scheme (@ref deepbsde).
+* [the Deep BSDE scheme](@ref deepbsde).
 
+* [NNKolmogorov](@ref nn_kolmogorov) and [NNParamKolmogorov](@ref nn_paramkolmogorov) schemes for Kolmogorov PDEs.
+
+* [NNStopping] scheme for solving optimal stopping time problem.
 
 To make the most out of **HighDimPDE.jl**, we advise to first have a look at the 
 

docs/src/tutorials/deepbsde.md

@@ -5,33 +5,46 @@ Hamilton-Jacobi-Bellman equation that takes a few minutes on a laptop:
 
 ```@example deepbsde
 using HighDimPDE
-using Flux, OptimizationOptimisers
-using DifferentialEquations
+using Flux
+using StochasticDiffEq
 using LinearAlgebra
+
 d = 100 # number of dimensions
-X0 = fill(0.0f0, d) # initial value of stochastic control process
+x0 = fill(0.0f0, d)
 tspan = (0.0f0, 1.0f0)
+dt = 0.2f0
 λ = 1.0f0
+#
+g(X) = log(0.5f0 + 0.5f0 * sum(X .^ 2))
+f(X, u, σᵀ∇u, p, t) = -λ * sum(σᵀ∇u .^ 2)
+μ_f(X, p, t) = zero(X)  #Vector d x 1 λ
+σ_f(X, p, t) = Diagonal(sqrt(2.0f0) * ones(Float32, d)) #Matrix d x d
+prob = ParabolicPDEProblem(μ_f, σ_f, x0, tspan; g, f)
 
-g(X) = log(0.5f0 + 0.5f0 * sum(X.^2))
-f(X,u,σᵀ∇u,p,t) = -λ * sum(σᵀ∇u.^2)
-μ_f(X,p,t) = zero(X)  # Vector d x 1 λ
-σ_f(X,p,t) = Diagonal(sqrt(2.0f0) * ones(Float32, d)) # Matrix d x d
-prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
-hls = 10 + d # hidden layer size
-opt = Optimisers.Adam(0.01)  # optimizer
-# sub-neural network approximating solutions at the desired point
+hls = 10 + d #hidden layer size
+opt = Flux.Optimise.Adam(0.1)  #optimizer
+#sub-neural network approximating solutions at the desired point
 u0 = Flux.Chain(Dense(d, hls, relu),
-                Dense(hls, hls, relu),
-                Dense(hls, 1))
+    Dense(hls, 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, hls, relu),
-                  Dense(hls, d))
-pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)
-@time ans = solve(prob, pdealg, verbose=true, maxiters=100, trajectories=100,
-                            alg=EM(), dt=1.2, pabstol=1f-2)
+    Dense(hls, hls, relu),
+    Dense(hls, hls, relu),
+    Dense(hls, hls, relu),
+    Dense(hls, d))
+pdealg = DeepBSDE(u0, σᵀ∇u, opt = opt)
+
+@time sol = solve(prob,
+    pdealg,
+    StochasticDiffEq.EM(),
+    verbose = true,
+    maxiters = 150,
+    trajectories = 30,
+    dt = 1.2f0,
+    pabstol = 1.0f-4)
+
 ```
 
 Now, let's explain the details!
@@ -61,9 +74,14 @@ with terminating condition $g(x) = \log(1/2 + 1/2 \|x\|^2))$.
 
 #### Define the Problem
 
-To get the solution above using the `PIDEProblem`, we write:
+To get the solution above using the [`ParabolicPDEProblem`](@ref), we write:
 
 ```@example deepbsde2
+using HighDimPDE
+using Flux
+using StochasticDiffEq
+using LinearAlgebra
+
 d = 100 # number of dimensions
 X0 = fill(0.0f0,d) # initial value of stochastic control process
 tspan = (0.0f0, 1.0f0)
@@ -73,12 +91,12 @@ g(X) = log(0.5f0 + 0.5f0*sum(X.^2))
 f(X,u,σᵀ∇u,p,t) = -λ*sum(σᵀ∇u.^2)
 μ_f(X,p,t) = zero(X)  #Vector d x 1 λ
 σ_f(X,p,t) = Diagonal(sqrt(2.0f0)*ones(Float32,d)) #Matrix d x d
-prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
+prob = ParabolicPDEProblem(μ_f, σ_f, X0, tspan; g, f)
 ```
 
 #### Define the Solver Algorithm
 
-As described in the API docs, we now need to define our `NNPDENS` algorithm
+As described in the API docs, we now need to define our `DeepBSDE` algorithm
 by giving it the Flux.jl chains we want it to use for the neural networks.
 `u0` needs to be a `d` dimensional -> 1 dimensional chain, while `σᵀ∇u`
 needs to be `d+1` dimensional to `d` dimensions. Thus we define the following:
@@ -95,14 +113,13 @@ u0 = Flux.Chain(Dense(d,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,d))
-pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)
+pdealg = DeepBSDE(u0, σᵀ∇u, opt = opt)
 ```
 
 #### Solving with Neural Nets
 
 ```@example deepbsde2
-@time ans = solve(prob, pdealg, verbose=true, maxiters=100, trajectories=100,
-                            alg=EM(), dt=0.2, pabstol = 1f-2)
+@time ans = solve(prob, pdealg, EM(), verbose=true, maxiters=100, trajectories=100, dt=0.2f0, pabstol = 1f-2)
 
 ```
 
@@ -121,7 +138,7 @@ equation, which models uncertain volatility and interest rates derived from the
 Black-Scholes equation. This model results in a nonlinear PDE whose dimension
 is the number of assets in the portfolio.
 
-To solve it using the `PIDEProblem`, we write:
+To solve it using the `ParabolicPDEProblem`, we write:
 
 ```julia
 d = 100 # number of dimensions
@@ -133,7 +150,7 @@ f(X,u,σᵀ∇u,p,t) = r * (u - sum(X.*σᵀ∇u))
 g(X) = sum(X.^2)
 μ_f(X,p,t) = zero(X) #Vector d x 1
 σ_f(X,p,t) = Diagonal(sigma*X) #Matrix d x d
-prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
+prob = ParabolicPDEProblem(μ_f, σ_f, X0, tspan; g, f)
 ```
 
 As described in the API docs, we now need to define our `NNPDENS` algorithm
@@ -151,7 +168,7 @@ u0 = Flux.Chain(Dense(d,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,d))
-pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)
+pdealg = DeepBSDE(u0, σᵀ∇u, opt=opt)
 ```
 
 And now we solve the PDE. Here, we say we want to solve the underlying neural
@@ -160,8 +177,7 @@ SDE using the Euler-Maruyama SDE solver with our chosen `dt=0.2`, do at most
 and stop if the loss ever goes below `1f-6`.
 
 ```julia
-ans = solve(prob, pdealg, verbose=true, maxiters=150, trajectories=100,
-                            alg=EM(), dt=0.2, pabstol = 1f-6)
+ans = solve(prob, pdealg, EM(), verbose=true, maxiters=150, trajectories=100, dt=0.2f0)
 ```
 
 ## References

docs/src/tutorials/deepsplitting.md

@@ -11,7 +11,7 @@ where $\Omega = [-1/2, 1/2]^d$, and let's assume Neumann Boundary condition on $
 Let's solve Eq. (1) with the [`DeepSplitting`](@ref deepsplitting) solver. 
 
 ```@example deepsplitting
-using HighDimPDE
+using HighDimPDE, Flux
 
 ## Definition of the problem
 d = 10 # dimension of the problem
@@ -22,10 +22,13 @@ g(x) = exp.(- sum(x.^2, dims=1) ) # initial condition
 σ(x, p, t) = 0.1f0 # diffusion coefficients
 x0_sample = UniformSampling(fill(-5f-1, d), fill(5f-1, d))
 f(x, y, v_x, v_y, ∇v_x, ∇v_y, p, t) = v_x .* (1f0 .- v_y)
+```
+Since this is a non-local equation, we will define our problem as a [`PIDEProblem`](@ref)
+```@example deepsplitting
 prob = PIDEProblem(μ, σ, x0, tspan, g, f; x0_sample = x0_sample)
-
+```
+```@example deepsplitting
 ## Definition of the neural network to use
-using Flux # needed to define the neural network
 
 hls = d + 50 #hidden layer size
 
@@ -53,7 +56,7 @@ sol = solve(prob,
 
 `DeepSplitting` can run on the GPU for (much) improved performance. To do so, just set `use_cuda = true`.
 
-```@example deepsplitting2
+```@example deepsplitting
 sol = solve(prob, 
             alg, 
             0.1, 
@@ -62,4 +65,50 @@ sol = solve(prob,
             maxiters = 1000,
             batch_size = 1000,
             use_cuda=true)
+```
+
+# Solving local Allen Cahn Equation with Neumann BC
+```@example deepsplitting2 
+using HighDimPDE, Flux, StochasticDiffEq
+batch_size = 1000
+train_steps = 1000
+
+tspan = (0.0f0, 5.0f-1)
+dt = 5.0f-2  # time step
+
+μ(x, p, t) = 0.0f0 # advection coefficients
+σ(x, p, t) = 1.0f0 # diffusion coefficients
+
+d = 10
+∂ = fill(5.0f-1, d)
+
+hls = d + 50 #hidden layer size
+
+nn = Flux.Chain(Dense(d, hls, relu),
+        Dense(hls, hls, relu),
+        Dense(hls, 1)) # Neural network used by the scheme
+
+opt = Flux.Optimise.Adam(1e-2) #optimiser
+alg = DeepSplitting(nn, opt = opt)
+
+X0 = fill(0.0f0, d)  # initial point
+g(X) = exp.(-0.25f0 * sum(X .^ 2, dims = 1))   # initial condition
+a(u) = u - u^3
+f(y, v_y, ∇v_y, p, t) = a.(v_y) # nonlocal nonlinear function
+```
+Since we are dealing with a local problem here, i.e. no integration term used, we use [`ParabolicPDEProblem`](@ref) to define the problem.
+```@example deepsplitting2
+# defining the problem
+prob = ParabolicPDEProblem(μ, σ, X0, tspan; g, f, neumann_bc = [-∂, ∂])
+```
+```@example deepsplitting2
+# solving
+@time sol = solve(prob,
+        alg,
+        dt,
+        verbose = false,
+        abstol = 1e-5,
+        use_cuda = false,
+        maxiters = train_steps,
+        batch_size = batch_size)
 ```

docs/src/tutorials/mlp.md

@@ -17,7 +17,7 @@ g(x) = exp(- sum(x.^2) ) # initial condition
 μ(x, p, t) = 0.0 # advection coefficients
 σ(x, p, t) = 0.1 # diffusion coefficients
 f(x, v_x, ∇v_x, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_x)) # nonlocal nonlinear part of the
-prob = ParabolicPDEProblem(μ, σ, x0, tspan, g, f) # defining the problem
+prob = ParabolicPDEProblem(μ, σ, x0, tspan; g, f) # defining the problem
 
 ## Definition of the algorithm
 alg = MLP() # defining the algorithm. We use the Multi Level Picard algorithm
@@ -43,7 +43,7 @@ g(x) = exp( -sum(x.^2) ) # initial condition
 μ(x, p, t) = 0.0 # advection coefficients
 σ(x, p, t) = 0.1 # diffusion coefficients
 mc_sample = UniformSampling(fill(-5f-1, d), fill(5f-1, d))
-f(x, y, v_x, v_y, ∇v_x, ∇v_y, t) = max(0.0, v_x) * (1 -  max(0.0, v_y)) 
+f(x, y, v_x, v_y, ∇v_x, ∇v_y, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_y)) 
 prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining x0_sample is sufficient to implement Neumann boundary conditions
 
 ## Definition of the algorithm

docs/src/tutorials/nnkolmogorov.md

@@ -2,7 +2,9 @@
 
 ## Solving high dimensional Rainbow European Options for a range of initial stock prices:
 
-```julia
+```@example nnkolmogorov
+using HighDimPDE, StochasticDiffEq, Flux, LinearAlgebra
+
 d = 10 # dims
 T = 1/12 
 sigma = 0.01 .+ 0.03.*Matrix(Diagonal(ones(d))) # volatility 
@@ -22,12 +24,12 @@ xspan = [(98.00, 102.00) for i in 1:d]
 
 g(x) = max(maximum(x) -K, 0)
 
-sdealg = EM()
+sdealg = SRIW1()
 # provide `x0` as nothing to the problem since we are provinding a range for `x0`.
 prob = ParabolicPDEProblem(μ_func, σ_func, nothing, tspan, g = g, xspan = xspan)
 opt = Flux.Optimisers.Adam(0.01)
-alg = NNKolmogorov(m, opt)
 m = Chain(Dense(d, 16, elu), Dense(16, 32, elu), Dense(32, 16, elu), Dense(16, 1))
+alg = NNKolmogorov(m, opt)
 sol = solve(prob, alg, sdealg, verbose = true, dt = 0.01,
     dx = 0.0001, trajectories = 1000, abstol = 1e-6, maxiters = 300)
 ```