#! /usr/bin/env python3
# -*- coding:utf8 -*-

import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
import numpy as np
from math import *

plt.rcParams['font.sans-serif'] = 'DejaVu Sans'
np.random.seed(8)

ntraces = 5
a = 0.
mu = 0. # long-term mean
tau = 1 # relaxation time scale
theta = 1. / tau # mean reversion rate
sigma = 1. # diffusion
D = sigma**2 / 2 # diffusion constant
sigmaX = sigma / sqrt(2 * theta) # standard deviation of X
t = np.linspace(0, 4, 1001)
dt = t[1:] - t[:-1]

fig, axes = plt.subplots(nrows=2, figsize=(520 / 90.0, 680 / 90.0), dpi=72,
    sharex=True, gridspec_kw={'height_ratios':(2, 1)})

for itrace in range(ntraces):
    # Wiener process
    randnorm = np.random.normal(0, 1, len(t))
    X = np.empty_like(t)
    X[0] = a

    for i in range(1, len(t)):
        X[i] = X[i-1] + sigma * sqrt(dt[i-1]) * randnorm[i]
    
    axes[0].plot(t, X, lw=1)
    
    # Ornstein-Uhlenbeck process
    X = np.empty_like(t)
    X[0] = mu + sigmaX * randnorm[0]

    for i in range(1, len(t)):
        X[i] = X[i-1] + (mu - X[i-1]) * dt[i-1] * theta + sigma * sqrt(dt[i-1]) * randnorm[i]
    
    axes[1].plot(t, X, lw=1)

plt.sca(axes[0])
plt.gca().set_title('Wiener process')
sigma0 = sigma * np.sqrt(t)
plt.fill_between(t, a - sigma0, a + sigma0, color='#dddddd')
plt.grid(True)
axes[0].xaxis.set_major_locator(MultipleLocator(1))
axes[0].yaxis.set_major_locator(MultipleLocator(1))
plt.ylim(-4.2, 4.2)
plt.ylabel('X')

plt.sca(axes[1])
plt.gca().set_title('Ornstein-Uhlenbeck process')
plt.gca().axhspan(mu - sigmaX, mu + sigmaX, color='#dddddd')
plt.grid(True)
axes[1].xaxis.set_major_locator(MultipleLocator(1))
axes[1].yaxis.set_major_locator(MultipleLocator(1))
plt.xlim(t[0], t[-1])
plt.ylim(-2.1, 2.1)
plt.xlabel('t')
plt.ylabel('X')
plt.tight_layout()
plt.savefig('Wiener-Ornstein-Uhlenbeck-5traces-samedata.svg')