feat: track variables evaluated in integrand for downstream analysis
A backward-compatible addition on top of v1.0.0: the adaptive quadrature engine can now track and return extra per-node quantities computed by the integrand, without integrating them.
✨ What's new
The integrand f may now return a 2-tuple (integrand, tracked_variables) instead of a bare integrand tensor. Tracked variables are evaluated at every quadrature node but not integrated — they ride alongside the integrand through the adaptive loop and are returned at the accepted nodes in the new IntegrationResult.tracked_variables field.
Use it to surface intermediate diagnostics (an energy, a label, a field, …) computed during integration without a second pass over the nodes.
def f(t):
return torch.sin(t), (t**2, torch.cos(t)) # integrand + 2 tracked vars
result = integrate(f, mesh_init=..., mesh_final=...)
result.integral # ∫ sin
result.tracked_variables # (t**2 at nodes, cos(t) at nodes), each [N, C, ...]result.tracked_variablesis a tuple of tensors, each shaped[N, C, *var_dims], aligned withresult.nodes([N, C, T]) andresult.y([N, C, D]) over the accepted nodes.- Values are detached (diagnostic-only; no autograd graph retained).
- Works across both sampling modes, both
take_gradientpaths, andfloat32/float64. - Non-float tracked variables (e.g.
int64/boollabels) are supported — the sorted-insert machinery is now dtype/device-aware.
🔁 Backward compatibility
Fully backward compatible — the feature is opt-in:
- An integrand returning a bare tensor behaves exactly as before, with
result.tracked_variables = None. - A single tracked tensor is wrapped into a 1-tuple, and
(integrand, None)is accepted.
No changes are required to existing code.
✅ Verification
- New tracked-variables test suite plus public-API coverage.
- Snapshot regression tests unchanged — numerics are preserved on the existing (bare-tensor) path.
- Full suite green;
ruffclean; no newmypyregressions.