stochastic.py

"""
Ornstein-Uhlenbeck (O-U) Continuous-Time Stochastic Process Module
Fits dS_t = theta * (mu - S_t) * dt + sigma * dW_t on stationary cointegrated spreads.
Calculates exact mean-reversion speed theta, long-term equilibrium mu, half-life t_{half},
and asymptotic equilibrium standard deviation sigma_{eq}.
"""

import math
import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import adfuller

class OrnsteinUhlenbeckProcess:
    def __init__(self, dt=1.0):
        self.dt = dt

    def fit(self, spread_series):
        """
        Fits O-U SDE parameters via Exact Maximum Likelihood / Discretized AR(1) on stationary spread series.
        S_t = a + b * S_{t-1} + eps
        b = exp(-theta * dt)
        mu = a / (1 - b)
        sigma = std(eps) * sqrt(-2*ln(b) / (dt * (1 - b^2)))
        half_life = ln(2) / theta
        """
        if len(spread_series) < 15:
            last_val = spread_series.iloc[-1] if hasattr(spread_series, 'iloc') else float(spread_series[-1])
            return {
                'theta': 0.1,
                'mu': last_val,
                'half_life': 10.0,
                'sigma_ou': 0.01,
                'sigma_eq': 0.01,
                'z_ou': 0.0,
                'is_stationary': True,
                'adf_pvalue': 0.01
            }

        y = np.asarray(spread_series, dtype=float)
        x_lag = y[:-1]
        y_curr = y[1:]

        # AR(1) OLS Regression: y_curr = a + b * x_lag
        b, a = np.polyfit(x_lag, y_curr, 1)

        # Check for mean-reversion: b must be strictly < 1.0 (and > 0.0)
        if 0.0001 < b < 0.9999:
            theta = -math.log(b) / self.dt
            mu = a / (1.0 - b)
            residuals = y_curr - (a + b * x_lag)
            sigma_eps = np.std(residuals)
            
            # Continuous SDE diffusion coefficient
            variance_factor = (1.0 - (b ** 2))
            if variance_factor > 1e-6:
                sigma_ou = sigma_eps * math.sqrt(-2.0 * math.log(b) / (self.dt * variance_factor))
                sigma_eq = math.sqrt(max(1e-8, (sigma_ou ** 2) / (2.0 * max(1e-5, theta))))
            else:
                sigma_ou = sigma_eps
                sigma_eq = max(1e-6, np.std(y))

            half_life = math.log(2.0) / max(1e-5, theta)
        else:
            # Unit root / non-mean-reverting fallback
            theta = 1e-4
            mu = float(np.mean(y))
            half_life = 999.0
            sigma_ou = float(np.std(y))
            sigma_eq = max(1e-6, sigma_ou)

        z_ou = (y[-1] - mu) / sigma_eq if sigma_eq > 0 else 0.0

        # Quick ADF Stationarity Check
        try:
            adf_res = adfuller(y, maxlag=1)
            adf_pvalue = float(adf_res[1])
            is_stationary = adf_pvalue < 0.05
        except Exception:
            adf_pvalue = 0.05
            is_stationary = half_life < 30.0

        return {
            'theta': float(theta),
            'mu': float(mu),
            'half_life': float(half_life),
            'sigma_ou': float(sigma_ou),
            'sigma_eq': float(sigma_eq),
            'z_ou': float(z_ou),
            'is_stationary': is_stationary,
            'adf_pvalue': float(adf_pvalue)
        }