[FEAT]: Add NNStopping

docs/pages.jl

@@ -3,11 +3,13 @@ pages = [
     "Getting started" => "getting_started.md",
     "Solver Algorithms" => ["MLP.md",
         "DeepSplitting.md",
-        "DeepBSDE.md"],
+        "DeepBSDE.md",
+        "NNStopping.md"],
     "Tutorials" => [
         "tutorials/deepsplitting.md",
         "tutorials/deepbsde.md",
         "tutorials/mlp.md",
+        "tutorials/nnstopping.md",
     ],
     "Feynman Kac formula" => "Feynman_Kac.md",
 ]

docs/src/DeepSplitting.md

@@ -62,14 +62,14 @@ In `HighDimPDE.jl` the right parameter combination $\theta$ is found by iterativ
 `DeepSplitting` allows obtaining $u(t,x)$ on a single point  $x \in \Omega$ with the keyword $x$.
 
 ```julia
-prob = PIDEProblem(g, f, μ, σ, x, tspan)
+prob = PIDEProblem(μ, σ, x, tspan, g, f,)
 ```
 
 ### Hypercube
 Yet more generally, one wants to solve Eq. (1) on a $d$-dimensional cube $[a,b]^d$. This is offered by `HighDimPDE.jl` with the keyword `x0_sample`.
 
 ```julia
-prob = PIDEProblem(g, f, μ, σ, x, tspan, x0_sample = x0_sample)
+prob = PIDEProblem(μ, σ, x, tspan, g, f; x0_sample = x0_sample)
 ```
 Internally, this is handled by assigning a random variable as the initial point of the particles, i.e.
 ```math

docs/src/NNStopping.md

@@ -0,0 +1,29 @@
+# [The `NNStopping` algorithm](@id nn_stopping)
+
+```@autodocs
+Modules = [HighDimPDE]
+Pages   = ["NNStopping.jl"]
+```
+## The general idea 💡
+
+Similar to DeepSplitting and DeepBSDE, NNStopping evaluates the PDE as a Stochastic Differential Equation. Consider an Obstacle PDE of the form:
+```math
+ max\lbrace\partial_t u(t,x) + \mu(t, x) \nabla_x u(t,x) + \frac{1}{2} \sigma^2(t, x) \Delta_x u(t,x) , g(t,x) - u(t,x)\rbrace
+```
+
+Such PDEs are commonly used as representations for the dynamics of stock prices that can be exercised before maturity, such as American Options.
+
+Using the Feynman-Kac formula, the underlying SDE will be:
+
+```math
+dX_{t}=\mu(X,t)dt + \sigma(X,t)\ dW_{t}^{Q}
+```
+
+The payoff of the option would then be:
+
+```math
+sup\lbrace\mathbb{E}[g(X_\tau, \tau)]\rbrace
+```
+Where τ is the stopping (exercising) time. The goal is to retrieve both the optimal exercising strategy (τ) and the payoff.
+
+We approximate each stopping decision with a neural network architecture, inorder to maximise the expected payoff.

docs/src/getting_started.md

@@ -33,7 +33,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, 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(g, f, μ, σ, x0, tspan) # defining the problem
+prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining the problem
 
 ## Definition of the algorithm
 alg = MLP() # defining the algorithm. We use the Multi Level Picard algorithm
@@ -62,7 +62,7 @@ g(x) = exp( -sum(x.^2) ) # initial condition
 σ(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, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_y)) 
-prob = PIDEProblem(g, f, μ, σ, x0, tspan) # defining x0_sample is sufficient to implement Neumann boundary conditions
+prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining x0_sample is sufficient to implement Neumann boundary conditions
 
 ## Definition of the algorithm
 alg = MLP(mc_sample = mc_sample)
@@ -87,7 +87,7 @@ 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)
-prob = PIDEProblem(g, f, μ, σ, x0, tspan,
+prob = PIDEProblem(μ, σ, x0, tspan, g, f;
                    x0_sample = x0_sample)
 
 ## Definition of the neural network to use

docs/src/tutorials/deepbsde.md

@@ -19,7 +19,7 @@ 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(g, f, μ_f, σ_f, X0, tspan)
+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
@@ -75,7 +75,7 @@ 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(g, f, μ_f, σ_f, X0, tspan)
+prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
 ```
 
 #### Define the Solver Algorithm
@@ -135,7 +135,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(g, f, μ_f, σ_f, X0, tspan)
+prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
 ```
 
 As described in the API docs, we now need to define our `NNPDENS` algorithm

docs/src/tutorials/deepsplitting.md

@@ -22,9 +22,7 @@ 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)
-prob = PIDEProblem(g, f, μ, 
-                    σ, x0, tspan, 
-                    x0_sample = x0_sample)
+prob = prob = PIDEProblem(μ, σ, x0, tspan, g, f; x0_sample = x0_sample)
 
 ## Definition of the neural network to use
 using Flux # needed to define the neural network

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, 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(g, f, μ, σ, x0, tspan) # defining the problem
+prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining the problem
 
 ## Definition of the algorithm
 alg = MLP() # defining the algorithm. We use the Multi Level Picard algorithm
@@ -44,8 +44,7 @@ g(x) = exp( -sum(x.^2) ) # initial condition
 σ(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)) 
-prob = PIDEProblem(g, f, μ, 
-                    σ, x0, tspan) # defining x0_sample is sufficient to implement Neumann boundary conditions
+prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining x0_sample is sufficient to implement Neumann boundary conditions
 
 ## Definition of the algorithm
 alg = MLP(mc_sample = mc_sample ) 

docs/src/tutorials/nnstopping.md

@@ -0,0 +1,50 @@
+## Solving for optimal strategy and expected payoff of a Bermudan Max-Call option
+
+We will calculate optimal strategy for Bermudan Max-Call option with following drift, diffusion and payoff:
+```math
+μ(x) =(r − δ) x, σ(x) = β diag(x1, ... , xd),\\
+g(t, x) =  e^{-rt}max\lbrace max\lbrace x1, ... , xd \rbrace − K, 0\rbrace
+```
+We define the parameters, drift function and the diffusion function for the dynamics of the option.
+```julia
+d = 3 # Number of assets in the stock
+r = 0.05 # interest rate
+beta = 0.2 # volatility
+T = 3 # maturity
+u0 = fill(90.0, d) # initial stock value
+delta = 0.1 # delta
+f(du, u, p, t) = du .= (r - delta) * u # drift
+sigma(du, u, p, t) = du .= beta * u # diffusion
+tspan = (0.0, T)
+N = 9 # discretization parameter
+dt = T / (N)
+K = 100.00 # strike price
+
+# payoff function
+function g(x, t)
+    return exp(-r * t) * (max(maximum(x) - K, 0))
+end
+
+```
+We then define a `PIDEProblem` with no non linear term:
+```julia
+prob = PIDEProblem(f, sigma, u0, tspan; payoff = g)
+```
+!!! note 
+    We provide the payoff function with a keyword argument `payoff` 
+
+And now we define our models:
+```julia
+models = [Chain(Dense(d + 1, 32, tanh), BatchNorm(32, tanh), Dense(32, 1, sigmoid))
+          for i in 1:N]
+```
+!!! note 
+    The number of models should be equal to the time discritization.
+
+And finally we define our optimizer and algorithm, and call `solve`:
+```julia
+opt = Flux.Optimisers.Adam(0.01)
+alg = NNStopping(models, opt)
+
+sol = solve(prob, alg, SRIW1(); dt = dt, trajectories = 1000, maxiters = 1000, verbose = true)
+```

paper/paper.md

@@ -127,7 +127,7 @@ vol = prod(x0_sample[2] - x0_sample[1])
 f(y, z, v_y, v_z, p, t) =  max.(v_y, 0f0) .* (m(y) .- vol *  max.(v_z, 0f0) .* m(z)) # nonlocal nonlinear part of the
 
 # defining the problem
-prob = PIDEProblem(g, f, μ, σ, tspan, 
+prob = PIDEProblem(μ, σ, tspan, g, f,
                     x0_sample = x0_sample
                     )
 # solving
@@ -162,8 +162,8 @@ g(x) = exp( -sum(x.^2) ) # initial condition
 σ(x, p, t) = 0.1 # diffusion coefficients
 x0_sample = [-1/2, 1/2]
 f(x, y, v_x, v_y, ∇v_x, ∇v_y, t) = max(0.0, v_x) * (1 -  max(0.0, v_y)) 
-prob = PIDEProblem(g, f, μ, 
-                    σ, x0, tspan, 
+prob = PIDEProblem(μ, 
+                    σ, x0, tspan, g, f, 
                     x0_sample = x0_sample) # defining x0_sample is sufficient to implement Neumann boundary conditions
 
 ## Definition of the algorithm

src/DeepBSDE.jl

@@ -26,7 +26,7 @@
 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(g, f, μ_f, σ_f, x0, tspan)
+prob = PIDEProblem(μ_f, σ_f, x0, tspan, g, f)
 
 hls  = 10 + d #hidden layer size
 opt = Flux.Optimise.Adam(0.001)
@@ -75,7 +75,7 @@
     [Deep Primal-Dual algorithm for BSDEs](https://papers.ssrn.com/sol3/papers.cfm?abstract_id=3071506).
 - Extra keyword arguments passed to `solve` will be further passed to the SDE solver.
 """
-function DiffEqBase.solve(prob::PIDEProblem,
+function DiffEqBase.solve(prob::ParabolicPDEProblem,
         pdealg::DeepBSDE,
         sdealg;
         verbose = false,

src/DeepBSDE_Han.jl

@@ -1,5 +1,5 @@
 # called whenever sdealg is not specified.
-function DiffEqBase.solve(prob::PIDEProblem,
+function DiffEqBase.solve(prob::ParabolicPDEProblem,
         alg::DeepBSDE;
         dt,
         abstol = 1.0f-6,

src/DeepSplitting.jl

@@ -64,7 +64,7 @@ Returns a `PIDESolution` object.
 - `use_cuda` : set to `true` to use CUDA.
 - `cuda_device` : integer, to set the CUDA device used in the training, if `use_cuda == true`.
 """
-function DiffEqBase.solve(prob::PIDEProblem,
+function DiffEqBase.solve(prob::Union{PIDEProblem, ParabolicPDEProblem},
         alg::DeepSplitting,
         dt;
         batch_size = 1,
@@ -98,7 +98,13 @@ function DiffEqBase.solve(prob::PIDEProblem,
     K = alg.K
     opt = alg.opt
     λs = alg.λs
-    g, f, μ, σ, p = prob.g, prob.f, prob.μ, prob.σ, prob.p
+    g, μ, σ, p = prob.g, prob.μ, prob.σ, prob.p
+
+    f = if isa(prob, ParabolicPDEProblem)
+        (y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) -> prob.f(y, v_y, ∇v_y, p, t )
+    else
+        prob.f
+    end
     T = eltype(x0)
 
     # neural network model