From d4d3dd783e6923b78ebc9c931a0c95ae7b4ea217 Mon Sep 17 00:00:00 2001 From: MoritzMaibaum Date: Sat, 25 Jul 2026 15:04:34 +0200 Subject: [PATCH] Optimize direct PZ two-jet norm integration --- .../affine_pz_twojet_adaquad_benchmarks.ipynb | 64 ++++-- src/intervalnets/pz_integration.py | 182 ++++++++++++++++++ src/intervalnets/pz_norms.py | 17 +- tests/test_pz_integration.py | 68 ++++++- tests/test_pz_norms.py | 39 +++- 5 files changed, 343 insertions(+), 27 deletions(-) diff --git a/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb b/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb index c62d181..f488e70 100644 --- a/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb +++ b/notebooks/affine_pz_twojet_adaquad_benchmarks.ipynb @@ -40,8 +40,9 @@ " sys.path.insert(0, str(repo_root / \"src\"))\n", "\n", "\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 import Interval, IntervalTensor, PZIntegrationCell, enable_interval_eval\n", + "from intervalnets.pz_norms import pz_twojet_l2_integrand, pz_twojet_l2_norm, pz_twojet_w12_integrand, pz_twojet_w12_norm, pz_twojet_w22_integrand, pz_twojet_w22_norm\n", + "from intervalnets.pz_integration import integrate_over_cell, integrate_pz_twojet_squared, _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", "enable_interval_eval(\"slope\")\n", @@ -105,7 +106,9 @@ "metadata": {}, "source": [ "cell = PZIntegrationCell.from_affine_box(domain)\n", + "t0 = time.perf_counter()\n", "jet = model.eval_pz_twojet(cell.domain)\n", + "twojet_construction_s = time.perf_counter() - t0\n", "\n", "from collections import Counter\n", "\n", @@ -143,26 +146,43 @@ "metadata": {}, "outputs": [], "source": [ - "## Direct PZ two-jet norm bounds\n", + "## Optimized PZ two-jet norm bounds and phase timings\n", "\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", + "# Keep the explicit squared PZ only as an optional regression/benchmark path.\n", + "EXPLICIT_COMPARISON = True\n", + "constructors = {\"L2\": pz_twojet_l2_integrand, \"W12\": pz_twojet_w12_integrand, \"W22\": pz_twojet_w22_integrand}\n", + "norms = {\"L2\": pz_twojet_l2_norm, \"W12\": pz_twojet_w12_norm, \"W22\": pz_twojet_w22_norm}\n", + "kinds = {\"L2\": \"l2\", \"W12\": \"w12\", \"W22\": \"w22\"}\n", "\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", - "\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", + "phase_rows = []\n", + "for quantity in [\"L2\", \"W12\", \"W22\"]:\n", + " explicit_integrand = None\n", + " explicit_integrand_s = explicit_integration_s = None\n", + " explicit_bounds = None\n", + " if EXPLICIT_COMPARISON:\n", + " t0 = time.perf_counter(); explicit_integrand = constructors[quantity](jet)\n", + " explicit_integrand_s = time.perf_counter() - t0\n", + " t0 = time.perf_counter(); explicit_integral = integrate_over_cell(explicit_integrand, cell, output=\"interval\")\n", + " explicit_integration_s = time.perf_counter() - t0\n", + " explicit_bounds = _sqrt_interval_nonnegative(explicit_integral)\n", + "\n", + " t0 = time.perf_counter(); direct_integral = integrate_pz_twojet_squared(jet, cell, kinds[quantity])\n", + " direct_integrated_square_s = time.perf_counter() - t0\n", + " # This is the optimized public API used by applications.\n", + " bounds = norms[quantity](jet, cell)\n", + " if explicit_bounds is not None:\n", + " assert math.isclose(float(bounds.lower), float(explicit_bounds.lower), rel_tol=1e-11, abs_tol=1e-11)\n", + " assert math.isclose(float(bounds.upper), float(explicit_bounds.upper), rel_tol=1e-11, abs_tol=1e-11)\n", + " phase_rows.append({\n", + " \"quantity\": quantity, \"bounds\": bounds,\n", + " \"twojet_construction_s\": twojet_construction_s,\n", + " \"explicit_integrand_construction_s\": explicit_integrand_s,\n", + " \"direct_integrated_square_s\": direct_integrated_square_s,\n", + " \"explicit_integrand_integration_s\": explicit_integration_s,\n", + " \"total_cell_s\": twojet_construction_s + direct_integrated_square_s,\n", + " })\n", + "\n", + "direct_pz_norms = pd.DataFrame(phase_rows)\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", @@ -285,7 +305,9 @@ "source": [ "## 6. Benchmark tables: width, runtime, cells, and refinement steps\n", "\n", - "The following cell performs one cached adaptive run per `(method, quantity)` pair up to `MAX_REFINEMENT_STEPS`. Each row is the partial certified result after that many refinement steps from the same run.\n" + "The following cell performs one cached adaptive run per `(method, quantity)` pair up to `MAX_REFINEMENT_STEPS`. Each row is the partial certified result after that many refinement steps from the same run.\n", + "\n", + "`runtime_s` below is the separately accumulated adaptive runtime on the same model and domain; the per-cell construction/integrand/integration phases are reported above.\n" ] }, { diff --git a/src/intervalnets/pz_integration.py b/src/intervalnets/pz_integration.py index 2051eb2..16fb57e 100644 --- a/src/intervalnets/pz_integration.py +++ b/src/intervalnets/pz_integration.py @@ -10,19 +10,24 @@ from __future__ import annotations from dataclasses import dataclass +from itertools import product from math import inf, isfinite, nextafter, prod, sqrt +from numbers import Real from typing import Any, Literal, Sequence from .interval import Interval from .polynomial_zonotope import ( Exponent, + PZTwoJet, PolynomialZonotope, _abs_coeff, _add_coeff, _mul_coeff, + _merge_noise_kinds, _to_fallback, _zero_like, box_monomial_moment, + torch, ) try: # pragma: no cover - optional dependency @@ -235,6 +240,183 @@ def integrate_over_cell(pz_expr: PolynomialZonotope, cell: PZIntegrationCell, *, return integrate_pz_over_domain(weighted, cell.domain_noise_indices, mode="pointwise_interval").interval_enclosure() +TwoJetIntegrandKind = Literal["l2", "w12", "w22"] + + +def _scalar_coordinates(zonotope: PolynomialZonotope) -> list[PolynomialZonotope]: + """Flatten a tensor-valued PZ without converting its scalar coefficients.""" + + if zonotope.shape == (): + return [zonotope] + return [zonotope[index] for index in product(*(range(size) for size in zonotope.shape))] + + +def _twojet_weighted_coordinates( + jet: PZTwoJet, integrand_kind: TwoJetIntegrandKind +) -> tuple[list[PolynomialZonotope], list[float]]: + if integrand_kind not in {"l2", "w12", "w22"}: + raise ValueError("integrand_kind must be 'l2', 'w12', or 'w22'.") + + coordinates = _scalar_coordinates(jet.Y) + weights = [1.0] * len(coordinates) + if integrand_kind in {"w12", "w22"}: + jacobian = _scalar_coordinates(jet.J) + coordinates.extend(jacobian) + weights.extend([1.0] * len(jacobian)) + if integrand_kind == "w22": + shape = jet.H.shape + if len(shape) == 3 and shape[1] == shape[2]: + output_indices: tuple[int | None, ...] = tuple(range(shape[0])) + input_dim = shape[1] + elif len(shape) == 2 and shape[0] == shape[1]: + output_indices = (None,) + input_dim = shape[0] + else: + raise ValueError("w22 direct integration requires a square stored Hessian.") + for output in output_indices: + for row in range(input_dim): + for column in range(row, input_dim): + index = (row, column) if output is None else (output, row, column) + coordinates.append(jet.H[index]) + weights.append(1.0 if row == column else 2.0) + return coordinates, weights + + +def _validate_twojet_metadata(jet: PZTwoJet) -> tuple[int, tuple[str, ...]]: + components = (jet.Y, jet.J, jet.H) + num_noise = components[0].num_noise + if any(component.num_noise != num_noise for component in components[1:]): + raise ValueError("Y, J, and H must have identical num_noise metadata.") + noise_kinds = components[0].noise_kinds + for component in components[1:]: + noise_kinds = _merge_noise_kinds(noise_kinds, component.noise_kinds) + return num_noise, noise_kinds + + +def _coefficient_matrix( + coordinates: Sequence[PolynomialZonotope], union_support: Sequence[Exponent] +) -> tuple[Any, Any]: + centers = [coordinate.center for coordinate in coordinates] + matrix = [ + [coordinate.terms.get(exponent, _zero_like(coordinate.center)) for coordinate in coordinates] + for exponent in union_support + ] + if torch is not None and isinstance(centers[0], torch.Tensor): + center_vector = torch.stack(centers) + if matrix: + return center_vector, torch.stack([torch.stack(row) for row in matrix]) + return center_vector, torch.empty((0, len(centers)), dtype=center_vector.dtype, device=center_vector.device) + return centers, matrix + + +def _weighted_dot(left: Sequence[Any], right: Sequence[Any], weights: Sequence[float]): + result = _zero_like(left[0]) + for lhs, rhs, weight in zip(left, right, weights): + result = _add_coeff(result, _mul_coeff(_mul_coeff(lhs, rhs), weight)) + return result + + +def integrate_pz_twojet_squared( + jet: PZTwoJet, cell: PZIntegrationCell, integrand_kind: TwoJetIntegrandKind +): + """Directly integrate a squared two-jet over a supported affine cell. + + Unsupported densities and Hessian layouts deliberately use the explicit + squared-integrand reference pipeline. + """ + + def explicit_fallback(): + from .pz_norms import ( + pz_twojet_l2_integrand, + pz_twojet_w12_integrand, + pz_twojet_w22_integrand, + ) + + constructor = { + "l2": pz_twojet_l2_integrand, + "w12": pz_twojet_w12_integrand, + "w22": pz_twojet_w22_integrand, + }.get(integrand_kind) + if constructor is None: + raise ValueError("integrand_kind must be 'l2', 'w12', or 'w22'.") + return integrate_over_cell(constructor(jet), cell, output="interval") + + density = cell.jacobian_density + if not isinstance(density, Real) or not isfinite(float(density)) or float(density) < 0.0: + return explicit_fallback() + + num_noise, noise_kinds = _validate_twojet_metadata(jet) + try: + coordinates, weights = _twojet_weighted_coordinates(jet, integrand_kind) + except ValueError as error: + if "Hessian" in str(error): + return explicit_fallback() + raise + if not coordinates: + return explicit_fallback() + + domain_indices = tuple(int(index) for index in cell.domain_noise_indices) + if len(set(domain_indices)) != len(domain_indices) or any(index < 0 or index >= num_noise for index in domain_indices): + raise ValueError("cell domain noise index out of range or duplicated.") + domain_set = set(domain_indices) + retained_indices = tuple(index for index in range(num_noise) if index not in domain_set) + retained_kinds = tuple(noise_kinds[index] for index in retained_indices) + pointwise_indices = tuple(index for index, kind in enumerate(noise_kinds) if kind in POINTWISE_RESIDUAL_KINDS) + measure = float(2 ** len(domain_indices)) + scale = float(density) + + support = sorted(set().union(*(coordinate.terms for coordinate in coordinates))) + centers, matrix = _coefficient_matrix(coordinates, support) + if torch is not None and isinstance(centers, torch.Tensor): + weight_vector = torch.tensor(weights, dtype=centers.dtype, device=centers.device) + weighted_matrix = matrix * weight_vector.unsqueeze(0) + center_cross = weighted_matrix @ centers + gram = weighted_matrix @ matrix.T + center_square = torch.dot(centers * weight_vector, centers) + else: + center_cross = [_weighted_dot(row, centers, weights) for row in matrix] + gram = [[_weighted_dot(left, right, weights) for right in matrix] for left in matrix] + center_square = _weighted_dot(centers, centers, weights) + + retained: dict[Exponent, Any] = {} + pointwise: dict[Exponent, Any] = {} + zero_retained = (0,) * len(retained_indices) + + def accumulate(target: dict[Exponent, Any], exponent: Exponent, coefficient: Any) -> None: + target[exponent] = _add_coeff(target[exponent], coefficient) if exponent in target else coefficient + + def route(exponent: Exponent, coefficient: Any) -> None: + scaled = _mul_coeff(coefficient, scale) + if any(exponent[index] for index in pointwise_indices): + accumulate(pointwise, exponent, scaled) + return + moment = box_monomial_moment(tuple(exponent[index] for index in domain_indices)) + if moment == 0.0: + return + retained_exponent = tuple(exponent[index] for index in retained_indices) + accumulate(retained, retained_exponent, _mul_coeff(scaled, moment)) + + route((0,) * num_noise, center_square) + for index, exponent in enumerate(support): + route(exponent, _mul_coeff(center_cross[index], 2.0)) + for other_index in range(index, len(support)): + pair_exponent = tuple(a + b for a, b in zip(exponent, support[other_index])) + factor = 1.0 if index == other_index else 2.0 + route(pair_exponent, _mul_coeff(gram[index][other_index], factor)) + + center = retained.pop(zero_retained, _zero_like(centers[0])) + radius = _zero_like(center) + for coefficient in pointwise.values(): + radius = _add_coeff(radius, _mul_coeff(_abs_coeff(coefficient), measure)) + result = IntegratedPZResult( + polynomial=PolynomialZonotope(center, retained, num_noise=len(retained_indices), noise_kinds=retained_kinds), + interval_radius=radius, + measure=measure, + metadata={"mode": "pointwise_interval", "direct_twojet_squared": True, "integrand_kind": integrand_kind}, + ) + return result.interval_enclosure() + + def _require_interval_tensor_domain(domain: Any): from .pytorch import IntervalTensor as RuntimeIntervalTensor diff --git a/src/intervalnets/pz_norms.py b/src/intervalnets/pz_norms.py index f7dc8af..45e2664 100644 --- a/src/intervalnets/pz_norms.py +++ b/src/intervalnets/pz_norms.py @@ -8,7 +8,7 @@ from .interval import Interval from .polynomial_zonotope import PZTwoJet, PolynomialZonotope, pz_to_latex, pz_to_markdown_code, twojet_to_latex -from .pz_integration import PZIntegrationCell, integrate_over_cell, integrate_pz_over_domain +from .pz_integration import PZIntegrationCell, integrate_over_cell, integrate_pz_over_domain, integrate_pz_twojet_squared try: # pragma: no cover - optional dependency import torch @@ -212,12 +212,21 @@ def pz_norm_from_integrand( def pz_twojet_l2_norm(jet: PZTwoJet, cell: PZIntegrationCell | None = None, *, p: float = 2.0) -> Interval: - return pz_norm_from_integrand(pz_twojet_l2_integrand(jet), cell, p=p) + _require_p2(p) + if cell is not None: + return _sqrt_interval_nonnegative(integrate_pz_twojet_squared(jet, cell, "l2")) + return pz_norm_from_integrand(pz_twojet_l2_integrand(jet), p=p) def pz_twojet_w12_norm(jet: PZTwoJet, cell: PZIntegrationCell | None = None, *, p: float = 2.0) -> Interval: - return pz_norm_from_integrand(pz_twojet_w12_integrand(jet), cell, p=p) + _require_p2(p) + if cell is not None: + return _sqrt_interval_nonnegative(integrate_pz_twojet_squared(jet, cell, "w12")) + return pz_norm_from_integrand(pz_twojet_w12_integrand(jet), p=p) def pz_twojet_w22_norm(jet: PZTwoJet, cell: PZIntegrationCell | None = None, *, p: float = 2.0) -> Interval: - return pz_norm_from_integrand(pz_twojet_w22_integrand(jet), cell, p=p) + _require_p2(p) + if cell is not None: + return _sqrt_interval_nonnegative(integrate_pz_twojet_squared(jet, cell, "w22")) + return pz_norm_from_integrand(pz_twojet_w22_integrand(jet), p=p) diff --git a/tests/test_pz_integration.py b/tests/test_pz_integration.py index f3082be..908fff6 100644 --- a/tests/test_pz_integration.py +++ b/tests/test_pz_integration.py @@ -1,12 +1,15 @@ +from dataclasses import replace + import pytest -from intervalnets import PolynomialZonotope +from intervalnets import PZTwoJet, PolynomialZonotope from intervalnets.pz_integration import ( IntegratedPZResult, POINTWISE_RESIDUAL_KINDS, PZIntegrationCell, integrate_over_cell, integrate_pz_over_domain, + integrate_pz_twojet_squared, ) from intervalnets.pz_tanh import ( affine_tanh_double_prime_enclosure, @@ -21,6 +24,69 @@ torch = None +def _assert_interval_close(left, right): + assert float(left.lower) == pytest.approx(float(right.lower), rel=1e-12, abs=1e-12) + assert float(left.upper) == pytest.approx(float(right.upper), rel=1e-12, abs=1e-12) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +@pytest.mark.parametrize("kind", ["l2", "w12", "w22"]) +def test_direct_twojet_square_matches_explicit_with_unequal_tensor_supports(kind): + from intervalnets.pz_norms import pz_twojet_l2_integrand, pz_twojet_w12_integrand, pz_twojet_w22_integrand + + torch.manual_seed(19) + kinds = ("domain", "approximation_symbolic", "approximation_pointwise") + y = PolynomialZonotope(torch.randn(2, dtype=torch.float64), {(1, 0, 0): torch.randn(2, dtype=torch.float64)}, num_noise=3, noise_kinds=kinds) + j = PolynomialZonotope(torch.randn(2, 2, dtype=torch.float64), {(0, 1, 0): torch.randn(2, 2, dtype=torch.float64), (1, 0, 1): torch.randn(2, 2, dtype=torch.float64)}, num_noise=3, noise_kinds=kinds) + h = PolynomialZonotope(torch.randn(2, 2, 2, dtype=torch.float64), {(2, 0, 0): torch.randn(2, 2, 2, dtype=torch.float64), (0, 0, 1): torch.randn(2, 2, 2, dtype=torch.float64)}, num_noise=3, noise_kinds=kinds) + jet = PZTwoJet(y, j, h) + cell = PZIntegrationCell.from_bounds((-2.0,), (2.0,)) + constructors = {"l2": pz_twojet_l2_integrand, "w12": pz_twojet_w12_integrand, "w22": pz_twojet_w22_integrand} + + direct = integrate_pz_twojet_squared(jet, cell, kind) + explicit = integrate_over_cell(constructors[kind](jet), cell, output="interval") + + _assert_interval_close(direct, explicit) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_direct_twojet_square_canonicalizes_pointwise_cancellation_and_keeps_odd_domain_terms(): + kinds = ("domain", "approximation_pointwise") + zero_j = PolynomialZonotope.constant(torch.zeros((2, 1), dtype=torch.float64), num_noise=2, noise_kinds=kinds) + zero_h = PolynomialZonotope.constant(torch.zeros((2, 1, 1), dtype=torch.float64), num_noise=2, noise_kinds=kinds) + cell = PZIntegrationCell.from_bounds((-1.0,), (1.0,)) + + cancelling_y = PolynomialZonotope(torch.tensor([1.0, 1.0], dtype=torch.float64), {(0, 1): torch.tensor([1.0, -1.0], dtype=torch.float64)}, num_noise=2, noise_kinds=kinds) + cancelling = integrate_pz_twojet_squared(PZTwoJet(cancelling_y, zero_j, zero_h), cell, "l2") + assert float(cancelling.lower) == pytest.approx(0.0, abs=1e-14) + assert float(cancelling.upper) == pytest.approx(8.0) + + odd_y = PolynomialZonotope(torch.tensor([0.0], dtype=torch.float64), {(1, 0): torch.tensor([1.0], dtype=torch.float64), (0, 1): torch.tensor([1.0], dtype=torch.float64)}, num_noise=2, noise_kinds=kinds) + odd_jet = PZTwoJet(odd_y, zero_j[:1], zero_h[:1]) + direct = integrate_pz_twojet_squared(odd_jet, cell, "l2") + from intervalnets.pz_norms import pz_twojet_l2_integrand + explicit = integrate_over_cell(pz_twojet_l2_integrand(odd_jet), cell, output="interval") + _assert_interval_close(direct, explicit) + assert float(direct.lower) < -5.0 # The odd alpha*eta term receives full measure. + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_direct_twojet_square_falls_back_for_polynomial_density(): + kinds = ("domain",) + cell = PZIntegrationCell.from_bounds((-1.0,), (1.0,)) + cell = replace(cell, jacobian_density=PolynomialZonotope.constant(1.0, num_noise=1, noise_kinds=kinds)) + jet = PZTwoJet( + PolynomialZonotope(torch.tensor([1.0], dtype=torch.float64), {(1,): torch.tensor([0.5], dtype=torch.float64)}, num_noise=1, noise_kinds=kinds), + PolynomialZonotope.constant(torch.zeros((1, 1), dtype=torch.float64), num_noise=1, noise_kinds=kinds), + PolynomialZonotope.constant(torch.zeros((1, 1, 1), dtype=torch.float64), num_noise=1, noise_kinds=kinds), + ) + from intervalnets.pz_norms import pz_twojet_l2_integrand + + direct = integrate_pz_twojet_squared(jet, cell, "l2") + explicit = integrate_over_cell(pz_twojet_l2_integrand(jet), cell, output="interval") + _assert_interval_close(direct, explicit) + + def test_integrating_pointwise_residual_adds_radius_not_symbolic_moment(): z = PolynomialZonotope( 1.0, diff --git a/tests/test_pz_norms.py b/tests/test_pz_norms.py index 3123dec..ad7dd62 100644 --- a/tests/test_pz_norms.py +++ b/tests/test_pz_norms.py @@ -2,7 +2,7 @@ from intervalnets import IntervalTensor, PZTwoJet, PolynomialZonotope, enable_interval_eval from intervalnets.pz_integration import PZIntegrationCell, integrate_over_cell -from intervalnets.pz_norms import build_pz_twojet_norm_diagnostics, pz_sum_squares, pz_symmetric_hessian_sum_squares, pz_twojet_l2_integrand, pz_twojet_w12_integrand, pz_twojet_w22_integrand +from intervalnets.pz_norms import build_pz_twojet_norm_diagnostics, pz_norm_from_integrand, pz_sum_squares, pz_symmetric_hessian_sum_squares, pz_twojet_l2_integrand, pz_twojet_l2_norm, pz_twojet_w12_integrand, pz_twojet_w12_norm, pz_twojet_w22_integrand, pz_twojet_w22_norm try: import torch @@ -113,6 +113,43 @@ def test_pz_twojet_w22_integrand_uses_symmetric_hessian_accumulation_equivalent_ _assert_same_pz(optimized, dense) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +@pytest.mark.parametrize( + ("norm", "integrand"), + [(pz_twojet_l2_norm, pz_twojet_l2_integrand), (pz_twojet_w12_norm, pz_twojet_w12_integrand), (pz_twojet_w22_norm, pz_twojet_w22_integrand)], +) +@pytest.mark.parametrize("constant", [0.0, 2.5]) +def test_public_twojet_norm_direct_path_matches_explicit_for_zero_and_constant_jets(norm, integrand, constant): + kinds = ("domain",) + jet = PZTwoJet( + PolynomialZonotope.constant(torch.tensor([constant], dtype=torch.float64), num_noise=1, noise_kinds=kinds), + PolynomialZonotope.constant(torch.zeros((1, 2), dtype=torch.float64), num_noise=1, noise_kinds=kinds), + PolynomialZonotope.constant(torch.zeros((1, 2, 2), dtype=torch.float64), num_noise=1, noise_kinds=kinds), + ) + cell = PZIntegrationCell.from_bounds((-1.0,), (1.0,)) + + direct = norm(jet, cell) + explicit = pz_norm_from_integrand(integrand(jet), cell) + + assert float(direct.lower) == pytest.approx(float(explicit.lower), abs=1e-12) + assert float(direct.upper) == pytest.approx(float(explicit.upper), abs=1e-12) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_direct_w22_uses_authoritative_upper_hessian_and_weight_two(): + kinds = ("domain",) + zero_y = PolynomialZonotope.constant(torch.zeros(1, dtype=torch.float64), num_noise=1, noise_kinds=kinds) + zero_j = PolynomialZonotope.constant(torch.zeros((1, 2), dtype=torch.float64), num_noise=1, noise_kinds=kinds) + # The deliberately different lower entry must be ignored. + h = PolynomialZonotope.constant(torch.tensor([[[0.0, 3.0], [100.0, 0.0]]], dtype=torch.float64), num_noise=1, noise_kinds=kinds) + jet = PZTwoJet(zero_y, zero_j, h) + result = pz_twojet_w22_norm(jet, PZIntegrationCell.from_bounds((-1.0,), (1.0,))) + + expected = (2.0 * 2.0 * 3.0**2) ** 0.5 + assert float(result.lower) == pytest.approx(expected) + assert float(result.upper) == pytest.approx(expected) + + def _small_tanh_model(input_dim=1, hidden_dim=2, output_dim=1): model = nn.Sequential( nn.Linear(input_dim, hidden_dim, dtype=torch.float64),