{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "118782fe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-16T15:57:28.465029Z",
     "iopub.status.busy": "2026-02-16T15:57:28.464741Z",
     "iopub.status.idle": "2026-02-16T15:57:30.666978Z",
     "shell.execute_reply": "2026-02-16T15:57:30.665798Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Libraries loaded successfully\n",
      "\n",
      "📚 This tutorial covers:\n",
      "  1. Regime Switching Systems\n",
      "  2. Jump Diffusion Processes\n",
      "  3. Combined MRSJD Models\n",
      "  4. Numerical Methods (Finite Differences, Upwind Schemes)\n",
      "  5. Practical Parameter Selection\n"
     ]
    }
   ],
   "source": [
    "# Import required libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import seaborn as sns\n",
    "\n",
    "# Set style\n",
    "sns.set_style('whitegrid')\n",
    "plt.rcParams['figure.figsize'] = (14, 6)\n",
    "plt.rcParams['font.size'] = 11\n",
    "\n",
    "print(\"✅ Libraries loaded successfully\")\n",
    "print(\"\\n📚 This tutorial covers:\")\n",
    "print(\"  1. Regime Switching Systems\")\n",
    "print(\"  2. Jump Diffusion Processes\")\n",
    "print(\"  3. Combined MRSJD Models\")\n",
    "print(\"  4. Numerical Methods (Finite Differences, Upwind Schemes)\")\n",
    "print(\"  5. Practical Parameter Selection\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcfec9d8",
   "metadata": {},
   "source": [
    "## 2. Mathematical Background <a id=\"math\"></a>\n",
    "\n",
    "### Stochastic Differential Equations (SDEs)\n",
    "\n",
    "A general SDE has the form:\n",
    "\n",
    "$$\n",
    "dX_t = \\mu(X_t)dt + \\sigma(X_t)dW_t\n",
    "$$\n",
    "\n",
    "where:\n",
    "- $\\mu(X_t)$ = **drift** (deterministic trend)\n",
    "- $\\sigma(X_t)$ = **diffusion** (volatility)\n",
    "- $dW_t$ = **Wiener process** increment: $dW_t \\sim \\mathcal{N}(0, dt)$\n",
    "\n",
    "### Key Properties\n",
    "\n",
    "**Itô's Lemma** (chain rule for SDEs):\n",
    "\n",
    "For $Y_t = f(X_t)$:\n",
    "\n",
    "$$\n",
    "dY_t = f'(X_t)dX_t + \\frac{1}{2}f''(X_t)\\sigma^2(X_t)dt\n",
    "$$\n",
    "\n",
    "**Feynman-Kac Formula** (connects PDEs to expectations):\n",
    "\n",
    "$$\n",
    "V(x,t) = \\mathbb{E}_x\\left[ \\int_t^T e^{-\\rho(s-t)} L(X_s)ds + e^{-\\rho(T-t)}\\Phi(X_T) \\right]\n",
    "$$\n",
    "\n",
    "satisfies the PDE:\n",
    "\n",
    "$$\n",
    "\\frac{\\partial V}{\\partial t} + \\mu(x)\\frac{\\partial V}{\\partial x} + \\frac{1}{2}\\sigma^2(x)\\frac{\\partial^2 V}{\\partial x^2} - \\rho V + L(x) = 0\n",
    "$$\n",
    "\n",
    "### Example: Ornstein-Uhlenbeck Process\n",
    "\n",
    "Mean-reverting process:\n",
    "\n",
    "$$\n",
    "dX_t = \\theta(\\mu - X_t)dt + \\sigma dW_t\n",
    "$$\n",
    "\n",
    "- $\\theta$ = speed of mean reversion\n",
    "- $\\mu$ = long-term mean\n",
    "- $\\sigma$ = volatility\n",
    "\n",
    "**Half-life**: $t_{1/2} = \\frac{\\ln 2}{\\theta}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b25e1746",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-16T15:57:30.670266Z",
     "iopub.status.busy": "2026-02-16T15:57:30.669891Z",
     "iopub.status.idle": "2026-02-16T15:57:31.331171Z",
     "shell.execute_reply": "2026-02-16T15:57:31.330002Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 117872 characters]",
      "text/plain": [
       "<Figure size 1400x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "📊 OU Process Analysis:\n",
      "  Half-life: 1.39 seconds\n",
      "  Theoretical equilibrium std: 2.00\n",
      "  Observed equilibrium std: 2.31\n"
     ]
    }
   ],
   "source": [
    "# Simulate Ornstein-Uhlenbeck process\n",
    "def simulate_ou(theta, mu, sigma, x0, T, dt):\n",
    "    \"\"\"\n",
    "    Simulate Ornstein-Uhlenbeck process using Euler-Maruyama method\n",
    "    \n",
    "    dX_t = θ(μ - X_t)dt + σ dW_t\n",
    "    \"\"\"\n",
    "    n_steps = int(T / dt)\n",
    "    t = np.linspace(0, T, n_steps)\n",
    "    X = np.zeros(n_steps)\n",
    "    X[0] = x0\n",
    "    \n",
    "    for i in range(1, n_steps):\n",
    "        dW = np.random.normal(0, np.sqrt(dt))\n",
    "        X[i] = X[i-1] + theta * (mu - X[i-1]) * dt + sigma * dW\n",
    "    \n",
    "    return t, X\n",
    "\n",
    "# Example: Temperature control\n",
    "theta = 0.5  # Mean reversion speed\n",
    "mu = 20.0    # Target temperature (°C)\n",
    "sigma = 2.0  # Noise level\n",
    "x0 = 10.0    # Initial temperature\n",
    "T = 10.0     # Time horizon (seconds)\n",
    "dt = 0.01    # Time step\n",
    "\n",
    "t, X = simulate_ou(theta, mu, sigma, x0, T, dt)\n",
    "\n",
    "# Plot\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "# Trajectory\n",
    "ax1.plot(t, X, linewidth=1.5, color='steelblue', label='Temperature')\n",
    "ax1.axhline(y=mu, color='red', linestyle='--', label=f'Target μ={mu}')\n",
    "ax1.fill_between(t, mu-sigma, mu+sigma, alpha=0.2, color='red', label='±σ band')\n",
    "ax1.set_xlabel('Time (s)')\n",
    "ax1.set_ylabel('Temperature (°C)')\n",
    "ax1.set_title('Ornstein-Uhlenbeck Process (Mean-Reverting System)')\n",
    "ax1.legend()\n",
    "ax1.grid(alpha=0.3)\n",
    "\n",
    "# Distribution at equilibrium\n",
    "equilibrium_samples = X[len(X)//2:]  # Second half (near equilibrium)\n",
    "ax2.hist(equilibrium_samples, bins=30, density=True, alpha=0.7, color='steelblue', edgecolor='black')\n",
    "\n",
    "# Theoretical distribution: N(μ, σ²/(2θ))\n",
    "x_range = np.linspace(X.min(), X.max(), 100)\n",
    "theoretical_std = sigma / np.sqrt(2 * theta)\n",
    "from scipy.stats import norm\n",
    "ax2.plot(x_range, norm.pdf(x_range, mu, theoretical_std), \n",
    "        'r-', linewidth=2, label=f'Theory: N({mu:.1f}, {theoretical_std:.2f}²)')\n",
    "\n",
    "ax2.set_xlabel('Temperature (°C)')\n",
    "ax2.set_ylabel('Probability Density')\n",
    "ax2.set_title('Equilibrium Distribution')\n",
    "ax2.legend()\n",
    "ax2.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"\\n📊 OU Process Analysis:\")\n",
    "print(f\"  Half-life: {np.log(2)/theta:.2f} seconds\")\n",
    "print(f\"  Theoretical equilibrium std: {theoretical_std:.2f}\")\n",
    "print(f\"  Observed equilibrium std: {equilibrium_samples.std():.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78636a0b",
   "metadata": {},
   "source": [
    "## 3. Regime Switching Systems <a id=\"regime\"></a>\n",
    "\n",
    "### Motivation\n",
    "\n",
    "Many real systems exhibit **multiple operating modes** or **regimes**:\n",
    "- Weather: sunny ↔ rainy ↔ stormy\n",
    "- Manufacturing: normal ↔ maintenance ↔ failure\n",
    "- Traffic: free-flow ↔ congested ↔ gridlock\n",
    "- Economic activity: expansion ↔ recession\n",
    "\n",
    "### Continuous-Time Markov Chain\n",
    "\n",
    "The regime $i_t \\in \\{1, 2, ..., N\\}$ follows a Markov chain with **transition rate matrix** $Q$:\n",
    "\n",
    "$$\n",
    "\\mathbb{P}(i_{t+dt} = j | i_t = i) = \n",
    "\\begin{cases}\n",
    "q_{ij} dt & \\text{if } i \\neq j \\\\\n",
    "1 + q_{ii} dt & \\text{if } i = j\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "where $q_{ii} = -\\sum_{j \\neq i} q_{ij}$ (rows sum to zero).\n",
    "\n",
    "### Coupled HJB System\n",
    "\n",
    "The value function $V^i(x)$ in regime $i$ satisfies:\n",
    "\n",
    "$$\n",
    "\\rho V^i(x) = \\sup_u \\left[ \\mu^i(x,u) (V^i)'(x) + \\frac{1}{2}(\\sigma^i)^2(x,u) (V^i)''(x) + L^i(x,u) + \\sum_{j \\neq i} q_{ij}[V^j(x) - V^i(x)] \\right]\n",
    "$$\n",
    "\n",
    "Key insight: The term $\\sum_{j \\neq i} q_{ij}[V^j(x) - V^i(x)]$ represents the **expected change in value due to regime switching**.\n",
    "\n",
    "### Stationary Distribution\n",
    "\n",
    "The long-run probability of being in each regime solves:\n",
    "\n",
    "$$\n",
    "Q^T \\pi = 0, \\quad \\sum_i \\pi_i = 1\n",
    "$$\n",
    "\n",
    "### Parameter Selection Tips\n",
    "\n",
    "| Parameter | Typical Range | Effect | How to Choose |\n",
    "|-----------|--------------|--------|---------------|\n",
    "| $q_{ij}$ | 0.1 - 10.0 | Regime persistence | Higher = faster switching. Set $q_{ij} = 1/\\text{expected duration}$ |\n",
    "| $\\mu^i$ | Problem-specific | Drift in regime $i$ | Estimate from data or physics |\n",
    "| $\\sigma^i$ | $> 0$ | Volatility in regime $i$ | Measure from observations or experiments |\n",
    "\n",
    "**Example**: If regime 1 typically lasts 5 time units, set $q_{12} \\approx 0.2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "79238099",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-16T15:57:31.335006Z",
     "iopub.status.busy": "2026-02-16T15:57:31.334669Z",
     "iopub.status.idle": "2026-02-16T15:57:35.959770Z",
     "shell.execute_reply": "2026-02-16T15:57:35.957422Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 92820 characters]",
      "text/plain": [
       "<Figure size 1400x1000 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "📊 Regime Switching Analysis:\n",
      "  Observed frequencies: [0.3002 0.3802 0.3196]\n",
      "  Theoretical stationary: [0.41573034 0.3258427  0.25842697]\n"
     ]
    }
   ],
   "source": [
    "# Simulate regime-switching process\n",
    "def simulate_regime_switching(Q, regime_params, x0, T, dt):\n",
    "    \"\"\"\n",
    "    Simulate regime-switching stochastic process\n",
    "    \n",
    "    Args:\n",
    "        Q: Transition rate matrix (N x N)\n",
    "        regime_params: List of (mu, sigma) for each regime\n",
    "        x0: Initial state\n",
    "        T: Time horizon\n",
    "        dt: Time step\n",
    "    \"\"\"\n",
    "    n_steps = int(T / dt)\n",
    "    n_regimes = Q.shape[0]\n",
    "    \n",
    "    t = np.linspace(0, T, n_steps)\n",
    "    X = np.zeros(n_steps)\n",
    "    regimes = np.zeros(n_steps, dtype=int)\n",
    "    \n",
    "    X[0] = x0\n",
    "    regimes[0] = 0  # Start in regime 0\n",
    "    \n",
    "    for i in range(1, n_steps):\n",
    "        current_regime = regimes[i-1]\n",
    "        \n",
    "        # Check for regime transition\n",
    "        for j in range(n_regimes):\n",
    "            if j != current_regime:\n",
    "                if np.random.rand() < Q[current_regime, j] * dt:\n",
    "                    current_regime = j\n",
    "                    break\n",
    "        \n",
    "        regimes[i] = current_regime\n",
    "        \n",
    "        # Evolve state according to current regime\n",
    "        mu, sigma = regime_params[current_regime]\n",
    "        dW = np.random.normal(0, np.sqrt(dt))\n",
    "        X[i] = X[i-1] + mu * dt + sigma * dW\n",
    "    \n",
    "    return t, X, regimes\n",
    "\n",
    "# Example: 3-regime system (Slow/Normal/Fast)\n",
    "Q = np.array([\n",
    "    [-0.5,  0.3,  0.2],  # Slow regime\n",
    "    [ 0.4, -0.7,  0.3],  # Normal regime\n",
    "    [ 0.3,  0.4, -0.7]   # Fast regime\n",
    "])\n",
    "\n",
    "regime_params = [\n",
    "    (0.1, 0.2),  # Slow: low drift, low vol\n",
    "    (0.3, 0.4),  # Normal: medium drift, medium vol\n",
    "    (0.5, 0.8)   # Fast: high drift, high vol\n",
    "]\n",
    "\n",
    "t, X, regimes = simulate_regime_switching(Q, regime_params, x0=0.0, T=50.0, dt=0.01)\n",
    "\n",
    "# Plot\n",
    "fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(14, 10), sharex=True)\n",
    "\n",
    "# State trajectory\n",
    "colors = ['blue', 'green', 'red']\n",
    "for i in range(len(t)-1):\n",
    "    ax1.plot(t[i:i+2], X[i:i+2], color=colors[regimes[i]], alpha=0.8, linewidth=0.8)\n",
    "\n",
    "ax1.set_ylabel('State X')\n",
    "ax1.set_title('Regime-Switching Process')\n",
    "ax1.grid(alpha=0.3)\n",
    "\n",
    "# Regime evolution\n",
    "ax2.step(t, regimes, where='post', linewidth=1.5, color='black')\n",
    "ax2.set_ylabel('Regime')\n",
    "ax2.set_yticks([0, 1, 2])\n",
    "ax2.set_yticklabels(['Slow', 'Normal', 'Fast'])\n",
    "ax2.set_title('Regime Evolution')\n",
    "ax2.grid(alpha=0.3)\n",
    "\n",
    "# Regime distribution\n",
    "regime_counts = np.bincount(regimes, minlength=3) / len(regimes)\n",
    "ax3.bar([0, 1, 2], regime_counts, color=colors, alpha=0.7, edgecolor='black')\n",
    "ax3.set_xlabel('Regime')\n",
    "ax3.set_ylabel('Frequency')\n",
    "ax3.set_xticks([0, 1, 2])\n",
    "ax3.set_xticklabels(['Slow', 'Normal', 'Fast'])\n",
    "ax3.set_title('Regime Distribution')\n",
    "ax3.grid(alpha=0.3, axis='y')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Compute stationary distribution\n",
    "from scipy.linalg import null_space\n",
    "pi_stationary = null_space(Q.T)\n",
    "pi_stationary = pi_stationary / pi_stationary.sum()\n",
    "\n",
    "print(\"\\n📊 Regime Switching Analysis:\")\n",
    "print(f\"  Observed frequencies: {regime_counts}\")\n",
    "print(f\"  Theoretical stationary: {pi_stationary.flatten()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3751412",
   "metadata": {},
   "source": [
    "## 4. Jump Diffusion Processes <a id=\"jumps\"></a>\n",
    "\n",
    "### Motivation\n",
    "\n",
    "Continuous diffusion models fail to capture **sudden, discrete events**:\n",
    "- Market crashes/rallies\n",
    "- Equipment failures\n",
    "- Policy changes\n",
    "- Natural disasters\n",
    "- Phase transitions\n",
    "\n",
    "### Lévy Processes and Compound Poisson\n",
    "\n",
    "A jump diffusion process combines:\n",
    "1. **Continuous diffusion**: $\\sigma dW_t$\n",
    "2. **Discrete jumps**: $dJ_t = \\sum_{i=1}^{N_t} Y_i$\n",
    "\n",
    "$$\n",
    "dX_t = \\mu dt + \\sigma dW_t + dJ_t\n",
    "$$\n",
    "\n",
    "where:\n",
    "- $N_t \\sim \\text{Poisson}(\\lambda t)$ = number of jumps by time $t$\n",
    "- $Y_i \\sim F$ = jump size distribution\n",
    "- $\\lambda$ = **jump intensity** (expected jumps per unit time)\n",
    "\n",
    "### HJB with Jump Integral\n",
    "\n",
    "$$\n",
    "\\rho V(x) = \\sup_u \\left[ \\mu(x,u) V'(x) + \\frac{1}{2}\\sigma^2(x,u) V''(x) + L(x,u) + \\lambda \\int [V(x+y) - V(x)] F(dy) \\right]\n",
    "$$\n",
    "\n",
    "The integral term $\\lambda \\mathbb{E}[V(x+Y) - V(x)]$ represents the **expected value change from jumps**.\n",
    "\n",
    "### Jump Size Distributions\n",
    "\n",
    "| Distribution | Density | Use Case |\n",
    "|--------------|---------|----------|\n",
    "| Normal | $\\mathcal{N}(\\mu_j, \\sigma_j^2)$ | Symmetric jumps (up/down equally likely) |\n",
    "| Exponential | $\\lambda e^{-\\lambda y}$ | One-sided jumps (failures, crashes) |\n",
    "| Laplace | $\\frac{1}{2b}e^{-|y-\\mu|/b}$ | Heavy-tailed jumps |\n",
    "| Uniform | $U(a, b)$ | Bounded jumps |\n",
    "\n",
    "### Parameter Selection\n",
    "\n",
    "| Parameter | Typical Range | Effect | How to Choose |\n",
    "|-----------|--------------|--------|---------------|\n",
    "| $\\lambda$ | 0.01 - 5.0 | Jump frequency | Count events per unit time from data |\n",
    "| $\\mu_j$ | Problem-specific | Average jump size | Measure typical event magnitude |\n",
    "| $\\sigma_j$ | $> 0$ | Jump size variability | Standard deviation of observed jumps |\n",
    "\n",
    "**Rule of thumb**: If you expect ~1 jump per 10 time units, set $\\lambda = 0.1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "733539e5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-16T15:57:35.963250Z",
     "iopub.status.busy": "2026-02-16T15:57:35.962947Z",
     "iopub.status.idle": "2026-02-16T15:57:36.683892Z",
     "shell.execute_reply": "2026-02-16T15:57:36.682687Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 115372 characters]",
      "text/plain": [
       "<Figure size 1400x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "📊 Jump Diffusion Analysis:\n",
      "  Expected jumps: 40.0\n",
      "  Observed jumps: 50\n",
      "  Average jump size: -0.536 (theoretical: -0.5)\n",
      "  Std of jumps: 0.178 (theoretical: 0.2)\n",
      "\n",
      "  Impact: Final value with jumps = -7.04 vs 17.97 without jumps\n"
     ]
    }
   ],
   "source": [
    "# Simulate jump diffusion process\n",
    "def simulate_jump_diffusion(mu, sigma, lambda_jump, jump_mean, jump_std, x0, T, dt):\n",
    "    \"\"\"\n",
    "    Simulate Merton jump diffusion model\n",
    "    \n",
    "    dX_t = μ dt + σ dW_t + dJ_t\n",
    "    \n",
    "    where J_t is compound Poisson with Normal jumps\n",
    "    \"\"\"\n",
    "    n_steps = int(T / dt)\n",
    "    t = np.linspace(0, T, n_steps)\n",
    "    X = np.zeros(n_steps)\n",
    "    jumps = np.zeros(n_steps)\n",
    "    \n",
    "    X[0] = x0\n",
    "    \n",
    "    for i in range(1, n_steps):\n",
    "        # Diffusion component\n",
    "        dW = np.random.normal(0, np.sqrt(dt))\n",
    "        dX = mu * dt + sigma * dW\n",
    "        \n",
    "        # Jump component\n",
    "        n_jumps = np.random.poisson(lambda_jump * dt)\n",
    "        if n_jumps > 0:\n",
    "            jump_sizes = np.random.normal(jump_mean, jump_std, n_jumps)\n",
    "            total_jump = jump_sizes.sum()\n",
    "            dX += total_jump\n",
    "            jumps[i] = total_jump\n",
    "        \n",
    "        X[i] = X[i-1] + dX\n",
    "    \n",
    "    return t, X, jumps\n",
    "\n",
    "# Example: System with occasional failures/shocks\n",
    "mu = 0.5        # Baseline drift\n",
    "sigma = 0.3     # Continuous volatility\n",
    "lambda_jump = 2.0  # 2 jumps per time unit (on average)\n",
    "jump_mean = -0.5   # Negative jumps (failures)\n",
    "jump_std = 0.2     # Jump size variability\n",
    "x0 = 10.0       # Initial state\n",
    "T = 20.0        # Time horizon\n",
    "dt = 0.01\n",
    "\n",
    "t, X, jumps = simulate_jump_diffusion(mu, sigma, lambda_jump, jump_mean, jump_std, x0, T, dt)\n",
    "\n",
    "# Also simulate without jumps for comparison\n",
    "t_nodiff, X_nodiff, _ = simulate_jump_diffusion(mu, sigma, 0.0, 0.0, 0.0, x0, T, dt)\n",
    "\n",
    "# Plot\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10), sharex=True)\n",
    "\n",
    "# Trajectories comparison\n",
    "ax1.plot(t, X, linewidth=1.5, color='red', label='With Jumps', alpha=0.8)\n",
    "ax1.plot(t_nodiff, X_nodiff, linewidth=1.5, color='blue', label='Pure Diffusion', alpha=0.6)\n",
    "\n",
    "# Mark jump times\n",
    "jump_times = t[jumps != 0]\n",
    "jump_values = X[jumps != 0]\n",
    "ax1.scatter(jump_times, jump_values, color='black', s=50, zorder=5, label='Jump Events', alpha=0.7)\n",
    "\n",
    "ax1.set_ylabel('State X')\n",
    "ax1.set_title('Jump Diffusion Process vs Pure Diffusion')\n",
    "ax1.legend(fontsize=11)\n",
    "ax1.grid(alpha=0.3)\n",
    "\n",
    "# Jump sizes over time\n",
    "ax2.stem(t, jumps, linefmt='red', markerfmt='ro', basefmt=' ', label='Jump Sizes')\n",
    "ax2.axhline(y=0, color='black', linewidth=0.8)\n",
    "ax2.set_xlabel('Time')\n",
    "ax2.set_ylabel('Jump Size')\n",
    "ax2.set_title('Jump Events')\n",
    "ax2.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Statistics\n",
    "n_observed_jumps = np.sum(jumps != 0)\n",
    "expected_jumps = lambda_jump * T\n",
    "avg_jump_size = jumps[jumps != 0].mean() if n_observed_jumps > 0 else 0\n",
    "\n",
    "print(\"\\n📊 Jump Diffusion Analysis:\")\n",
    "print(f\"  Expected jumps: {expected_jumps:.1f}\")\n",
    "print(f\"  Observed jumps: {n_observed_jumps}\")\n",
    "print(f\"  Average jump size: {avg_jump_size:.3f} (theoretical: {jump_mean})\")\n",
    "print(f\"  Std of jumps: {jumps[jumps != 0].std() if n_observed_jumps > 0 else 0:.3f} (theoretical: {jump_std})\")\n",
    "print(f\"\\n  Impact: Final value with jumps = {X[-1]:.2f} vs {X_nodiff[-1]:.2f} without jumps\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca7aba3",
   "metadata": {},
   "source": [
    "## 5. Combined MRSJD Models <a id=\"mrsjd\"></a>\n",
    "\n",
    "### Why Combine Regime Switching and Jumps?\n",
    "\n",
    "Real systems often exhibit **both**:\n",
    "1. **State-dependent behavior** (regimes)\n",
    "2. **Sudden shocks** (jumps)\n",
    "\n",
    "Examples:\n",
    "- **Manufacturing**: Normal/maintenance regimes + equipment failures (jumps)\n",
    "- **Power grid**: Low/high demand regimes + blackout events (jumps)\n",
    "- **Epidemic**: Endemic/outbreak regimes + super-spreader events (jumps)\n",
    "\n",
    "### Full MRSJD Dynamics\n",
    "\n",
    "$$\n",
    "dX_t = \\mu^{i_t}(X_t)dt + \\sigma^{i_t}(X_t)dW_t + dJ_t^{i_t}\n",
    "$$\n",
    "\n",
    "where:\n",
    "- Drift $\\mu^i$ and volatility $\\sigma^i$ depend on current regime $i_t$\n",
    "- Jump intensity $\\lambda^i$ and distribution $F^i$ also regime-dependent\n",
    "- Regime switches according to $Q$\n",
    "\n",
    "### Coupled HJB with Both Effects\n",
    "\n",
    "$$\n",
    "\\boxed{\n",
    "\\rho V^i(x) = \\sup_u \\left[ \\mu^i V^i_x + \\frac{(\\sigma^i)^2}{2} V^i_{xx} + L^i(x,u) + \\lambda^i \\int [V^i(x+y) - V^i(x)] F^i(dy) + \\sum_{j \\neq i} q_{ij}[V^j(x) - V^i(x)] \\right]\n",
    "}\n",
    "$$\n",
    "\n",
    "This is the **most general** formulation combining:\n",
    "1. ✅ Diffusion: $(\\sigma^i)^2 V^i_{xx}$\n",
    "2. ✅ Jumps: $\\lambda^i \\int [V^i(x+y) - V^i(x)] F^i(dy)$\n",
    "3. ✅ Regime Switching: $\\sum_{j \\neq i} q_{ij}[V^j(x) - V^i(x)]$\n",
    "4. ✅ Optimal Control: $\\sup_u$\n",
    "\n",
    "### Numerical Solution: Finite Differences with Upwind Schemes\n",
    "\n",
    "**Grid discretization**: $x_k = x_{\\min} + k \\Delta x$, $k = 0, ..., N$\n",
    "\n",
    "**Value function approximation**: $V^i(x_k) \\approx V^i_k$\n",
    "\n",
    "**Derivatives**:\n",
    "- Forward: $V_x \\approx (V_{k+1} - V_k) / \\Delta x$\n",
    "- Backward: $V_x \\approx (V_k - V_{k-1}) / \\Delta x$\n",
    "- Central: $V_{xx} \\approx (V_{k+1} - 2V_k + V_{k-1}) / (\\Delta x)^2$\n",
    "\n",
    "**Upwind scheme** (for stability when $\\mu \\neq 0$):\n",
    "$$\n",
    "V_x \\approx \n",
    "\\begin{cases}\n",
    "(V_{k+1} - V_k) / \\Delta x & \\text{if } \\mu > 0 \\text{ (forward)} \\\\\n",
    "(V_k - V_{k-1}) / \\Delta x & \\text{if } \\mu < 0 \\text{ (backward)}\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "**Why upwind?** Prevents numerical oscillations when advection dominates diffusion.\n",
    "\n",
    "### Algorithm: Value Iteration\n",
    "\n",
    "```\n",
    "1. Initialize V^i_k = 0 for all regimes i and grid points k\n",
    "2. Repeat until convergence:\n",
    "   For each regime i:\n",
    "     For each grid point k:\n",
    "       a. Compute derivatives V_x, V_xx\n",
    "       b. Compute jump integral ∫[V(x+y) - V(x)]F(dy)\n",
    "       c. Compute regime switching term Σ q_ij[V^j - V^i]\n",
    "       d. Optimize over control: u* = argmax_u RHS(u)\n",
    "       e. Update: V^i_k ← RHS(u*) / ρ\n",
    "3. Convergence check: ||V_new - V_old|| < tol\n",
    "```\n",
    "\n",
    "### Computational Complexity\n",
    "\n",
    "- **Per iteration**: $O(N \\cdot M \\cdot K)$\n",
    "  - $N$ = number of regimes\n",
    "  - $M$ = grid points\n",
    "  - $K$ = control discretization\n",
    "- **Iterations**: Typically 100-1000\n",
    "- **Total**: $O(10^5 - 10^7)$ operations\n",
    "\n",
    "**Speedup techniques**:\n",
    "- Parallel computation across regimes (Rayon)\n",
    "- Adaptive grid refinement\n",
    "- Policy iteration instead of value iteration\n",
    "- Sparse matrix operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f104f849",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-16T15:57:36.687408Z",
     "iopub.status.busy": "2026-02-16T15:57:36.687121Z",
     "iopub.status.idle": "2026-02-16T15:57:41.069722Z",
     "shell.execute_reply": "2026-02-16T15:57:41.068651Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 134576 characters]",
      "text/plain": [
       "<Figure size 1400x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "📊 MRSJD Analysis:\n",
      "  Total jumps: 32\n",
      "  Time in Stable regime: 18.2 (90.8%)\n",
      "  Time in Volatile regime: 11.8 (59.2%)\n",
      "  Average jump size (Stable): -0.106\n",
      "  Average jump size (Volatile): -0.290\n"
     ]
    }
   ],
   "source": [
    "# Simplified MRSJD simulation (for illustration)\n",
    "def simulate_mrsjd(Q, regime_params_list, x0, T, dt):\n",
    "    \"\"\"\n",
    "    Simulate Markov Regime Switching Jump Diffusion\n",
    "    \n",
    "    Each regime has: (mu, sigma, lambda_jump, jump_mean, jump_std)\n",
    "    \"\"\"\n",
    "    n_steps = int(T / dt)\n",
    "    n_regimes = Q.shape[0]\n",
    "    \n",
    "    t = np.linspace(0, T, n_steps)\n",
    "    X = np.zeros(n_steps)\n",
    "    regimes = np.zeros(n_steps, dtype=int)\n",
    "    jump_events = []\n",
    "    \n",
    "    X[0] = x0\n",
    "    regimes[0] = 0\n",
    "    \n",
    "    for i in range(1, n_steps):\n",
    "        current_regime = regimes[i-1]\n",
    "        mu, sigma, lam, jmu, jsig = regime_params_list[current_regime]\n",
    "        \n",
    "        # Check regime transition\n",
    "        for j in range(n_regimes):\n",
    "            if j != current_regime and np.random.rand() < Q[current_regime, j] * dt:\n",
    "                current_regime = j\n",
    "                break\n",
    "        \n",
    "        regimes[i] = current_regime\n",
    "        \n",
    "        # Diffusion\n",
    "        dW = np.random.normal(0, np.sqrt(dt))\n",
    "        dX = mu * dt + sigma * dW\n",
    "        \n",
    "        # Jumps (regime-dependent)\n",
    "        n_jumps = np.random.poisson(lam * dt)\n",
    "        if n_jumps > 0:\n",
    "            jump_size = np.sum(np.random.normal(jmu, jsig, n_jumps))\n",
    "            dX += jump_size\n",
    "            jump_events.append((t[i], jump_size, current_regime))\n",
    "        \n",
    "        X[i] = X[i-1] + dX\n",
    "    \n",
    "    return t, X, regimes, jump_events\n",
    "\n",
    "# Example: 2-regime system with regime-dependent jumps\n",
    "Q = np.array([\n",
    "    [-0.3, 0.3],\n",
    "    [0.5, -0.5]\n",
    "])\n",
    "\n",
    "# Regime 0: Stable (low vol, rare small jumps)\n",
    "# Regime 1: Volatile (high vol, frequent large jumps)\n",
    "regime_params_list = [\n",
    "    (0.2, 0.3, 0.5, -0.1, 0.05),  # Stable: mu, sigma, lambda, jump_mu, jump_sigma\n",
    "    (0.1, 0.8, 2.0, -0.3, 0.15)   # Volatile\n",
    "]\n",
    "\n",
    "t, X, regimes, jump_events = simulate_mrsjd(Q, regime_params_list, x0=5.0, T=30.0, dt=0.01)\n",
    "\n",
    "# Plot\n",
    "fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=True)\n",
    "\n",
    "# State trajectory\n",
    "regime_colors = ['blue', 'red']\n",
    "for i in range(len(t)-1):\n",
    "    axes[0].plot(t[i:i+2], X[i:i+2], color=regime_colors[regimes[i]], alpha=0.8, linewidth=1.0)\n",
    "\n",
    "# Mark jumps\n",
    "if jump_events:\n",
    "    jump_t = [j[0] for j in jump_events]\n",
    "    jump_idx = [np.argmin(np.abs(t - jt)) for jt in jump_t]\n",
    "    axes[0].scatter([t[i] for i in jump_idx], [X[i] for i in jump_idx], \n",
    "                   color='black', s=60, zorder=5, marker='x', label='Jumps')\n",
    "\n",
    "axes[0].set_ylabel('State X')\n",
    "axes[0].set_title('MRSJD: Combined Regime Switching + Jump Diffusion')\n",
    "axes[0].legend()\n",
    "axes[0].grid(alpha=0.3)\n",
    "\n",
    "# Regime evolution\n",
    "axes[1].step(t, regimes, where='post', linewidth=1.5, color='black')\n",
    "axes[1].fill_between(t, regimes, alpha=0.3, step='post', \n",
    "                     color=['blue' if r==0 else 'red' for r in regimes])\n",
    "axes[1].set_ylabel('Regime')\n",
    "axes[1].set_yticks([0, 1])\n",
    "axes[1].set_yticklabels(['Stable', 'Volatile'])\n",
    "axes[1].set_title('Regime Transitions')\n",
    "axes[1].grid(alpha=0.3)\n",
    "\n",
    "# Jump events by regime\n",
    "if jump_events:\n",
    "    regime_0_jumps = [j for j in jump_events if j[2] == 0]\n",
    "    regime_1_jumps = [j for j in jump_events if j[2] == 1]\n",
    "    \n",
    "    if regime_0_jumps:\n",
    "        axes[2].scatter([j[0] for j in regime_0_jumps], [j[1] for j in regime_0_jumps],\n",
    "                       color='blue', s=50, alpha=0.7, label='Stable Regime Jumps')\n",
    "    if regime_1_jumps:\n",
    "        axes[2].scatter([j[0] for j in regime_1_jumps], [j[1] for j in regime_1_jumps],\n",
    "                       color='red', s=50, alpha=0.7, label='Volatile Regime Jumps')\n",
    "\n",
    "axes[2].axhline(y=0, color='black', linewidth=0.8)\n",
    "axes[2].set_xlabel('Time')\n",
    "axes[2].set_ylabel('Jump Size')\n",
    "axes[2].set_title('Jump Events by Regime')\n",
    "axes[2].legend()\n",
    "axes[2].grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Statistics\n",
    "print(\"\\n📊 MRSJD Analysis:\")\n",
    "print(f\"  Total jumps: {len(jump_events)}\")\n",
    "regime_times = [np.sum(regimes == i) * dt for i in range(2)]\n",
    "print(f\"  Time in Stable regime: {regime_times[0]:.1f} ({regime_times[0]/T*100:.1f}%)\")\n",
    "print(f\"  Time in Volatile regime: {regime_times[1]:.1f} ({regime_times[1]/T*100:.1f}%)\")\n",
    "if jump_events:\n",
    "    avg_jump_0 = np.mean([j[1] for j in jump_events if j[2] == 0]) if len([j for j in jump_events if j[2] == 0]) > 0 else 0\n",
    "    avg_jump_1 = np.mean([j[1] for j in jump_events if j[2] == 1]) if len([j for j in jump_events if j[2] == 1]) > 0 else 0\n",
    "    print(f\"  Average jump size (Stable): {avg_jump_0:.3f}\")\n",
    "    print(f\"  Average jump size (Volatile): {avg_jump_1:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "822b77f8",
   "metadata": {},
   "source": [
    "## 6. Practical Parameter Selection Guide <a id=\"params\"></a>\n",
    "\n",
    "### How to Choose Parameters for Your Problem\n",
    "\n",
    "#### Step 1: Identify Regimes\n",
    "\n",
    "Ask: Does the system have distinct \"modes\" or \"states\"?\n",
    "\n",
    "**Examples**:\n",
    "- Manufacturing: Normal / Degraded / Failed\n",
    "- Weather: Clear / Cloudy / Storm\n",
    "- Network: Low / Medium / High traffic\n",
    "\n",
    "**Tip**: Start with 2-3 regimes. More regimes = more parameters to estimate.\n",
    "\n",
    "#### Step 2: Estimate Regime Persistence\n",
    "\n",
    "**Question**: How long does each regime typically last?\n",
    "\n",
    "**Formula**: $q_{ij} = \\frac{1}{\\text{expected duration in regime } i}$\n",
    "\n",
    "**Example**: If \"Normal\" regime lasts ~10 time units:\n",
    "- Total exit rate from Normal: $q_{01} + q_{02} = 0.1$\n",
    "- Split based on transition probabilities\n",
    "\n",
    "#### Step 3: Characterize Within-Regime Dynamics\n",
    "\n",
    "For each regime $i$:\n",
    "\n",
    "| Parameter | Method | Example |\n",
    "|-----------|--------|----------|\n",
    "| $\\mu^i$ | Sample mean of increments | $\\bar{\\Delta X} / \\Delta t$ |\n",
    "| $\\sigma^i$ | Sample std of increments | $\\text{std}(\\Delta X) / \\sqrt{\\Delta t}$ |\n",
    "| $\\lambda^i$ | Count events per time | $N_{\\text{jumps}} / T$ |\n",
    "| Jump mean | Average jump size | $\\bar{Y}$ |\n",
    "| Jump std | Std of jump sizes | $\\text{std}(Y)$ |\n",
    "\n",
    "#### Step 4: Validate with Simulations\n",
    "\n",
    "Before solving the HJB:\n",
    "1. Simulate the process with chosen parameters\n",
    "2. Check if trajectories \"look right\"\n",
    "3. Compare summary statistics to data\n",
    "4. Adjust and iterate\n",
    "\n",
    "### Common Pitfalls and Solutions\n",
    "\n",
    "| Problem | Symptom | Solution |\n",
    "|---------|---------|----------|\n",
    "| Too many regimes | Overfitting, unstable estimates | Use 2-3 regimes; combine similar ones |\n",
    "| Wrong time scale | Unrealistic dynamics | Match $q_{ij}$ to actual durations |\n",
    "| Numerical instability | Oscillations, divergence | Reduce grid spacing, use upwind scheme |\n",
    "| Slow convergence | Many iterations needed | Better initial guess, increase tolerance |\n",
    "| High dimensionality | Curse of dimensionality | Reduce state space, use approximations |\n",
    "\n",
    "### Sensitivity Analysis\n",
    "\n",
    "Always check how results change with parameters:\n",
    "1. Vary each parameter by ±20%\n",
    "2. Observe impact on optimal policy and value function\n",
    "3. Identify which parameters matter most\n",
    "4. Focus calibration efforts on sensitive parameters\n",
    "\n",
    "### When to Use vs. Not Use\n",
    "\n",
    "✅ **Good fit for HJB optimal control**:\n",
    "- Continuous state space (position, temperature, concentration)\n",
    "- Known or learnable dynamics\n",
    "- Quantifiable objectives\n",
    "- Medium-dimensional problems (1-3 state variables)\n",
    "- Offline planning acceptable\n",
    "\n",
    "❌ **Not recommended**:\n",
    "- Purely discrete decisions (use dynamic programming)\n",
    "- Unknown dynamics (use reinforcement learning)\n",
    "- High-dimensional state (>5 variables)\n",
    "- Real-time requirements (<1ms response)\n",
    "- Purely deterministic problems (use calculus of variations)\n",
    "\n",
    "### Alternative Approaches\n",
    "\n",
    "| Method | When to Use | Pros | Cons |\n",
    "|--------|-------------|------|------|\n",
    "| **LQR/LQG** | Linear dynamics, quadratic cost | Fast, analytical solution | Limited to LQ problems |\n",
    "| **MPC** | Need real-time receding horizon | Handles constraints well | Computational cost |\n",
    "| **RL (DQN, PPO)** | Unknown dynamics | Model-free, flexible | Sample inefficient |\n",
    "| **PID Control** | Simple SISO systems | Easy to tune | No optimality guarantee |\n",
    "| **Bang-Bang** | Hard constraints | Simple implementation | Non-smooth control |\n",
    "\n",
    "### Further Reading\n",
    "\n",
    "1. **Books**:\n",
    "   - Fleming & Rishel: \"Deterministic and Stochastic Optimal Control\"\n",
    "   - Øksendal & Sulem: \"Applied Stochastic Control of Jump Diffusions\"\n",
    "   - Bertsekas: \"Dynamic Programming and Optimal Control\"\n",
    "\n",
    "2. **Papers**:\n",
    "   - Guo & Hernandez-Lerma: \"Continuous-Time Markov Decision Processes\"\n",
    "   - Pham: \"Continuous-time Stochastic Control and Optimization with Financial Applications\"\n",
    "\n",
    "3. **Software**:\n",
    "   - This library (optimizr): Generic optimal control solvers\n",
    "   - PROPT: MATLAB optimal control toolbox\n",
    "   - CasADi: Nonlinear optimization and optimal control"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "432ef4af",
   "metadata": {},
   "source": [
    "## Summary and Next Steps\n",
    "\n",
    "### What We Covered\n",
    "\n",
    "1. ✅ **Optimal Control Theory**: HJB equations, value functions\n",
    "2. ✅ **Regime Switching**: Markov chains, coupled HJB systems\n",
    "3. ✅ **Jump Diffusion**: Lévy processes, compound Poisson\n",
    "4. ✅ **MRSJD Models**: Combined framework for complex systems\n",
    "5. ✅ **Numerical Methods**: Finite differences, upwind schemes, value iteration\n",
    "6. ✅ **Parameter Selection**: Practical guidance and sensitivity analysis\n",
    "\n",
    "### Key Takeaways\n",
    "\n",
    "- Optimal control finds the **best** policy, not just a good one\n",
    "- HJB equations require **solving PDEs** (computational cost)\n",
    "- Regime switching captures **state-dependent behavior**\n",
    "- Jumps model **sudden events** and tail risk\n",
    "- Start simple (2 regimes, pure diffusion) then add complexity\n",
    "\n",
    "### Exercises for Practice\n",
    "\n",
    "1. **Temperature Control**: Design an optimal heating/cooling policy to maintain room temperature near 20°C while minimizing energy cost\n",
    "\n",
    "2. **Inventory Management**: Optimize reorder policy for warehouse with regime-switching demand (normal/holiday)\n",
    "\n",
    "3. **Robot Navigation**: Find optimal path for robot avoiding obstacles with uncertain dynamics and occasional sensor failures (jumps)\n",
    "\n",
    "### Next Tutorial\n",
    "\n",
    "- **Hidden Markov Models (HMM)**: When regime is not directly observable\n",
    "- **MCMC Sampling**: Bayesian inference for parameter estimation\n",
    "- **Sparse Optimization**: High-dimensional problems with sparsity\n",
    "\n",
    "---\n",
    "\n",
    "**Questions?** Open an issue on the repository or consult the API documentation.\n",
    "\n",
    "**Happy Optimizing! 🚀**"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
