Add mesh LSQ Hessian functional

[loliverhennigh]

PhysicsNeMo Pull Request

Description

Adds mesh_lsq_hessian, a FunctionSpec-backed Hessian reconstruction for unstructured CSR neighborhoods.

• Fits linear and symmetric quadratic Taylor coefficients directly, avoiding the noise amplification of applying a first-derivative reconstruction twice.
• Supports 1D, 2D, and 3D scalar or tensor-valued fields, ragged neighborhoods, coordinate normalization, and rank/duplicate-neighbor safeguards.
• Provides a Torch reference and a preferred Warp column-pivoted QR backend with explicit coordinate and value adjoints, FakeTensor support, and AOT-compiled backward support.
• Exposes the public API, registers forward/backward benchmark cases, and adds focused tests, API documentation, a documentation figure, and changelog coverage.

Related derivative task: https://github.com/NVIDIA/physicsnemo-roadmap/issues/2358

Related functional roadmap: https://github.com/NVIDIA/physicsnemo-roadmap/issues/2662

Validation

• 53 passed in the focused Mesh LSQ Hessian suite.
• 247 passed across the derivative-functional suite.
• 2 passed in the Hessian package doctests.
• All 12 Torch/Warp forward/backward benchmark-plan cases completed successfully.
• Functional benchmark harness passed.
• Changed-file pre-commit hooks and git diff --check passed.

Local Warp validation used CPU. CUDA execution remains delegated to CI.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Model implementation standards are not applicable; this PR does not add or modify a model.

Dependencies

None. Warp remains an optional backend through the existing FunctionSpec dependency dispatch.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds mesh_lsq_hessian, a differentiable quadratic least-squares Hessian reconstruction for unstructured CSR neighborhoods. It fits linear and symmetric-quadratic Taylor coefficients jointly, avoiding the noise amplification of a twice-applied gradient reconstruction, and exposes both a PyTorch reference and a Warp CPQR backend with explicit coordinate and value adjoints.

• Math is correct end-to-end: the coefficient ordering in design_value, the quadratic_column_scale (0.5 for diagonals, 1.0 for cross-terms), the Hessian-to-coefficient mapping in forward_kernel, and the adjoint formulas in both backward kernels are all internally consistent across 1D, 2D, and 3D.
• Edge cases are handled properly: coincident neighbors (distance=0) receive zero weight; rank-deficient entities return zero Hessians; empty entity batches are guarded; scale_squared=0 is replaced by 1 in the Torch path to avoid rsqrt(0).
• Design pattern matches existing Warp backends (mesh_lsq_gradient): the same wp.init() at import time, the same FunctionSpec.register(required_imports=...) dispatch, and the same unconditional import of _warp_impl are already established patterns in this codebase.

Important Files Changed

Filename	Overview
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/mesh_lsq_hessian.py	Introduces MeshLSQHessian FunctionSpec class with Warp (rank-0) and Torch (rank-1/baseline) implementations, benchmark scaffolding, and the public mesh_lsq_hessian function alias.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/_torch_impl.py	Pure-PyTorch reference implementation: groups entities by neighbor count, builds normalized quadratic Taylor design matrices, applies a detached SVD rank gate, then solves full-rank groups with gels. Math is internally consistent.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/_warp_impl/op.py	Warp custom-op integration: forward, backward, and FakeTensor registrations. The QR factorization is intentionally re-run during backward (for AOT compatibility), which doubles geometry-kernel cost on the backward pass.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/_warp_impl/_kernels/factorization.py	Column-pivoted QR factorization kernel per entity; correctly handles rank detection via pivot thresholds. After factorization_valid is set to 0, the inner loop continues running but is guarded — wasted work in failing cases but not a correctness issue.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/_warp_impl/_kernels/backward_points.py	Coordinate adjoint kernel; adjoint math (residual, weight, and scale cotangents) is correct. The _row_coordinate_cotangent helper is called twice per neighbor edge (once for scale accumulation, once for edge gradient), which is a performance trade-off to avoid dynamic local storage.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/_warp_impl/_kernels/backward_values.py	Value adjoint kernel; correctly applies the transposed LSQ map to cotangents, with atomic accumulation for thread safety.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/_warp_impl/_kernels/utils.py	Shared fixed-width 9x9 algebra utilities; coefficient ordering, design values, back-substitution, and gradient helpers are internally consistent across all dimensionalities.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/_warp_impl/_kernels/forward.py	Forward solve kernel; correctly assembles Q^T W b, back-substitutes through stored R, unpermutes, and writes symmetric Hessian entries with appropriate inverse_scale_squared factors.
physicsnemo/nn/functional/derivatives/mesh_lsq_hessian/utils.py	Input validation and parameter-resolution helpers; CSR structural checks, dtype validation, and default coefficient-count logic are correct.
test/nn/functional/derivatives/test_mesh_lsq_hessian.py	Comprehensive 1025-line test suite covering correctness, edge cases, gradient checks, and Warp/Torch parity; tests are well-structured and mark Warp cases as optional.

_{Reviews (1): Last reviewed commit: "Add mesh LSQ Hessian functional" | Re-trigger Greptile}