functional: lazy functional algebra (mirror LinOp)

[Pavlo3P]

Summary

Gives Functional the same lazy additive/scalar algebra LinOp already has (0.4.2 W4), so objectives compose naturally:

2.0 * F        # ScaledFunctional
F + G          # SumFunctional
F - G          # SumFunctional(F, -G)
-F             # ScaledFunctional(-1, F)
F + 3.0        # ShiftedFunctional (affine — the one place functionals extend beyond LinOp)

The key correctness invariant: a functional's grad is a metric (Riesz) gradient — a domain element — so gradients combine through X.add/X.scale (not raw +/*), which is correct on tree/stacked/weighted domains.

Changes (four gradual commits)

1. Scalar algebra — ScaledFunctional (grad = X.scale(a, F.grad)), make_scaled_functional (unit passthrough, nested-scalar fold), __neg__/__mul__/__rmul__.
2. Additive algebra — SumFunctional (grad = term grads folded via X.add), make_functional_sum (flatten nested sums, single-unwrap, shared-domain validation), __add__/__radd__/__sub__/__rsub__ (0 + F enables builtin sum).
3. Affine + identity — ShiftedFunctional (F + c, grad unchanged), ZeroFunctional (additive identity; make_scaled returns it for a zero scalar, make_functional_sum drops it). Scalar operands route to the affine shift. Registered all four in functional_cases() so the registry-completeness check + generic functional laws (value / gradient / directional-derivative / conversion) exercise them.
4. Verify fixes (below).

Adversarial verify — two bugs caught and fixed pre-merge

A multi-agent review (gradient correctness, factory laws, overload dispatch, jit integration) found:

• Complex-scalar conjugation: the Riesz gradient of a·F is conj(a)·grad(F) (the domain inner product conjugates its first argument), matching ScaledLinOp.rapply. ScaledFunctional.grad was scaling by a, not conj(a) — wrong for complex scalars. Fixed; value still scales by a.
• Swallowed domain mismatch: make_functional_sum dropped ZeroFunctional terms before validating domains, so F(X) + Zero(Y) silently returned F. Now validates all flattened terms' domains (folding in backend/dtype) before dropping zeros, mirroring linop.make_sum.

Both are covered by regression tests (complex Riesz identity; mismatched-domain Zero raising). Two jit-edge findings were reviewed and dismissed as unreachable.

Testing

• tests/functional/test_algebra.py (31 tests): scalar/additive/affine/zero algebra, factory simplifications, tree-domain grad via X.add/X.scale, complex Riesz identity, domain-mismatch guards, pytree round-trips, and an end-to-end minimize_optax(0.5F + 0.5F).
• Full suite: 3455 passed, 170 skipped; ruff / ruff-D / jit_audit / docstring_audit all clean.

🤖 Generated with Claude Code