PDE — Fokker–Planck, HJB, elliptic Poisson ========================================== Three CPU-only finite-difference solvers covering the two canonical PDE pillars of stochastic analysis: the **forward** equation for the marginal density of a diffusion (Fokker–Planck), the **backward** equation for an optimally controlled diffusion (Hamilton–Jacobi–Bellman), and a static **elliptic** boundary-value problem (Poisson). Mathematical background ----------------------- **Fokker–Planck (Kolmogorov forward).** For a 1-D Itô diffusion :math:`dX_t = \mu(t, x)\, dt + \sigma(t, x)\, dW_t`, the marginal density :math:`\rho(t, x)` of :math:`X_t` satisfies the parabolic PDE .. math:: \partial_t \rho \;+\; \partial_x\!\bigl(\mu(t,x)\, \rho\bigr) \;=\; \tfrac12\, \partial^2_{xx}\!\bigl(\sigma^2(t,x)\, \rho\bigr), \qquad \rho(0, \cdot) = \rho_0 . For the *pure-diffusion* test (:math:`\mu \equiv 0`, :math:`\sigma^2 \equiv 1`, :math:`\rho_0 = \mathcal{N}(0, 1)`) the analytic Gaussian heat kernel gives :math:`\rho(t, x) = \frac{1}{\sqrt{2\pi(1+t)}}\exp\!\bigl(-\frac{x^2}{2(1+t)}\bigr)`, so the variance grows linearly: :math:`\mathrm{Var}(X_t) = 1 + t`. The conservative Lax–Wendroff / centred-flux scheme implemented by `fokker_planck_constant` preserves total mass (checked in the notebook to machine precision). **Hamilton–Jacobi–Bellman.** Consider the controlled diffusion :math:`dX_t = \mu(X_t, \alpha_t)\, dt + \sigma(X_t)\, dW_t` and the value function :math:`v(t, x) = \sup_\alpha \mathbb{E}_{t,x}\!\bigl[\int_t^T r(X_s, \alpha_s)\, ds + g(X_T)\bigr]`. Dynamic programming produces .. math:: \partial_t v \;+\; \sup_{a \in \mathcal{A}}\Bigl\{ \mu(x, a) \cdot \nabla v \;+\; \tfrac12\, \mathrm{tr}\!\bigl(\sigma\sigma^\top(x)\, \nabla^2 v\bigr) \;+\; r(x, a) \Bigr\} \;=\; 0, \qquad v(T, x) = g(x). `hjb_quadratic_2d` discretises this in 2-D by an explicit finite-difference scheme; the simple heat-only relaxation case (:math:`H \equiv 0`, :math:`\sigma^2 > 0`) preserves a constant value while a quadratic terminal :math:`g(x) = \tfrac12 \lVert x \rVert^2` smooths into a Gaussian-shaped value surface. **Elliptic Poisson with zero Dirichlet boundary.** On the unit square :math:`\Omega = (0,1)^2`, .. math:: -\Delta u(x, y) = f(x, y) \text{ in } \Omega, \qquad u\!\restriction_{\partial\Omega} = 0 . The Laplace eigenfunctions :math:`\phi_{m,n}(x, y) = \sin(m\pi x)\sin(n\pi y)` form an orthonormal basis with eigenvalues :math:`\lambda_{m,n} = (m^2 + n^2)\pi^2`, so for :math:`f = 2\pi^2 \sin(\pi x)\sin(\pi y)` the *exact* solution is :math:`u(x, y) = \sin(\pi x)\sin(\pi y)`. `poisson_2d_zero_boundary` solves the 5-point stencil by **Successive Over-Relaxation** with optimal relaxation parameter :math:`\omega^* = 2 / (1 + \sin(\pi h))` for grid spacing :math:`h = 1/(N-1)`, achieving spectral radius :math:`\rho \sim 1 - 2\pi h` — i.e. :math:`O(h^{-1})` iterations to reach a fixed tolerance, against :math:`O(h^{-2})` for plain Gauss–Seidel. **Probabilistic representation (Feynman–Kac).** Both the parabolic HJB and the elliptic Poisson PDE admit stochastic representations: :math:`u(x) = \mathbb{E}_x\!\bigl[\int_0^{\tau_\Omega} f(X_s)\, ds\bigr]` for the latter, where :math:`\tau_\Omega` is the first exit time of the diffusion from :math:`\Omega`. This links the PDE solvers above to the BSDE primitives of :doc:`bsde`. Why it matters -------------- * **Density estimation under controlled noise.** Fokker–Planck is the workhorse of non-equilibrium statistical physics, plasma transport, calibration of stochastic-volatility models, and Langevin-based MCMC convergence diagnostics. * **Optimal control & inverse problems.** HJB is the cornerstone of dynamic programming, reinforcement learning (continuous-time policy iteration), and stochastic-control routing. * **Mean-field games.** The MFG fixed point is exactly the coupled system *(backward HJB + forward Fokker–Planck)* with cost depending on the density — building this loop on top of the two solvers above is one of the v2.0 milestones. * **Image processing & PDE-constrained optimisation.** Poisson editing, electric-potential reconstruction, gravitational-potential inversion all reduce to the same elliptic stencil. .. note:: 📓 **Companion notebook** — `view on GitHub `_ · `download .ipynb `_ 11 — PDE solvers ================ Fokker–Planck, HJB, Poisson. .. code-block:: python import numpy as np import matplotlib.pyplot as plt from optimizr import _core as opt plt.rcParams['figure.figsize'] = (7, 4) plt.rcParams['figure.dpi'] = 110 Pure-diffusion Fokker–Planck ---------------------------- :math:`\partial_t m = \tfrac12 \partial_{xx} m` with Gaussian initial density should remain centred and approximately Gaussian. .. code-block:: python res = opt.fokker_planck_constant( mu=0.0, sigma_sq=1.0, init_sigma=1.0, x_min=-8.0, x_max=8.0, n_x=401, t_horizon=0.5, n_t=8000, ) x = np.array(res['x_grid']) t = np.array(res['time_grid']) nx = res['n_x']; nt = res['n_t'] M = np.array(res['density']).reshape(nt + 1, nx) print('total mass at t=0:', np.trapezoid(M[0], x)) print('total mass at t=T:', np.trapezoid(M[-1], x)) print('mean at t=T:', np.trapezoid(x * M[-1], x)) .. code-block:: python fig, ax = plt.subplots() for k in [0, nt // 4, nt // 2, 3 * nt // 4, nt]: ax.plot(x, M[k], label=f't = {t[k]:.2f}') ax.set_xlim(-5, 5); ax.set_xlabel('x'); ax.set_ylabel('m(x, t)') ax.set_title('Pure-diffusion Fokker–Planck'); ax.grid(alpha=0.3); ax.legend() fig.tight_layout(); plt.show() .. AUTO-PLOT-BEGIN .. image:: ../_static/auto/algorithms__pde/block_03_fig_01.png :align: center :width: 80% .. AUTO-PLOT-END .. image:: ../_static/v2/pde/plot_01.png :align: center :width: 80% 2-D Poisson eigenfunction ------------------------- :math:`-\Delta u = 2\pi^2 \sin(\pi x)\sin(\pi y)` on the unit square with zero Dirichlet boundary admits the exact solution :math:`u(x,y) = \sin(\pi x)\sin(\pi y)`. .. code-block:: python n = 65 xs = np.linspace(0, 1, n); ys = np.linspace(0, 1, n) X, Y = np.meshgrid(xs, ys, indexing='ij') F = 2 * np.pi ** 2 * np.sin(np.pi * X) * np.sin(np.pi * Y) res = opt.poisson_2d_zero_boundary(F.flatten().tolist(), n, n) U = np.array(res['u']).reshape(n, n) U_exact = np.sin(np.pi * X) * np.sin(np.pi * Y) print('iterations =', res['iterations']) print('residual =', res['residual']) print('max error =', float(np.max(np.abs(U - U_exact)))) .. code-block:: python fig, axes = plt.subplots(1, 2, figsize=(11, 4)) im0 = axes[0].imshow(U.T, origin='lower', extent=(0, 1, 0, 1), cmap='viridis') axes[0].set_title('SOR solution'); plt.colorbar(im0, ax=axes[0]) im1 = axes[1].imshow((U - U_exact).T, origin='lower', extent=(0, 1, 0, 1), cmap='RdBu_r') axes[1].set_title('error vs analytic'); plt.colorbar(im1, ax=axes[1]) fig.tight_layout(); plt.show() .. AUTO-PLOT-BEGIN .. image:: ../_static/auto/algorithms__pde/block_05_fig_01.png :align: center :width: 80% .. AUTO-PLOT-END .. image:: ../_static/v2/pde/plot_02.png :align: center :width: 80% 2-D HJB with quadratic terminal ------------------------------- Heat-only relaxation (:math:`H = 0`, σ² > 0) preserves a constant value, while a quadratic terminal :math:`g(x) = ½(x²+y²)` smooths. .. code-block:: python res = opt.hjb_quadratic_2d(n_per_dim=21, x_min=-1.0, x_max=1.0, n_t=200, t_horizon=0.2, sigma_sq=0.1) ax_x = np.array(res['axis']); npd = res['n_per_dim'] V = np.array(res['value']).reshape(npd, npd) print('V(0,0) =', V[npd // 2, npd // 2]) print('V(±1,±1) =', V[0, 0], V[-1, -1]) .. code-block:: python fig, ax = plt.subplots() im = ax.imshow(V.T, origin='lower', extent=(-1, 1, -1, 1), cmap='magma') ax.set_title('HJB value V(0, x, y) — quadratic terminal') plt.colorbar(im, ax=ax) fig.tight_layout(); plt.show() .. AUTO-PLOT-BEGIN .. image:: ../_static/auto/algorithms__pde/block_07_fig_01.png :align: center :width: 80% .. AUTO-PLOT-END .. image:: ../_static/v2/pde/plot_03.png :align: center :width: 80% **Verified:** Poisson max-error vs analytic eigenfunction below `5e-3`; Fokker–Planck mean stays at 0 within `0.05`. API --- .. code-block:: rust pub fn solve_fokker_planck_1d(drift: F, diffusion_sq: G, initial_density: H, cfg: &FokkerPlanckConfig) -> Result where F: Fn(f64) -> f64, G: Fn(f64) -> f64, H: Fn(f64) -> f64; pub fn solve_hjb_multid(hamiltonian: H, terminal: G, cfg: &HjbMultidConfig) -> Result where H: Fn(&[f64], &[f64]) -> f64, G: Fn(&[f64]) -> f64; pub fn solve_poisson_2d(rhs: F, boundary: G, cfg: &EllipticFdConfig) -> Result where F: Fn(f64, f64) -> f64, G: Fn(f64, f64) -> f64;