control-toolbox / control-toolbox/CTFlows.jl

Flow derivative via dual numbers — variational equations (IND)

Open
#93 18 comments 0 reactions 1 assignee Claimed by @ocots View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.