drifted_bm_example.py

# import libraries
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# ignore warnings
import warnings
warnings.filterwarnings("ignore")

def brownian_motion(N, T, h, seed=42):
    """Simulates a Brownian motion.
    
    :param int N : number of discrete steps
    :param int T: number of continuous time steps
    :param float h: variance of the increments
    :param int seed: initial seed of the random generator
    :returns tuplpe: the brownian motion and its increments
    """
    # set the seed
    np.random.seed(seed)
    # the normalizing constant
    dt = 1. * T/N
    # the epsilon values
    random_increments = np.random.normal(0.0, 1.0 * h, N)*np.sqrt(dt)
    # calculate the brownian motion
    brownian_motion = np.cumsum(random_increments)
    # insert the initial condition
    brownian_motion = np.insert(brownian_motion, 0, 0.0)
    
    return brownian_motion, random_increments

def drifted_brownian_motion(mu, sigma, N, T, seed=42):
    """Simulates a Brownian Motion with drift.
    
    :param float mu: drift coefficient
    :param float sigma: volatility coefficient
    :param int N : number of discrete steps
    :param int T: number of continuous time steps
    :param int seed: initial seed of the random generator
    :returns list: drifted Brownian motion
    """
    # set the seed
    np.random.seed(seed)
    # standard brownian motion
    W, _ = brownian_motion(N, T ,1.0)
    # the normalizing constant
    dt = 1. * T/N
    # generate the time steps
    time_steps = np.linspace(0.0, N*dt, N+1)
    # calculate the Brownian Motion with drift
    X = mu * time_steps + sigma * W
    return X

seed = 42 # the seed to use
mu = 1.45 # the drift
sigma = 1.0 # the diffusial

N = 10000 # number of discret points
T = 10 # number of time units
dt = 1.0 * T/N  # total number of time steps

W, _ = brownian_motion(N, T, 1.0, seed)  # standard Brownian Motion
min_W = np.min(W)  # min of W

X = drifted_brownian_motion(mu, sigma, N, T, seed)  # drifted version
max_X = np.max(X)  # max of X

t = np.linspace(0.0, N*dt, N+1)

fig = plt.figure(figsize=(15, 7))  # instantiate a figure
ax = plt.axes(xlim=(0, T), ylim=(min_W, max_X))  # create an axes object

line_w, = ax.step([], [], where='mid', lw=1, color='#0492c2', alpha=0.8, label='without drift')  # line for W
line_x, = ax.step([], [], where='mid', lw=1, color='#ff4500', alpha=0.8, label='with drift')  # line for X
diff_line, = ax.plot([], [], 'ko-', lw=2)  # line for the difference
text = ax.text(0, 0, '', fontsize=18)

# formatting options
ax.set_title('Drifted Brownian Motion without volatility', fontsize=30)
ax.set_xticks(np.linspace(0, T, 2*T + 1))
ax.set_xlabel('Time', fontsize=18)
ax.set_ylabel('Value', fontsize=18)
ax.tick_params(labelsize=22)
ax.grid(True, which='major', linestyle='--', color='black', alpha=0.6)
ax.legend(loc=2)

frames = 400
factor = N // frames
text_offset = 10*dt

def animate(i):
    upper_bound = (i + 1)*factor
    t_i = t[:upper_bound]
    
    line_w.set_data(list(t_i), list(W[:upper_bound]))
    line_x.set_data(list(t_i), list(X[:upper_bound]))
    diff_line.set_data([t[upper_bound], t[upper_bound]], [W[upper_bound], X[upper_bound]])
    text.set_position((t[upper_bound] + text_offset, (W[upper_bound] + X[upper_bound])/2))
    text.set_text('Diff. $\mu t = $ {:.2f}'.format(np.abs(X[upper_bound] - W[upper_bound])))
    
    return line_w, line_x, diff_line, text,
    
# call the animator	 
anim = animation.FuncAnimation(fig, animate, frames=frames, interval=25, blit=True)
# save the animation as mp4 video file 
anim.save('drift_no_vol_bm.gif',writer='imagemagick')