Approximate and exact sampling of an OU noise process

ornstein_uhlenbeck.py

import numpy as np


def ou_approx(t, theta=0.15, sigma=0.2, x0=0, xm=0):
    """
    Ornstein-Uhlenbeck process sampled using an approximate updating formula, first-order in the time step.
    
    The approximation gets worse as the time steps get larger.
    
    Parameters
    ----------
    t: 1D array (n,)
        Time points to sample. Must be monotonic increasing.
    theta, sigma : float
        OU parameters as described in https://en.wikipedia.org/wiki/Ornstein%E2%80%93Uhlenbeck_process
    x0 : float, optional, default: 0
        Initial value.
    xm : float, optional, default: 0
        Equilibrium value.
        
    Returns
    -------
    x : 1D array (n,)
        Values of the sampled OU trajectory.
    
    """
    noise = np.random.randn(len(t) - 1)
    tsteps = np.diff(t)
    
    x = np.zeros(len(t))
    x[0] = x0
    for i in range(len(t) - 1):
        dt = tsteps[i]
        
        x[i + 1] = (
            x[i] + 
            theta * (xm - x[i]) * dt + 
            sigma * noise[i] * np.sqrt(dt)
        )
    
    return x


def ou_exact(t, theta=0.15, sigma=0.2, x0=0, xm=0):
    """
    Ornstein-Uhlenbeck process sampled using the exact updating formula derived in Gillespie, 1996 [1].
    
    This method permits time steps of any size without loss of accuracy.
    
    Parameters
    ----------
    t: 1D array (n,)
        Time points to sample. Must be monotonic increasing.
    theta, sigma : float
        OU parameters as described in https://en.wikipedia.org/wiki/Ornstein%E2%80%93Uhlenbeck_process
    x0 : float, optional, default: 0
        Initial value.
    xm : float, optional, default: 0
        Equilibrium value.
        
    Returns
    -------
    x : 1D array (n,)
        Values of the sampled OU trajectory.
        
    Notes
    -----
    Gillespie uses a different parameterization in terms of:
    
    * The "relaxation time" tau = 1/theta
    * The "diffusion constant" c = sigma^2
    
    References
    ----------
    [1] Gillespie, Daniel T. (1996). Phys. Rev. E. doi:10.1103/PhysRevE.54.2084
    
    """
    noise = np.random.randn(len(t) - 1)
    tsteps = np.diff(t)
    
    x = np.zeros(len(t))
    x[0] = x0
    for i in range(len(t) - 1):
        k = -theta * tsteps[i]
        MU = np.exp(k)
        SIG = sigma * np.sqrt((1 - np.exp(2*k)) / 2 / theta)
        
        x[i + 1] = (
            xm +
            MU * (x[i] - xm) + 
            SIG * noise[i]
        )
    
    return x