From 0ec216f1c42f8a407d5f87d3cd0403f1968df7f7 Mon Sep 17 00:00:00 2001 From: MoritzMaibaum Date: Fri, 24 Jul 2026 18:33:06 +0200 Subject: [PATCH] Clean up affine PZ two-jet notebook diagnostics --- .../affine_pz_twojet_adaquad_benchmarks.ipynb | 277 +++--------------- 1 file changed, 48 insertions(+), 229 deletions(-) diff --git a/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb b/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb index c445452..c62d181 100644 --- a/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb +++ b/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb @@ -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", @@ -96,7 +95,9 @@ "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" ] }, { @@ -104,131 +105,34 @@ "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 @@ -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" ] }, { @@ -570,4 +389,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +}