{ "cells": [ { "cell_type": "markdown", "id": "intro", "metadata": {}, "source": "# Jacobian diagnostics\n\nWhen `JacobianAutograd` and `JacobianNumerical` disagree on a model, the next question is *why*. This notebook is a runnable companion to [Differentiability and Jacobian Tools](diff.rst). It walks through three progressively trickier models, shows how `compare_jacobian_dicts` localises the disagreement to specific parameters, and closes with a decision tree for reading a report.\n\n* **LINEAR2D**: the test model with an analytical Jacobian and a zero parameter. Demonstrates `parameter_space=\"direct\"`.\n* **GL06SIMEX**: a small SFC model whose autograd and numerical Jacobians agree cleanly under the default log-space setting.\n* **GL06INSOUT**: the production Tobin portfolio model, which uses a step-function `torch.where` inside `portfolio_switches`. The diagnostic shows up as a subset of parameters with visibly worse agreement than the others, and testing different values of `epsilon` in the numerical differentiation confirms that the disagreement is caused by the step function." }, { "cell_type": "code", "execution_count": 1, "id": "version", "metadata": { "execution": { "iopub.execute_input": "2026-04-10T08:12:13.163602Z", "iopub.status.busy": "2026-04-10T08:12:13.163472Z", "iopub.status.idle": "2026-04-10T08:12:14.034842Z", "shell.execute_reply": "2026-04-10T08:12:14.034428Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "macrostat version: 0.5.7.post1.dev7+g35c4d7655.d20260409\n", "torch version: 2.11.0+cu130\n" ] } ], "source": [ "import sys, os\n", "sys.path.insert(0, os.path.abspath('.'))\n", "\n", "import macrostat\n", "import torch\n", "\n", "from macrostat.models import get_model\n", "from macrostat.diff import (\n", " JacobianAutograd,\n", " JacobianNumerical,\n", " compare_jacobian_dicts,\n", ")\n", "from _doc_helpers import (\n", " mean_nominal_output,\n", " mean_national_income,\n", " final_state_norm,\n", ")\n", "\n", "print('macrostat version:', macrostat.__version__)\n", "print('torch version: ', torch.__version__)" ] }, { "cell_type": "markdown", "id": "linear2d-section", "metadata": {}, "source": "## LINEAR2D: direct-space differences\n\nLINEAR2D is a 2x2 linear map $x_{t+1} = A x_t$ with six parameters: four matrix entries and two initial conditions. The initial condition `x0_2` defaults to zero, so log-space perturbation is undefined for it. We therefore set `parameter_space=\"direct\"` with a small `epsilon=1e-5`, which is the standard central-difference setup. With an analytical Jacobian available on this model, we expect low but non-zero relative error dominated by truncation of the finite-difference formula." }, { "cell_type": "code", "execution_count": 2, "id": "linear2d-code", "metadata": { "execution": { "iopub.execute_input": "2026-04-10T08:12:14.035986Z", "iopub.status.busy": "2026-04-10T08:12:14.035812Z", "iopub.status.idle": "2026-04-10T08:12:15.600294Z", "shell.execute_reply": "2026-04-10T08:12:15.599747Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Jacobian comparison: autograd vs numerical (direct, eps=1e-5)\n", "Parameters compared: 6\n", "Overall max abs diff: 5.045e-03 max rel diff: 5.045e-03\n", "\n", "Parameter max_abs mean_abs max_rel mean_rel close nan/inf\n", "---------------------------------------------------------------------------------------------------------\n", "x0_2 5.045e-03 5.045e-03 5.045e-03 5.045e-03 0/1 no\n", "x0_1 4.689e-03 4.689e-03 4.689e-03 4.689e-03 0/1 no\n", "a12 3.227e-04 3.227e-04 2.999e-04 2.999e-04 1/1 no\n", "a22 1.304e-04 1.304e-04 1.304e-04 1.304e-04 1/1 no\n", "a11 2.871e-04 2.871e-04 9.710e-05 9.710e-05 1/1 no\n", "a21 1.221e-04 1.221e-04 9.525e-05 9.525e-05 1/1 no\n" ] } ], "source": [ "model = get_model('LINEAR2D')()\n", "\n", "jac_auto = JacobianAutograd(model).compute(\n", " loss_fn=final_state_norm, mode='rev',\n", ")\n", "jac_num = JacobianNumerical(\n", " model, parameter_space='direct', epsilon=1e-5,\n", ").compute(loss_fn=final_state_norm, mode='central', num_workers=1)\n", "\n", "report = compare_jacobian_dicts(\n", " jac_auto, jac_num,\n", " method_a='autograd', method_b='numerical (direct, eps=1e-5)',\n", " rtol=1e-3, atol=1e-6,\n", ")\n", "print(report.summary())" ] }, { "cell_type": "markdown", "id": "linear2d-interp", "metadata": {}, "source": [ "All six parameters land in the ballpark of `1e-4` to `1e-3` relative error. The two initial conditions `x0_1` and `x0_2` sit at the top of the list because they enter linearly and so pick up the full truncation error of the central-difference formula; the matrix entries roll into the propagation and average down. There is no outlier. This is the reference shape of a clean report." ] }, { "cell_type": "markdown", "id": "simex-section", "metadata": {}, "source": [ "## GL06SIMEX: log-space on a smooth SFC model\n", "\n", "GL06SIMEX is the smallest SFC model in the suite: three free parameters (`TaxRate`, `PropensityToConsumeIncome`, `PropensityToConsumeSavings`) and no portfolio switch. All behaviour is smooth, so log-space perturbation with the default `epsilon=1e-3` should agree with autograd to roughly `1e-4` relative error." ] }, { "cell_type": "code", "execution_count": 3, "id": "simex-code", "metadata": { "execution": { "iopub.execute_input": "2026-04-10T08:12:15.601417Z", "iopub.status.busy": "2026-04-10T08:12:15.601262Z", "iopub.status.idle": "2026-04-10T08:12:16.935164Z", "shell.execute_reply": "2026-04-10T08:12:16.934785Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Jacobian comparison: autograd vs numerical (log, eps=1e-3)\n", "Parameters compared: 3\n", "Overall max abs diff: 2.362e-02 max rel diff: 4.386e-04\n", "\n", "Parameter max_abs mean_abs max_rel mean_rel close nan/inf\n", "---------------------------------------------------------------------------------------------------------\n", "PropensityToConsumeIncome 4.344e-03 4.344e-03 4.386e-04 4.386e-04 1/1 no\n", "PropensityToConsumeSavings 2.014e-03 2.014e-03 2.034e-04 2.034e-04 1/1 no\n", "TaxRate 2.362e-02 2.362e-02 5.301e-05 5.301e-05 1/1 no\n" ] } ], "source": [ "model = get_model('GL06SIMEX')()\n", "\n", "jac_auto = JacobianAutograd(model).compute(\n", " loss_fn=mean_national_income, mode='rev',\n", ")\n", "jac_num = JacobianNumerical(model).compute(\n", " loss_fn=mean_national_income, mode='central', num_workers=1,\n", ")\n", "\n", "report = compare_jacobian_dicts(\n", " jac_auto, jac_num,\n", " method_a='autograd', method_b='numerical (log, eps=1e-3)',\n", " rtol=1e-3, atol=1e-6,\n", ")\n", "print(report.summary())" ] }, { "cell_type": "markdown", "id": "simex-interp", "metadata": {}, "source": [ "All three parameters agree to better than `5e-4` relative error. The max absolute differences look larger (`2e-2` on `TaxRate`) but that reflects the scale of the Jacobian entry, not a disagreement: once normalised by the magnitude, the report is clean. This is the pattern to expect on a smooth model: relative error sits near the log-space truncation floor." ] }, { "cell_type": "markdown", "id": "insout-section", "metadata": {}, "source": "## GL06INSOUT: localising a step-function disagreement\n\nThe Tobin portfolio model behaves differently. Some parameters agree well and others show visibly worse agreement. The goal here is not to fix the disagreement but to use `compare_jacobian_dicts` to identify which parameters are affected, and then to confirm, by testing different values of `epsilon` in the numerical differentiation, that the cause is a step-function operation in the model code rather than a smooth numerical artefact." }, { "cell_type": "code", "execution_count": 4, "id": "insout-code", "metadata": { "execution": { "iopub.execute_input": "2026-04-10T08:12:16.936106Z", "iopub.status.busy": "2026-04-10T08:12:16.936010Z", "iopub.status.idle": "2026-04-10T08:12:22.848639Z", "shell.execute_reply": "2026-04-10T08:12:22.848306Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Jacobian comparison: autograd vs numerical (log, eps=1e-3)\n", "Parameters compared: 5\n", "Overall max abs diff: 1.899e+00 max rel diff: 2.077e-01\n", "\n", "Parameter max_abs mean_abs max_rel mean_rel close nan/inf\n", "---------------------------------------------------------------------------------------------------------\n", "WealthShareBills_Constant 1.899e+00 1.899e+00 2.077e-01 2.077e-01 0/1 no\n", "WealthShareM2_Constant 5.841e-02 5.841e-02 5.841e-02 5.841e-02 0/1 no\n", "WealthShareBonds_Constant 1.534e-01 1.534e-01 7.729e-03 7.729e-03 0/1 no\n", "InventorySalesRatioBaseline 9.693e-02 9.693e-02 3.941e-03 3.941e-03 0/1 no\n", "TaxRate 3.356e-01 3.356e-01 1.868e-04 1.868e-04 1/1 no\n" ] } ], "source": [ "model = get_model('GL06INSOUT')()\n", "\n", "sample_params = [\n", " 'WealthShareM2_Constant',\n", " 'WealthShareBills_Constant',\n", " 'WealthShareBonds_Constant',\n", " 'TaxRate',\n", " 'InventorySalesRatioBaseline',\n", "]\n", "\n", "jac_auto = JacobianAutograd(model).compute(\n", " loss_fn=mean_nominal_output, mode='rev',\n", ")\n", "jac_num = JacobianNumerical(model).compute(\n", " loss_fn=mean_nominal_output, mode='central', num_workers=1,\n", " param_names=sample_params,\n", ")\n", "\n", "jac_auto_subset = {k: jac_auto[k] for k in sample_params}\n", "\n", "report = compare_jacobian_dicts(\n", " jac_auto_subset, jac_num,\n", " method_a='autograd', method_b='numerical (log, eps=1e-3)',\n", " rtol=1e-3, atol=1e-6,\n", ")\n", "print(report.summary())" ] }, { "cell_type": "markdown", "id": "insout-interp", "metadata": {}, "source": "The report is sorted by descending relative disagreement. `TaxRate` and `InventorySalesRatioBaseline` sit at the bottom with relative errors in the `1e-4` to `1e-3` range, consistent with the clean-model baseline from GL06SIMEX. The three Tobin constants are two to three orders of magnitude worse: `WealthShareBonds_Constant` at `8e-3`, `WealthShareM2_Constant` at `6e-2`, and `WealthShareBills_Constant` at `2e-1`.\n\nThe disagreement is concentrated on the three Tobin portfolio constants and not on the unrelated parameters. This tells us the affected parameters share a structural role in the model, and that role must involve a non-smooth operation. In GL06INSOUT, the `portfolio_switches` method uses a step-function `torch.where` to select between the M1- and M2-dominated regimes of the Tobin allocation. Parameters that shift the baseline portfolio composition can push the economy across that switch for some timesteps. When a central-difference perturbation of size `epsilon` straddles the discontinuity, the numerical Jacobian sees a finite jump while autograd reports zero gradient at the switch." }, { "cell_type": "markdown", "id": "eps-sweep-section", "metadata": {}, "source": "## Testing different values of epsilon in the numerical differentiation\n\nTesting different values of `epsilon` in the numerical differentiation will reveal whether the disagreement comes from a step-function operation or from a smooth numerical artefact. A step-function operation produces a non-monotone relationship between `epsilon` and the relative disagreement: at large `epsilon` the perturbation spans a region much wider than the discontinuity and the numerical derivative is dominated by whichever side of the switch each perturbed trajectory lands on, producing a large disagreement with autograd; at very small `epsilon` roundoff and cancellation in the central-difference formula erode the agreement from the other direction; and between these two regimes sits a value where the disagreement is smallest. A smooth numerical artefact, by contrast, produces a monotone curve.\n\nThe following cell computes the numerical Jacobian at three values of `epsilon` on the parameter with the largest disagreement in the previous report, `WealthShareBills_Constant`." }, { "cell_type": "code", "execution_count": 5, "id": "eps-sweep-code", "metadata": { "execution": { "iopub.execute_input": "2026-04-10T08:12:22.849726Z", "iopub.status.busy": "2026-04-10T08:12:22.849630Z", "iopub.status.idle": "2026-04-10T08:12:28.684101Z", "shell.execute_reply": "2026-04-10T08:12:28.683570Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Relative disagreement on WealthShareBills_Constant vs epsilon:\n", " eps=1e-02 max_rel_diff=1.186e+00\n", " eps=1e-03 max_rel_diff=2.077e-01\n", " eps=1e-05 max_rel_diff=2.503e-01\n" ] } ], "source": [ "worst = 'WealthShareBills_Constant'\n", "eps_values = [1e-2, 1e-3, 1e-5]\n", "rel_diffs = {}\n", "\n", "for eps in eps_values:\n", " jn = JacobianNumerical(model, epsilon=eps).compute(\n", " loss_fn=mean_nominal_output, mode='central', num_workers=1,\n", " param_names=[worst],\n", " )\n", " a = jac_auto[worst]\n", " n = jn[worst]\n", " denom = torch.maximum(a.abs(), n.abs()).clamp(min=1e-12)\n", " rel_diffs[eps] = ((a - n).abs() / denom).max().item()\n", "\n", "print(f'Relative disagreement on {worst} vs epsilon:')\n", "for eps, r in rel_diffs.items():\n", " print(f' eps={eps:>.0e} max_rel_diff={r:.3e}')" ] }, { "cell_type": "markdown", "id": "eps-sweep-interp", "metadata": {}, "source": "The three values show the expected non-monotone pattern. At `eps=1e-2` the perturbation is large enough that the relative disagreement exceeds 100%, well above the other two values. At `eps=1e-3` the disagreement is smallest, around 20%. At `eps=1e-5` it increases again as floating-point cancellation erodes the finite-difference signal. A smooth numerical source of error would produce a monotone curve in `epsilon`; the non-monotone pattern observed here is consistent with a step-function operation, and matches the `portfolio_switches` diagnosis." }, { "cell_type": "markdown", "id": "decision-tree", "metadata": {}, "source": "## Decision tree: reading a `compare_jacobian_dicts` report\n\nWhen an autograd vs numerical comparison fails on a MacroStat model, use the following checklist before reaching for deeper debugging.\n\n1. **Every parameter at `max_rel_diff` below roughly `1e-3`.** Nothing is wrong. The report is clean.\n2. **One or more parameters with `nan/inf`.** Autograd found a `NaN` or `Inf` in the gradient. Typical causes are divide-by-zero or `log(0)` inside the model behaviour; instrument with `torch.autograd.set_detect_anomaly(True)` and rerun.\n3. **A derived parameter (residual of a `LinearConstraint`) reports `max_rel_diff` of order unity, but autograd gives exactly zero.** Expected. The constraint system absorbs all sensitivity into the free parameters of that group. See the constraint-system notes in [Differentiability and Jacobian Tools](diff.rst) and the worked example in [Parameter constraints](constraints_example.ipynb). If you are numerically perturbing the *derived* parameter directly, stop: that perturbation violates the constraint. Perturb the free parameters instead.\n4. **A small subset of parameters shows disagreement two or more orders of magnitude worse than the rest, and the worst offenders share a structural role in the model.** Suspect a step-function operation. Recompute the numerical Jacobian at several values of `epsilon` on the worst offender as shown in the previous section; a non-monotone relationship between `epsilon` and the relative disagreement confirms the diagnosis. Locate the switch with a `grep` for `torch.where`, `torch.clamp`, or explicit `if` branches in the behaviour code, and replace it with one of the smooth alternatives provided by `Behavior.diffwhere`, `Behavior.tanhmask`, `Behavior.diffmin`, or `Behavior.diffmax`.\n5. **All parameters disagree uniformly, around `1e-2` or worse.** Likely a bug in the numerical configuration rather than the model. Check whether `parameter_space=\"log\"` is being used with zero parameters (should raise, but verify), whether `epsilon` is appropriate for the parameter scale, and whether the loss function itself is well-conditioned.\n\nFor the constraint system specifically, see [Constraints](core/constraints.rst) and the [constraints example notebook](constraints_example.ipynb)." } ], "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.13.5" } }, "nbformat": 4, "nbformat_minor": 5 }