feat(pde): Deep-BSDE equations

PARITY.md

@@ -444,4 +444,25 @@ by a functional API over the `DecisionTree.jl` random forest; the tunable grid i
 `{:n_trees, :n_subfeatures, :max_depth}` and scoring is `:accuracy` /
 `:neg_log_loss`.
 
+## Pde — equations (Deep BSDE)  — PR (wired)
+
+Port of `pde.equation`: the forward-SDE sampler and BSDE generator functions for
+four financial PDEs (Han, Jentzen & E, 2018). The neural Deep-BSDE solver follows
+with a deep-learning backend (Lux.jl).
+
+| Concept | Python | Julia | Notes |
+|---|---|---|---|
+| Equation base | `Equation` | `Pde.Equation` (abstract) | dispatch over the four concrete types |
+| HJB-LQ | `HJBLQ` | `Pde.HJBLQ` | **exact** generators (`pde_driver`/`pde_hamiltonian`/`pde_terminal`) |
+| Black–Scholes–Barenblatt | `BlackScholesBarenblatt` | `Pde.BlackScholesBarenblatt` | **exact** |
+| Default-risk pricing | `PricingDefaultRisk` | `Pde.PricingDefaultRisk` | **exact** (piecewise-linear hazard) |
+| Different-rate pricing | `PricingDiffRate` | `Pde.PricingDiffRate` | **exact** |
+| Forward-SDE sample | `Equation.sample` | `Pde.pde_sample` | **behavioural** (Euler–Maruyama; stochastic) |
+
+**Deliberate divergence:** torch tensors → `Matrix`es (`batch × dim`); the method
+names are `pde_driver` (`r_u`), `pde_hamiltonian` (`h_z`), `pde_terminal`,
+`pde_sigma`. The neural solver (`FBSDESolver`/`DeepBSDE` + the PyTorch
+set-transformer architectures) is deferred to the Lux.jl slice; the transformer
+variants are research scaffolding.
+
 _(further submodules appended as they are wired)_

src/Pde/Equations.jl

@@ -0,0 +1,199 @@
+"""
+PDE equations — native Julia port mirroring the Python `RiskLabAI.pde.equation`
+module: the forward-SDE sampler and the BSDE generator functions (driver `f`,
+Hamiltonian `H(z)`, terminal payoff) for four financial PDEs solved by the Deep
+BSDE method (Han, Jentzen & E, 2018). The neural Deep-BSDE solver itself is wired
+separately (it needs a deep-learning backend).
+
+Parity notes: the generator functions (`pde_driver`, `pde_hamiltonian`,
+`pde_terminal`, `pde_sigma`) are **deterministic** and match Python exactly
+(verified in `test/runtests.jl`). `pde_sample` is **stochastic** (Euler–Maruyama
+forward simulation) and validated structurally.
+
+Representation note (deliberate divergence): torch tensors become `Matrix`es with
+shape `batch × dim`; `pde_sample` returns dense `num_sample × dim × steps` arrays.
+
+Reference: Han, J., Jentzen, A., E, W. (2018), *Solving high-dimensional PDEs
+using deep learning*, PNAS. Based on https://github.com/frankhan91/DeepBSDE.
+"""
+
+using Random: default_rng
+
+abstract type Equation end
+
+_relu(v) = max.(v, 0.0)
+
+# Common config accessors.
+delta_t(eq::Equation) = eq.delta_t
+sqrt_delta_t(eq::Equation) = eq.sqrt_delta_t
+
+"""
+    HJBLQ(dim, total_time, num_time_interval)
+
+Hamilton–Jacobi–Bellman equation with linear-quadratic control.
+"""
+struct HJBLQ <: Equation
+    dim::Int
+    total_time::Float64
+    num_time_interval::Int
+    delta_t::Float64
+    sqrt_delta_t::Float64
+    x_init::Vector{Float64}
+    sigma::Float64
+    lambd::Float64
+end
+
+function HJBLQ(dim::Integer, total_time::Real, num_time_interval::Integer)
+    dt = total_time / num_time_interval
+    return HJBLQ(dim, total_time, num_time_interval, dt, sqrt(dt), zeros(dim), sqrt(2.0), 1.0)
+end
+
+pde_sigma(eq::HJBLQ, x) = eq.sigma
+pde_driver(eq::HJBLQ, t, x, y, z) = zeros(size(x, 1), 1)
+pde_hamiltonian(eq::HJBLQ, t, x, y, z) = sum(z .^ 2; dims = 2) ./ eq.sigma^2
+pde_terminal(eq::HJBLQ, t, x) = log.(0.5 .* (1 .+ sum(x .^ 2; dims = 2)))
+
+"""
+    BlackScholesBarenblatt(dim, total_time, num_time_interval)
+
+Black–Scholes–Barenblatt equation.
+"""
+struct BlackScholesBarenblatt <: Equation
+    dim::Int
+    total_time::Float64
+    num_time_interval::Int
+    delta_t::Float64
+    sqrt_delta_t::Float64
+    x_init::Vector{Float64}
+    sigma::Float64
+    rate::Float64
+    mu_bar::Float64
+end
+
+function BlackScholesBarenblatt(dim::Integer, total_time::Real, num_time_interval::Integer)
+    dt = total_time / num_time_interval
+    x_init = [1.0 / (1.0 + (i - 1) % 2) for i = 1:dim]
+    return BlackScholesBarenblatt(
+        dim, total_time, num_time_interval, dt, sqrt(dt), x_init, 0.4, 0.05, 0.0,
+    )
+end
+
+pde_sigma(eq::BlackScholesBarenblatt, x) = eq.sigma .* x
+pde_driver(eq::BlackScholesBarenblatt, t, x, y, z) = fill(eq.rate, size(x, 1), 1)
+pde_hamiltonian(eq::BlackScholesBarenblatt, t, x, y, z) =
+    -sum(z; dims = 2) .* eq.rate ./ eq.sigma
+pde_terminal(eq::BlackScholesBarenblatt, t, x) = sum(x .^ 2; dims = 2)
+
+"""
+    PricingDefaultRisk(dim, total_time, num_time_interval)
+
+Nonlinear pricing PDE with default risk (piecewise-linear hazard rate).
+"""
+struct PricingDefaultRisk <: Equation
+    dim::Int
+    total_time::Float64
+    num_time_interval::Int
+    delta_t::Float64
+    sqrt_delta_t::Float64
+    x_init::Vector{Float64}
+    sigma::Float64
+    rate::Float64
+    delta::Float64
+    gammah::Float64
+    gammal::Float64
+    mu_bar::Float64
+    vh::Float64
+    vl::Float64
+    slope::Float64
+end
+
+function PricingDefaultRisk(dim::Integer, total_time::Real, num_time_interval::Integer)
+    dt = total_time / num_time_interval
+    gammah, gammal, vh, vl = 0.2, 0.02, 50.0, 70.0
+    return PricingDefaultRisk(
+        dim, total_time, num_time_interval, dt, sqrt(dt), ones(dim) .* 100.0,
+        0.2, 0.02, 2.0 / 3, gammah, gammal, 0.02, vh, vl, (gammah - gammal) / (vh - vl),
+    )
+end
+
+pde_sigma(eq::PricingDefaultRisk, x) = eq.sigma .* x
+function pde_driver(eq::PricingDefaultRisk, t, x, y, z)
+    piecewise = _relu(_relu(y .- eq.vh) .* eq.slope .+ eq.gammah .- eq.gammal) .+ eq.gammal
+    return (1 - eq.delta) .* piecewise .+ eq.rate
+end
+pde_hamiltonian(eq::PricingDefaultRisk, t, x, y, z) = zeros(size(x, 1), 1)
+pde_terminal(eq::PricingDefaultRisk, t, x) = _relu(minimum(x; dims = 2))
+
+"""
+    PricingDiffRate(dim, total_time, num_time_interval)
+
+Nonlinear Black–Scholes with different borrowing/lending rates.
+"""
+struct PricingDiffRate <: Equation
+    dim::Int
+    total_time::Float64
+    num_time_interval::Int
+    delta_t::Float64
+    sqrt_delta_t::Float64
+    x_init::Vector{Float64}
+    sigma::Float64
+    mu_bar::Float64
+    rl::Float64
+    rb::Float64
+    alpha::Float64
+end
+
+function PricingDiffRate(dim::Integer, total_time::Real, num_time_interval::Integer)
+    dt = total_time / num_time_interval
+    return PricingDiffRate(
+        dim, total_time, num_time_interval, dt, sqrt(dt), ones(dim) .* 100.0,
+        0.2, 0.06, 0.04, 0.06, 1.0 / dim,
+    )
+end
+
+pde_sigma(eq::PricingDiffRate, x) = eq.sigma .* x
+function pde_driver(eq::PricingDiffRate, t, x, y, z)
+    temp = sum(z; dims = 2) ./ eq.sigma .- y
+    return ifelse.(temp .> 0, eq.rb, eq.rl)
+end
+function pde_hamiltonian(eq::PricingDiffRate, t, x, y, z)
+    temp = sum(z; dims = 2) ./ eq.sigma .- y
+    sum_z = sum(z; dims = 2)
+    return ifelse.(
+        temp .> 0,
+        (eq.mu_bar - eq.rb) .* sum_z ./ eq.sigma,
+        (eq.mu_bar - eq.rl) .* sum_z ./ eq.sigma,
+    )
+end
+function pde_terminal(eq::PricingDiffRate, t, x)
+    temp = maximum(x; dims = 2)
+    return _relu(temp .- 120) .- 2 .* _relu(temp .- 150)
+end
+
+"""
+    pde_sample(eq, num_sample; rng=default_rng()) -> (dw, x)
+
+Euler–Maruyama simulation of the forward SDE: returns Wiener increments
+`dw` (`num_sample × dim × num_time_interval`) and the asset paths `x`
+(`num_sample × dim × num_time_interval+1`). Stochastic. Mirrors Python's
+`Equation.sample`.
+"""
+function pde_sample(eq::Equation, num_sample::Integer; rng = default_rng())
+    dw = randn(rng, num_sample, eq.dim, eq.num_time_interval) .* eq.sqrt_delta_t
+    x = zeros(num_sample, eq.dim, eq.num_time_interval + 1)
+    x[:, :, 1] .= reshape(eq.x_init, 1, eq.dim)
+    for i = 1:eq.num_time_interval
+        x[:, :, i+1] = _sample_step(eq, x[:, :, i], dw[:, :, i])
+    end
+    return dw, x
+end
+
+_sample_step(eq::HJBLQ, xi, dwi) = xi .+ eq.sigma .* dwi
+_sample_step(eq::BlackScholesBarenblatt, xi, dwi) =
+    (1 + eq.mu_bar * eq.delta_t) .* xi .+ (eq.sigma .* xi) .* dwi
+_sample_step(eq::PricingDefaultRisk, xi, dwi) =
+    (1 + eq.mu_bar * eq.delta_t) .* xi .+ (eq.sigma .* xi) .* dwi
+function _sample_step(eq::PricingDiffRate, xi, dwi)
+    factor = exp((eq.mu_bar - (eq.sigma^2) / 2) * eq.delta_t)
+    return (factor .* exp.(eq.sigma .* dwi)) .* xi
+end

src/Pde/Pde.jl

@@ -0,0 +1,27 @@
+"""
+    RiskLabAI.Pde
+
+PDE submodule, mirroring the Python `RiskLabAI.pde` sub-package: financial PDEs
+solved by the Deep BSDE method (Han, Jentzen & E, 2018).
+
+This slice wires the **equations** — the forward-SDE sampler and the BSDE
+generator functions. The neural Deep-BSDE solver (which needs a deep-learning
+backend) is wired in a follow-up.
+"""
+module Pde
+
+include("Equations.jl")
+
+export
+    Equation,
+    HJBLQ,
+    BlackScholesBarenblatt,
+    PricingDefaultRisk,
+    PricingDiffRate,
+    pde_sample,
+    pde_driver,
+    pde_hamiltonian,
+    pde_terminal,
+    pde_sigma
+
+end # module Pde

src/RiskLabAI.jl

@@ -60,6 +60,10 @@ include("Ensemble/Ensemble.jl")
 using .Ensemble: bagging_classifier_accuracy, fit_bagging, bagging_evaluate_schemes,
     calculate_bootstrap_accuracy
 
+include("Pde/Pde.jl")
+using .Pde: Equation, HJBLQ, BlackScholesBarenblatt, PricingDefaultRisk, PricingDiffRate,
+    pde_sample, pde_driver, pde_hamiltonian, pde_terminal, pde_sigma
+
 # --------------------------------------------------------------------------- #
 # Top-level exports.
 # --------------------------------------------------------------------------- #
@@ -116,6 +120,9 @@ export
     # Ensemble — bagging accuracy
     bagging_classifier_accuracy, fit_bagging, bagging_evaluate_schemes,
     calculate_bootstrap_accuracy,
+    # Pde — equations (Deep BSDE)
+    HJBLQ, BlackScholesBarenblatt, PricingDefaultRisk, PricingDiffRate,
+    pde_sample, pde_driver, pde_hamiltonian, pde_terminal, pde_sigma,
     # Backtest (legacy)
     probabilityOfBacktestOverfitting,
     # BetSize

test/runtests.jl

@@ -1194,3 +1194,35 @@ end
     @test length(rs.results) == 3
     @test rs.best_score > 0.7
 end
+
+@testset "Pde — equations (parity with Python)" begin
+    P = RiskLabAI.Pde
+    x = [100.0 110.0; 90.0 120.0; 100.0 100.0]
+    y = reshape([1.0, 2.0, 0.5], 3, 1)
+    z = [0.1 0.2; 0.3 -0.1; 0.0 0.5]
+
+    hjb = P.HJBLQ(2, 1.0, 4)
+    @test vec(P.pde_driver(hjb, 0.0, x, y, z)) == zeros(3)
+    @test vec(P.pde_hamiltonian(hjb, 0.0, x, y, z)) ≈ [0.025, 0.05, 0.125]
+    @test vec(P.pde_terminal(hjb, 1.0, x)) ≈ [9.3102309548, 9.3281678511, 9.2103903707]
+
+    bsb = P.BlackScholesBarenblatt(2, 1.0, 4)
+    @test vec(P.pde_driver(bsb, 0.0, x, y, z)) ≈ [0.05, 0.05, 0.05]
+    @test vec(P.pde_hamiltonian(bsb, 0.0, x, y, z)) ≈ [-0.0375, -0.025, -0.0625]
+    @test vec(P.pde_terminal(bsb, 1.0, x)) ≈ [22100.0, 22500.0, 20000.0]
+
+    dr = P.PricingDefaultRisk(2, 1.0, 4)
+    @test vec(P.pde_driver(dr, 0.0, x, y, z)) ≈ [0.0866666667, 0.0866666667, 0.0866666667]
+    @test vec(P.pde_terminal(dr, 1.0, x)) ≈ [100.0, 90.0, 100.0]
+
+    dfr = P.PricingDiffRate(2, 1.0, 4)
+    @test vec(P.pde_driver(dfr, 0.0, x, y, z)) ≈ [0.06, 0.04, 0.06]
+    @test vec(P.pde_hamiltonian(dfr, 0.0, x, y, z)) ≈ [0.0, 0.02, 0.0]
+    @test vec(P.pde_terminal(dfr, 1.0, x)) ≈ [0.0, 0.0, 0.0]
+
+    # Forward-SDE sampling: shapes and the initial condition.
+    dw, xs = P.pde_sample(bsb, 16; rng = MersenneTwister(1))
+    @test size(dw) == (16, 2, 4)
+    @test size(xs) == (16, 2, 5)
+    @test xs[:, :, 1] == repeat([1.0 0.5], 16, 1)
+end