{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d555a3fc",
   "metadata": {},
   "source": [
    "# 12 — Stochastic control (Pontryagin LQR)\n",
    "\n",
    "Doc page: [stochastic_control.rst](../../docs/source/algorithms/stochastic_control.rst).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "630ebe94",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T14:05:29.702153Z",
     "iopub.status.busy": "2026-05-12T14:05:29.701881Z",
     "iopub.status.idle": "2026-05-12T14:05:30.298994Z",
     "shell.execute_reply": "2026-05-12T14:05:30.297851Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from optimizr import _core as opt\n",
    "plt.rcParams['figure.figsize'] = (8.5, 4.5)\n",
    "plt.rcParams['figure.dpi'] = 110\n",
    "plt.rcParams['axes.grid'] = True\n",
    "plt.rcParams['grid.alpha'] = 0.3\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ece7571",
   "metadata": {},
   "source": [
    "## Cellule 1 — LQR scalaire et équation de Riccati\n",
    "\n",
    "**Théorème (Kalman, LQR).** Pour le système $\\dot x = ax + bu$ et le\n",
    "coût\n",
    "$$J = \\int_0^T (q x^2 + r u^2)\\,dt + s\\,x(T)^2,$$\n",
    "le contrôle optimal est $u^*(t) = -(b/r)\\,P(t)\\,x(t)$ où $P(t)$\n",
    "satisfait la Riccati rétrograde\n",
    "$$-\\dot P = 2aP - (b^2/r) P^2 + q,\\qquad P(T) = s.$$\n",
    "\n",
    "**Équation pivot (cas $a = b = q = r = 1$).** Le point fixe\n",
    "stationnaire est $P^* = 1 + \\sqrt{2} \\approx 2.414$.\n",
    "\n",
    "**Ce que la cellule vérifie.** Le primitive `pontryagin_lqr` retourne\n",
    "$P(0)$ proche du point fixe pour $T = 5$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b3ac43bf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T14:05:30.302507Z",
     "iopub.status.busy": "2026-05-12T14:05:30.302154Z",
     "iopub.status.idle": "2026-05-12T14:05:30.867738Z",
     "shell.execute_reply": "2026-05-12T14:05:30.866267Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P(0) numérique  = 2.7321\n",
      "P* analytique   = 2.4142\n",
      "P(T) (terminal) = 0.5000\n",
      "Coût optimal    = 2.4701\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 63648 characters]",
      "text/plain": [
       "<Figure size 1430x418 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "res = opt.pontryagin_lqr(\n",
    "    a=1.0, b=1.0, q=1.0, r=1.0, s_terminal=0.5,\n",
    "    x0=1.0, t_horizon=5.0, n_steps=400,\n",
    ")\n",
    "ts = np.array(res['time_grid'])\n",
    "P = np.array(res['riccati'])\n",
    "state = np.array(res['state'])\n",
    "control = np.array(res['control'])\n",
    "\n",
    "P_star = 1.0 + np.sqrt(2.0)\n",
    "print(f\"P(0) numérique  = {P[0]:.4f}\")\n",
    "print(f\"P* analytique   = {P_star:.4f}\")\n",
    "print(f\"P(T) (terminal) = {P[-1]:.4f}\")\n",
    "print(f\"Coût optimal    = {res['cost']:.4f}\")\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(13, 3.8))\n",
    "axes[0].plot(ts, P, lw=2)\n",
    "axes[0].axhline(P_star, ls='--', color='gray',\n",
    "                label=f'P* = {P_star:.3f}')\n",
    "axes[0].set_xlabel('t'); axes[0].set_ylabel('P(t)')\n",
    "axes[0].set_title(\"Riccati\"); axes[0].legend()\n",
    "axes[1].plot(ts, state, lw=2, color='C2')\n",
    "axes[1].set_xlabel('t'); axes[1].set_ylabel('x(t)')\n",
    "axes[1].set_title(\"État optimal\")\n",
    "ts_u = ts[:len(control)]\n",
    "axes[2].plot(ts_u, control, lw=2, color='C3')\n",
    "axes[2].set_xlabel('t'); axes[2].set_ylabel('u(t)')\n",
    "axes[2].set_title(\"Commande optimale\")\n",
    "fig.tight_layout(); plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "274eb1aa",
   "metadata": {},
   "source": [
    "**Résultat attendu.** $P(0) \\approx 1 + \\sqrt{2}$ ; l'état $x$\n",
    "décroît rapidement vers $0$ ; la commande est proportionnelle à\n",
    "$-P(t)\\,x(t)/r$.\n",
    "\n",
    "**Lecture du graphique.** $P(t)$ plate sur $[0, T-\\epsilon]$ puis\n",
    "descend vers $s = 0.5$ ; trajectoires d'état/commande lisses.\n",
    "\n",
    "**Conclusion.** Le LQR scalaire est validé contre le point fixe\n",
    "analytique.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "790043ec",
   "metadata": {},
   "source": [
    "## Cellule 2 — Étude paramétrique : sensibilité au coût terminal\n",
    "\n",
    "**Théorème.** Lorsque $T \\to \\infty$, $P(0)$ tend vers le point fixe\n",
    "indépendamment de $s$.\n",
    "\n",
    "**Équation pivot.** $P^* = (a + \\sqrt{a^2 + b^2 q / r}) \\cdot r / b^2$.\n",
    "Pour $a = b = q = r = 1$, $P^* = 1 + \\sqrt{2}$.\n",
    "\n",
    "**Ce que la cellule vérifie.** Sweep sur plusieurs valeurs de $s$ et\n",
    "de $T$ : $P(0)$ converge vers $P^*$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0e6af886",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T14:05:30.870857Z",
     "iopub.status.busy": "2026-05-12T14:05:30.870588Z",
     "iopub.status.idle": "2026-05-12T14:05:31.158164Z",
     "shell.execute_reply": "2026-05-12T14:05:31.155824Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s =  0.1 : P(0) à T=10 = 2.7321\n",
      "s =  0.5 : P(0) à T=10 = 2.7321\n",
      "s =  2.0 : P(0) à T=10 = 2.7321\n",
      "s =  5.0 : P(0) à T=10 = 2.7321\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 51508 characters]",
      "text/plain": [
       "<Figure size 935x495 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "s_values = [0.1, 0.5, 2.0, 5.0]\n",
    "T_values = [0.5, 1.0, 2.0, 5.0, 10.0]\n",
    "P_star = 1.0 + np.sqrt(2.0)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for s in s_values:\n",
    "    p0s = []\n",
    "    for T_ in T_values:\n",
    "        r = opt.pontryagin_lqr(1.0, 1.0, 1.0, 1.0, s, 1.0, T_, 200)\n",
    "        p0s.append(r['riccati'][0])\n",
    "    ax.plot(T_values, p0s, 'o-', lw=2, label=f's = {s}')\n",
    "    print(f\"s = {s:4.1f} : P(0) à T=10 = {p0s[-1]:.4f}\")\n",
    "ax.axhline(P_star, ls='--', color='black',\n",
    "           label=f'P* = {P_star:.3f}')\n",
    "ax.set_xlabel('horizon T'); ax.set_ylabel('P(0)')\n",
    "ax.set_title(\"Convergence vers le point fixe Riccati\")\n",
    "ax.legend()\n",
    "fig.tight_layout(); plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cd43961",
   "metadata": {},
   "source": [
    "**Résultat attendu.** Toutes les courbes convergent vers\n",
    "$P^* \\approx 2.414$ pour $T$ grand.\n",
    "\n",
    "**Lecture du graphique.** L'effet du terminal $s$ s'estompe quand $T$\n",
    "augmente.\n",
    "\n",
    "**Conclusion.** En horizon long, le coût terminal devient négligeable.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65d4c5e1",
   "metadata": {},
   "source": [
    "## Cellule 3 — Exemple concret : stabilisation d'un pendule inversé\n",
    "\n",
    "**Modèle physique.** Pendule inversé linéarisé autour de la verticale\n",
    "$\\ddot\\theta = (g/\\ell)\\theta + (1/m\\ell^2) u$.  Pour le mode\n",
    "sur-amorti $(\\dot\\theta \\equiv 0)$ on retient l'équation scalaire\n",
    "$\\dot\\theta = a\\theta + bu$ avec $a = g/\\ell$, $b = 1/(m\\ell^2)$.\n",
    "\n",
    "**Équation pivot (Riccati).** $P$ vérifie $-\\dot P = 2aP - (b^2/r)P^2 + q$.\n",
    "\n",
    "**Ce que la cellule vérifie.** Le contrôleur LQR ramène l'angle\n",
    "initial $0.3$ rad vers $0$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7aabb9b7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T14:05:31.162225Z",
     "iopub.status.busy": "2026-05-12T14:05:31.161917Z",
     "iopub.status.idle": "2026-05-12T14:05:31.762984Z",
     "shell.execute_reply": "2026-05-12T14:05:31.761885Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Angle initial : 0.300 rad (17.2°)\n",
      "Angle final   : 5.955e-25 rad\n",
      "Effort max    : 6.177\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 52572 characters]",
      "text/plain": [
       "<Figure size 1210x440 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g, ell, m = 9.81, 1.0, 1.0\n",
    "a, b = g / ell, 1.0 / (m * ell ** 2)\n",
    "res = opt.pontryagin_lqr(a, b, q=10.0, r=1.0, s_terminal=1.0,\n",
    "                          x0=0.3, t_horizon=5.0, n_steps=500)\n",
    "ts = np.array(res['time_grid'])\n",
    "theta = np.array(res['state'])\n",
    "u = np.array(res['control'])\n",
    "\n",
    "print(f\"Angle initial : {theta[0]:.3f} rad ({np.degrees(theta[0]):.1f}°)\")\n",
    "print(f\"Angle final   : {theta[-1]:.3e} rad\")\n",
    "print(f\"Effort max    : {np.abs(u).max():.3f}\")\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n",
    "axes[0].plot(ts, theta, lw=2, color='C3', label=r'$\\theta(t)$')\n",
    "axes[0].axhline(0, ls='--', color='gray', alpha=0.6)\n",
    "axes[0].set_xlabel('t (s)'); axes[0].set_ylabel(r'$\\theta$ (rad)')\n",
    "axes[0].set_title(\"Stabilisation du pendule inversé\")\n",
    "axes[0].legend()\n",
    "ts_u = ts[:len(u)]\n",
    "axes[1].plot(ts_u, u, lw=2, color='C2', label='u(t)')\n",
    "axes[1].set_xlabel('t (s)'); axes[1].set_ylabel('u (couple)')\n",
    "axes[1].set_title(\"Couple appliqué (LQR)\")\n",
    "axes[1].legend()\n",
    "fig.tight_layout(); plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6c1b078",
   "metadata": {},
   "source": [
    "**Résultat attendu.** $\\theta \\to 0$ rapidement, effort borné.\n",
    "\n",
    "**Lecture du graphique.** Décroissance exponentielle de l'angle ;\n",
    "couple initial fort puis amorti.\n",
    "\n",
    "**Conclusion.** Le primitive transforme un système naturellement\n",
    "instable ($a = 9.81 > 0$) en système stable, illustration directe\n",
    "du contrôle optimal.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rhftlab",
   "language": "python",
   "name": "rhftlab"
  },
  "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
