{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f7ec988d",
   "metadata": {},
   "source": [
    "# DCT Laboratory — Volume II, Chapter 6\n",
    "## Optimal Control of Enterprise Systems\n",
    "**Seed `26206`** · Companion to the chapter and AXIOM Module **AXIOM-06 (Vol. II)**\n",
    "\n",
    "Pontryagin, solved by hand twice. **Problem 1**: quadratic effort cost,\n",
    "terminal capital reward — the adjoint equation gives every control in closed\n",
    "form ($u_k = \\lambda_{k+1}$), and the costate is verified as **the marginal\n",
    "value of capital**. **Problem 2**: linear cost with $u \\in [0,1]$ — the\n",
    "switching function delivers genuine **bang-bang**, with the switch at $k = 1$\n",
    "computed from a sign change. Mirrored in `DCT_V2_Ch06_Lab.xlsx`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25ec4129",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-14T00:00:13.368159Z",
     "iopub.status.busy": "2026-07-14T00:00:13.367403Z",
     "iopub.status.idle": "2026-07-14T00:00:13.878761Z",
     "shell.execute_reply": "2026-07-14T00:00:13.877442Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams['figure.dpi']=110\n",
    "\n",
    "import numpy as np\n",
    "SEED = 26206\n",
    "A, Q, N = 0.9, 2.0, 6\n",
    "K0 = 4.0\n",
    "# Problem 1: max Q*K_N - sum u^2/2 ; K' = A*K + u\n",
    "# adjoint: lam_k = A*lam_{k+1}, lam_N = Q  =>  lam_k = Q*A^(N-k) ; u_k = lam_{k+1}\n",
    "def costate(k): return Q*A**(N-k)\n",
    "def u_star(k): return costate(k+1)\n",
    "def solve_p1():\n",
    "    K, J = K0, 0.0\n",
    "    for k in range(N):\n",
    "        u = u_star(k); J -= u*u/2\n",
    "        K = A*K + u\n",
    "    return J + Q*K, K\n",
    "# Problem 2: max Q*K_N - C*sum(u) ; u in [0,1] ; switching sigma_k = lam_{k+1} - C\n",
    "C = 1.2\n",
    "def sigma(k): return costate(k+1) - C\n",
    "def solve_p2():\n",
    "    K, J = K0, 0.0\n",
    "    us = []\n",
    "    for k in range(N):\n",
    "        u = 1.0 if sigma(k) > 0 else 0.0\n",
    "        us.append(u); J -= C*u\n",
    "        K = A*K + u\n",
    "    return J + Q*K, K, us\n",
    "\n",
    "def reference_values():\n",
    "    J1, K1 = solve_p1()\n",
    "    J2, K2, us = solve_p2()\n",
    "    switch = next(k for k,u in enumerate(us) if u > 0)\n",
    "    return {\n",
    "        \"lambda_0\": round(costate(0),4),\n",
    "        \"u0_star\": round(u_star(0),4), \"u5_star\": round(u_star(5),4),\n",
    "        \"K_N_p1\": round(K1,4), \"J_star_p1\": round(J1,4),\n",
    "        \"marginal_dJ_dK0\": round(Q*A**N,4),      # = lambda_0, exactly (linear dynamics)\n",
    "        \"sigma_0\": round(sigma(0),4), \"sigma_1\": round(sigma(1),4),\n",
    "        \"switch_k\": switch, \"n_active\": int(sum(us)),\n",
    "        \"J_star_p2\": round(J2,4), \"K_N_p2\": round(K2,4),\n",
    "    }\n",
    "if __name__ == \"__main__\":\n",
    "    [print(f\"{k:18s} {v}\") for k,v in reference_values().items()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5d154bb",
   "metadata": {},
   "source": [
    "## Panel 1 — The maximum principle, in closed form\n",
    "Maximize $Q K_N - \\sum u_k^2/2$ with $K_{k+1} = 0.9K_k + u_k$. The Hamiltonian\n",
    "$H = -u^2/2 + \\lambda_{k+1}(0.9K + u)$; the Costate Evolution Theorem gives\n",
    "$\\lambda_k = 0.9\\lambda_{k+1}$ with transversality $\\lambda_N = Q = 2$, so\n",
    "$\\lambda_k = 2 \\cdot 0.9^{6-k}$ — and maximizing $H$ in $u$ gives\n",
    "$u_k^* = \\lambda_{k+1}$. **Controls rise toward the horizon** (1.18 → 2.00):\n",
    "late capital decays less before payday, so late effort is worth more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db4f5720",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-14T00:00:13.880895Z",
     "iopub.status.busy": "2026-07-14T00:00:13.880592Z",
     "iopub.status.idle": "2026-07-14T00:00:14.196162Z",
     "shell.execute_reply": "2026-07-14T00:00:14.195049Z"
    }
   },
   "outputs": [],
   "source": [
    "ks = np.arange(N)\n",
    "lams = [costate(k) for k in range(N+1)]\n",
    "us = [u_star(k) for k in ks]\n",
    "fig, axes = plt.subplots(1,2, figsize=(10,3.9))\n",
    "axes[0].plot(range(N+1), lams, \"o-\", c=\"#0B3D2E\", lw=2.2, ms=5)\n",
    "axes[0].set(xlabel=\"k\", ylabel=\"λ_k\", title=\"Costate: marginal value of capital, λ_k = 2·0.9^(6−k)\")\n",
    "axes[0].grid(alpha=.25)\n",
    "axes[1].bar(ks, us, color=\"#C8A24B\", width=.6)\n",
    "axes[1].set(xlabel=\"k\", ylabel=\"u_k*\", title=\"Optimal controls: u_k = λ_{k+1}, rising to the horizon\")\n",
    "axes[1].grid(alpha=.25, axis=\"y\")\n",
    "plt.tight_layout(); plt.show()\n",
    "J1, K1 = solve_p1()\n",
    "print(f\"u0* = {u_star(0):.4f}   u5* = {u_star(5):.4f}   K_N = {K1:.4f}   J* = {J1:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e4e0ea8",
   "metadata": {},
   "source": [
    "## Panel 2 — The costate is a derivative\n",
    "Adjoint Variables Quantify the Marginal Value of Enterprise States (Prop.):\n",
    "$\\partial J^*/\\partial K_0 = \\lambda_0$. With linear dynamics this is exact,\n",
    "no finite difference needed: one extra unit of $K_0$ compounds to $0.9^6$ at\n",
    "the horizon and earns $Q \\cdot 0.9^6 = 1.0629 = \\lambda_0$. Chapter 3's\n",
    "shadow price, Chapter 4's envelope derivative, and now the costate: **the same\n",
    "economic object, its third appearance** — this time indexed by time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b68e48f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-14T00:00:14.198138Z",
     "iopub.status.busy": "2026-07-14T00:00:14.197927Z",
     "iopub.status.idle": "2026-07-14T00:00:14.203447Z",
     "shell.execute_reply": "2026-07-14T00:00:14.202474Z"
    }
   },
   "outputs": [],
   "source": [
    "print(f\"lambda_0 (adjoint at k=0):        {costate(0):.4f}\")\n",
    "print(f\"dJ*/dK0 (exact, linear dynamics): {Q*A**N:.4f}\")\n",
    "print(\"The multiplier family tree: shadow price (Ch.3) = envelope derivative (Ch.4) = costate (Ch.6)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87addd47",
   "metadata": {},
   "source": [
    "## Panel 3 — Bang-bang: when the Hamiltonian is linear in control\n",
    "Maximize $Q K_N - 1.2\\sum u_k$ with $u_k \\in [0,1]$: $H$ is linear in $u$, so\n",
    "the maximum principle pushes $u$ to a **bound**, chosen by the switching\n",
    "function $\\sigma_k = \\lambda_{k+1} - 1.2$. At $k = 0$: $\\sigma = -0.019$\n",
    "(barely negative — don't invest); from $k = 1$: positive — invest at full\n",
    "throttle. One sign change, one switch, five active periods: the structure of\n",
    "the optimal control read directly off the costate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e47a6e43",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-14T00:00:14.205237Z",
     "iopub.status.busy": "2026-07-14T00:00:14.205046Z",
     "iopub.status.idle": "2026-07-14T00:00:14.456971Z",
     "shell.execute_reply": "2026-07-14T00:00:14.455762Z"
    }
   },
   "outputs": [],
   "source": [
    "J2, K2, us2 = solve_p2()\n",
    "sigs = [sigma(k) for k in range(N)]\n",
    "fig, ax = plt.subplots(figsize=(7.8,4.0))\n",
    "ax.bar(range(N), us2, color=\"#C8A24B\", width=.55, label=\"u_k* (bang-bang)\")\n",
    "ax2 = ax.twinx()\n",
    "ax2.plot(range(N), sigs, \"o-\", c=\"#0B3D2E\", lw=2, label=\"switching σ_k = λ_{k+1} − c\")\n",
    "ax2.axhline(0, c=\"#8A8F8B\", lw=1, ls=\":\")\n",
    "ax.set(xlabel=\"k\", ylabel=\"u_k*\", title=\"σ crosses zero once — the control switches once (seed 26206)\")\n",
    "ax2.set_ylabel(\"σ_k\")\n",
    "ax.legend(loc=\"upper left\", frameon=False, fontsize=9); ax2.legend(loc=\"lower right\", frameon=False, fontsize=9)\n",
    "plt.tight_layout(); plt.show()\n",
    "print(f\"sigma_0 = {sigma(0):+.4f}  sigma_1 = {sigma(1):+.4f}   switch at k = {next(k for k,u in enumerate(us2) if u>0)}\")\n",
    "print(f\"J* = {J2:.4f}   K_N = {K2:.4f}   active periods: {int(sum(us2))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bb8fad9",
   "metadata": {},
   "source": [
    "## Validation — agrees with `DCT_V2_Ch06_Lab.xlsx`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d89515bd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-14T00:00:14.459096Z",
     "iopub.status.busy": "2026-07-14T00:00:14.458862Z",
     "iopub.status.idle": "2026-07-14T00:00:14.466642Z",
     "shell.execute_reply": "2026-07-14T00:00:14.465231Z"
    }
   },
   "outputs": [],
   "source": [
    "ref = reference_values()\n",
    "expected = {\"lambda_0\":1.0629,\"u0_star\":1.181,\"u5_star\":2.0,\"K_N_p1\":9.6791,\"J_star_p1\":11.8049,\n",
    " \"marginal_dJ_dK0\":1.0629,\"sigma_0\":-0.019,\"sigma_1\":0.1122,\"switch_k\":1,\"n_active\":5,\n",
    " \"J_star_p2\":6.4417,\"K_N_p2\":6.2209}\n",
    "for k,v in expected.items():\n",
    "    assert abs(ref[k]-v)<5e-4, f\"MISMATCH {k}\"\n",
    "    print(f\"PASS  {k:18s} {ref[k]}\")\n",
    "print(\"\\nAll checkpoints agree — seed 26206.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bcb00ea",
   "metadata": {},
   "source": [
    "**Next**: Exercises 6.5–6.9 (Part C) move the cost c through the switching threshold; AXIOM-06's Hamiltonian console animates costate, control, and trajectory together. Chapter 7 solves the same problems backward: dynamic programming. Solutions: IM Vol. II, Ch. 6."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}