control-toolbox / control-toolbox/CTFlows.jl
Flow derivative via dual numbers — variational equations (IND)
- Dominant language
- Julia
- Stars
- 2
- Forks
- 3
- Avg merge
- 22h 16m
- Merged PRs (30d)
- 9
Description
## Goal
When the flow is evaluated on `ForwardDiff.Dual` inputs — typically to compute the Jacobian of a shooting function via AD — the naive approach pushes duals through the black-box ODE integrator. This is correct but differentiates through every internal integrator step, which is expensive and can be numerically fragile for long integration intervals or high-order solvers.
The correct approach is **Internal Numerical Differentiation (IND)**: detect dual-typed inputs and instead integrate the **variational (sensitivity) equations** alongside the state, giving `∂φ/∂x₀`, `∂φ/∂p₀`, `∂φ/∂t₀`, `∂φ/∂tf` from the sensitivity ODE.
Roadmap: [v4 §4, Discussion #299](https://github.com/control-toolbox/CTFlows.jl/discussions/299). Refs: [CTBase #25](https://github.com/control-toolbox/CTBase.jl/issues/25).
---
## Why this matters
The main use case is Newton-based shooting, where the Jacobian of the shooting function is computed by AD:
```julia
# Today — ForwardDiff pushes duals through the integrator (black-box, potentially slow/fragile)
J = ForwardDiff.jacobian(ξ -> shoot(ξ), ξ0)
# After IND — same call, but duals are detected and routed to the variational equations (fast, accurate)
J = ForwardDiff.jacobian(ξ -> shoot(ξ), ξ0)
```
The user call is identical; only the internal path changes.
---
## Design
### Variational system
Augment `(x, p)` with `(δx, δp)` propagated by the Jacobian of the Hamiltonian vector field `X_H`:
```
δż = J · δz, z = (x, p), J = DX_H(z)
```
For a Hamiltonian flow `X_H = (∂H/∂p, -∂H/∂x)`, the Jacobian block is:
```
J = [ ∂²H/∂p∂x ∂²H/∂p² ]
[ -∂²H/∂x² -∂²H/∂x∂p ]
```
The augmented state `(x, p, δx, δp)` is integrated as a single ODE. `J` is computed by AD (ForwardDiff on `X_H`).
### Tagged Dual dispatch
Overload the flow call for `Dual`-typed inputs with a **dedicated tag** (e.g. `CTFlowsVariationalTag`) so we can recognise "our" duals and route to the variational path, rather than colliding with an outer AD pass (e.g. an enclosing `ForwardDiff.jacobian` call). The tag disambiguates nested/foreign differentiation.
### Activation: option vs. strategy
Two choices:
- **Keyword option** on the flow call: `variational=true` or `derivative=true`
- **`CTBase.Strategies` strategy** — same family as the existing `:di` (backend) and `:sciml` (integrator) strategies, making the derivative-propagation method (IND / plain dual push-forward / future adjoint) a pluggable, dispatchable component
➡️ Decision pending. **Decide together with #46** (implicit control) so both features use the same option-or-strategy convention.
---
## Implementation steps
1. **Variational RHS functor** — augment the Hamiltonian system with `δż = J·δz` (Hessian of `H` via AD on `X_H`)
2. **Tagged Dual dispatch** — overload the flow call: Dual inputs with the CTFlows tag → variational path; foreign Dual inputs → pass through unchanged
3. **Derivative extraction** — extract `∂φ/∂x₀`, `∂φ/∂p₀`, etc. from the dual part of the output
4. **Validation** — cross-check against:
- finite differences
- naive dual push-forward (black-box AD)
---
## Work items
- [ ] Decide option vs. strategy (together with #46)
- [ ] Design and register `CTFlowsVariationalTag`
- [ ] Implement `VariationalRHS` functor (augmented `(x, p, δx, δp)` ODE, `J` via AD)
- [ ] Overload flow call for tagged Dual inputs → variational path
- [ ] Derivative extraction helpers (unpack dual part of `(xf, pf)`)
- [ ] Tests:
- `@inferred` on the variational call
- `J_ind ≈ J_fd` (finite-difference cross-check) for double integrator energy
- `J_ind ≈ J_ad` (naive AD cross-check) for small problem
Contributor guide
Assessment
This issue has not been assessed yet.