Skip to content

Clean up affine PZ two-jet notebook diagnostics #188

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MoritzMaibaum merged 1 commit into from
Jul 24, 2026
Merged

Clean up affine PZ two-jet notebook diagnostics #188

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
277 changes: 48 additions & 229 deletions notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,8 @@
"if str(repo_root / \"src\") not in sys.path:\n",
" sys.path.insert(0, str(repo_root / \"src\"))\n",
"\n",
"from IPython.display import Markdown, display\n",
"\n",
"from intervalnets import Interval, IntervalTensor, PZIntegrationCell, build_pz_twojet_norm_diagnostics, collect_pz_diagnostics, enable_interval_eval, pz_sum_squares, pz_to_latex\n",
"from intervalnets import Interval, IntervalTensor, PZIntegrationCell, enable_interval_eval, pz_sum_squares\n",
"from intervalnets.pz_integration import integrate_over_cell, _dorfler_marking as _pz_dorfler_marking, _evaluate_squared_contribution_cache, _interval_add, _interval_width, _split_box, _sqrt_interval_nonnegative\n",
"from intervalnets.pytorch import _box_volume, _choose_split_dim, _dorfler_marking, _hessian_is_exact_zero, _interval_tensor_is_exact_constant, _jacobian_is_exact_zero, _lp_pointwise_power_bounds_refined, _sobolev_pointwise_power_bounds_refined\n",
"\n",
Expand Down Expand Up @@ -96,139 +95,44 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The next cell uses `model.eval_pz_twojet(..., return_trace=True)` to capture the actual propagated polynomial zonotopes. The generated Markdown/LaTeX trace is size-bounded by `TRACE_MAX_TERMS`, `TRACE_PRECISION`, and `TRACE_COMPONENT_LIMIT`, and it displays the domain polynomial plus `Y`, `J`, and `H` before they are squared into norm integrands and before integration over domain noise or Jacobian-determinant factors.\n"
"## Final two-jet monomial diagnostics\n",
"\n",
"Summarize the final propagated polynomial-zonotope two-jet before constructing the norm integrands.\n"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"cell = PZIntegrationCell.from_affine_box(domain)\n",
"jet = model.eval_pz_twojet(cell.domain)\n",
"\n",
"# Opt-in, size-bounded rendering of the actual propagated PZ two-jet polynomials.\n",
"# These limits are intentionally conservative so the notebook remains usable even\n",
"# as term counts grow. Set TRACE_DISPLAY_MARKDOWN=False to write the trace file\n",
"# without rendering it inline.\n",
"TRACE_MAX_TERMS = 12\n",
"TRACE_PRECISION = 5\n",
"TRACE_COMPONENT_LIMIT = 8\n",
"TRACE_DISPLAY_MARKDOWN = True\n",
"TRACE_OUTPUT_PATH = repo_root / \"notebooks\" / \"generated\" / \"pz_twojet_trace.md\"\n",
"\n",
"trace_result = model.eval_pz_twojet(cell.domain, return_trace=True)\n",
"jet = trace_result.final\n",
"trace_records = trace_result.records\n",
"\n",
"\n",
"def _pz_component_count(z):\n",
" if z.shape == ():\n",
" return 1\n",
" return int(math.prod(z.shape))\n",
"\n",
"\n",
"def _fallback_scalar_indices(value, prefix=()):\n",
" if isinstance(value, tuple):\n",
" for idx, item in enumerate(value):\n",
" yield from _fallback_scalar_indices(item, prefix + (idx,))\n",
" else:\n",
" yield prefix\n",
"\n",
"\n",
"def _pz_limited_components(z, limit):\n",
" if z.shape == ():\n",
" yield (), z\n",
" return\n",
" total = _pz_component_count(z)\n",
" displayed = min(total, limit)\n",
" if isinstance(z.center, torch.Tensor):\n",
" for flat_idx in range(displayed):\n",
" flat_tensor = torch.tensor(flat_idx, device=z.center.device)\n",
" multi = tuple(int(i) for i in torch.unravel_index(flat_tensor, z.center.shape))\n",
" yield multi, z[multi]\n",
" return\n",
" for flat_idx, multi in enumerate(_fallback_scalar_indices(z.center)):\n",
" if flat_idx >= displayed:\n",
" break\n",
" yield multi, z[multi]\n",
"\n",
"\n",
"def _pz_stats(z):\n",
" return f\"shape={tuple(z.shape)}, terms={len(z.terms)}, max_degree={max((sum(exp) for exp in z.terms), default=0)}, noise={z.num_noise}\"\n",
"\n",
"\n",
"def _render_pz_components(label, z):\n",
" lines = [f\"#### {label} ({_pz_stats(z)})\"]\n",
" for index, scalar in _pz_limited_components(z, TRACE_COMPONENT_LIMIT):\n",
" suffix = \"\" if index == () else \"[\" + \", \".join(str(i) for i in index) + \"]\"\n",
" lines.append(f\"**{label}{suffix}**: terms={len(scalar.terms)}, max_degree={max((sum(exp) for exp in scalar.terms), default=0)}\")\n",
" lines.append(\"```latex\")\n",
" lines.append(pz_to_latex(scalar, max_terms=TRACE_MAX_TERMS, precision=TRACE_PRECISION))\n",
" lines.append(\"```\")\n",
" omitted = _pz_component_count(z) - min(_pz_component_count(z), TRACE_COMPONENT_LIMIT)\n",
" if omitted > 0:\n",
" lines.append(f\"_Omitted {omitted} additional scalar components because TRACE_COMPONENT_LIMIT={TRACE_COMPONENT_LIMIT}._\")\n",
" return \"\\n\\n\".join(lines)\n",
"\n",
"\n",
"def _render_trace_markdown():\n",
" lines = [\n",
" \"# PZ two-jet propagated polynomial trace\",\n",
" \"\",\n",
" \"This file is generated by `notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb`.\",\n",
" \"\",\n",
" \"The displayed polynomials are the propagated input/domain and two-jet components `Y`, `J`, and `H` after each recorded hidden layer / `nn.Sequential` child. They are shown before squaring into the norm integrands (`y_sq`, `j_sq`, `h_sq`) and before integration over domain-noise moments or Jacobian-determinant factors.\",\n",
" \"\",\n",
" f\"Rendering limits: TRACE_MAX_TERMS={TRACE_MAX_TERMS}, TRACE_PRECISION={TRACE_PRECISION}, TRACE_COMPONENT_LIMIT={TRACE_COMPONENT_LIMIT}.\",\n",
" \"\",\n",
" \"## Input/domain polynomial\",\n",
" _render_pz_components(\"X\", cell.domain),\n",
" \"\",\n",
" \"## Propagated two-jet trace\",\n",
" ]\n",
" for record in trace_records:\n",
" lines.extend([\n",
" \"\",\n",
" f\"### Layer {record.layer_index}: `{record.layer_name}` ({record.layer_type})\",\n",
" f\"Summary: Y {record.summary['Y']}; J {record.summary['J']}; H {record.summary['H']}.\",\n",
" _render_pz_components(\"Y\", record.jet.Y),\n",
" _render_pz_components(\"J\", record.jet.J),\n",
" _render_pz_components(\"H\", record.jet.H),\n",
" ])\n",
" return \"\\n\\n\".join(lines) + \"\\n\"\n",
"\n",
"\n",
"TRACE_OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)\n",
"TRACE_OUTPUT_PATH.write_text(_render_trace_markdown(), encoding=\"utf-8\")\n",
"print(f\"Wrote PZ two-jet trace to {TRACE_OUTPUT_PATH.relative_to(repo_root)}\")\n",
"if TRACE_DISPLAY_MARKDOWN:\n",
" display(Markdown(TRACE_OUTPUT_PATH.read_text(encoding=\"utf-8\")))\n",
"\n",
"# Final two-jet diagnostics are built before norm integrand construction so the\n",
"# raw pre-norm PZ objects, rendered snippets, and metadata can be inspected.\n",
"norm_diagnostics = build_pz_twojet_norm_diagnostics(\n",
" jet,\n",
" max_terms=TRACE_MAX_TERMS,\n",
" precision=TRACE_PRECISION,\n",
")\n",
"norm_diagnostics[\"metadata\"]\n",
"\n",
"l2_integrand = norm_diagnostics[\"integrands\"][\"l2_integrand\"]\n",
"w12_integrand = norm_diagnostics[\"integrands\"][\"w12_integrand\"]\n",
"w22_integrand = norm_diagnostics[\"integrands\"][\"w22_integrand\"]\n",
"from collections import Counter\n",
"\n",
"def _cell_norm_from_cached_integrand(integrand):\n",
" integral = integrate_over_cell(integrand, cell, output=\"interval\")\n",
" return _sqrt_interval_nonnegative(integral)\n",
"\n",
"direct_pz_norms = pd.DataFrame([\n",
" {\"quantity\": \"L2\", \"bounds\": _cell_norm_from_cached_integrand(l2_integrand)},\n",
" {\"quantity\": \"W12\", \"bounds\": _cell_norm_from_cached_integrand(w12_integrand)},\n",
" {\"quantity\": \"W22\", \"bounds\": _cell_norm_from_cached_integrand(w22_integrand)},\n",
"def summarize_final_twojet_component(component, z):\n",
" kind_counts = Counter(z.noise_kinds)\n",
" approximation_noise = sum(\n",
" count for kind, count in kind_counts.items()\n",
" if \"approximation\" in str(kind)\n",
" )\n",
" return {\n",
" \"component\": component,\n",
" \"shape\": tuple(z.shape),\n",
" \"distinct_monomials\": len(z.terms),\n",
" \"max_degree\": max((sum(exp) for exp in z.terms), default=0),\n",
" \"domain_noise\": kind_counts.get(\"domain\", 0),\n",
" \"approximation_noise\": approximation_noise,\n",
" \"total_noise\": z.num_noise,\n",
" }\n",
"\n",
"\n",
"final_twojet_monomial_diagnostics = pd.DataFrame([\n",
" summarize_final_twojet_component(\"value Y\", jet.Y),\n",
" summarize_final_twojet_component(\"first derivative J\", jet.J),\n",
" summarize_final_twojet_component(\"second derivative H\", jet.H),\n",
"])\n",
"direct_pz_norms[\"lower\"] = direct_pz_norms[\"bounds\"].map(lambda z: float(z.lower))\n",
"direct_pz_norms[\"upper\"] = direct_pz_norms[\"bounds\"].map(lambda z: float(z.upper))\n",
"direct_pz_norms[\"width\"] = direct_pz_norms[\"upper\"] - direct_pz_norms[\"lower\"]\n",
"direct_pz_norms.drop(columns=\"bounds\")\n"
"final_twojet_monomial_diagnostics\n"
],
"outputs": [],
"execution_count": null
Expand All @@ -239,115 +143,30 @@
"metadata": {},
"outputs": [],
"source": [
"# Optional polynomial-zonotope multiplication diagnostics for the direct two-jet path.\n",
"# This does not change default behavior; diagnostics are collected only inside\n",
"# collect_pz_diagnostics(...) contexts.\n",
"_pz_diag_records = []\n",
"with collect_pz_diagnostics(\"eval_pz_twojet\") as records:\n",
" diag_jet = model.eval_pz_twojet(cell.domain)\n",
"_pz_diag_records.extend(records)\n",
"\n",
"with collect_pz_diagnostics(\"squared_components\") as records:\n",
" diag_y_sq = pz_sum_squares(diag_jet.Y)\n",
" diag_j_sq = pz_sum_squares(diag_jet.J)\n",
" diag_h_sq = pz_sum_squares(diag_jet.H)\n",
"_pz_diag_records.extend(records)\n",
"\n",
"with collect_pz_diagnostics(\"cumulative_integrands\") as records:\n",
" diag_l2_integrand = diag_y_sq\n",
" diag_w12_integrand = diag_y_sq + diag_j_sq\n",
" diag_w22_integrand = diag_w12_integrand + diag_h_sq\n",
"_pz_diag_records.extend(records)\n",
"\n",
"with collect_pz_diagnostics(\"integrate_cached_integrands\") as records:\n",
" _ = integrate_over_cell(diag_l2_integrand, cell, output=\"interval\")\n",
" _ = integrate_over_cell(diag_w12_integrand, cell, output=\"interval\")\n",
" _ = integrate_over_cell(diag_w22_integrand, cell, output=\"interval\")\n",
"_pz_diag_records.extend(records)\n",
"\n",
"pz_diag_df = pd.DataFrame(_pz_diag_records)\n",
"if pz_diag_df.empty:\n",
" pz_diag_summary = pd.DataFrame(columns=[\"multiplications\", \"raw_pair_count\", \"output_term_count\", \"max_output_degree\"])\n",
"else:\n",
" pz_diag_summary = (\n",
" pz_diag_df.groupby(\"phase\", dropna=False)\n",
" .agg(\n",
" multiplications=(\"raw_pair_count\", \"size\"),\n",
" raw_pair_count=(\"raw_pair_count\", \"sum\"),\n",
" output_term_count=(\"output_term_count\", \"sum\"),\n",
" max_output_degree=(\"max_output_degree\", \"max\"),\n",
" )\n",
" .reset_index()\n",
" )\n",
"pz_diag_summary\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Monomial growth diagnostics\n",
"\n",
"Statically summarize polynomial-zonotope term growth for the propagated two-jet and norm integrands.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from collections import Counter\n",
"\n",
"## Direct PZ two-jet norm bounds\n",
"\n",
"def _pz_scalar_entries_for_diagnostics(z):\n",
" if z.shape == ():\n",
" yield z\n",
" return\n",
" if isinstance(z.center, torch.Tensor):\n",
" for flat_idx in range(z.center.numel()):\n",
" multi = tuple(int(i) for i in torch.unravel_index(torch.tensor(flat_idx, device=z.center.device), z.center.shape))\n",
" yield z[multi]\n",
" return\n",
"y_sq = pz_sum_squares(jet.Y)\n",
"j_sq = pz_sum_squares(jet.J)\n",
"h_sq = pz_sum_squares(jet.H)\n",
"l2_integrand = y_sq\n",
"w12_integrand = y_sq + j_sq\n",
"w22_integrand = w12_integrand + h_sq\n",
"\n",
" def _fallback_scalar_indices(value, prefix=()):\n",
" if isinstance(value, tuple):\n",
" for idx, item in enumerate(value):\n",
" yield from _fallback_scalar_indices(item, prefix + (idx,))\n",
" else:\n",
" yield prefix\n",
"\n",
" for index in _fallback_scalar_indices(z.center):\n",
" yield z[index]\n",
"def _cell_norm_from_cached_integrand(integrand):\n",
" integral = integrate_over_cell(integrand, cell, output=\"interval\")\n",
" return _sqrt_interval_nonnegative(integral)\n",
"\n",
"\n",
"def summarize_pz_monomials(component, z):\n",
" scalar_term_counts = [len(entry.terms) for entry in _pz_scalar_entries_for_diagnostics(z)]\n",
" kind_counts = Counter(z.noise_kinds)\n",
" row = {\n",
" \"component\": component,\n",
" \"shape\": tuple(z.shape),\n",
" \"num_noise\": z.num_noise,\n",
" \"num_terms\": len(z.terms),\n",
" \"max_degree\": max((sum(exp) for exp in z.terms), default=0),\n",
" \"scalar_entries\": len(scalar_term_counts),\n",
" \"mean_terms_per_scalar\": float(np.mean(scalar_term_counts)) if scalar_term_counts else 0.0,\n",
" \"max_terms_per_scalar\": max(scalar_term_counts, default=0),\n",
" \"square_pair_work_estimate\": sum(term_count ** 2 for term_count in scalar_term_counts),\n",
" }\n",
" row.update({f\"noise_kind:{kind}\": count for kind, count in sorted(kind_counts.items())})\n",
" return row\n",
"\n",
"\n",
"monomial_diagnostics = pd.DataFrame([\n",
" summarize_pz_monomials(\"jet.Y\", jet.Y),\n",
" summarize_pz_monomials(\"jet.J\", jet.J),\n",
" summarize_pz_monomials(\"jet.H\", jet.H),\n",
" summarize_pz_monomials(\"l2_integrand\", l2_integrand),\n",
" summarize_pz_monomials(\"w12_integrand\", w12_integrand),\n",
" summarize_pz_monomials(\"w22_integrand\", w22_integrand),\n",
"]).fillna(0)\n",
"monomial_diagnostics\n"
"direct_pz_norms = pd.DataFrame([\n",
" {\"quantity\": \"L2\", \"bounds\": _cell_norm_from_cached_integrand(l2_integrand)},\n",
" {\"quantity\": \"W12\", \"bounds\": _cell_norm_from_cached_integrand(w12_integrand)},\n",
" {\"quantity\": \"W22\", \"bounds\": _cell_norm_from_cached_integrand(w22_integrand)},\n",
"])\n",
"direct_pz_norms[\"lower\"] = direct_pz_norms[\"bounds\"].map(lambda z: float(z.lower))\n",
"direct_pz_norms[\"upper\"] = direct_pz_norms[\"bounds\"].map(lambda z: float(z.upper))\n",
"direct_pz_norms[\"width\"] = direct_pz_norms[\"upper\"] - direct_pz_norms[\"lower\"]\n",
"direct_pz_norms.drop(columns=\"bounds\")\n"
]
},
{
Expand Down Expand Up @@ -570,4 +389,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
}
}
Loading