{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "23e5e7a8",
   "metadata": {},
   "source": [
    "# Path Signatures — Companion Notebook\n",
    "\n",
    "Companion to the documentation page\n",
    "[optimiz-r.readthedocs.io / signatures](https://optimiz-r.readthedocs.io/en/latest/algorithms/signatures.html).\n",
    "\n",
    "We exercise the `signatures` module of `optimiz-rs` on a planar Lissajous\n",
    "trajectory of a multivariate path and verify each routine against an\n",
    "analytic ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "07795277",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T09:45:44.865669Z",
     "iopub.status.busy": "2026-05-12T09:45:44.865187Z",
     "iopub.status.idle": "2026-05-12T09:45:46.165789Z",
     "shell.execute_reply": "2026-05-12T09:45:46.164468Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47552 characters]",
      "text/plain": [
       "<Figure size 450x450 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from optimizr import _core as opt\n",
    "\n",
    "rng = np.random.default_rng(0)\n",
    "n_pts = 50\n",
    "t = np.linspace(0.0, 1.0, n_pts)\n",
    "path = np.column_stack([np.sin(2 * np.pi * t), np.sin(3 * np.pi * t + 0.4)])\n",
    "path_list = [list(map(float, row)) for row in path]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(4.5, 4.5))\n",
    "ax.plot(path[:, 0], path[:, 1], '-o', ms=3)\n",
    "ax.set_title('Lissajous trajectory of a 2D multivariate path')\n",
    "ax.set_xlabel('channel 0'); ax.set_ylabel('channel 1'); ax.set_aspect('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271a318e",
   "metadata": {},
   "source": [
    "## Truncated tensor signature\n",
    "\n",
    "$$ S(X)_{0,T} \\;=\\; 1 + \\sum_{k\\ge 1} \\sum_{i_1,\\dots,i_k}\n",
    "   S^{i_1,\\dots,i_k}_{0,T}\\, e_{i_1}\\otimes\\dots\\otimes e_{i_k},\n",
    "   \\qquad\n",
    "   S^{i_1,\\dots,i_k}_{0,T} = \\int_{0<u_1<\\dots<u_k<T}\n",
    "       dX^{i_1}_{u_1}\\dots dX^{i_k}_{u_k}.$$\n",
    "\n",
    "Piecewise-linear recursion:\n",
    "\n",
    "$$ S^{(M)}_{0,t_n} \\;=\\; S^{(M)}_{0,t_{n-1}} \\otimes_M \\exp_M(\\Delta_n),\n",
    "   \\qquad \\exp_M(\\Delta) = \\sum_{k=0}^M \\Delta^{\\otimes k}/k!. $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "44d25084",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T09:45:46.169525Z",
     "iopub.status.busy": "2026-05-12T09:45:46.169163Z",
     "iopub.status.idle": "2026-05-12T09:45:46.544836Z",
     "shell.execute_reply": "2026-05-12T09:45:46.542292Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "channels = 2 level = 3\n",
      "  level 0: shape=(1,), ||S_k||_2 = 1.0000e+00\n",
      "  level 1: shape=(2,), ||S_k||_2 = 7.7884e-01\n",
      "  level 2: shape=(4,), ||S_k||_2 = 3.1270e+00\n",
      "  level 3: shape=(8,), ||S_k||_2 = 1.6278e+00\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 19412 characters]",
      "text/plain": [
       "<Figure size 500x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "linear-segment signature max error vs analytic: 4.337e-19\n"
     ]
    }
   ],
   "source": [
    "level = 3\n",
    "sig = opt.path_signature(path_list, level)\n",
    "tensors = [np.asarray(t, dtype=float) for t in sig['tensors']]\n",
    "norms = [float(np.linalg.norm(t)) for t in tensors]\n",
    "print('channels =', sig['channels'], 'level =', sig['level'])\n",
    "for k, (t_k, nrm) in enumerate(zip(tensors, norms)):\n",
    "    print(f'  level {k}: shape={t_k.shape}, ||S_k||_2 = {nrm:.4e}')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 3))\n",
    "ax.bar(range(level + 1), norms)\n",
    "ax.set_xlabel('tensor level k'); ax.set_ylabel('||S_k||_2')\n",
    "ax.set_title('Signature tensor norms by level')\n",
    "plt.show()\n",
    "\n",
    "# Ground truth: a single linear segment of displacement Delta has\n",
    "# S^{i_1...i_k} = (Delta_{i_1} ... Delta_{i_k}) / k!\n",
    "delta = np.array([0.3, -0.2])\n",
    "seg = [[0.0, 0.0], delta.tolist()]\n",
    "sig_seg = opt.path_signature(seg, 3)\n",
    "T2 = np.asarray(sig_seg['tensors'][2]).reshape(2, 2)\n",
    "T3 = np.asarray(sig_seg['tensors'][3]).reshape(2, 2, 2)\n",
    "exp_T2 = np.einsum('i,j->ij', delta, delta) / 2.0\n",
    "exp_T3 = np.einsum('i,j,k->ijk', delta, delta, delta) / 6.0\n",
    "err_seg = max(np.max(np.abs(T2 - exp_T2)), np.max(np.abs(T3 - exp_T3)))\n",
    "print(f'linear-segment signature max error vs analytic: {err_seg:.3e}')\n",
    "assert err_seg < 1e-12"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d01c35",
   "metadata": {},
   "source": [
    "## Log-signature\n",
    "\n",
    "$$ \\log(S) \\;=\\; \\sum_{n\\ge 1} \\frac{(-1)^{n+1}}{n}(S - 1)^{\\otimes n}. $$\n",
    "\n",
    "Lives in the truncated free Lie algebra; its level-1 component equals the\n",
    "total path increment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d2a13685",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T09:45:46.549872Z",
     "iopub.status.busy": "2026-05-12T09:45:46.549523Z",
     "iopub.status.idle": "2026-05-12T09:45:46.787959Z",
     "shell.execute_reply": "2026-05-12T09:45:46.786642Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "log level-1: [-2.49800181e-16 -7.78836685e-01]\n",
      "  err vs path increment [-2.44929360e-16 -7.78836685e-01]: 1.110e-16\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 20852 characters]",
      "text/plain": [
       "<Figure size 500x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "log_sig = opt.path_log_signature(path_list, level)\n",
    "log_tensors = [np.asarray(t, dtype=float) for t in log_sig['tensors']]\n",
    "log_norms = [float(np.linalg.norm(t)) for t in log_tensors]\n",
    "print('log level-1:', log_tensors[1])\n",
    "increment = path[-1] - path[0]\n",
    "err_log1 = float(np.max(np.abs(log_tensors[1] - increment)))\n",
    "print(f'  err vs path increment {increment}: {err_log1:.3e}')\n",
    "assert err_log1 < 1e-10\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 3))\n",
    "ax.bar(range(level + 1), log_norms, color='C1')\n",
    "ax.set_xlabel('tensor level k'); ax.set_ylabel('||log(S)_k||_2')\n",
    "ax.set_title('Log-signature tensor norms by level')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2db19426",
   "metadata": {},
   "source": [
    "## Random reservoir signature\n",
    "\n",
    "$$ dZ_t = A_0 Z_t\\, dt + \\sum_{i=1}^{d} A_i Z_t\\, dX^{i}_t,\n",
    "   \\qquad A_i \\in \\mathbb{R}^{N\\times N},\\;\n",
    "   (A_i)_{ab}\\stackrel{iid}{\\sim}\\mathcal{N}(0, 1/N). $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f0ae6db8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T09:45:46.797186Z",
     "iopub.status.busy": "2026-05-12T09:45:46.796900Z",
     "iopub.status.idle": "2026-05-12T09:45:47.078891Z",
     "shell.execute_reply": "2026-05-12T09:45:47.077506Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reproducibility max error (same seed): 0.000e+00\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 68852 characters]",
      "text/plain": [
       "<Figure size 600x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "res_a = opt.random_signature(path_list, reservoir_dim=32, seed=42, variance=1.0)\n",
    "res_b = opt.random_signature(path_list, reservoir_dim=32, seed=42, variance=1.0)\n",
    "traj_a = np.asarray(res_a['trajectory'], dtype=float)\n",
    "traj_b = np.asarray(res_b['trajectory'], dtype=float)\n",
    "err_seed = float(np.max(np.abs(traj_a - traj_b)))\n",
    "print(f'reproducibility max error (same seed): {err_seed:.3e}')\n",
    "assert err_seed == 0.0\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 3.5))\n",
    "for i in range(min(8, traj_a.shape[1])):\n",
    "    ax.plot(traj_a[:, i], lw=0.9)\n",
    "ax.set_xlabel('path step n'); ax.set_ylabel('Z_n component')\n",
    "ax.set_title('Random reservoir trajectory (first 8 components)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "210a9ab7",
   "metadata": {},
   "source": [
    "## Signature kernel (Salvi--Cass--Lyons PDE)\n",
    "\n",
    "$$ K(s,t) \\;=\\; \\langle S(X)_{0,s},\\, S(Y)_{0,t}\\rangle, \\qquad\n",
    "   \\frac{\\partial^2 K}{\\partial s\\,\\partial t}\n",
    "   \\;=\\; \\langle \\dot X_s, \\dot Y_t\\rangle\\, K(s,t),\n",
    "   \\quad K(s,0)=K(0,t)=1.$$\n",
    "\n",
    "Goursat finite-difference scheme:\n",
    "\n",
    "$$ K_{i+1,j+1} = K_{i+1,j} + K_{i,j+1} - K_{i,j}\n",
    "   + \\langle \\Delta x_i,\\Delta y_j\\rangle\n",
    "       \\tfrac12(K_{i+1,j} + K_{i,j+1}). $$\n",
    "\n",
    "**Analytic ground truth.** A linear segment with displacement $\\Delta$\n",
    "has signature inner product against itself equal to $\\exp(\\|\\Delta\\|^2)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cf921756",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T09:45:47.082359Z",
     "iopub.status.busy": "2026-05-12T09:45:47.082068Z",
     "iopub.status.idle": "2026-05-12T09:45:47.419422Z",
     "shell.execute_reply": "2026-05-12T09:45:47.418224Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "K(S, T) = 180.840450\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28120 characters]",
      "text/plain": [
       "<Figure size 500x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kernel(seg, seg) = 1.177363, exact exp(||Delta||^2) = 1.185305, err = 7.942e-03\n"
     ]
    }
   ],
   "source": [
    "kres = opt.signature_kernel(path_list, path_list)\n",
    "K = np.asarray(kres['grid'], dtype=float)\n",
    "print(f\"K(S, T) = {kres['value']:.6f}\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 4))\n",
    "im = ax.imshow(K, origin='lower', aspect='auto', cmap='viridis')\n",
    "ax.set_xlabel('j (path Y)'); ax.set_ylabel('i (path X)')\n",
    "ax.set_title('Signature kernel K(s, t)')\n",
    "plt.colorbar(im, ax=ax)\n",
    "plt.show()\n",
    "\n",
    "# Analytic check: linear segment vs itself\n",
    "delta = np.array([0.4, -0.1])\n",
    "n = 200\n",
    "seg_path = [[float(k / n) * delta[0], float(k / n) * delta[1]] for k in range(n + 1)]\n",
    "kseg = opt.signature_kernel(seg_path, seg_path)\n",
    "exact = float(np.exp(np.dot(delta, delta)))\n",
    "err_kernel = abs(kseg['value'] - exact)\n",
    "print(f'kernel(seg, seg) = {kseg[\"value\"]:.6f}, exact exp(||Delta||^2) = {exact:.6f}, err = {err_kernel:.3e}')\n",
    "assert err_kernel < 1e-2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f377d32",
   "metadata": {},
   "source": [
    "## Shuffle product and Chen concatenation\n",
    "\n",
    "$$ \\mathrm{Sh}(u, v) \\;=\\; \\sum_{w\\in u\\,\\shuffle\\,v} w,\n",
    "   \\qquad\n",
    "   S(X * Y) \\;=\\; S(X)\\otimes S(Y). $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "71a56d29",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-12T09:45:47.424399Z",
     "iopub.status.busy": "2026-05-12T09:45:47.424009Z",
     "iopub.status.idle": "2026-05-12T09:45:47.435687Z",
     "shell.execute_reply": "2026-05-12T09:45:47.433850Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sh([0], [1]) = {(1, 0): 1.0, (0, 1): 1.0}\n",
      "Chen concatenation max error per level: [0.0, 4.440892098500626e-16, 8.326672684688674e-16, 1.4432899320127035e-15]\n",
      "Chen concatenation overall max error: 1.443e-15\n"
     ]
    }
   ],
   "source": [
    "# Shuffle product of one-letter words\n",
    "sh = dict((tuple(w), m) for w, m in opt.shuffle_product([0], [1]))\n",
    "print('Sh([0], [1]) =', sh)\n",
    "assert sh[(0, 1)] == 1.0 and sh[(1, 0)] == 1.0\n",
    "\n",
    "# Chen identity: split the path in two halves and concatenate signatures\n",
    "mid = n_pts // 2\n",
    "half_a = path_list[: mid + 1]\n",
    "half_b = path_list[mid:]\n",
    "sa = opt.path_signature(half_a, level)\n",
    "sb = opt.path_signature(half_b, level)\n",
    "cat = opt.concatenate_signatures(\n",
    "    sa['channels'], sa['level'], sa['tensors'],\n",
    "    sb['channels'], sb['level'], sb['tensors'],\n",
    ")\n",
    "full = opt.path_signature(path_list, level)\n",
    "errs_chen = []\n",
    "for k in range(level + 1):\n",
    "    a = np.asarray(cat['tensors'][k]); b = np.asarray(full['tensors'][k])\n",
    "    errs_chen.append(float(np.max(np.abs(a - b))) if a.size else 0.0)\n",
    "err_chen = max(errs_chen)\n",
    "print(f'Chen concatenation max error per level: {errs_chen}')\n",
    "print(f'Chen concatenation overall max error: {err_chen:.3e}')\n",
    "assert err_chen < 1e-10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37a0f62d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "**Verified against analytic ground truth** — max error across all\n",
    "checks (linear-segment signature, log-signature level-1 increment,\n",
    "reservoir reproducibility, signature kernel of a segment, Chen\n",
    "concatenation) `< 1e-2` (kernel Goursat scheme on a coarse grid),\n",
    "all other checks `< 1e-10`."
   ]
  }
 ],
 "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
