{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-000",
   "metadata": {},
   "source": [
    "# Lesson 3 — Stochastic Simulation & Performance\n",
    "\n",
    "**Companion notebook** for the Lesson 3 slides. You'll build a Monte Carlo ALM model of\n",
    "5,000 immediate annuities under stochastic interest rates and longevity, check it\n",
    "against a textbook value, profile a slow version, make it fast with NumPy and numba,\n",
    "and compute economic capital for two asset strategies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-001",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check(label, got, expected, tol=1e-6):\n",
    "    \"\"\"Compare your answer with the expected value; prints ✅ or ❌ (never raises).\"\"\"\n",
    "    import math\n",
    "    try:\n",
    "        if isinstance(expected, float) or isinstance(got, float):\n",
    "            ok = math.isclose(float(got), float(expected), rel_tol=tol, abs_tol=tol)\n",
    "        elif isinstance(expected, (list, tuple)) and expected and isinstance(expected[0], str):\n",
    "            ok = list(got) == list(expected)\n",
    "        elif hasattr(expected, \"__len__\") and not isinstance(expected, str):\n",
    "            import numpy as np\n",
    "            ok = np.allclose(np.asarray(got, dtype=float), np.asarray(expected, dtype=float), rtol=tol, atol=tol)\n",
    "        else:\n",
    "            ok = got == expected\n",
    "    except Exception:\n",
    "        ok = False\n",
    "    print((\"✅ \" if ok else \"❌ \") + f\"{label}: got {got!r}, expected {expected!r}\")\n",
    "import cProfile\n",
    "import io\n",
    "import math\n",
    "import pstats\n",
    "import random\n",
    "import time\n",
    "import tomllib\n",
    "from pathlib import Path\n",
    "\n",
    "import numba as nb\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "LIFE = Path(\"../data/life\")\n",
    "a = tomllib.loads((LIFE / \"assumptions.toml\").read_text())\n",
    "ann = pd.read_csv(LIFE / \"annuities.csv\")\n",
    "ages = ann[\"age\"].to_numpy()\n",
    "payments = ann[\"annual_payment\"].to_numpy(dtype=float)\n",
    "T = 60   # projection years: the youngest annuitant (60) reaches 120\n",
    "print(a)\n",
    "print(f\"{len(ann):,} annuities, {payments.sum():,.0f} a year\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-002",
   "metadata": {},
   "source": [
    "---\n",
    "## Your turn 1 — mortality, and a textbook check\n",
    "\n",
    "The Standard Ultimate Survival Model has force of mortality μ_x = A + B·c^x, so\n",
    "\n",
    "q_x = 1 − exp(−A − B·c^x·(c − 1) / ln c).\n",
    "\n",
    "`q_table` builds q_x for ages 0 to 120 (q = 1 at 120), and `survival(q, ages, g)`\n",
    "returns, in row i, the t-year survival probabilities of a life aged `ages[i]`, t = 0..60,\n",
    "with mortality improvement g.\n",
    "\n",
    "1. Compute ä_65 at 5% (annuity-**due**: payments at t = 0, 1, 2, …). The textbook\n",
    "   value is 13.5498.\n",
    "2. Compute the reserve for the book at 4%: each annuity pays at the **end** of each year\n",
    "   (a_x, annuity-immediate), so reserve = Σ payment × a_x."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-003",
   "metadata": {},
   "outputs": [],
   "source": [
    "def q_table(A, B, c, omega=121):\n",
    "    x = np.arange(omega)\n",
    "    q = 1 - np.exp(-A - B * c ** x * (c - 1) / np.log(c))\n",
    "    q[-1] = 1.0\n",
    "    return q\n",
    "\n",
    "\n",
    "def survival(q, ages, g=0.0, horizon=T):\n",
    "    t = np.arange(1, horizon + 1)\n",
    "    att = np.minimum(ages[:, None] + t[None, :] - 1, len(q) - 1)\n",
    "    qt = q[att] * (1 - g) ** t[None, :]\n",
    "    qt[att == len(q) - 1] = 1.0\n",
    "    return np.hstack([np.ones((len(ages), 1)), np.cumprod(1 - qt, axis=1)])\n",
    "\n",
    "\n",
    "m = a[\"mortality\"]\n",
    "q = q_table(m[\"A\"], m[\"B\"], m[\"c\"])\n",
    "a_due_65 = reserve = None\n",
    "...  # your code here\n",
    "check(\"a-due(65) at 5%\", None if a_due_65 is None else round(a_due_65, 4), 13.5498)\n",
    "check(\"reserve at 4%\", None if reserve is None else round(reserve), 762202760)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-004",
   "metadata": {},
   "source": [
    "## 2. Random numbers you can reproduce\n",
    "\n",
    "One `Generator`, one seed, every draw reproducible. For parallel work, `SeedSequence`\n",
    "spawns independent child streams from one parent seed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-005",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng1, rng2 = np.random.default_rng(2026), np.random.default_rng(2026)\n",
    "print(\"same seed, same numbers:\", np.array_equal(rng1.random(5), rng2.random(5)))\n",
    "children = np.random.SeedSequence(2026).spawn(3)\n",
    "print(\"three independent streams:\", [np.random.default_rng(s).random() for s in children])\n",
    "\n",
    "\n",
    "def draws(n, n_lives, seed):\n",
    "    rng = np.random.default_rng(seed)\n",
    "    return rng.standard_normal((n, T)), rng.standard_normal(n), rng.random((n, n_lives))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-006",
   "metadata": {},
   "source": [
    "## 3. Vasicek short rates, vectorized over scenarios\n",
    "\n",
    "Exact annual steps: r_{t+1} = r_t·e^{−a} + b(1 − e^{−a}) + σ·√((1 − e^{−2a}) / 2a)·z_{t+1},\n",
    "with r_0 given (column 0 of `z` is unused). The loop runs over the 60 years, never over\n",
    "scenarios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-007",
   "metadata": {},
   "outputs": [],
   "source": [
    "im, lv = a[\"interest_model\"], a[\"longevity\"]\n",
    "\n",
    "\n",
    "def vasicek(z, r0, a_, b, sigma):\n",
    "    n, horizon = z.shape\n",
    "    r = np.empty((n, horizon))\n",
    "    r[:, 0] = r0\n",
    "    e = math.exp(-a_)\n",
    "    sd = sigma * math.sqrt((1 - e * e) / (2 * a_))\n",
    "    for t in range(1, horizon):\n",
    "        r[:, t] = r[:, t - 1] * e + b * (1 - e) + sd * z[:, t]\n",
    "    return r\n",
    "\n",
    "\n",
    "z, zg, u = draws(10_000, 1, seed=1)\n",
    "r = vasicek(z, im[\"r0\"], im[\"a\"], im[\"b\"], im[\"sigma\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-008",
   "metadata": {},
   "source": [
    "## Your turn 2 — discount factors and bond prices\n",
    "\n",
    "The discount factor for a payment at the end of year t is D_t = exp(−(r_0 + … + r_{t−1})),\n",
    "one row per scenario (use `np.cumsum` along axis 1). The model price of a zero-coupon\n",
    "bond paying 1 in 10 years is the mean of D_10 over scenarios. Store `D` and `p10`, and\n",
    "compare `p10` with 1.04^−10: why is the model price higher? (Look at r_0 and b.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-009",
   "metadata": {},
   "outputs": [],
   "source": [
    "D = p10 = None\n",
    "...  # your code here\n",
    "check(\"mean short rate in year 60\", round(float(r[:, -1].mean()), 5), 0.03975)\n",
    "check(\"10-year zero-coupon price\", None if p10 is None else round(p10, 4), 0.7221)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-010",
   "metadata": {},
   "source": [
    "## 4. The same projection, three ways\n",
    "\n",
    "Each scenario draws a longevity improvement g, a rate path, and one uniform U per life;\n",
    "the life receives payments at t = 1..K, where K = max{t : S(t) > U}. The scenario's PV\n",
    "is Σ payment × C(K) with C(k) = D_1 + … + D_k. First, plain loops:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-011",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pv_loops(z, zg, u):\n",
    "    r = vasicek(z, im[\"r0\"], im[\"a\"], im[\"b\"], im[\"sigma\"])\n",
    "    out = np.empty(len(zg))\n",
    "    for s in range(len(zg)):\n",
    "        S = survival(q, ages, lv[\"mean\"] + lv[\"sd\"] * zg[s])\n",
    "        C = np.r_[0.0, np.cumsum(np.exp(-np.cumsum(r[s])))]\n",
    "        total = 0.0\n",
    "        for i in range(len(ages)):\n",
    "            k = 0\n",
    "            while k < T and S[i, k + 1] > u[s, i]:\n",
    "                k += 1\n",
    "            total += payments[i] * C[k]\n",
    "        out[s] = total\n",
    "    return out\n",
    "\n",
    "\n",
    "z, zg, u = draws(40, len(ages), seed=2026)\n",
    "t0 = time.perf_counter()\n",
    "pv1 = pv_loops(z, zg, u)\n",
    "t_loops = time.perf_counter() - t0\n",
    "print(f\"{len(pv1)} scenarios in {t_loops:.2f} s; scenario 0 PV {pv1[0]:,.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-012",
   "metadata": {},
   "source": [
    "## Your turn 3 — vectorize\n",
    "\n",
    "Write `pv_numpy` with no Python loop over lives: for each scenario (or a chunk of\n",
    "scenarios), compare the survival curves of all lives with their uniforms at once,\n",
    "`K = (S[idx, 1:] > u[s, :, None]).sum(axis=1)`, where `uniq, idx = np.unique(ages,\n",
    "return_inverse=True)` lets you compute one survival curve per **age**, not per life.\n",
    "It must agree with `pv_loops`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-013",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pv_numpy(z, zg, u):\n",
    "    ...  # your code here\n",
    "\n",
    "\n",
    "t0 = time.perf_counter()\n",
    "pv2 = pv_numpy(z, zg, u)\n",
    "t_numpy = time.perf_counter() - t0\n",
    "if pv2 is not None:\n",
    "    print(f\"NumPy: {t_numpy:.2f} s, {t_loops / t_numpy:.0f}x faster\")\n",
    "check(\"NumPy agrees with loops\", None if pv2 is None else bool(np.allclose(pv2, pv1, rtol=1e-12)), True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-014",
   "metadata": {},
   "source": [
    "## Your turn 4 — profile before you optimize\n",
    "\n",
    "`legacy_pv` is the model as first written. Profile two scenarios with `cProfile`, sort\n",
    "by `tottime`, and store the name of the function that takes the most time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-015",
   "metadata": {},
   "outputs": [],
   "source": [
    "def legacy_q(age, year, g):\n",
    "    if age >= 120:\n",
    "        return 1.0\n",
    "    return (1 - math.exp(-m[\"A\"] - m[\"B\"] * m[\"c\"] ** age * (m[\"c\"] - 1) / math.log(m[\"c\"]))) * (1 - g) ** year\n",
    "\n",
    "\n",
    "def legacy_pv(n_scen, seed):\n",
    "    rnd = random.Random(seed)\n",
    "    out = []\n",
    "    for _ in range(n_scen):\n",
    "        g = rnd.gauss(lv[\"mean\"], lv[\"sd\"])\n",
    "        rt, rates = im[\"r0\"], []\n",
    "        for _t in range(T):\n",
    "            rates.append(rt)\n",
    "            e = math.exp(-im[\"a\"])\n",
    "            rt = rt * e + im[\"b\"] * (1 - e) + im[\"sigma\"] * math.sqrt((1 - e * e) / (2 * im[\"a\"])) * rnd.gauss(0, 1)\n",
    "        pv = 0.0\n",
    "        for age, pay in zip(ages, payments):\n",
    "            df = 1.0\n",
    "            for t in range(1, T + 1):\n",
    "                df *= math.exp(-rates[t - 1])\n",
    "                if rnd.random() < legacy_q(age + t - 1, t, g):\n",
    "                    break\n",
    "                pv += pay * df\n",
    "        out.append(pv)\n",
    "    return out\n",
    "\n",
    "\n",
    "slowest = None\n",
    "...  # your code here\n",
    "check(\"the hot spot\", slowest, \"legacy_q\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-016",
   "metadata": {},
   "source": [
    "## Your turn 5 — compile the loops with numba\n",
    "\n",
    "Loops are slow in Python, not in machine code. Write `pv_numba` with `@nb.njit(parallel=True)`,\n",
    "a `nb.prange` loop over scenarios, and inside it the same logic as `pv_loops` (a\n",
    "`_survival_nb` helper is written for you). Pass in precomputed rates `r` and improvements\n",
    "`g`. It must agree with the other versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-017",
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def _survival_nb(q, uniq, g):\n",
    "    S = np.empty((len(uniq), T + 1))\n",
    "    for j in range(len(uniq)):\n",
    "        S[j, 0] = 1.0\n",
    "        for t in range(1, T + 1):\n",
    "            age = min(uniq[j] + t - 1, len(q) - 1)\n",
    "            qt = 1.0 if age == len(q) - 1 else q[age] * (1 - g) ** t\n",
    "            S[j, t] = S[j, t - 1] * (1 - qt)\n",
    "    return S\n",
    "\n",
    "\n",
    "...  # your code here\n",
    "\n",
    "pv3 = None\n",
    "if \"pv_numba\" in globals():\n",
    "    uniq, idx = np.unique(ages, return_inverse=True)\n",
    "    r = vasicek(z, im[\"r0\"], im[\"a\"], im[\"b\"], im[\"sigma\"])\n",
    "    g = lv[\"mean\"] + lv[\"sd\"] * zg\n",
    "    pv_numba(q, uniq, idx, payments, r, g, u)            # first call compiles\n",
    "    t0 = time.perf_counter()\n",
    "    pv3 = pv_numba(q, uniq, idx, payments, r, g, u)\n",
    "    t_numba = time.perf_counter() - t0\n",
    "    print(f\"numba: {t_numba * 1000:.1f} ms, {t_loops / t_numba:,.0f}x faster than loops\")\n",
    "check(\"numba agrees with loops\", None if pv3 is None else bool(np.allclose(pv3, pv1, rtol=1e-12)), True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-018",
   "metadata": {},
   "source": [
    "## 5. Economic capital\n",
    "\n",
    "`simulate` draws its own random numbers inside numba, **one seed per scenario**, so the\n",
    "answer doesn't depend on how many threads run it. It also returns the PV of the\n",
    "best-estimate cash flows along each rate path (what cash-flow-matched bonds would pay)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-019",
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit(parallel=True)\n",
    "def simulate(q, uniq, idx, payments, n, seed, r0, a_, b, sigma, g_mean, g_sd, cf_expected):\n",
    "    pv, pv_expected = np.empty(n), np.empty(n)\n",
    "    e = math.exp(-a_)\n",
    "    sd = sigma * math.sqrt((1 - e * e) / (2 * a_))\n",
    "    for s in nb.prange(n):\n",
    "        np.random.seed(seed + s)\n",
    "        g = g_mean + g_sd * np.random.standard_normal()\n",
    "        S = _survival_nb(q, uniq, g)\n",
    "        C = np.zeros(T + 1)\n",
    "        acc, df, r, pe = 0.0, 1.0, r0, 0.0\n",
    "        for t in range(T):\n",
    "            df *= math.exp(-r)\n",
    "            acc += df\n",
    "            C[t + 1] = acc\n",
    "            pe += df * cf_expected[t]\n",
    "            r = r * e + b * (1 - e) + sd * np.random.standard_normal()\n",
    "        total = 0.0\n",
    "        for i in range(len(idx)):\n",
    "            uu = np.random.random()\n",
    "            k = 0\n",
    "            while k < T and S[idx[i], k + 1] > uu:\n",
    "                k += 1\n",
    "            total += payments[i] * C[k]\n",
    "        pv[s], pv_expected[s] = total, pe\n",
    "    return pv, pv_expected\n",
    "\n",
    "\n",
    "uniq, idx = np.unique(ages, return_inverse=True)\n",
    "cf_expected = (payments[:, None] * survival(q, ages, lv[\"mean\"])[:, 1:]).sum(axis=0)\n",
    "t0 = time.perf_counter()\n",
    "pv, pv_e = simulate(q, uniq, idx, payments, 20_000, 20260101, im[\"r0\"], im[\"a\"], im[\"b\"], im[\"sigma\"],\n",
    "                    lv[\"mean\"], lv[\"sd\"], cf_expected)\n",
    "print(f\"20,000 scenarios in {time.perf_counter() - t0:.1f} s (including compilation)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-020",
   "metadata": {},
   "source": [
    "## Your turn 6 — capital for two asset strategies\n",
    "\n",
    "- **cash**: the reserve earns the short rate, so the cost of the liabilities in a\n",
    "  scenario is `pv`.\n",
    "- **matched**: at t = 0 buy bonds paying the best-estimate flows; they cost\n",
    "  `pv_e.mean()`. Only the differences roll at the short rate: cost = `pv_e.mean() + pv − pv_e`.\n",
    "\n",
    "Economic capital = 99.5th percentile of the cost minus its mean. Store both.\n",
    "\n",
    "With 20,000 scenarios the 99.5th percentile rests on the worst 100 of them, so expect it\n",
    "to move by several million with another seed; the slides use a million scenarios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-021",
   "metadata": {},
   "outputs": [],
   "source": [
    "ec_cash = ec_matched = None\n",
    "...  # your code here\n",
    "check(\"EC, reserve in cash ($m)\", None if ec_cash is None else round(ec_cash / 1e6, 1), 306.7)\n",
    "check(\"EC, cash-flow matched ($m)\", None if ec_matched is None else round(ec_matched / 1e6, 1), 90.5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}