Merge pull request #2100 from SciML/fb/frame_connector

RFC: handle multibody frames in `connection2set!`

Author: baggepinnen

src/systems/connectors.jl

@@ -181,14 +181,42 @@ end
     error("Different types of connectors are in one connection statement: <$(map(nameof, ss))>")
 end
 
+"Return true if the system is a 3D multibody frame, otherwise return false."
+function isframe(sys)
+    sys.metadata isa Dict || return false
+    get(sys.metadata, :frame, false)
+end
+
+"Return orienation object of a multibody frame."
+function ori(sys)
+    if sys.metadata isa Dict && (O = get(sys.metadata, :orientation, nothing)) !== nothing
+        return O
+    else
+        error("System $(sys.name) does not have an orientation object.")
+    end
+end
+
 function connection2set!(connectionsets, namespace, ss, isouter)
     nn = map(nameof, ss)
-    sts1 = Set(states(first(ss)))
+    s1 = first(ss)
+    sts1v = states(s1)
+    if isframe(s1) # Multibody
+        O = ori(s1)
+        orientation_vars = Symbolics.unwrap.(collect(vec(O.R)))
+        sts1v = [sts1v; orientation_vars]
+    end
+    sts1 = Set(sts1v)
     T = ConnectionElement
-    csets = [T[] for _ in 1:length(sts1)]
+    num_statevars = length(sts1)
+    csets = [T[] for _ in 1:num_statevars] # Add 9 orientation variables if connection is between multibody frames
     for (i, s) in enumerate(ss)
         sts = states(s)
-        i != 1 && ((length(sts1) == length(sts) && all(Base.Fix2(in, sts1), sts)) ||
+        if isframe(s) # Multibody
+            O = ori(s)
+            orientation_vars = Symbolics.unwrap.(vec(O.R))
+            sts = [sts; orientation_vars]
+        end
+        i != 1 && ((num_statevars == length(sts) && all(Base.Fix2(in, sts1), sts)) ||
          connection_error(ss))
         io = isouter(s)
         for (j, v) in enumerate(sts)