I am trying to numerically solve the following Hamilton-Jacobi-Bellman (HJB) equation applied to a specifically easy value function with the implicit finite difference method (FDM). However, I get unsatisfying results but I don't know at which point the error occurs.
In my case I am considering the value function and the gain function $v, f:[0,T] \times \mathbb{R} \to \mathbb{R}$. I will leave out the optimal control as by now I am just trying to get my code to work. I want to implement a solution to the following HJB:
$$-\frac{\partial}{\partial t}v(t,x)-(b_t-x)\frac{\partial}{\partial x}v(t,x)-\frac{\sigma^2}{2}\frac{\partial^2}{\partial x^2}v(t,x)-f(t,x)=0$$
with terminal condition $v(T,x)=0$, where
$$dX_t= (b_t-X_t)dt + \sigma dW_t, \qquad W_t\sim\mathcal{N}(0,1)$$
is a mean-reverting process with forecast $b_t$. Setting $v(t,x)=tx^2$ this leads to
$$f(t,x)=-x^2-(b_{t}-x)2tx-t\sigma^2$$
The implicit method is obtained by the forward method as I am solving the value function backwards in time. Therefore, I should not worry about stability concerns. Also I use the backward finite difference instead of the central finite difference method for the $\frac{\partial}{\partial x}$ term. The HJB would therefore translate into
\begin{align} -\frac{v_i^{t_n}-v_{i}^{t_{n-1}}}{\Delta t}&=\frac{\sigma^2}{2}\frac{v_{i+1}^{t_{n-1}}-2v_{i}^{t_{n-1}}+v_{i-1}^{t_{n-1}}}{h^2}+(b_{t_{n-1}}-x_i)\frac{v_{i}^{t_{n-1}}-v_{i-1}^{t_{n-1}}}{h}+f(t_{n-1}, x_i)\\ \iff v_i^{t_n} &= -\Delta t \big((\frac{\sigma^2}{2h^2}-\frac{b_{t_{n-1}}-x_i}{h})v_{i-1}^{t_{n-1}}-(\frac{\sigma^2}{h^2}+\frac{1}{\Delta t}-\frac{b_{t_{n-1}}-x_i}{h})v_{i}^{t_{n-1}}+\frac{\sigma^2}{2h^2}v_{i+1}^{t_{n-1}}+f(t_{n-1}, x_i)\big) \end{align}
with discretization of the domain of interest $[0,T]\times[a,b], a<b \in \mathbb{R}$ with $a=x_0,\dots,x_N=b, x_i = x_0 + ih$ and $0=t_0, t_1, \dots, t_M=T, t_i=t_0+i\Delta t$.
As I want to check consistency for the true value function, e.g. $v(t,x)=tx^2$ the boundary values for $v(t_n, a)$ and $v(t_n, b)$ are known. My Pseudo goes as follows:
Store true boundary values of $v$ into arrays $v_a$ and $v_b$ (thus being of length $[0,\dots,T]$)\
for $t_n$ in $[T,\dots,t_1]$: # $t_1$ instead of $t_0=0$ in order to not overstep last iteration\
Initiate $v$ to be zero matrix with corresponding rows and columns\
Solve $$\left( {\begin{array}{ccc} -\frac{\sigma^2}{h^2}-\frac{1}{\Delta t}+\frac{b_{t_{n-1}}-x_1}{h} & \frac{\sigma^2}{2h^2} &\\ \frac{\sigma^2}{2h^2}-\frac{b_{t_{n-1}}-x_2}{h} & \ddots & \frac{\sigma^2}{2h^2} \\ & \frac{\sigma^2}{2h^2}-\frac{b_{t_{n-1}}-x_{N-1}}{h} & -\frac{\sigma^2}{h^2}-\frac{1}{\Delta t}+\frac{b_{t_{n-1}}-x_{N-1}}{h} \end{array} } \right) \left({\begin{array}{c} v_1^{t_{n-1}}\\ \vdots \\ v_{N-1}^{t_{n-1}}\\ \end{array}}\right)= -\frac{1}{\Delta t} \left({\begin{array}{c} v_1^{t_n}\\ \vdots \\ v_{N-1}^{t_n}\\ \end{array}}\right)- \left({\begin{array}{c} f(t_{n-1}, x_1)+(\frac{\sigma^2}{2h^2}-\frac{b_{t_{n-1}}-x_1}{h})v_a^{t_{n-1}}\\ \vdots \\ f(t_{n-1}, x_{N-1})+\frac{\sigma^2}{2h^2}v_b^{t_{n-1}}\\ \end{array}}\right)$$ Glue $v_a, v, v_b$ together
This is my code so far
import numpy as np
import scipy.stats as st
import scipy.sparse as spa
import matplotlib.pyplot as plt
import plotly.graph_objs as go
import numpy as np
T = 3
d_t = 0.25
t_grid = np.arange(0, T+d_t, d_t)
prodmax = 6
consmax = 6
h = 0.2
x_grid = np.arange(-prodmax, consmax+h, h)
b = np.zeros(len(t_grid))
sig = 0.1
for i, t in enumerate(t_grid):
if i==0:
b[0] = 0
else:
b[i] = b[i-1]+(6*np.sin(np.pi*t_grid[i-1])-b[i-1])*d_t
v_true = np.zeros((len(t_grid), len(x_grid)))
for l, t in enumerate(t_grid):
for i, x in enumerate(x_grid):
v_true[l,i] = t*x**2
va_store = v_true[:, 0]
vb_store = v_true[:, -1]
va_store[-1], vb_store[-1] = 0, 0
v_fdm = np.zeros((len(t_grid), len(x_grid[1:-1])))
for l, t in enumerate(t_grid[1:]):
# iteration backwards in time, attention between l and 1
l_back = len(t_grid)-1-l
# Create finite difference matrix with truncated domain as we now have boundary conditions
up_diag = sig**2/h**2/2*np.ones(len(x_grid[1:-2]))
diag = -sig**2/h**2 - 1/d_t + (b[l_back-1]-x_grid[1:-1])/h
low_diag = sig**2/h**2/2+(b[l_back-1]-x_grid[2:-1])/h
# Evaluate and update gain function
f = np.array([-x**2-(b[l_back-1]-x)*2*t*x-t*sig**2 for x in x_grid[1:-1]])
f[0] = f[0] + (sig**2/h**2/2-(b[l_back-1]-x_grid[1])/h)*va_store[l_back-1]
f[-1] = f[-1] + sig**2/h**2/2*vb_store[l_back-1]
A = np.diag(up_diag, k=1) + np.diag(diag, k=0) + np.diag(low_diag, k=-1)
v_fdm[l_back-1, :] = np.linalg.solve(A, -v_fdm[l_back, :]/d_t-f)
v_fdm = np.column_stack((np.column_stack((va_store, v_fdm)), vb_store))
# Create 3D surface plot
T, X = np.meshgrid(t_grid, x_grid)
fig = go.Figure(data=[go.Surface(x=X, y=T, z=v_fdm.T)])
# Customize plot
fig.update_traces(colorscale='Viridis')
fig.update_layout(scene=dict(xaxis_title='x', yaxis_title='t', zaxis_title='v_fdm(t,x)'))
# Show plot
fig.show()
However, running this code leads to 
instead of something close to 
I hope you can help me out on this one.
Just in case someone has the same issue, I implemented the second part of the algorithm correctly. In this case I went for a central finite difference method for the $\frac{\partial}{\partial x}$ term