I am trying to write a python implementation of Euler-Maruyama and Milstein schemes for numerically solving stochastic differential equations. The pseudo-code for the algorithms is in the Wikipedia links.
Would you review my code to see if it can be improved and is closer to production-quality? I've tried to use snake_style convention for all variable names.
from dataclasses import dataclass
from abc import ABC, abstractmethod
from typing import Callable, Optional
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style("whitegrid")
@dataclass
class SDE:
"""
An abstraction for a stochastic differential equation
"""
initial_condition: float
drift: Callable[[float, float], float]
vol: Callable[[float, float], float]
dvol_dx: Optional[Callable[[float, float], float]]
@dataclass
class Solver(ABC):
"""
An abstract base class for all numerical schemes
"""
t_start: float = 0.0
t_end: float = 1.0
step_size: float = 0.001
num_paths: int = 10
def __post_init__(self):
self.iter = 0
self.x_curr = np.zeros((self.num_paths,))
self.num_steps = int((self.t_end - self.t_start) / self.step_size)
self.times = np.linspace(self.t_start, self.t_end, self.num_steps + 1)
self.brownian_increments = np.sqrt(self.step_size) * np.random.standard_normal(
size=(self.num_paths, self.num_steps)
)
self.brownian = np.cumsum(self.brownian_increments, axis=1)
self.brownian = np.concatenate(
[np.zeros(shape=(self.num_paths, 1)), self.brownian], axis=1
)
self.solution = []
@abstractmethod
def iterate(self):
"""
Compute the next iterate X(n+1)
"""
def solve(self, sde: SDE):
"""
Solve the SDE
"""
self.x_curr = np.full(shape=(self.num_paths,), fill_value=sde.initial_condition)
self.solution = [self.x_curr.copy()]
while self.iter < self.num_steps:
self.solution.append(self.iterate(sde))
self.iter += 1
return np.array(self.solution).transpose()
def reset(self):
self.__post_init__()
@dataclass
class Milstein(Solver):
"""
Numerical solver for a stochastic differential equation(SDE) using
the Euler-Maruyama method.
Consider an SDE of the form :
dX_t = mu(t,X_t) dt + sigma(t,X_t) dB_t
with initial condition X_0 = x_0
The solution to the SDE can be computed using the increments
X_{n+1} - X_n = mu(n,X_n)(t_{n+1}-t_n) + sigma(n,X_n)(B(n+1)-B(n))
+ 0.5 * sigma(n,X_n) * sigma'(n,X_n) * ((B(n+1) - B(n))**2 - (t_{n+1} - t_n))
"""
def iterate(self, sde: SDE):
"""
Generate the next iterate X(n+1)
"""
mu_n = np.array([sde.drift(self.times[self.iter], x) for x in self.x_curr])
sigma_n = np.array([sde.vol(self.times[self.iter], x) for x in self.x_curr])
dvol_dx_n = np.array(
[sde.dvol_dx(self.times[self.iter], x) for x in self.x_curr]
)
d_brownian = self.brownian[:, self.iter + 1] - self.brownian[:, self.iter]
self.x_curr += (
mu_n * self.step_size
+ sigma_n * d_brownian
+ 0.5 * sigma_n * dvol_dx_n * (d_brownian**2 - self.step_size)
)
return self.x_curr.copy()
@dataclass
class EulerMaruyama(Solver):
"""
Numerical solver for a stochastic differential equation(SDE) using
the Euler-Maruyama method.
Consider an SDE of the form :
dX_t = mu(t,X_t) dt + sigma(t,X_t) dB_t
with initial condition X_0 = x_0
The solution to the SDE can be computed using the increments
X_{n+1} - X_n = mu(n,X_n)(t_{n+1}-t_n) + sigma(n,X_n)(B(n+1)-B(n))
"""
def iterate(self, sde: SDE):
"""
Generate the next iterate X(n+1)
"""
mu_n = np.array([sde.drift(self.times[self.iter], x) for x in self.x_curr])
sigma_n = np.array([sde.vol(self.times[self.iter], x) for x in self.x_curr])
d_brownian = self.brownian[:, self.iter + 1] - self.brownian[:, self.iter]
self.x_curr += mu_n * self.step_size + sigma_n * d_brownian
return self.x_curr.copy()
if __name__ == "__main__":
# Let's solve the following SDE
# dS_t = S_t dB_t + (mu + 1/2) S_t dt
# where mu = 0. This has the known solution S_t = exp(B_t - t/2)
gbm_sde = SDE(
initial_condition=1.0,
drift = lambda t, s_t : 0.50 * s_t,
vol = lambda t, s_t : s_t,
dvol_dx = lambda t, s_t : 1
)
t = np.linspace(start=0.0, stop=1.0, num=1001)
euler = EulerMaruyama()
solution = euler.solve(gbm_sde)
plt.xlabel(r'Time $t$')
plt.ylabel(r'$S(t)$')
plt.title(r'10 sample paths of Geometric Brownian Motion, $\mu=0.0$, Euler method')
plt.grid(True)
for i in range(10):
plt.plot(t, solution[i])
plt.show()
milstein = Milstein()
solution = milstein.solve(gbm_sde)
plt.xlabel(r'Time $t$')
plt.ylabel(r'$S(t)$')
plt.title(r'10 sample paths of Geometric Brownian Motion, $\mu=0.0$, Milstein method')
plt.grid(True)
for i in range(10):
plt.plot(t, solution[i])
plt.show()
dataclassas well.dataclassis mainly used to represent data, and a long__post_init__()usually means code smell. Some of my classes are workers. \$\endgroup\$