Merge pull request #3748 from AayushSabharwal/as/loop-opening-defs

feat: better handle loop opening defaults

Author: ChrisRackauckas

src/linearization.jl

@@ -45,6 +45,7 @@ function linearization_function(sys::AbstractSystem, inputs,
         warn_initialize_determined = true,
         guesses = Dict(),
         warn_empty_op = true,
+        t = 0.0,
         kwargs...)
     op = Dict(op)
     if isempty(op) && warn_empty_op
@@ -73,7 +74,7 @@ function linearization_function(sys::AbstractSystem, inputs,
     end
 
     prob = ODEProblem{true, SciMLBase.FullSpecialize}(
-        sys, merge(op, anydict(p)), (nothing, nothing); allow_incomplete = true,
+        sys, merge(op, anydict(p)), (t, t); allow_incomplete = true,
         algebraic_only = true, guesses)
     u0 = state_values(prob)
 
@@ -753,7 +754,7 @@ function linearize(sys, inputs, outputs; op = Dict(), t = 0.0,
         inputs,
         outputs;
         zero_dummy_der,
-        op,
+        op, t,
         kwargs...)
     mats, extras = linearize(ssys, lin_fun; op, t, allow_input_derivatives)
     mats, ssys, extras

src/systems/analysis_points.jl

@@ -486,14 +486,19 @@ struct Break <: AnalysisPointTransformation
     Whether to add a new input variable connected to all the outputs of `ap`.
     """
     add_input::Bool
+    """
+    Whether the default of the added input variable should be the input of `ap`. Only
+    applicable if `add_input == true`.
+    """
+    default_outputs_to_input::Bool
 end
 
 """
     $(TYPEDSIGNATURES)
 
-`Break` the given analysis point `ap` without adding an input.
+`Break` the given analysis point `ap`.
 """
-Break(ap::AnalysisPoint) = Break(ap, false)
+Break(ap::AnalysisPoint, add_input::Bool = false) = Break(ap, add_input, false)
 
 function apply_transformation(tf::Break, sys::AbstractSystem)
     modify_nested_subsystem(sys, tf.ap) do breaksys
@@ -517,7 +522,11 @@ function apply_transformation(tf::Break, sys::AbstractSystem)
             push!(breaksys_eqs, ap_var(outsys) ~ new_var)
         end
         defs = copy(get_defaults(breaksys))
-        defs[new_var] = new_def
+        defs[new_var] = if tf.default_outputs_to_input
+            ap_ivar
+        else
+            new_def
+        end
         @set! breaksys.defaults = defs
         unks = copy(get_unknowns(breaksys))
         push!(unks, new_var)
@@ -803,7 +812,7 @@ Given a list of analysis points, break the connection for each and set the outpu
 """
 function handle_loop_openings(sys::AbstractSystem, aps)
     for ap in canonicalize_ap(sys, aps)
-        sys, (outvar,) = apply_transformation(Break(ap, true), sys)
+        sys, (outvar,) = apply_transformation(Break(ap, true, true), sys)
         if Symbolics.isarraysymbolic(outvar)
             push!(get_eqs(sys), outvar ~ zeros(size(outvar)))
         else
@@ -815,7 +824,7 @@ end
 
 const DOC_LOOP_OPENINGS = """
     - `loop_openings`: A list of analysis points whose connections should be removed and
-      the outputs set to zero as a part of the linear analysis.
+      the outputs set to the input as a part of the linear analysis.
 """
 
 const DOC_SYS_MODIFIER = """
@@ -952,7 +961,7 @@ function linearization_ap_transform(sys,
     for input in inputs
         if nameof(input) in loop_openings
             delete!(loop_openings, nameof(input))
-            sys, (input_var,) = apply_transformation(Break(input, true), sys)
+            sys, (input_var,) = apply_transformation(Break(input, true, true), sys)
         else
             sys, (input_var,) = apply_transformation(PerturbOutput(input), sys)
         end

src/systems/connectiongraph.jl

@@ -414,10 +414,6 @@ function remove_negative_connections!(
                 # if there is any other variable, start removing
                 push!(idxs_to_rm[edge_j], var_j)
             end
-            if should_rm
-                # if there was any other variable, also remove `input_j`
-                push!(idxs_to_rm, input_j)
-            end
         end
     end
 
@@ -459,6 +455,7 @@ function connectionsets(graph::ConnectionGraph)
     disjoint_sets = IntDisjointSets(length(invmap))
     for edge_i in 𝑠vertices(bigraph)
         hyperedge = 𝑠neighbors(bigraph, edge_i)
+        isempty(hyperedge) && continue
         root, rest = Iterators.peel(hyperedge)
         for vert in rest
             union!(disjoint_sets, root, vert)

src/systems/connectors.jl

@@ -468,6 +468,16 @@ function _generate_connectionsets!(connection_state::AbstractConnectionState,
                 [namespace; var_ns], length(var_ns) == 1 || isouter(var_ns[1]), type)
         end
         add_connection_edge!(connection_state, hyperedge)
+
+        # Removed analysis points generate causal connections in the negative graph. These
+        # should also remove `Equality` connections involving the same variables, so also
+        # add an `Equality` variant of the edge.
+        if connection_state isa NegativeConnectionState
+            hyperedge = map(hyperedge) do cvert
+                ConnectionVertex(cvert.name, cvert.isouter, Equality)
+            end
+            add_connection_edge!(connection_state, hyperedge)
+        end
     end
 end