Merge pull request #4702 from SebastianM-C/smc/codegen_input_defaults

[AI] Use declared inputs for input-aware codegen on compiled systems

Author: AayushSabharwal

lib/ModelingToolkitBase/ext/MTKCasADiDynamicOptExt.jl

@@ -71,12 +71,15 @@ struct CasADiDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
     wrapped_model::CasADiModel
     kwargs::K
 
-    function CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
+    function CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs)
         return new{
             typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5),
             typeof(p), typeof(f), typeof(kwargs),
         }(f, u0, tspan, p, model, kwargs)
     end
+    function CasADiDynamicOptProblem(f, u0, tspan, p, model; kwargs...)
+        return CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs)
+    end
 end
 
 function (M::MXLinearInterpolation)(τ)

lib/ModelingToolkitBase/ext/MTKInfiniteOptExt.jl

@@ -56,12 +56,15 @@ struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
     wrapped_model::InfiniteOptModel
     kwargs::K
 
-    function JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
+    function JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs)
         return new{
             typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5),
             typeof(p), typeof(f), typeof(kwargs),
         }(f, u0, tspan, p, model, kwargs)
     end
+    function JuMPDynamicOptProblem(f, u0, tspan, p, model; kwargs...)
+        return JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs)
+    end
 end
 
 struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
@@ -73,12 +76,15 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
     wrapped_model::InfiniteOptModel
     kwargs::K
 
-    function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
+    function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs)
         return new{
             typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
             typeof(p), typeof(f), typeof(kwargs),
         }(f, u0, tspan, p, model, kwargs)
     end
+    function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model; kwargs...)
+        return InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs)
+    end
 end
 
 MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel()

lib/ModelingToolkitBase/ext/MTKPyomoDynamicOptExt.jl

@@ -83,12 +83,15 @@ struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
     wrapped_model::PyomoDynamicOptModel
     kwargs::K
 
-    function PyomoDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
+    function PyomoDynamicOptProblem(f, u0, tspan, p, model, kwargs)
         return new{
             typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5),
             typeof(p), typeof(f), typeof(kwargs),
         }(f, u0, tspan, p, model, kwargs)
     end
+    function PyomoDynamicOptProblem(f, u0, tspan, p, model; kwargs...)
+        return PyomoDynamicOptProblem(f, u0, tspan, p, model, kwargs)
+    end
 end
 
 function pysym_getproperty(s::Union{Num, SymbolicT}, name::Symbol)

lib/ModelingToolkitBase/src/inputoutput.jl

@@ -33,6 +33,23 @@ See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref), [`bound_outputs`](@re
 """
 unbound_inputs(sys) = filter(x -> !is_bound(sys, x), inputs(sys))
 
+"""
+    default_codegen_inputs(sys)
+
+The inputs to use by default for input-aware code generation.
+
+For scheduled (compiled) systems this is `inputs(sys)`: the inputs declared to
+`mtkcompile` are the contract the simplified system was built around. The
+[`unbound_inputs`](@ref) heuristic cannot be used there — after flattening and
+simplification an effective input appears in equations together with variables
+from other namespaces and is therefore classified as bound, so `unbound_inputs`
+is empty for compiled hierarchical systems.
+
+For unscheduled systems this is [`unbound_inputs`](@ref), which inspects the
+connection structure of the hierarchy to find external inputs.
+"""
+default_codegen_inputs(sys) = isscheduled(sys) ? inputs(sys) : unbound_inputs(sys)
+
 """
     bound_outputs(sys)
 
@@ -184,7 +201,7 @@ has_var(ex, x) = x ∈ Set(get_variables(ex))
 """
     (f_oop, f_ip), x_sym, p_sym, io_sys = generate_control_function(
             sys::System,
-            inputs                   = unbound_inputs(sys),
+            inputs                   = default_codegen_inputs(sys),
             disturbance_inputs       = disturbances(sys);
             known_disturbance_inputs = nothing,
             implicit_dae             = false,
@@ -222,7 +239,7 @@ f[1](x, inputs, p, t)
 ```
 """
 function generate_control_function(
-        sys::AbstractSystem, inputs = unbound_inputs(sys),
+        sys::AbstractSystem, inputs = default_codegen_inputs(sys),
         disturbance_inputs = disturbances(sys);
         known_disturbance_inputs = nothing,
         disturbance_argument = false,

lib/ModelingToolkitBase/src/problems/bvproblem.jl

@@ -36,7 +36,7 @@
         wrap_gfw = Val{true}, cse, checkbounds
     )
 
-    n_controls = length(unbound_inputs(sys))
+    n_controls = length(default_codegen_inputs(sys))
     f_prototype = n_controls > 0 ? zeros(eltype(u0), length(dvs) - n_controls) : nothing
     bcresid_prototype = zeros(eltype(u0), length(u0_idxs) + length(constraints(sys)))
 

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -1024,7 +1024,7 @@ function calculate_control_jacobian(
         sparse = false, simplify = false
     )
     rhs = [eq.rhs for eq in full_equations(sys)]
-    ctrls = unbound_inputs(sys)
+    ctrls = default_codegen_inputs(sys)
 
     if sparse
         jac = sparsejacobian(rhs, ctrls, simplify = simplify)

lib/ModelingToolkitBase/src/systems/optimal_control_interface.jl

@@ -120,7 +120,7 @@ is_explicit(tableau) = tableau isa DiffEqBase.ExplicitRKTableau
 
 @fallback_iip_specialize function SciMLBase.ODEInputFunction{iip, specialize}(
         sys::System;
-        inputs = inputs(sys),
+        inputs = default_codegen_inputs(sys),
         disturbance_inputs = disturbances(sys),
         u0 = nothing, tgrad = false,
         jac = false, controljac = false,
@@ -363,7 +363,7 @@ function process_DynamicOptProblem(
     add_user_constraints!(fullmodel, sys, tspan, pmap)
     add_initial_constraints!(fullmodel, u0, u0_idxs, model_tspan[1])
 
-    return prob_type(f, u0, tspan, p, fullmodel, kwargs...), pmap
+    return prob_type(f, u0, tspan, p, fullmodel; kwargs...), pmap
 end
 
 function generate_time_variable! end

lib/ModelingToolkitBase/test/input_output_handling.jl

@@ -391,6 +391,44 @@ eqs = [D(x) ~ u]
     @test isequal(ModelingToolkitBase.outputs(ss1), [x[1], x[2], x[3]])
 end
 
+@testset "default_codegen_inputs: declared inputs on compiled hierarchical systems" begin
+    # After mtkcompile flattens a hierarchical model, an effective input appears
+    # in equations together with variables from other namespaces and is therefore
+    # classified as bound, so `unbound_inputs` is empty. Input-aware codegen must
+    # fall back to the inputs declared to `mtkcompile` instead of silently
+    # generating input-free dynamics.
+    function TestActuator(; name)
+        @variables u(t) [input = true] o(t) [output = true]
+        return System([o ~ 2u], t; name)
+    end
+    function TestPlant(; name)
+        @variables y(t) i(t)
+        return System([D(y) ~ -y + i], t; name)
+    end
+    @named act = TestActuator()
+    @named plant = TestPlant()
+    @named hier = System([plant.i ~ act.o], t; systems = [act, plant])
+    hier = complete(hier)
+    ss = mtkcompile(hier; inputs = [hier.act.u])
+
+    @test isempty(unbound_inputs(ss))
+    @test length(ModelingToolkitBase.default_codegen_inputs(ss)) == 1
+    @test isequal(
+        collect(ModelingToolkitBase.default_codegen_inputs(ss)),
+        ModelingToolkitBase.inputs(ss)
+    )
+
+    # The `ODEInputFunction` default must pick up the declared input rather than
+    # generating dynamics with the input bound to its operating-point value.
+    f = ModelingToolkitBase.SciMLBase.ODEInputFunction(ss)
+    p = ModelingToolkitBase.get_p(ss, Dict(hier.act.u => 0.0, hier.plant.y => 1.0))
+    @test f([1.0], [0.0], p, 0.0) == [-1.0]
+    @test f([1.0], [5.0], p, 0.0) == [9.0]
+
+    # The control jacobian wrt the declared inputs is non-empty.
+    @test size(ModelingToolkitBase.calculate_control_jacobian(ss)) == (1, 1)
+end
+
 using ModelingToolkitStandardLibrary.Blocks
 
 if @isdefined(ModelingToolkit)

lib/ModelingToolkitBase/test/optimization/Project.toml

@@ -3,7 +3,6 @@ AmplNLWriter = "7c4d4715-977e-5154-bfe0-e096adeac482"
 CasADi = "c49709b8-5c63-11e9-2fb2-69db5844192f"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
-DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 Ipopt_jll = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"
@@ -14,7 +13,9 @@ Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb"
 OptimizationMOI = "fd9f6733-72f4-499f-8506-86b2bdd0dea1"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+OrdinaryDiffEqExplicitTableaus = "3278f1b1-0f5c-4cde-98e0-ba5eb00db955"
 OrdinaryDiffEqFIRK = "5960d6e9-dd7a-4743-88e7-cf307b64f125"
+OrdinaryDiffEqImplicitTableaus = "75f66a49-58fc-43e3-9173-2340726368f7"
 OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
@@ -32,5 +33,7 @@ ModelingToolkitBase = {path = "../.."}
 [compat]
 CasADi = "1.0.7"
 DataInterpolations = "8.8"
+OrdinaryDiffEqExplicitTableaus = "2"
+OrdinaryDiffEqImplicitTableaus = "2"
 SafeTestsets = "0.1, 1"
 SciMLTesting = "1"