Optimize direct PZ two-jet squared integration

[MoritzMaibaum]

Motivation

• Avoid constructing the full squared polynomial zonotope when integrating squared two-jet quantities over affine cells with scalar non-negative Jacobian density to reduce expensive sparse convolution and improve performance.
• Preserve existing semantics: pointwise-residual handling, canonical cancellation, outward rounding/enclosure behavior, tensor-shaped coefficients, and authoritative upper-triangular Hessian convention.

Description

• Add integrate_pz_twojet_squared(jet, cell, integrand_kind) and private helpers in src/intervalnets/pz_integration.py that select/flatten two-jet coordinates, build a deterministically sorted union support, form a weighted coefficient matrix, compute the weighted Gram and center cross terms, traverse unordered exponent pairs, canonicalize contributions before absolute-value collapse, route pointwise and retained terms by domain moments, and return the result via IntegratedPZResult.interval_enclosure(); fallback to the explicit-squared pipeline when cell.jacobian_density is unsupported.
• Implement coordinate extraction that assigns weight 2 to upper off-diagonal Hessian coordinates while treating the stored upper triangle as authoritative, plus metadata validation and deterministic union-support construction.
• Make pz_twojet_l2_norm, pz_twojet_w12_norm, and pz_twojet_w22_norm in src/intervalnets/pz_norms.py use the new direct integration path when a supported PZIntegrationCell is supplied while retaining the explicit integrand helpers for diagnostics and output="pz".
• Add regression tests in tests/test_pz_integration.py and tests/test_pz_norms.py covering direct vs explicit equivalence for l2, w12, w22, unequal tensor supports, cancellation examples, odd-domain pointwise terms, Hessian upper-triangle weighting, constants/zero jets, and the polynomial-density fallback; update the affine two-jet benchmark notebook to call the optimized public norm API and record separate phase timings.

Testing

• Ran targeted tests: pytest -q tests/test_polynomial_zonotope.py tests/test_pz_integration.py tests/test_pz_norms.py, which passed.
• Ran full test suite: pytest -q, which completed successfully (all automated tests passed).
• Ran additional targeted runs and compilation checks: pytest -q tests/test_pz_integration.py tests/test_pz_norms.py and python -m compileall -q src tests, all succeeded; test runs emitted a benign PyTorch/NumPy initialization warning but no failures.

---

Codex Task

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d4d3dd783e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Route adaptive norm cells through the direct integrator

When applications call model.pz_l2norm(...) or model.pz_sobolev_norm(...), _pz_adaptive_squared_integral still reaches _evaluate_squared_contribution_cache, which constructs _squared_twojet_integrand(jet, integrand_kind) and passes it to integrate_over_cell; it never calls this new helper. Consequently, the primary adaptive API—and the notebook's PZ API/trace benchmarks—continues materializing the expensive squared PZ for every active cell, so the intended optimization is absent from the path where repeated Hessian convolution matters most.

AGENTS.md reference: AGENTS.md:L45-L47

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Validate only components used by the selected norm

When a manually constructed PZTwoJet has valid Y metadata but unused J or H components with different num_noise metadata, pz_twojet_l2_norm(jet, cell) now raises here even though the previous explicit L2 path only squared Y and succeeded; the same problem affects W12 when only H differs. Validate metadata after selecting the coordinates, or limit validation to the components required by integrand_kind, so supplying a cell does not change otherwise supported public behavior.

AGENTS.md reference: AGENTS.md:L30-L32

Useful? React with 👍 / 👎.