Layers

Project.toml

@@ -14,20 +14,19 @@ OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
 [weakdeps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
-ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 
 [extensions]
+CorleoneComponentArraysExtension = "ComponentArrays"
 CorleoneMakieExtension = "Makie"
-CorleoneComponentArraysExtension = ["ComponentArrays"]
-CorleoneModelingToolkitExtension = "ModelingToolkit"
 
 [compat]
 Aqua = "0.8"
@@ -51,6 +50,7 @@ OrdinaryDiffEqTsit5 = "1"
 Random = "1"
 RecursiveArrayTools = "3.30"
 Reexport = "1.2"
+RuntimeGeneratedFunctions = "0.5"
 SafeTestsets = "0.1"
 SciMLBase = "2.110"
 SciMLSensitivity = "7.87"

docs/src/api.md

@@ -1,6 +1,34 @@
 # API Reference
 
-```@autodocs
-Modules = [Corleone]
-Order = [:type, :function]
+## Types
+
+```@docs
+Corleone.ControlParameter
+Corleone.ControlParameters
+Corleone.ControlSignal
+Corleone.InitialCondition
+Corleone.SingleShootingLayer
+Corleone.Trajectory
+```
+
+## Functions
+
+```@docs
+Corleone.default_controls
+Corleone.find_idx
+Corleone.get_block_structure
+Corleone.get_bounds
+Corleone.get_lower_bound
+Corleone.get_new_system
+Corleone.get_timegrid
+Corleone.get_tspan
+Corleone.get_upper_bound
+Corleone.is_shooted
+Corleone.is_shooting_solution
+Corleone.mythreadmap
+Corleone.remake_system
+Corleone.shooting_violations
+Corleone.to_val
+Corleone.ttype
+Corleone.utype
 ```

ext/CorleoneComponentArraysExtension.jl

@@ -2,24 +2,18 @@ module CorleoneComponentArraysExtension
 using Corleone
 using ComponentArrays
 
-struct CAFunctionWrapper{A, F} <: Corleone.AbstractCorleoneFunctionWrapper
-    "The axes of the componentarray"
-    ax::A
-    "The original function"
-    f::F
+Corleone.to_vec(::Val{:ComponentArrays}, u) = begin
+    collect(ComponentArray(u))
 end
 
-(f::CAFunctionWrapper)(ps, st) = f.f(ComponentArray(ps, f.ax), st)
-(f::CAFunctionWrapper)(res, ps, st) = f.f(res, ComponentArray(ps, f.ax), st)
-
-Corleone.to_vec(::CAFunctionWrapper, x...) = map(x -> isnothing(x) ? x : (collect ∘ ComponentArray)(x), x)
-
-function Corleone.wrap_functions(::Val{:ComponentArrays}, u0::NamedTuple, f...)
-    u0 = ComponentVector(u0)
-    ax = getaxes(u0)
-    return map(f) do fi
-        isnothing(fi) ? fi : CAFunctionWrapper{typeof(ax), typeof(fi)}(ax, fi)
+function Corleone.WrappedFunction(::Val{:ComponentArrays}, f, p, st; kwargs...)
+    u0 = ComponentVector(p)
+    pre = let ax = getaxes(u0)
+        (p) -> ComponentArray(p, ax)
     end
+    return Corleone.WrappedFunction{
+        typeof(f), typeof(pre),
+    }(f, pre)
 end
 
 end

ext/CorleoneMakieExtension.jl

@@ -1,5 +1,6 @@
 module CorleoneMakieExtension
 using Corleone
+using SymbolicIndexingInterface
 using Makie
 
 Makie.plottype(sol::Trajectory) = Makie.Lines
@@ -8,34 +9,59 @@ function Makie.used_attributes(::Type{<:Plot}, sol::Trajectory)
     return (:vars, :idxs)
 end
 
+maybevec(x::AbstractArray) = eachrow(reduce(hcat, x))
+maybevec(x) = x
+
+_getindex(x, i) = getindex(x, i)
+_getindex(x::Symbol, i) = x
+
 function Makie.convert_arguments(
         PT::Type{<:Plot},
         sol::Trajectory;
-        vars = nothing,
-        idxs = nothing
+        idxs::AbstractVector{<:Int} = Int64[],
+        vars::AbstractVector = [],
+        kwargs...
     )
-
-    if !isnothing(vars)
-        (!isnothing(idxs)) && error("Can't simultaneously provide vars and idxs!")
-        idxs = vars
+    if !isempty(idxs)
+        append!(
+            vars,
+            variable_symbols(sol)[idxs]
+        )
+    end
+    if isempty(vars)
+        for v in variable_symbols(sol)
+            push!(vars, variable_index(sol, v))
+        end
+    end
+    ts = []
+    xs = []
+    labels = String[]
+    foreach(vars) do var
+        if is_timeseries_parameter(sol, var)
+            x_current = maybevec(getp(sol, var)(sol))
+            append!(xs, x_current)
+            for i in eachindex(x_current)
+                push!(ts, sol.controls.collection[1].t)
+                push!(labels, string(_getindex(var, i)))
+            end
+        else
+            x_current = maybevec(getsym(sol, var)(sol))
+            append!(xs, x_current)
+            for i in eachindex(x_current)
+                push!(ts, sol.t)
+                push!(labels, string(_getindex(var, i)))
+            end
+        end
     end
-
-    idxs = isnothing(idxs) ? eachindex(sol.u[1]) : idxs
-    idxs = eltype(idxs) == Symbol ? [sol.sys.variables[i] for i in idxs] : idxs
-
-    plot_vecs = reduce(hcat, sol.u)[idxs, :]
 
     plot_type_sym = Makie.plotsym(PT)
 
-    inv_sys_variables = Dict(value => key for (key, value) in sol.sys.variables)
-    labels = string.([inv_sys_variables[i] for i in idxs])
-
     return map(
-        (x, y, label, i) -> PlotSpec(plot_type_sym, Point2f.(x, y); label, color = Cycled(i)),
-        [sol.t for _ in eachindex(idxs)],
-        eachrow(plot_vecs),
+        (x, y, label, i) -> PlotSpec(plot_type_sym, Point2f.(x, y); label, color = Cycled(i), kwargs...),
+        ts,
+        xs,
         labels,
-        eachindex(idxs)
+        eachindex(labels)
     )
 end
 end

ext/CorleoneModelingToolkitExtension.jl

@@ -16,10 +16,92 @@ using ModelingToolkit.Setfield
 using ModelingToolkit.SymbolicIndexingInterface
 using ModelingToolkit.Symbolics.RuntimeGeneratedFunctions
 
+using ModelingToolkit: IndexCache, DiscreteIndex, BufferTemplate, ParameterTimeseriesIndex
+
+@info "Loading MTK Extension..."
+
 RuntimeGeneratedFunctions.init(@__MODULE__)
 
+function _control_timeseries_pairs(controls)
+    pairs = Pair[]
+    for (i, c) in enumerate(values(controls.controls))
+        name = c.name
+        # Whole control timeseries, e.g. `u(t)`.
+        push!(pairs, name => ParameterTimeseriesIndex(i, :))
+        # Per-component access for array-valued controls, e.g. `u[1]`.
+        try
+            for j in eachindex(name)
+                push!(
+                    pairs, name[j] => ParameterTimeseriesIndex(i, j)
+                )
+            end
+        catch
+            nothing
+        end
+    end
+    return pairs
+end
+
+function _mtk_parameter_index_map(sys, controls)
+    syms = collect(parameter_symbols(sys))
+
+    # Add component symbols for array-valued parameters so `traj.ps[α[1]]` works.
+    for sym in copy(syms)
+        try
+            for j in eachindex(sym)
+                subsym = sym[j]
+                if !any(isequal(subsym), syms)
+                    push!(syms, subsym)
+                end
+            end
+        catch
+            nothing
+        end
+    end
+
+    for c in values(controls.controls)
+        name = c.name
+        if !any(isequal(name), syms)
+            push!(syms, name)
+        end
+        try
+            for j in eachindex(name)
+                sym = name[j]
+                if !any(isequal(sym), syms)
+                    push!(syms, sym)
+                end
+            end
+        catch
+            nothing
+        end
+    end
+
+    idxmap = Dict{Any, Any}()
+    for sym in syms
+        idxmap[sym] = parameter_index(sys, sym)
+    end
+    return idxmap
+end
+
+function Corleone.remake_system(sys::ModelingToolkit.AbstractSystem, controls)
+    control_symbols = map(values(controls.controls)) do ci
+        csym = ci.name
+        csym => ParameterTimeseriesIndex(1, parameter_index(sys, csym))
+    end
+    ps = map(parameters(sys)) do xi
+        xi => parameter_index(sys, xi)
+    end
+    return SymbolCache(
+        unknowns(sys),
+        ps,
+        independent_variable_symbols(sys);
+        timeseries_parameters = Dict(reduce(vcat, control_symbols)),
+    )
+end
+
+#= 
 include("MTKExtension/utils.jl")
 
 include("MTKExtension/optimal_control.jl")
-
+ =#
 end