{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "brightcon-01",
   "metadata": {},
   "source": [
    "# From 10 minutes to 4 seconds\n",
    "## Solving large systems with Jacobi + GMRES\n",
    "\n",
    "**Romain Sacchi (PSI) · Brightcon 2026 · 25 September 2026 · 8 minutes**<br> Direct solvers are a good default. But their factors can contain many more nonzeros than the original matrix. This extra storage is called **fill-in**. Jacobi + GMRES can avoid it when convergence is fast enough.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6452918-aff3-497e-a4e6-569e38eba969",
   "metadata": {},
   "source": [
    "# How it started"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffa52b94-e108-4068-86fc-043a8e410312",
   "metadata": {},
   "source": [
    "![UMFPACK vs. JacobiGMRED](assets/figure_1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-02",
   "metadata": {},
   "source": [
    "## Goals\n",
    "\n",
    "1. See how matrix size, density, and structure affect solve time and memory.\n",
    "2. Compare direct solvers with Jacobi + GMRES.\n",
    "3. Reuse a factorization for many demands on one matrix.\n",
    "4. Solve 500 changing matrices, with and without a previous solution as a starting guess.\n",
    "\n",
    "Run this notebook from the repository folder. The main comparisons use synthetic data.\n",
    "The final optional examples use an existing local Brightway project.\n",
    "The larger grids can take several minutes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-03",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dev.benchmark_synthetic import banded_matrix, constant_degree_matrix, solve\n",
    "from dev.run_synthetic_suite import run_benchmark\n",
    "from scipy.sparse.linalg import LinearOperator, gmres"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-04",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib.metadata\n",
    "import sys\n",
    "from time import perf_counter\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import PercentFormatter\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import psutil\n",
    "import scipy\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "SEED = 2026\n",
    "RTOL = 1e-4\n",
    "available_mib = psutil.virtual_memory().available / 2**20\n",
    "memory_guard_mib = int(min(8_192, available_mib * 0.25))\n",
    "try:\n",
    "    import pypardiso\n",
    "\n",
    "    PARDISO_AVAILABLE = True\n",
    "except ImportError:\n",
    "    PARDISO_AVAILABLE = False\n",
    "DIRECT_SOLVERS = [\"umfpack\"] + ([\"pardiso\"] if PARDISO_AVAILABLE else [])\n",
    "colors = {\n",
    "    \"numpy-dense\": \"#7F8C8D\",\n",
    "    \"superlu\": \"#C44E52\",\n",
    "    \"umfpack\": \"#DD8452\",\n",
    "    \"pardiso\": \"#937860\",\n",
    "    \"gmres\": \"#4C72B0\",\n",
    "    \"jacobi-gmres\": \"#12A594\",\n",
    "    \"jacobi-gmres-no-guess\": \"#8172B2\",\n",
    "}\n",
    "solver_names = {\n",
    "    \"numpy-dense\": \"NumPy dense\",\n",
    "    \"superlu\": \"SuperLU\",\n",
    "    \"umfpack\": \"UMFPACK\",\n",
    "    \"pardiso\": \"Pardiso\",\n",
    "    \"gmres\": \"GMRES\",\n",
    "    \"jacobi-gmres\": \"Jacobi + GMRES\",\n",
    "    \"jacobi-gmres-no-guess\": \"Jacobi + GMRES (no warm start)\",\n",
    "}\n",
    "\n",
    "\n",
    "def observed_log_ratio_scale(values):\n",
    "    finite = np.asarray(values, dtype=float)\n",
    "    finite = finite[np.isfinite(finite)]\n",
    "    limit = max(float(np.max(np.abs(finite))), 0.25) if finite.size else 1.0\n",
    "    ticks = np.linspace(-limit, limit, 5)\n",
    "    labels = [\n",
    "        f\"{10 ** tick:,.0f}×\" if 10**tick >= 1 else f\"{10 ** tick:.2g}×\"\n",
    "        for tick in ticks\n",
    "    ]\n",
    "    return limit, ticks, labels\n",
    "\n",
    "\n",
    "plt.style.use(\"seaborn-v0_8-whitegrid\")\n",
    "plt.rcParams.update(\n",
    "    {\n",
    "        \"axes.spines.top\": False,\n",
    "        \"axes.spines.right\": False,\n",
    "        \"axes.titleweight\": \"bold\",\n",
    "        \"figure.dpi\": 120,\n",
    "    }\n",
    ")\n",
    "\n",
    "\n",
    "solver_order = [\"numpy-dense\", \"superlu\", *DIRECT_SOLVERS, \"gmres\", \"jacobi-gmres\"]\n",
    "\n",
    "{\n",
    "    \"benchmark mode\": \"live synthetic benchmarks\",\n",
    "    \"worker Python\": sys.executable,\n",
    "    \"NumPy\": np.__version__,\n",
    "    \"SciPy\": scipy.__version__,\n",
    "    \"scikit-umfpack\": importlib.metadata.version(\"scikit-umfpack\"),\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-05",
   "metadata": {},
   "source": [
    "## Solver toolbox\n",
    "\n",
    "- **NumPy dense** stores all entries, including zeros. It is useful as a small reference case.\n",
    "- **SuperLU**, **UMFPACK**, and optional **Pardiso** factor a sparse matrix, then solve using those factors.\n",
    "- **GMRES** improves an approximate solution through repeated matrix-vector products.\n",
    "- **Jacobi + GMRES** uses the inverse diagonal to rescale the equations before iterating.\n",
    "\n",
    "We use `rtol=1e-4`, `restart=50`, and `maxiter=300`.\n",
    "Check both the GMRES return code (`info == 0`) and the relative residual\n",
    "$\\|Ax-b\\|/\\|b\\|$.  \n",
    "Note that a small residual alone does not guarantee a small error in the solution.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-06",
   "metadata": {},
   "source": [
    "### Jacobi + GMRES: try one solve\n",
    "\n",
    "For $Ax=b$, Jacobi applies $D^{-1}$, where $D$ is the diagonal of $A$.  \n",
    "A `LinearOperator` applies this scaling without creating a dense matrix.  \n",
    "The callback counts inner GMRES steps. `maxiter` limits restart cycles.  \n",
    "Compare the solution with a direct solve as well as checking the residual.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-07",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = constant_degree_matrix(500, degree=8, seed=SEED, diagonal_span=4.0)\n",
    "b = np.zeros(A.shape[0])\n",
    "b[0] = 1.0\n",
    "inverse_diagonal = 1.0 / A.diagonal()\n",
    "M = LinearOperator(A.shape, matvec=lambda vector: inverse_diagonal * vector)\n",
    "residual_history = []\n",
    "x, info = gmres(\n",
    "    A,\n",
    "    b,\n",
    "    M=M,\n",
    "    rtol=RTOL,\n",
    "    atol=0.0,\n",
    "    restart=50,\n",
    "    maxiter=300,\n",
    "    callback=residual_history.append,\n",
    "    callback_type=\"pr_norm\",\n",
    ")\n",
    "reference, _, _ = solve(A, b, \"superlu\", RTOL, 50, 300)\n",
    "relative_residual = np.linalg.norm(A @ x - b) / np.linalg.norm(b)\n",
    "relative_difference = np.linalg.norm(x - reference) / np.linalg.norm(reference)\n",
    "assert info == 0 and relative_residual <= RTOL * 1.01\n",
    "print(f\"{len(residual_history)} GMRES steps\")\n",
    "print(f\"Relative residual: {relative_residual:.2e}\")\n",
    "print(f\"Relative solution difference: {relative_difference:.2e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-08",
   "metadata": {},
   "source": [
    "# 1 · Matrix size and solve cost\n",
    "\n",
    "Start with roughly eight inputs per activity, from 500 to 50,000 activities.  \n",
    "Each solver runs in a separate process so previous allocations do not affect its memory measurement.  \n",
    "The helper samples resident memory (RSS) every 2 ms; very short peaks may be missed.  \n",
    "Additional memory is measured from the point just before solving.  \n",
    "\n",
    "The process runner is ordinary Python in `dev/run_synthetic_suite.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-09",
   "metadata": {},
   "outputs": [],
   "source": [
    "synthetic_payload = run_benchmark(\n",
    "    \"scaling\",\n",
    "    sizes=[500, 1_000, 5_000, 10_000, 20_000],\n",
    "    solvers=[\"numpy-dense\", \"jacobi-gmres\", *DIRECT_SOLVERS, \"gmres\", \"superlu\"],\n",
    "    rtol=RTOL,\n",
    "    worker_timeout=12, # we limit each case to 12 sec, otherwise we skip it\n",
    "    total_budget=65, # we limit the entire benchmark to 65 sec\n",
    ")\n",
    "synthetic = pd.json_normalize(synthetic_payload[\"results\"])\n",
    "synthetic[\"memory_mib\"] = synthetic[\"incremental_peak_rss_bytes\"] / 2**20\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-11",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))\n",
    "valid = synthetic[synthetic[\"status\"] == \"COMPLETED\"]\n",
    "for solver in solver_order:\n",
    "    subset = valid[valid.solver == solver].sort_values(\"size\")\n",
    "    if subset.empty:\n",
    "        continue\n",
    "    label = solver_names[solver]\n",
    "    axes[0].plot(\n",
    "        subset[\"size\"], subset[\"solve_seconds\"], \"o-\", label=label, color=colors[solver]\n",
    "    )\n",
    "    axes[1].plot(\n",
    "        subset[\"size\"], subset[\"memory_mib\"], \"o-\", label=label, color=colors[solver]\n",
    "    )\n",
    "\n",
    "axes[0].set_ylabel(\"solve time [s]\")\n",
    "axes[1].set_ylabel(\"additional peak memory [MiB]\")\n",
    "for axis in axes:\n",
    "    axis.set_xlabel(\"matrix rows / columns\")\n",
    "    axis.grid(alpha=0.25, which=\"both\")\n",
    "axes[0].set_title(\"Runtime separates as the system grows\")\n",
    "axes[1].set_title(\"Direct factorization can increase memory\")\n",
    "handles, labels = axes[0].get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc=\"lower center\", ncol=3, frameon=False, fontsize=8)\n",
    "fig.tight_layout(rect=[0, 0.12, 1, 1])\n",
    "plt.show()\n",
    "\n",
    "not_completed_scaling = synthetic[synthetic.status != \"COMPLETED\"].copy()\n",
    "if not not_completed_scaling.empty:\n",
    "    oom = (\n",
    "        not_completed_scaling[\"error\"]\n",
    "        .fillna(\"\")\n",
    "        .str.contains(\"out_of_memory\", case=False)\n",
    "    )\n",
    "    not_completed_scaling.loc[oom, \"status\"] = \"OUT OF MEMORY\"\n",
    "    display(\n",
    "        not_completed_scaling[[\"solver\", \"size\", \"status\"]].style.hide(axis=\"index\")\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-12",
   "metadata": {},
   "source": [
    "### Density versus size\n",
    "\n",
    "**Matrix density** is the share of matrix entries that are nonzero.\n",
    "This sweep compares 1%, 3%, and 5% density at three matrix sizes.\n",
    "It uses **Pardiso if installed**, otherwise **UMFPACK**, against Jacobi + GMRES.\n",
    "\n",
    "Each case has up to **30 seconds**, within a **120-second budget** for the sweep.\n",
    "The heatmap shows direct solve time divided by Jacobi + GMRES time: values above 1 favour Jacobi.\n",
    "Both solves must complete and pass their convergence checks to show a ratio.\n",
    "Missing results are labelled, with the reason listed below the plot.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-13",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "density_direct_solver = \"pardiso\" if PARDISO_AVAILABLE else \"umfpack\"\n",
    "\n",
    "density_size_payload = run_benchmark(\n",
    "    \"density versus size\",\n",
    "    sizes=[1_000, 5_000, 7_500],\n",
    "    solvers=[\"jacobi-gmres\", density_direct_solver],\n",
    "    topology=\"fixed-density\",\n",
    "    densities=[0.01, 0.03, 0.05],\n",
    "    rtol=RTOL,\n",
    "    worker_timeout=30,\n",
    "    total_budget=120,\n",
    "    memory_guard_mib=memory_guard_mib,\n",
    ")\n",
    "density_size = pd.json_normalize(density_size_payload[\"results\"])\n",
    "density_size[\"speedup\"] = np.nan\n",
    "for (size, density), subset in density_size[density_size.status == \"COMPLETED\"].groupby(\n",
    "    [\"size\", \"target_density\"]\n",
    "):\n",
    "    direct = subset.loc[subset.solver == density_direct_solver, \"solve_seconds\"]\n",
    "    iterative = subset.loc[subset.solver == \"jacobi-gmres\", \"solve_seconds\"]\n",
    "    if len(direct) == len(iterative) == 1:\n",
    "        density_size.loc[subset.index, \"speedup\"] = direct.iloc[0] / iterative.iloc[0]\n",
    "\n",
    "heat = density_size[density_size.solver == \"jacobi-gmres\"].pivot(\n",
    "    index=\"target_density\", columns=\"size\", values=\"speedup\"\n",
    ")\n",
    "log_heat = np.log10(heat)\n",
    "heat_limit, heat_ticks, heat_tick_labels = observed_log_ratio_scale(log_heat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-14",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axis = plt.subplots(figsize=(8.5, 4.5))\n",
    "image = axis.imshow(\n",
    "    log_heat,\n",
    "    aspect=\"auto\",\n",
    "    origin=\"lower\",\n",
    "    cmap=\"PiYG\",\n",
    "    vmin=-heat_limit,\n",
    "    vmax=heat_limit,\n",
    ")\n",
    "axis.set_xticks(range(len(heat.columns)), [f\"{n:,}\" for n in heat.columns])\n",
    "axis.set_yticks(range(len(heat.index)), [f\"{density:.1%}\" for density in heat.index])\n",
    "axis.set_xlabel(\"matrix size\")\n",
    "axis.set_ylabel(\"matrix density\")\n",
    "axis.set_title(f\"{solver_names[density_direct_solver]} time ÷ Jacobi + GMRES time (>1× favors Jacobi)\")\n",
    "for row, density in enumerate(heat.index):\n",
    "    for column, size in enumerate(heat.columns):\n",
    "        value = heat.loc[density, size]\n",
    "        if np.isfinite(value):\n",
    "            label = f\"{value:.0f}×\" if value >= 1 else \"<1×\"\n",
    "            text_color = (\n",
    "                \"white\" if abs(np.log10(value)) > 0.45 * heat_limit else \"black\"\n",
    "            )\n",
    "            axis.text(\n",
    "                column,\n",
    "                row,\n",
    "                label,\n",
    "                ha=\"center\",\n",
    "                va=\"center\",\n",
    "                fontsize=8,\n",
    "                color=text_color,\n",
    "                fontweight=\"bold\",\n",
    "            )\n",
    "        else:\n",
    "            axis.text(column, row, \"no result\", ha=\"center\", va=\"center\", fontsize=8)\n",
    "bar = fig.colorbar(image, ax=axis, label=\"time ratio (scaled to observed maximum)\")\n",
    "bar.set_ticks(heat_ticks)\n",
    "bar.set_ticklabels(heat_tick_labels)\n",
    "fig.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "missing = density_size[density_size.status != \"COMPLETED\"]\n",
    "if not missing.empty:\n",
    "    display(missing[[\"solver\", \"size\", \"target_density\", \"status\", \"skip_reason\", \"error\"]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-18",
   "metadata": {},
   "source": [
    "### Large synthetic systems\n",
    "\n",
    "Here we measure **Jacobi + GMRES only** at 10,000, 20,000, 30,000, and 60,000 rows.\n",
    "The earlier comparisons already show the cost of direct solving.\n",
    "The left plot shows solve time at **1% density**; the right shows peak memory at\n",
    "**0.5%, 1%, 5%, and 10% density**.\n",
    "\n",
    "Each sweep has a **60-second budget**, with at most **30 seconds per case**.\n",
    "Cases that exceed the construction-memory estimate or remaining time budget are skipped.\n",
    "A slower server may show fewer completed points.\n",
    "\n",
    "The intermediate sizes give the curves more points without raising the memory limit.\n",
    "More time cannot make a case fit in memory; the largest cases may still be skipped.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-19",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "large_sizes = [10_000, 20_000, 30_000, 60_000]\n",
    "large_memory_guard_mib = min(4_096, memory_guard_mib)\n",
    "large_jacobi_scale_payload = run_benchmark(\n",
    "    \"large Jacobi systems at 1 percent\",\n",
    "    sizes=large_sizes,\n",
    "    solvers=[\"jacobi-gmres\"],\n",
    "    topology=\"fixed-density\",\n",
    "    densities=[0.01],\n",
    "    rtol=RTOL,\n",
    "    worker_timeout=30,\n",
    "    total_budget=60,\n",
    "    memory_guard_mib=large_memory_guard_mib,\n",
    "    construction_memory_multiplier=3.0,\n",
    ")\n",
    "large_memory_payload = run_benchmark(\n",
    "    \"large systems memory by density\",\n",
    "    sizes=large_sizes,\n",
    "    solvers=[\"jacobi-gmres\"],\n",
    "    topology=\"fixed-density\",\n",
    "    densities=[0.005, 0.05, 0.1],\n",
    "    rtol=RTOL,\n",
    "    worker_timeout=30,\n",
    "    total_budget=60,\n",
    "    memory_guard_mib=large_memory_guard_mib,\n",
    "    construction_memory_multiplier=3.0,\n",
    ")\n",
    "large_jacobi_scale = pd.json_normalize(large_jacobi_scale_payload[\"results\"])\n",
    "large_memory = pd.json_normalize(large_memory_payload[\"results\"])\n",
    "completed_scale = large_jacobi_scale[large_jacobi_scale.status == \"COMPLETED\"].copy()\n",
    "completed_memory = large_memory[large_memory.status == \"COMPLETED\"].copy()\n",
    "memory_plot = pd.concat(\n",
    "    [\n",
    "        completed_memory,\n",
    "        completed_scale[\n",
    "            (completed_scale[\"solver\"] == \"jacobi-gmres\")\n",
    "            & np.isclose(completed_scale[\"target_density\"], 0.01)\n",
    "        ],\n",
    "    ],\n",
    "    ignore_index=True,\n",
    ").drop_duplicates([\"size\", \"target_density\"], keep=\"last\")\n",
    "memory_plot[\"peak_rss_mib\"] = memory_plot[\"peak_rss_bytes\"] / 2**20\n",
    "size_position = {size: index for index, size in enumerate(large_sizes)}\n",
    "density_colors = {0.005: \"#4C72B0\", 0.01: \"#55A868\", 0.05: \"#C44E52\", 0.1: \"#8172B2\"}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-20",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n",
    "for solver, subset in completed_scale.groupby(\"solver\"):\n",
    "    subset = subset.sort_values(\"size\")\n",
    "    label = solver_names[solver]\n",
    "    axes[0].plot(\n",
    "        [size_position[size] for size in subset[\"size\"]],\n",
    "        subset[\"solve_seconds\"],\n",
    "        \"o-\",\n",
    "        label=label,\n",
    "        color=colors[solver],\n",
    "    )\n",
    "for density, subset in memory_plot.groupby(\"target_density\"):\n",
    "    subset = subset.sort_values(\"size\")\n",
    "    axes[1].plot(\n",
    "        [size_position[size] for size in subset[\"size\"]],\n",
    "        subset[\"peak_rss_mib\"],\n",
    "        \"o-\",\n",
    "        label=f\"{density:.1%}\",\n",
    "        color=density_colors[density],\n",
    "    )\n",
    "for axis in axes:\n",
    "    axis.set_xlabel(\"matrix size\")\n",
    "    axis.grid(alpha=0.25)\n",
    "    axis.set_xticks(\n",
    "        range(len(large_sizes)), [f\"{size:,}\" for size in large_sizes], rotation=25\n",
    "    )\n",
    "axes[0].set_ylabel(\"solve time [s]\")\n",
    "axes[1].set_ylabel(\"worker peak RSS [MiB]\")\n",
    "axes[0].set_title(\"Jacobi solve time at 1% target density\")\n",
    "axes[1].set_title(\"Jacobi peak memory across densities\")\n",
    "axes[0].legend(frameon=False, fontsize=8)\n",
    "axes[1].legend(frameon=False, fontsize=8, ncol=2, title=\"target density\")\n",
    "fig.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "large_status = pd.concat(\n",
    "    [\n",
    "        large_jacobi_scale.assign(suite=\"1% Jacobi comparison\"),\n",
    "        large_memory.assign(suite=\"Jacobi density sweep\"),\n",
    "    ],\n",
    "    ignore_index=True,\n",
    ")\n",
    "not_completed = large_status[large_status.status != \"COMPLETED\"].copy()\n",
    "if not not_completed.empty:\n",
    "    error_text = (\n",
    "        not_completed[\"error\"].fillna(\"\")\n",
    "        if \"error\" in not_completed\n",
    "        else pd.Series(\"\", index=not_completed.index)\n",
    "    )\n",
    "    oom = error_text.str.contains(\"out_of_memory\", case=False)\n",
    "    not_completed.loc[oom, \"status\"] = \"OUT OF MEMORY\"\n",
    "    failed = not_completed[not_completed[\"status\"] != \"SKIPPED\"]\n",
    "    if not failed.empty:\n",
    "        display(\n",
    "            failed[[\"suite\", \"solver\", \"size\", \"target_density\", \"status\"]].style.hide(\n",
    "                axis=\"index\"\n",
    "            )\n",
    "        )\n",
    "    skipped = not_completed[not_completed[\"status\"] == \"SKIPPED\"]\n",
    "    if not skipped.empty:\n",
    "        display(\n",
    "            skipped[[\"suite\", \"size\", \"target_density\", \"skip_reason\"]].style.hide(\n",
    "                axis=\"index\"\n",
    "            )\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "banded-0",
   "metadata": {},
   "source": [
    "### A large matrix where direct solving works well\n",
    "\n",
    "A banded matrix connects nearby activities. Its factors have little fill-in, so a direct solver can be fast even for a large system.\n",
    "Compare `lu_fill_ratio` (nonzeros in the factors divided by nonzeros in the matrix), solve time, and residual.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "banded-1",
   "metadata": {},
   "outputs": [],
   "source": [
    "banded_payload = run_benchmark(\n",
    "    \"low-fill counterexample\",\n",
    "    sizes=[50_000],\n",
    "    solvers=[\"superlu\", *DIRECT_SOLVERS, \"jacobi-gmres\"],\n",
    "    topology=\"banded\",\n",
    "    degree=4,\n",
    "    rtol=RTOL,\n",
    "    worker_timeout=15,\n",
    "    total_budget=45,\n",
    ")\n",
    "banded = pd.json_normalize(banded_payload[\"results\"])\n",
    "display(\n",
    "    banded.reindex(\n",
    "        columns=[\n",
    "            \"solver\",\n",
    "            \"size\",\n",
    "            \"nnz\",\n",
    "            \"status\",\n",
    "            \"solve_seconds\",\n",
    "            \"lu_fill_ratio\",\n",
    "            \"relative_residual\",\n",
    "        ]\n",
    "    )\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-21",
   "metadata": {},
   "source": [
    "# 2 · Fixed matrix, many right-hand sides\n",
    "\n",
    "One fixed 40,000 × 40,000 matrix can serve many demands.  \n",
    "The left plot separates the one-time factorization from the first solve.  \n",
    "The right plot then shows total time as more demands reuse that same matrix.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-22",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "rhs_payload = run_benchmark(\n",
    "    \"fixed matrix demands\",\n",
    "    sizes=[40_000],\n",
    "    solvers=[*DIRECT_SOLVERS, \"jacobi-gmres\"],\n",
    "    topology=\"banded\",\n",
    "    degree=30,\n",
    "    rhs_counts=[1, 3, 10, 25, 50, 100],\n",
    "    rtol=RTOL,\n",
    "    worker_timeout=15,\n",
    "    total_budget=50,\n",
    ")\n",
    "rhs_all = pd.json_normalize(rhs_payload[\"results\"])\n",
    "display(rhs_all[[\"solver\", \"rhs_count\", \"status\", \"skip_reason\"]])\n",
    "rhs = rhs_all[rhs_all.status == \"COMPLETED\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-23",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n",
    "\n",
    "first_demand = rhs[\n",
    "    (rhs.rhs_count == 1) & rhs.solver.isin([\"umfpack\", \"jacobi-gmres\"])\n",
    "].set_index(\"solver\")\n",
    "bar_solvers = [\n",
    "    solver for solver in [\"umfpack\", \"jacobi-gmres\"] if solver in first_demand.index\n",
    "]\n",
    "x = np.arange(len(bar_solvers))\n",
    "setup_time = [\n",
    "    first_demand.loc[solver, \"factorization_seconds\"] for solver in bar_solvers\n",
    "]\n",
    "solve_time = [first_demand.loc[solver, \"rhs_solve_seconds\"] for solver in bar_solvers]\n",
    "axes[0].bar(x, setup_time, label=\"one-time factorization\", color=\"#B0B0B0\")\n",
    "axes[0].bar(\n",
    "    x,\n",
    "    solve_time,\n",
    "    bottom=setup_time,\n",
    "    label=\"solve first demand\",\n",
    "    color=[colors[solver] for solver in bar_solvers],\n",
    ")\n",
    "axes[0].set_xticks(x, [solver_names[solver] for solver in bar_solvers])\n",
    "axes[0].set_ylabel(\"time [s]\")\n",
    "axes[0].set_title(\"Cost of the first demand\")\n",
    "axes[0].legend(frameon=False, fontsize=8)\n",
    "\n",
    "for solver, subset in rhs.groupby(\"solver\"):\n",
    "    axes[1].plot(\n",
    "        subset[\"rhs_count\"],\n",
    "        subset[\"solve_seconds\"],\n",
    "        \"o-\",\n",
    "        label=solver_names[solver],\n",
    "        color=colors[solver],\n",
    "    )\n",
    "\n",
    "comparison = rhs[rhs.solver.isin([\"umfpack\", \"jacobi-gmres\"])].pivot(\n",
    "    index=\"rhs_count\", columns=\"solver\", values=\"solve_seconds\"\n",
    ")\n",
    "if {\"umfpack\", \"jacobi-gmres\"}.issubset(comparison.columns):\n",
    "    direct_wins = comparison.index[comparison[\"umfpack\"] < comparison[\"jacobi-gmres\"]]\n",
    "    if len(direct_wins):\n",
    "        crossover = direct_wins.min()\n",
    "        axes[1].axvline(crossover, color=\"black\", linestyle=\"--\", linewidth=1)\n",
    "        axes[1].text(\n",
    "            crossover,\n",
    "            axes[1].get_ylim()[1],\n",
    "            f\"  direct faster from ~{crossover:,}\",\n",
    "            va=\"top\",\n",
    "            fontsize=8,\n",
    "        )\n",
    "\n",
    "axes[1].set_xlabel(\"demands solved with one fixed matrix\")\n",
    "axes[1].set_ylabel(\"total solve time [s]\")\n",
    "axes[1].set_title(\"Factorization reuse changes the winner\")\n",
    "axes[1].legend(frameon=False, fontsize=8)\n",
    "for axis in axes:\n",
    "    axis.grid(alpha=0.25)\n",
    "fig.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-24",
   "metadata": {},
   "source": [
    "# 3 · Changing matrix, repeated solves\n",
    "\n",
    "This is a 500-iteration synthetic **Monte Carlo** analogue: the technosphere matrix changes slightly at every sample.  \n",
    "UMFPACK must refactorize each matrix. Jacobi + GMRES is shown both with a warm start from the previous solution and without one, matching the role of `use_guess` in `bw2calc`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-25",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "demand = np.zeros(10_000)\n",
    "demand[0] = 1.0\n",
    "records = []\n",
    "coupling_rng = np.random.default_rng(SEED)\n",
    "sample_couplings = np.clip(0.8 + coupling_rng.normal(0, 0.015, 500), 0.7, 0.9)\n",
    "solver_cases = [\n",
    "    *((solver, solver, False) for solver in DIRECT_SOLVERS),\n",
    "    (\"jacobi-gmres\", \"jacobi-gmres\", True),\n",
    "    (\"jacobi-gmres-no-guess\", \"jacobi-gmres\", False),\n",
    "]\n",
    "\n",
    "for label, solver, reuse_guess in solver_cases:\n",
    "    previous_solution = None\n",
    "    for sample in range(500):\n",
    "        matrix = banded_matrix(\n",
    "            10_000,\n",
    "            4,\n",
    "            SEED,\n",
    "            4.0,\n",
    "            coupling=float(sample_couplings[sample]),\n",
    "        )\n",
    "        solve_started = perf_counter()\n",
    "        solution, info, count = solve(\n",
    "            matrix,\n",
    "            demand,\n",
    "            solver,\n",
    "            RTOL,\n",
    "            50,\n",
    "            300,\n",
    "            x0=previous_solution if reuse_guess else None,\n",
    "        )\n",
    "        elapsed = perf_counter() - solve_started\n",
    "        residual = float(\n",
    "            np.linalg.norm(matrix @ solution - demand) / np.linalg.norm(demand)\n",
    "        )\n",
    "        assert info in (None, 0) and residual <= RTOL * 1.01\n",
    "        if reuse_guess:\n",
    "            previous_solution = solution\n",
    "        records.append(\n",
    "            {\n",
    "                \"solver\": label,\n",
    "                \"sample\": sample,\n",
    "                \"solve_seconds\": elapsed,\n",
    "                \"iterations\": count,\n",
    "                \"relative_residual\": residual,\n",
    "                \"info\": info,\n",
    "            }\n",
    "        )\n",
    "changing = pd.DataFrame(records)\n",
    "per_sample = changing.copy()\n",
    "per_sample[\"cumulative_seconds\"] = per_sample.groupby(\"solver\")[\n",
    "    \"solve_seconds\"\n",
    "].cumsum()\n",
    "per_sample[\"monte_carlo_iteration\"] = per_sample[\"sample\"].astype(int) + 1\n",
    "mc_summary = (\n",
    "    per_sample.groupby(\"solver\")\n",
    "    .agg(\n",
    "        total_seconds=(\"solve_seconds\", \"sum\"),\n",
    "        median_iterations=(\"iterations\", \"median\"),\n",
    "        max_relative_residual=(\"relative_residual\", \"max\"),\n",
    "        failed_solves=(\"info\", lambda values: int(values.dropna().ne(0).sum())),\n",
    "    )\n",
    "    .reset_index()\n",
    ")\n",
    "assert mc_summary[\"failed_solves\"].eq(0).all()\n",
    "assert mc_summary[\"max_relative_residual\"].le(RTOL * 1.01).all()\n",
    "display(\n",
    "    mc_summary.style.format(\n",
    "        {\n",
    "            \"total_seconds\": \"{:.2f}\",\n",
    "            \"median_iterations\": \"{:.0f}\",\n",
    "            \"max_relative_residual\": \"{:.2e}\",\n",
    "        }\n",
    "    ).hide(axis=\"index\")\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brightcon-26",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axis = plt.subplots(figsize=(8.5, 4.5))\n",
    "for solver, subset in per_sample.groupby(\"solver\"):\n",
    "    label = (\n",
    "        \"Jacobi + GMRES (warm start)\"\n",
    "        if solver == \"jacobi-gmres\"\n",
    "        else solver_names[solver]\n",
    "    )\n",
    "    axis.plot(\n",
    "        subset[\"monte_carlo_iteration\"],\n",
    "        subset[\"cumulative_seconds\"],\n",
    "        \"o-\",\n",
    "        label=label,\n",
    "        color=colors[solver],\n",
    "        markersize=3,\n",
    "    )\n",
    "    final = subset.iloc[-1]\n",
    "    axis.annotate(\n",
    "        f\"{final['cumulative_seconds']:.2f} s\",\n",
    "        (final[\"monte_carlo_iteration\"], final[\"cumulative_seconds\"]),\n",
    "        xytext=(5, 0),\n",
    "        textcoords=\"offset points\",\n",
    "        va=\"center\",\n",
    "        fontsize=8,\n",
    "    )\n",
    "axis.set_xlabel(\"Monte Carlo iteration\")\n",
    "axis.set_ylabel(\"cumulative solve time [s]\")\n",
    "axis.set_title(\"Cumulative cost across changing matrices\")\n",
    "axis.set_xlim(1, per_sample[\"monte_carlo_iteration\"].max() * 1.05)\n",
    "axis.grid(alpha=0.25)\n",
    "axis.legend(frameon=False)\n",
    "fig.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-27",
   "metadata": {},
   "source": [
    "# 4 · JacobiGMRESLCA in bw2calc\n",
    "\n",
    "In Brightway, `bw2calc.JacobiGMRESLCA` changes the technosphere solve while leaving biosphere and\n",
    "characterization processing unchanged. Main arguments are `demand`, `method`, `rtol`, `use_guess`,\n",
    "`restart`, and `maxiter`.\n",
    "\n",
    "For repeated solves, `use_guess=True` starts from the previous supply array; `False` starts each\n",
    "solve without that warm start.\n",
    "\n",
    "```python\n",
    "from bw2calc import JacobiGMRESLCA\n",
    "lca = JacobiGMRESLCA({activity.id: 1.0}, method, rtol=1e-4, use_guess=True)\n",
    "lca.lci()\n",
    "lca.lcia()\n",
    "print(lca.score)\n",
    "```\n",
    "\n",
    "Check convergence, residual, and agreement before interpreting an iterative score.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b780c208-dcc7-4f18-aa81-10e4f9c03fd6",
   "metadata": {},
   "source": [
    "# How it ended"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "941ca8c3-a10c-44b4-9f6a-6c8fb5a02db9",
   "metadata": {},
   "source": [
    "![UMFPACK vs. JacobiGMRED](assets/figure_2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brightcon-28",
   "metadata": {},
   "source": [
    "# When to use **Jacobi + GMRES**\n",
    "\n",
    "**Good fit when:**\n",
    "\n",
    "- The system is **very large and sparse** and not **narrowly banded**.\n",
    "- The matrix **changes repeatedly**.\n",
    "- A direct factorization is close to the machine's **memory limit**.\n",
    "- A **controlled tolerance** is acceptable.\n",
    "\n",
    "**Choose UMFPACK or Pardiso when:**\n",
    "\n",
    "- The matrix is small or moderate, or **narrowly banded**.\n",
    "- Many demands reuse one fixed matrix.\n",
    "- You need a direct answer and factorization fits comfortably in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87391e5c-aebd-4df5-921b-8070d61eb0ec",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}