Layers by AlCap23 · Pull Request #43 · SciML/Corleone.jl

AlCap23 · 2026-03-16T11:50:50Z

Checklist

Appropriate tests were added
Any code changes were done in a way that does not break public API
All documentation related to code changes were updated
The new code follows the
contributor guidelines, in particular the SciML Style Guide and
COLPRAC.
Any new documentation only uses public API

Additional context

This effectively replaces #38.

Add tests

Draft initialconditions

Easier and type stable

Copilot

Pull request overview

This PR is a substantial refactor of Corleone’s optimal-control “shooting” stack: controls and initial conditions are re-modeled as Lux layers/containers, single/multiple shooting evaluation is rewritten around binned time intervals, and new parallel/observed layers plus updated tests/docs are added.

Changes:

Replace the previous “local controls” implementation with ControlParameter/ControlParameters (+ fixed controls) and an InitialCondition layer.
Rework SingleShootingLayer/MultipleShootingLayer trajectory evaluation and introduce ParallelShootingLayer.
Add/adjust tests, docs, and extensions; add RuntimeGeneratedFunctions dependency.

Reviewed changes

Copilot reviewed 23 out of 24 changed files in this pull request and generated 14 comments.

Show a summary per file

File	Description
`src/Corleone.jl`	Wire up new layers/utilities; add `RuntimeGeneratedFunctions` init; update exports.
`src/controls.jl`	New control parameter and control container implementations (incl. fixed controls).
`src/initializers.jl`	Introduce `InitialCondition` Lux layer for tunable ICs and bounds.
`src/single_shooting.jl`	Rewrite `SingleShootingLayer` as a Lux container and implement binned interval evaluation.
`src/parallel_shooting.jl`	Add `ParallelShootingLayer` wrapper for ensemble-based evaluation of multiple layers.
`src/multiple_shooting.jl`	Reimplement multiple shooting as a wrapper over `ParallelShootingLayer` + matching constraints.
`src/trajectory.jl`	Extend `Trajectory` to carry control time series; adjust indexing/constraints helpers.
`src/observed.jl`	Add `ObservedLayer` that builds runtime-generated observation functions over trajectories.
`src/dynprob.jl`	Partially stub/disable dynamic opt problem machinery (mostly commented out).
`src/local_controls.jl`	Removed legacy local-control implementation.
`test/controls.jl`	New tests for control parameter/container behavior and remake semantics.
`test/initializers.jl`	New tests for `InitialCondition` bounds/setup/call/remake.
`test/single_shooting.jl`	New single-shooting tests: constructors, binning, evaluation, symbolic access fallback.
`test/parallel_shooting.jl`	New tests for `ParallelShootingLayer` construction/evaluation/remake.
`test/multiple_shooting.jl`	Replace old Lotka multiple-shooting tests with new layer-level tests.
`test/observed.jl`	Add observed-layer tests (currently not wired into `runtests.jl`).
`test/runtests.jl`	Reorganize testsets and add new shooting tests; comment out some example testsets.
`test/local_controls.jl`	Removed legacy tests for old control implementation.
`test/examples/lotka_oc.jl`	Update example to new controls/symbolic setup and revised optimization wiring.
`Project.toml`	Add `RuntimeGeneratedFunctions` dependency; reorder extensions.
`ext/CorleoneModelingToolkitExtension.jl`	Update MTK extension for control time-series indexing and system remaking.
`ext/CorleoneMakieExtension.jl`	Update plotting to handle state vs parameter time series.
`ext/CorleoneComponentArraysExtension.jl`	Comment out prior ComponentArrays wrappers (extension becomes mostly inert).
`docs/src/api.md`	Replace autodocs with explicit public API list.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/single_shooting.jl

src/controls.jl

src/observed.jl

src/parallel_shooting.jl

src/trajectory.jl

test/runtests.jl

src/parallel_shooting.jl

src/controls.jl

src/trajectory.jl

src/multiple_shooting.jl

Copilot

Pull request overview

This PR refactors Corleone’s shooting/optimization “layers” API, introducing structured control and initial-condition layers, adding parallel/multiple shooting composition, and updating the trajectory representation and dynamic optimization wrapper accordingly.

Changes:

Introduces ControlParameter/ControlParameters and InitialCondition Lux layers, and rewires SingleShootingLayer to compose them.
Adds ParallelShootingLayer and updates MultipleShootingLayer to wrap parallel solves and emit unified Trajectory objects with shooting violations.
Reworks Trajectory to expose time-series controls/parameters via SymbolicIndexingInterface; updates tests, docs, and project dependencies (RuntimeGeneratedFunctions).

Reviewed changes

Copilot reviewed 21 out of 22 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
test/single_shooting.jl	Adds coverage for new `SingleShootingLayer` constructors/state/time-binning and symbolic access.
test/runtests.jl	Reorganizes/renames testsets and adds new shooting/initializer/controls test files.
test/parallel_shooting.jl	Adds tests for `ParallelShootingLayer` block structure, evaluation, and remake behavior.
test/multiple_shooting.jl	Replaces older multi-shooting tests with new API exercising matching constraints and remake.
test/local_controls.jl	Removes legacy local-controls tests (replaced by `test/controls.jl`).
test/initializers.jl	Adds tests for `InitialCondition` construction/bounds/setup/call/remake.
test/examples/lotka_oc.jl	Updates example tests to new control/system wiring and `DynamicOptimizationLayer`.
test/controls.jl	Adds comprehensive tests for `ControlParameter`/`ControlParameters`, bounds, evaluation, and remake.
src/trajectory.jl	Extends `Trajectory` with control time series and revises symbolic indexing + shooting constraints utilities.
src/single_shooting.jl	Refactors single shooting to compose initial conditions + controls and rebuild symbolic systems.
src/parallel_shooting.jl	Introduces a wrapper to solve multiple layers via SciML ensemble mapping.
src/node_initialization.jl	Adds `InitialCondition` Lux layer implementation (and leaves legacy node init commented).
src/multiple_shooting.jl	Reimplements multiple shooting on top of `ParallelShootingLayer`, adds matching calculations.
src/local_controls.jl	Removes legacy local-controls implementation (replaced by `src/controls.jl`).
src/dynprob.jl	Replaces old dynamic opt problem wrapper with `DynamicOptimizationLayer` using runtime-generated getters.
src/controls.jl	Adds new controls layer system: tunable/fixed controls, containers, bounds/remake/eval utilities.
src/Corleone.jl	Wires new modules/exports, introduces `MAXBINSIZE`, and adds RuntimeGeneratedFunctions init.
ext/CorleoneModelingToolkitExtension.jl	Updates MTK extension logic for timeseries parameter indexing (but see Project.toml note).
ext/CorleoneMakieExtension.jl	Updates plotting conversion to support variables and time-series parameters/controls.
ext/CorleoneComponentArraysExtension.jl	Updates ComponentArrays vectorization integration to the new `WrappedFunction` approach.
docs/src/api.md	Replaces broad autodocs with curated API docs list.
Project.toml	Adds RuntimeGeneratedFunctions; changes extension registration and removes ModelingToolkit weakdep/extension.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/controls.jl

src/trajectory.jl

src/dynprob.jl

ext/CorleoneModelingToolkitExtension.jl

src/parallel_shooting.jl

src/dynprob.jl

src/controls.jl

Fix docstring Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

length(params) == length(t)

AlCap23 · 2026-03-18T05:02:42Z

┌──────────────────────────┬───────┬─────┬──────┬─────┐
│ Filename                 │ Lines │ Hit │ Miss │   % │
├──────────────────────────┼───────┼─────┼──────┼─────┤
│ src/Corleone.jl          │    33 │  21 │   12 │ 64% │
│ src/controls.jl          │   132 │ 124 │    8 │ 94% │
│ src/dynprob.jl           │   108 │  85 │   23 │ 79% │
│ src/initializers.jl      │    43 │  42 │    1 │ 98% │
│ src/multiple_shooting.jl │    62 │  61 │    1 │ 98% │
│ src/parallel_shooting.jl │    29 │  25 │    4 │ 86% │
│ src/single_shooting.jl   │    97 │  94 │    3 │ 97% │
│ src/trajectory.jl        │    35 │  24 │   11 │ 69% │
├──────────────────────────┼───────┼─────┼──────┼─────┤
│ TOTAL                    │   539 │ 476 │   63 │ 88% │
└──────────────────────────┴───────┴─────┴──────┴─────┘
 Target coverage was met (80%)

AlCap23 · 2026-03-18T05:07:29Z

@chplate Let's open up a PR into this branch to fix OED. I do not think we need a dev branch. Its just a naming convention here.

Copilot

Pull request overview

This PR refactors Corleone’s shooting/optimal-control stack into more composable Lux layers (initial conditions + controls), adds parallel/modernized multiple-shooting support, and updates the dynamic-optimization wrapper plus docs/tests accordingly.

Changes:

Introduces new control layer APIs (ControlParameter, FixedControlParameter, ControlParameters) and an InitialCondition layer.
Reworks SingleShootingLayer/MultipleShootingLayer internals and adds ParallelShootingLayer.
Replaces the prior dynamic optimization problem wrapper with DynamicOptimizationLayer, updating docs and tests.

Reviewed changes

Copilot reviewed 21 out of 22 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
src/controls.jl	New control parameter and container layer implementations.
src/initializers.jl	Adds `InitialCondition` layer and remaking behavior.
src/single_shooting.jl	Refactors single-shooting to compose IC + controls, adds binning-based evaluation.
src/parallel_shooting.jl	Adds parallel wrapper for solving multiple shooting layers concurrently.
src/multiple_shooting.jl	Re-implements multiple shooting as wrapper over `ParallelShootingLayer` + matching constraints.
src/trajectory.jl	Extends `Trajectory` to include control time series and updates indexing behavior.
src/dynprob.jl	Introduces `DynamicOptimizationLayer` and new expression-driven objective/constraints.
src/Corleone.jl	Wires in new modules/exports and adds `RuntimeGeneratedFunctions`.
ext/CorleoneComponentArraysExtension.jl	Updates ComponentArrays integration for the new optimization wrappers.
ext/CorleoneMakieExtension.jl	Updates plotting to support parameter time-series controls.
ext/CorleoneModelingToolkitExtension.jl	Updates MTK extension code (still present as an extension module).
docs/src/api.md	Updates API reference to list new types/functions explicitly.
Project.toml	Adds `RuntimeGeneratedFunctions` and adjusts extensions/weak deps.
test/*	Adds/reworks tests for controls, initializers, and shooting layers; trims older tests.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/parallel_shooting.jl

+    layers = map(keys(layer.layers)) do k
+        layer_kwargs = get(kwargs, k, kwargs)
+        k, remake(layer.layers[k]; layer_kwargs...)
+    end |> NamedTuple
+    ensemble_algorithm = get(kwargs, :ensemble_algorithm, layer.ensemble_algorithm)
+    return ParallelShootingLayer(layer.name, layers, ensemble_algorithm)


ext/CorleoneModelingToolkitExtension.jl

 module CorleoneModelingToolkitExtension

 @info "Loading MTK Extension"

 using Corleone
 using ModelingToolkit

 using Corleone.LuxCore
 using Corleone.Random
 using Corleone.DocStringExtensions
 using Corleone.SciMLBase

 using ModelingToolkit.Symbolics
 using ModelingToolkit.SymbolicUtils
 using ModelingToolkit.Setfield
 using ModelingToolkit.SymbolicIndexingInterface
 using ModelingToolkit.Symbolics.RuntimeGeneratedFunctions

+using ModelingToolkit: IndexCache, DiscreteIndex, BufferTemplate, ParameterTimeseriesIndex
+
+@info "Loading MTK Extension..."
+
 RuntimeGeneratedFunctions.init(@__MODULE__)



src/dynprob.jl

+"""
+$(SIGNATURES)
+
+Construct a SciML `OptimizationProblem` from a [`CorleoneDynamicOptProblem`](@ref).
+"""
 function SciMLBase.OptimizationProblem(
-        prob::CorleoneDynamicOptProblem, ad::SciMLBase.ADTypes.AbstractADType, vectorizer;
+        prob::DynamicOptimizationLayer, ad::SciMLBase.ADTypes.AbstractADType,
+        ps = nothing,
+        st = nothing;
        rng::Random.AbstractRNG = Random.default_rng(),
+        vectorizer,
        sense = nothing,
        kwargs...
    )
-    p0, st = LuxCore.setup(rng, prob.layer)
-    objective, cons = wrap_functions(vectorizer, p0, prob.objective, prob.constraints)
-    optf = SciMLBase.OptimizationFunction{true}(objective, ad; cons, kwargs...)
-    u0_, lb, ub = to_vec(objective, p0, Corleone.get_bounds(prob.layer)...)
-    return SciMLBase.OptimizationProblem(optf, u0_, st; lb, ub, lcons = prob.lcons, ucons = prob.ucons, sense = sense)
+    ps = something(ps, LuxCore.initialparameters(rng, prob))
+    st = something(st, LuxCore.initialstates(rng, prob))
+    optf = SciMLBase.OptimizationFunction(prob, ad; vectorizer, rng, ps, st, kwargs...)
+    u0, lb, ub = map(Base.Fix1(to_vec, vectorizer), (ps, Corleone.get_bounds(prob)...))
+    return SciMLBase.OptimizationProblem(optf, u0, st; lb, ub, lcons = prob.lcons, ucons = prob.ucons, sense = sense)
 end

-function SciMLBase.OptimizationProblem(
-        layer::Union{SingleShootingLayer, MultipleShootingLayer},
-        ad::SciMLBase.ADTypes.AbstractADType, vectorizer;
-        loss::Union{Symbol, Expr},
-        constraints = [],
+"""
+$(SIGNATURES)
+
+Construct a SciML `OptimizationFunction` from a [`CorleoneDynamicOptProblem`](@ref).
+"""
+function SciMLBase.OptimizationFunction(


src/dynprob.jl

-                reduce(vcat, _collect_tspans(st_.tspans...))
-            end
-        )
+        throw(error("The operator $(op) is not supported to define constraints. Only ==, <=, and >= are supported."))


src/trajectory.jl

+Parameter indexing through `getindex` is deprecated; use `A.ps[sym]` instead.
+"""
+function Base.getindex(A::Trajectory, sym)
+    if is_parameter(A, sym)
+        error("Indexing with parameters is deprecated. Use `sol.ps[$sym]` for parameter indexing.")


AlCap23 added 30 commits March 6, 2026 09:04

ControlParameter Rework

d1023d1

Remake works;

3ef70b2

Add tests

Rename test file;

29b4e7d

Draft initialconditions

Add tests for initial condition; fixed typos

450f544

Works for symbols typestable

e8f09f5

Indexing and type stable

4d7fd8a

Rework MTK approach

690a75f

Easier and type stable

Switch to tspans

55972fd

Update plots

45359dc

Format

7c589cb

Update docstrings

1c0c897

Minor fixes for plotting and trajectory

8752038

Rework controls and introduce fixed controls

223dd2b

Fix Zygote

507bc82

Add parallel and multiple shooting

8e2bfb9

Reiterating tests I

a882ba9

Fix control tests

4eb811b

Initializer tests

b908b2e

Add LQR testcases

4fb9a45

Fix typo

d6dcf88

Add more tests

86e8f16

Adapt lotka optimal control

4908fc2

Fix test errors

babb0d4

Update methods for dynprob

f1848cb

Start working on exprs

be46e87

Timegrid handling

120bbf6

Add expression based API for evaluation of objective and constraints

db0c587

Adapted observed

b83ad88

Switch to Runtimegenerated

b4302f8

Rework AI slop

e744fcf

Copilot started reviewing on behalf of AlCap23 March 16, 2026 11:51 View session

Copilot AI reviewed Mar 16, 2026

View reviewed changes

AlCap23 added 14 commits March 16, 2026 14:55

Update Dynprob and extension

1575561

Update dynprob

5b854ed

Fix bounds

a210e46

Fix bounds of FixedControl

a7ad6b0

Update tests and remove prints

661508b

Update tests and format

6f538c4

Rmv observed

2377639

Update docstrings

a8a1461

Format

c3fa160

Remove MTK Extensions for now

5275620

Fix tests

e1ccaae

Update minor fixes

0bdae97

Rmv legacy comments

9d7a74d

Minor fix

8835fa8

AlCap23 requested a review from Copilot March 17, 2026 09:10

Fix typos and format

35bd1df

Copilot AI reviewed Mar 17, 2026

View reviewed changes

AlCap23 and others added 5 commits March 17, 2026 10:25

Fix wrong index of lcons and ucons

b24a35b

Potential fix for pull request finding

4eabfcc

Fix docstring Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Fix Copilot code

3fbc8a0

Adapt parameter handling

b1dfa37

Update MS and tests

914f5db

length(params) == length(t)

AlCap23 requested a review from Copilot March 18, 2026 05:06

Copilot started reviewing on behalf of AlCap23 March 18, 2026 05:06 View session

Copilot AI reviewed Mar 18, 2026

View reviewed changes

Adapt to parameter observed interface.

899109f

Uh oh!

Conversation

AlCap23 commented Mar 16, 2026

Checklist

Additional context

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AlCap23 commented Mar 18, 2026

Uh oh!

AlCap23 commented Mar 18, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants