Skip to content

feat: track variables evaluated in integrand for downstream analysis

Choose a tag to compare

@khegazy khegazy released this 31 May 06:47
· 13 commits to main since this release
81f01ba

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_variables is a tuple of tensors, each shaped [N, C, *var_dims], aligned with result.nodes ([N, C, T]) and result.y ([N, C, D]) over the accepted nodes.
  • Values are detached (diagnostic-only; no autograd graph retained).
  • Works across both sampling modes, both take_gradient paths, and float32/float64.
  • Non-float tracked variables (e.g. int64/bool labels) 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; ruff clean; no new mypy regressions.