Adding controls kwarg and controls jacobian to AbstractODESystem

Author: cadojo

src/ModelingToolkit.jl

@@ -161,11 +161,13 @@ export Differential, expand_derivatives, @derivatives
 export Equation, ConstrainedEquation
 export Term, Sym
 export SymScope, LocalScope, ParentScope, GlobalScope
-export independent_variable, states, parameters, equations, controls, observed, structure
+export independent_variable, states, parameters, equations, controls, observed, structure, defaults
+export ssmodel, linearize
 export structural_simplify
 export DiscreteSystem, DiscreteProblem
 
 export calculate_jacobian, generate_jacobian, generate_function
+export calculate_control_jacobian
 export calculate_tgrad, generate_tgrad
 export calculate_gradient, generate_gradient
 export calculate_factorized_W, generate_factorized_W

src/systems/abstractsystem.jl

@@ -34,6 +34,18 @@ call will be cached in the system object.
 """
 function calculate_jacobian end
 
+"""
+```julia
+calculate_control_jacobian(sys::AbstractSystem)
+```
+
+Calculate the jacobian matrix of a system with respect to the system's controls.
+
+Returns a matrix of [`Num`](@ref) instances. The result from the first
+call will be cached in the system object.
+"""
+function calculate_control_jacobian end
+
 """
 ```julia
 calculate_factorized_W(sys::AbstractSystem)
@@ -140,10 +152,12 @@ for prop in [
              :iv
              :states
              :ps
+             :ctrl
              :defaults
              :observed
              :tgrad
              :jac
+             :ctrl_jac
              :Wfact
              :Wfact_t
              :systems
@@ -346,11 +360,17 @@ function states(sys::AbstractSystem)
            sts :
            [sts;reduce(vcat,namespace_variables.(systems))])
 end
+
 function parameters(sys::AbstractSystem)
     ps = get_ps(sys)
     systems = get_systems(sys)
     isempty(systems) ? ps : [ps;reduce(vcat,namespace_parameters.(systems))]
 end
+
+function controls(sys::AbstractSystem)
+    get_ctrl(sys)
+end
+
 function observed(sys::AbstractSystem)
     iv = independent_variable(sys)
     obs = get_observed(sys)

src/systems/diffeqs/abstractodesystem.jl

@@ -38,6 +38,28 @@ function calculate_jacobian(sys::AbstractODESystem;
     return jac
 end
 
+function calculate_control_jacobian(sys::AbstractODESystem;
+                                    sparse=false, simplify=false)
+    cache = get_ctrl_jac(sys)[]
+    if cache isa Tuple && cache[2] == (sparse, simplify)
+        return cache[1]
+    end
+
+    rhs = [eq.rhs for eq ∈ equations(sys)]
+
+    iv = get_iv(sys)
+    ctrls = controls(sys)
+
+    if sparse
+        jac = sparsejacobian(rhs, ctrls, simplify=simplify)
+    else
+        jac = jacobian(rhs, ctrls, simplify=simplify)
+    end
+
+    get_ctrl_jac(sys)[] = jac, (sparse, simplify) # cache Jacobian
+    return jac
+end
+
 function generate_tgrad(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);
                         simplify=false, kwargs...)
     tgrad = calculate_tgrad(sys,simplify=simplify)
@@ -50,6 +72,25 @@ function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = param
     return build_function(jac, dvs, ps, get_iv(sys); kwargs...)
 end
 
+"""
+```julia
+generate_linearization(sys::ODESystem, dvs = states(sys), ps = parameters(sys), ctrls = controls(sys), point = defaults(sys), expression = Val{true}; sparse = false, kwargs...)
+```
+
+Generates a function for the linearized state space model of the system. Extra arguments
+control the arguments to the internal [`build_function`](@ref) call.
+"""
+function generate_linearization(sys::AbstractSystem, dvs = states(sys), ps = parameters(sys), ctrls = controls(sys);
+                                simplify=false, sparse=false, kwargs...) 
+    ops = map(eq -> eq.rhs, equations(sys)) 
+    
+    jac = calculate_jacobian(sys;simplify=simplify,sparse=sparse)
+
+    A = @views J[1:length(states(sys)), 1:length(states(sys))]
+    B = @views J[1:length(states(sys)), length(states(sys))+1:end]
+
+    return A, B
+end
 @noinline function throw_invalid_derivative(dervar, eq)
     msg = "The derivative variable must be isolated to the left-hand " *
     "side of the equation like `$dervar ~ ...`.\n Got $eq."

src/systems/diffeqs/odesystem.jl

@@ -31,6 +31,7 @@ struct ODESystem <: AbstractODESystem
     states::Vector
     """Parameter variables. Must not contain the independent variable."""
     ps::Vector
+    ctrl::Vector
     observed::Vector{Equation}
     """
     Time-derivative matrix. Note: this field will not be defined until
@@ -43,6 +44,11 @@ struct ODESystem <: AbstractODESystem
     """
     jac::RefValue{Any}
     """
+    Control Jacobian matrix. Note: this field will not be defined until
+    [`calculate_control_jacobian`](@ref) is called on the system.
+    """
+    ctrl_jac::RefValue{Any}
+    """
     `Wfact` matrix. Note: this field will not be defined until
     [`generate_factorized_W`](@ref) is called on the system.
     """
@@ -84,6 +90,7 @@ end
 
 function ODESystem(
                    deqs::AbstractVector{<:Equation}, iv, dvs, ps;
+                   controls  = Num[],
                    observed = Num[],
                    systems = ODESystem[],
                    name=gensym(:ODESystem),
@@ -92,9 +99,13 @@ function ODESystem(
                    defaults=_merge(Dict(default_u0), Dict(default_p)),
                    connection_type=nothing,
                   )
+    
+    @assert all(control -> any(isequal.(control, ps)), controls) "All controls must also be parameters."
+
     iv′ = value(scalarize(iv))
     dvs′ = value.(scalarize(dvs))
     ps′ = value.(scalarize(ps))
+    ctrl′ = value.(scalarize(controls))
 
     if !(isempty(default_u0) && isempty(default_p))
         Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.", :ODESystem, force=true)
@@ -110,7 +121,7 @@ function ODESystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    ODESystem(deqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type)
+    ODESystem(deqs, iv′, dvs′, ps′, ctrl′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type)
 end
 
 vars(x::Sym) = Set([x])
@@ -349,4 +360,4 @@ function convert_system(::Type{<:ODESystem}, sys, t; name=nameof(sys))
     neweqs = map(sub, equations(sys))
     defs = Dict(sub(k) => sub(v) for (k, v) in defaults(sys))
     return ODESystem(neweqs, t, newsts, parameters(sys); defaults=defs, name=name)
-end
+end