feat: update PIDEProblem definition to accomodate all kinds of equations

src/HighDimPDE.jl

@@ -28,30 +28,8 @@ function Base.show(io::IO, A::AbstractPDEProblem)
     show(io, A.tspan)
 end
 
-"""
-$(SIGNATURES)
-
-Defines a Partial Integro Differential Problem, of the form
-```math
-\\begin{aligned}
-    \\frac{du}{dt} &= \\tfrac{1}{2} \\text{Tr}(\\sigma \\sigma^T) \\Delta u(x, t) + \\mu \\nabla u(x, t) \\\\
-    &\\quad + \\int f(x, y, u(x, t), u(y, t), ( \\nabla_x u )(x, t), ( \\nabla_x u )(y, t), p, t) dy,
-\\end{aligned}
-```
-with `` u(x,0) = g(x)``.
-
-## Arguments
+include("MCSample.jl")
 
-* `g` : initial condition, of the form `g(x, p, t)`.
-* `f` : nonlinear function, of the form `f(x, y, u(x, t), u(y, t), ∇u(x, t), ∇u(y, t), p, t)`.
-* `μ` : drift function, of the form `μ(x, p, t)`.
-* `σ` : diffusion function `σ(x, p, t)`.
-* `x`: point where `u(x,t)` is approximated. Is required even in the case where `x0_sample` is provided. Determines the dimensionality of the PDE.
-* `tspan`: timespan of the problem.
-* `p`: the parameter vector.
-* `x0_sample` : sampling method for `x0`. Can be `UniformSampling(a,b)`, `NormalSampling(σ_sampling, shifted)`, or `NoSampling` (by default). If `NoSampling`, only solution at the single point `x` is evaluated.
-* `neumann_bc`: if provided, Neumann boundary conditions on the hypercube `neumann_bc[1] × neumann_bc[2]`.
-"""
 struct PIDEProblem{uType, G, F, Mu, Sigma, xType, tType, P, UD, NBC, K} <:
        DiffEqBase.AbstractODEProblem{uType, tType, false}
     u0::uType
@@ -67,43 +45,107 @@ struct PIDEProblem{uType, G, F, Mu, Sigma, xType, tType, P, UD, NBC, K} <:
     kwargs::K
 end
 
-function PIDEProblem(g, f, μ, σ, x::Vector{X}, tspan;
-        p = nothing,
-        x0_sample = NoSampling(),
-        neumann_bc::NBC = nothing,
-        kwargs...) where {X <: AbstractFloat, NBC <: Union{Nothing, AbstractVector}}
-    @assert eltype(tspan)<:AbstractFloat "`tspan` should be a tuple of Float"
-
-    isnothing(neumann_bc) ? nothing : @assert eltype(eltype(neumann_bc)) <: eltype(x)
-    @assert eltype(g(x))==eltype(x) "Type of `g(x)` must match type of x"
-    try
-        @assert(eltype(f(x, x, g(x), g(x), x, x, p, tspan[1]))==eltype(x),
-            "Type of non linear function `f(x)` must type of x")
-    catch e
-        if isa(e, MethodError)
-            @assert(eltype(f(x, eltype(x)(0.0), x, p, tspan[1]))==eltype(x),
+"""
+$(SIGNATURES)
+
+Defines one of the following equations: 
+- Partial Integro Differential Equation
+    * f -> f(x, y, u(x, t), u(y, t), ∇u(x, t), ∇u(y, t), p, t)
+- Semilinear Parabolic Partial Differential Equation 
+    * f -> f(X, u, σᵀ∇u, p, t)
+- Kolmogorov Differential Equation
+    * f -> `nothing`
+    * x0 -> nothing, xspan must be provided.
+- Obstacle Partial Differential Equation 
+    * f -> `nothing`
+    * g -> `nothing`
+    * discounted payoff function provided.
+
+## Arguments
+
+* `μ` : drift function, of the form `μ(x, p, t)`.
+* `σ` : diffusion function `σ(x, p, t)`.
+* `x`: point where `u(x,t)` is approximated. Is required even in the case where `x0_sample` is provided. Determines the dimensionality of the PDE.
+* `tspan`: timespan of the problem.
+* `g` : initial condition, of the form `g(x, p, t)`.
+* `f` : when defining PIDE : nonlinear function, of the form `f(x, y, u(x, t), u(y, t), ∇u(x, t), ∇u(y, t), p, t)`. When defining Semilinear PDE: `f(X, u, σᵀ∇u, p, t)`
+
+## Optional Arguments 
+* `p`: the parameter vector.
+* `x0_sample` : sampling method for `x0`. Can be `UniformSampling(a,b)`, `NormalSampling(σ_sampling, shifted)`, or `NoSampling` (by default). If `NoSampling`, only solution at the single point `x` is evaluated.
+* `neumann_bc`: if provided, Neumann boundary conditions on the hypercube `neumann_bc[1] × neumann_bc[2]`.
+* `xspan`: The domain of the independent variable `x`
+* `payoff`: The discounted payoff function. Required when solving for optimal stopping problem (Obstacle PDEs).
+"""
+function PIDEProblem(μ,
+        σ,
+        x0::Union{Nothing, AbstractArray},
+        tspan::TF,
+        g = nothing,
+        f = nothing;
+        p::Union{Nothing, AbstractVector} = nothing,
+        xspan::Union{Nothing, TF, AbstractVector{<:TF}} = nothing,
+        x0_sample::Union{Nothing, AbstractSampling} = NoSampling(),
+        neumann_bc::Union{Nothing, AbstractVector} = nothing,
+        payoff = nothing,
+        kw...) where {TF <: Tuple{AbstractFloat, AbstractFloat}}
+
+    # Check the Initial Condition Function returns correct types.
+    isnothing(g) && @assert !isnothing(payoff) "Either of `g` or `payoff` must be provided."
+
+    isnothing(neumann_bc) ? nothing : @assert eltype(eltype(neumann_bc)) <: eltype(x0)
+
+    @assert !isnothing(x0)||!isnothing(xspan) "Either of `x0` or `xspan` must be provided."
+
+    # Check if the Non Linear Function `f` returns correct values. 
+    if !isnothing(f)
+        try
+            @assert(eltype(f(x0, x0, g(x0), g(x0), x0, x0, p, tspan[1]))==eltype(x0),
                 "Type of non linear function `f(x)` must type of x")
-        else
-            throw(e)
+        catch e
+            if isa(e, MethodError)
+                @assert(eltype(f(x0, eltype(x0)(0.0), x0, p, tspan[1]))==eltype(x0),
+                    "Type of non linear function `f(x)` must type of x")
+            else
+                throw(e)
+            end
         end
     end
 
-    PIDEProblem{typeof(g(x)),
+    # Wrap kwargs : 
+    kw = NamedTuple(kw)
+    prob_kw = (p = p, xspan = xspan, payoff = payoff)
+    kwargs = merge(prob_kw, kw)
+
+    # If xspan isa Tuple, then convert it as a Vector{Tuple} with single element
+    xspan = isa(xspan, Tuple) ? [xspan] : xspan
+
+    # if `x0` is not provided, pick up the lower-bound of `xspan`.
+    x0 = isnothing(x0) ? first.(xspan) : x0
+
+    # Initial Condition 
+    u0 = if haskey(kw, :p_prototype)
+        u0 = g(x0, kw.p_prototype.p_phi)
+    else
+        !isnothing(g) ? g(x0) : payoff(x0, 0.0)
+    end
+    @assert eltype(u0)==eltype(x0) "Type of `g(x)` must match the Type of x"
+    PIDEProblem{typeof(u0),
         typeof(g),
         typeof(f),
         typeof(μ),
         typeof(σ),
-        typeof(x),
+        typeof(x0),
         eltype(tspan),
         typeof(p),
         typeof(x0_sample),
         typeof(neumann_bc),
-        typeof(kwargs)}(g(x),
+        typeof(kwargs)}(u0,
         g,
         f,
         μ,
         σ,
-        x,
+        x0,
         tspan,
         p,
         x0_sample,
@@ -133,61 +175,6 @@ function PIDESolution(x0, ts, losses, usols, ufuns, limits = nothing)
         limits)
 end
 
-"""
-Defines a Partial Differential Equation without the non linear term `f` : 
-## Arguments : 
-* `μ` : drift function, of the form `μ(x, p, t)`.
-* `σ` : diffusion function `σ(x, p, t)`.
-* `x0`: initial condition for the `x`
-* `tspan`: timespan of the problem.
-* `g` : terminal condition, of the form `g(x)`. Defaults to `nothing`
-
-## Keyword Arguments:
-* `payoff`: While defining an optimal stopping problem provide the `payoff` function `pay(x,t)``
-* `xspan`: The domain of system state. This can be a tuple of floats for single dimension, and a vector of tuples for multiple dimensions. Where each tuple corresponds to a dimension of state vector.
-* `noise_rate_prototype` : Incase of a non diagonal noise, the prototype of `dx` in `σ`
-"""
-function PIDEProblem(μ,
-        σ,
-        x0,
-        tspan,
-        g = nothing;
-        xspan = nothing,
-        payoff = nothing,
-        p = nothing,
-        x0_sample = NoSampling(),
-        kwargs...)
-    kwargs = merge(NamedTuple(kwargs), (xspan = xspan, payoff = payoff))
-
-    @assert !isnothing(x0)||!isnothing(xspan) "Both `x0` and `xspan` cannot be nothing"
-    xspan = !isnothing(xspan) && !isa(xspan, AbstractVector) ? [xspan] : xspan
-
-    # For Optimal stopping problem 
-    u_est = isnothing(x0) && !isnothing(xspan) ? g(first.(xspan)) : payoff(x0, tspan[1])
-
-    PIDEProblem{typeof(u_est),
-        typeof(g),
-        Nothing,
-        typeof(μ),
-        typeof(σ),
-        typeof(x0),
-        eltype(tspan),
-        typeof(p),
-        typeof(x0_sample),
-        Nothing,
-        typeof(kwargs)}(u_est,
-        g,
-        nothing,
-        μ,
-        σ,
-        x0,
-        tspan,
-        p,
-        x0_sample,
-        nothing,
-        kwargs)
-end
-
 Base.summary(prob::PIDESolution) = string(nameof(typeof(prob)))
 
 function Base.show(io::IO, A::PIDESolution)
@@ -198,15 +185,14 @@ function Base.show(io::IO, A::PIDESolution)
     show(io, A.us)
 end
 
-include("MCSample.jl")
 include("reflect.jl")
 include("DeepSplitting.jl")
 include("DeepBSDE.jl")
 include("DeepBSDE_Han.jl")
 include("MLP.jl")
 include("NNStopping.jl")
 
-export PIDEProblem, PIDEProblem, PIDESolution, DeepSplitting, DeepBSDE, MLP, NNStopping
+export PIDEProblem, PIDESolution, DeepSplitting, DeepBSDE, MLP, NNStopping
 
 export NormalSampling, UniformSampling, NoSampling, solve
 end