control-toolbox / control-toolbox/CTFlows.jl

Ensemble flows — batch API for multiple shooting, globalization, and ensemble optimal control

Open
#341 1 comment 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

## Context

GPU support work (PR #326, issues #249 #111) validated two GPU execution models end-to-end on real H200 hardware and **explicitly deferred** the production ensemble/batch API on `Flow` to its own design note. This issue carries that design work forward. Roadmap: [v4 §4 (ensemble), Discussion #299, comment Jul 18 2026](https://github.com/control-toolbox/CTFlows.jl/discussions/299).

Three architecturally distinct use cases have been identified. They map to different execution models and will need different call shapes on `Flow`.

---

## Use case 1 — Multiple shooting (time-domain splitting)

Split `[t0, tf]` into `K` subintervals. Each leg `[tₖ, tₖ₊₁]` is integrated **independently in parallel**. Then a **synchronization/matching phase** solves the junction conditions between adjacent legs — the Newton step on the continuity residual.

```
parallel integration → CPU-side linear solve (matching conditions) → iterate
```

- **Parallel integration phase:** `K` independent ODEs → `DiffEqGPU.EnsembleProblem` (GPU model 2)
- **Synchronization phase:** small, inherently coupled/sequential linear system; stays CPU-side (or GPU-resident) between rounds of parallel integration

**GPU fit:** model 2 (`EnsembleProblem`) for the integration phase.

---

## Use case 2 — Multiple shooting over many initial guesses (globalization)

No time-splitting. Instead, `N` independent initial guesses (`p0`, or the full shooting unknown `ξ`) are integrated in parallel. Whichever guess(es) converge are kept — a standard globalization strategy for Newton on the shooting equation.

This is the most direct fit for `DiffEqGPU.EnsembleProblem` — genuinely independent trajectories, no coupling between them. **Already probed and validated on H200** (design report §6bis.1, run 11).

**GPU fit:** model 2 (`EnsembleProblem`), directly.

---

## Use case 3 — Ensemble optimal control (coupled via shared control and/or criterion)

A different problem class: **one OCP** whose state is the concatenation of `N` small, structurally-identical subsystems coupled through a **shared control** `u(t)` and/or a **pooled criterion** (average or worst-case over the ensemble). Classic "robust/ensemble control" setting.

**Concrete example:** MRI saturation for a single spin ([MagneticResonanceImaging.jl](https://control-toolbox.org/MagneticResonanceImaging.jl/stable/saturation.html)) generalises to multi-ε saturation — `N` isochromats with different resonance offsets `εᵢ`, all driven by the **same** control, optimised jointly (average or worst-case saturation error over the ensemble).

**Structure:** block-diagonal RHS (each subsystem depends only on its own state and the shared control, not on other subsystems) → naturally a matrix-batch RHS call `u::AbstractMatrix`, one column per subsystem. `HamIpAugRHS` already supports this mode today.

**GPU fit:** model 1 (array-level GPU) — the concatenated state can be large enough (N in the hundreds to thousands of isochromats) that the "large state" threshold from the design report is actually met. This is the rare case in optimal control where model 1 genuinely pays off.

---

## Architecture summary

| Use case | Trajectories | Coupling | GPU model | DiffEqGPU API |
|----------|-------------|----------|-----------|---------------|
| 1. Time-domain splitting | K legs, same (x₀,p₀), disjoint time windows | Matching conditions | Model 2 | `EnsembleProblem` |
| 2. Globalization (N guesses) | N independent guesses, same time window | None | Model 2 | `EnsembleProblem` |
| 3. Ensemble OCP (shared control) | N subsystems, same `u(t)` | Via `u(t)` + pooled criterion | Model 1 | Array-level GPU |

Cases 1 and 2 are independent-trajectory ensembles → `EnsembleProblem`. Case 3 is a single coupled system → matrix-batch RHS + array-level GPU.

---

## Scope of this issue

This issue covers the **design and implementation** of the batch/ensemble `Flow` API. GPU hardware validation is tracked in #249.

**Out of scope for the Phase 4b work already done (PR #326):**
- Production user-facing batch API
- `ext/CTFlowsDiffEqGPU` extension

---

## Open design questions

- **Call shape:** should `Flow` expose a `batch(flow, guesses)` / `ensemble(flow, subproblems)` surface, or a new `EnsembleFlow` type?
- **Unified vs. separate API:** cases 1/2 (EnsembleProblem) and case 3 (matrix-batch) are structurally different; they may warrant separate constructors rather than a unified overload.
- **CPU fallback:** cases 1/2 should degrade gracefully to `EnsembleThreads()` when no GPU is available.
- **Result layout:** how are per-trajectory results returned? Cf. the accessor footgun documented in the GPU design report (use `.u[i].u[end]`, never `sol[i][end]`).

---

## Work items

### Design (prerequisite)
- [ ] Write design note: call shapes for cases 1, 2, 3 — unified surface or separate APIs?
- [ ] Decide interaction with `ext/CTFlowsSciMLFlows` (already used for single-trajectory flows)

### Use case 2 — globalization (easiest, already probed)
- [ ] Implement `Flow` batch call for `N` initial guesses (wraps `EnsembleProblem`)
- [ ] CPU fallback via `EnsembleThreads()`
- [ ] Test: double integrator energy, `N=16` independent guesses → all converge to same `(xf, pf)`

### Use case 1 — time-domain splitting
- [ ] `K`-leg parallel integration + matching-condition API
- [ ] CPU fallback
- [ ] Test: double integrator time-optimal, 2-arc splitting

### Use case 3 — ensemble OCP (shared control)
- [ ] Confirm `HamIpAugRHS` matrix-batch path runs end-to-end on GPU array (under `allowscalar(false)`)
- [ ] Define `N`-subsystem OCP builder (concatenated state, shared control)
- [ ] MRI-inspired test case (`N` isochromats, shared saturation pulse)

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.