From 6b60b1aaf428f3d6adb1ff68782e6ad7579dde2a Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Tue, 2 Sep 2025 14:15:56 +0200
Subject: [PATCH 1/6] add input-affine equation factorization

---
 src/ModelingToolkit.jl    |   2 +
 src/input_affine_form.jl  |  54 +++++++++++++++
 test/input_affine_form.jl | 135 ++++++++++++++++++++++++++++++++++++++
 test/runtests.jl          |   1 +
 4 files changed, 192 insertions(+)
 create mode 100644 src/input_affine_form.jl
 create mode 100644 test/input_affine_form.jl

diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 8b3c4084c7..36b20aca4c 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -224,6 +224,7 @@ include("structural_transformation/StructuralTransformations.jl")
 
 @reexport using .StructuralTransformations
 include("inputoutput.jl")
+include("input_affine_form.jl")
 
 include("adjoints.jl")
 include("deprecations.jl")
@@ -259,6 +260,7 @@ export JumpProblem
 export alias_elimination, flatten
 export connect, domain_connect, @connector, Connection, AnalysisPoint, Flow, Stream,
        instream
+export input_affine_form
 export initial_state, transition, activeState, entry, ticksInState, timeInState
 export @component, @mtkmodel, @mtkcompile, @mtkbuild
 export isinput, isoutput, getbounds, hasbounds, getguess, hasguess, isdisturbance,
diff --git a/src/input_affine_form.jl b/src/input_affine_form.jl
new file mode 100644
index 0000000000..b0f7ac5062
--- /dev/null
+++ b/src/input_affine_form.jl
@@ -0,0 +1,54 @@
+"""
+    input_affine_form(eqs, inputs)
+
+Extract control-affine (input-affine) form from symbolic equations.
+
+Given a system of equations of the form `ẋ = F(x, u)` where `x` is the state
+and `u` are the inputs, this function extracts the control-affine representation:
+`ẋ = f(x) + g(x)*u`
+
+where:
+- `f(x)` is the drift term (dynamics without inputs)
+- `g(x)` is the input matrix
+
+# Arguments
+- `eqs`: Vector of symbolic equations representing the system dynamics
+- `inputs`: Vector of input/control variables
+
+# Returns
+- `f`: The drift vector f(x)
+- `g`: The input matrix g(x)
+
+# Example
+```julia
+using ModelingToolkit, Symbolics
+
+@variables x1 x2 u1 u2
+state  = [x1, x2]
+inputs = [u1, u2]
+
+# Define system equations: ẋ = F(x, u)
+eqs = [
+    -x1 + 2*x2 + u1,
+    x1*x2 - x2 + u1 + 2*u2
+]
+
+# Extract control-affine form
+f, g = input_affine_form(eqs, inputs)
+```
+
+# Notes
+The function assumes that the equations are affine in the inputs. If the equations
+are nonlinear in the inputs, the result may not be meaningful.
+"""
+function input_affine_form(eqs, inputs)
+    # Extract the input matrix g(x) by taking coefficients of each input
+    g = [Symbolics.coeff(Symbolics.simplify(eq, expand=true), u) for eq in eqs, u in inputs]
+    g = Symbolics.simplify.(g, expand=true)
+    
+    # Extract the drift term f(x) by substituting inputs = 0
+    f = Symbolics.substitute.(eqs, Ref(Dict(inputs .=> 0)))
+    f = Symbolics.simplify.(f, expand=true)
+    
+    return f, g
+end
\ No newline at end of file
diff --git a/test/input_affine_form.jl b/test/input_affine_form.jl
new file mode 100644
index 0000000000..fd01e22b85
--- /dev/null
+++ b/test/input_affine_form.jl
@@ -0,0 +1,135 @@
+using ModelingToolkit, Test
+using Symbolics
+using StaticArrays
+
+@testset "input_affine_form" begin
+    # Test with simple linear system
+    @testset "Simple linear system" begin
+        @variables x1 x2 u1 u2
+        state = [x1, x2]
+        inputs = [u1, u2]
+        
+        eqs = [
+            -x1 + 2*x2 + u1,
+            x1*x2 - x2 + u1 + 2*u2
+        ]
+        
+        f, g = input_affine_form(eqs, inputs)
+        
+        # Verify reconstruction
+        eqs_reconstructed = f + g * inputs
+        @test isequal(Symbolics.simplify.(eqs_reconstructed), Symbolics.simplify.(eqs))
+        
+        # Check dimensions
+        @test length(f) == length(eqs)
+        @test size(g) == (length(eqs), length(inputs))
+    end
+
+    # Test with Segway dynamics example
+    @testset "Segway dynamics" begin
+        # Segway parameters
+        grav = 9.81
+        R = 0.195
+        M = 2 * 2.485
+        Jc = 2 * 0.0559
+        L = 0.169
+        m = 44.798
+        Jg = 3.836
+        m0 = 52.710
+        J0 = 5.108
+        Km = 2 * 1.262
+        bt = 2 * 1.225
+
+        # Dynamics of Segway in Euler-Lagrange form
+        D(q) = [m0 m*L*cos(q[2]); m*L*cos(q[2]) J0]
+        function H(q, q̇)
+            return SA[
+                -m * L * sin(q[2]) * q̇[2] + bt * (q̇[1] - R * q̇[2]) / R,
+                -m * grav * L * sin(q[2]) - bt * (q̇[1] - R * q̇[2]),
+            ]
+        end
+        B(q) = SA[Km / R, -Km]
+
+        # Convert to control affine form
+        function f_seg(x)
+            q, q̇ = x[SA[1,2]], x[SA[3,4]]
+            return [q̇; -D(q) \ H(q, q̇)]
+        end
+        function g_seg(x)
+            q, q̇ = x[SA[1,2]], x[SA[3,4]]
+            return [SA[0,0]; D(q) \ B(q)]
+        end
+
+        # Trace dynamics symbolically
+        @variables q1 q2 qd1 qd2 u
+        x = [q1; q2; qd1; qd2]
+        inputs = [u]
+        eqs = f_seg(x) + g_seg(x) * u
+
+        # Extract control-affine form
+        fe, ge = input_affine_form(eqs, inputs)
+        
+        # Test reconstruction
+        eqs2 = fe + ge * inputs
+        diff = Symbolics.simplify.(eqs2 - eqs, expand=true)
+        
+        # The difference should be zero or very close to zero symbolically
+        # We test numerically since symbolic simplification might not be perfect
+        f2, _ = build_function(fe, x, expression=false)
+        g2, _ = build_function(ge, x, expression=false)
+        
+        for i = 1:10
+            x_val = rand(length(x))
+            @test f2(x_val) ≈ f_seg(x_val) rtol=1e-10
+            @test g2(x_val) ≈ g_seg(x_val) rtol=1e-10
+        end
+    end
+
+    # Test with multiple inputs
+    @testset "Multiple inputs" begin
+        @variables x1 x2 x3 u1 u2
+        state = [x1, x2, x3]
+        inputs = [u1, u2]
+        
+        eqs = [
+            x2,
+            x3,
+            -x1 - 2*x2 - x3 + u1 + 3*u2
+        ]
+        
+        f, g = input_affine_form(eqs, inputs)
+        
+        # Expected results
+        f_expected = [x2, x3, -x1 - 2*x2 - x3]
+        g_expected = [0 0; 0 0; 1 3]
+        
+        @test isequal(Symbolics.simplify.(f), Symbolics.simplify.(f_expected))
+        
+        # Test g matrix elements
+        for i in 1:size(g, 1), j in 1:size(g, 2)
+            @test isequal(Symbolics.simplify(g[i,j]), g_expected[i,j])
+        end
+    end
+
+    # Test with nonlinear state dynamics
+    @testset "Nonlinear state dynamics" begin
+        @variables x1 x2 u
+        state = [x1, x2]
+        inputs = [u]
+        
+        eqs = [
+            x2,
+            -sin(x1) - x2 + u
+        ]
+        
+        f, g = input_affine_form(eqs, inputs)
+        
+        # Expected results
+        f_expected = [x2, -sin(x1) - x2]
+        g_expected = reshape([0, 1], 2, 1)
+        
+        @test isequal(Symbolics.simplify.(f), Symbolics.simplify.(f_expected))
+        @test isequal(g, g_expected)
+    end
+
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 47230c9539..71b3318b91 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -34,6 +34,7 @@ end
             @safetestset "Direct Usage Test" include("direct.jl")
             @safetestset "System Linearity Test" include("linearity.jl")
             @safetestset "Input Output Test" include("input_output_handling.jl")
+            @safetestset "Input Affine Form Test" include("input_affine_form.jl")
             @safetestset "Clock Test" include("clock.jl")
             @safetestset "ODESystem Test" include("odesystem.jl")
             @safetestset "Dynamic Quantities Test" include("dq_units.jl")

From bb55cb574e5624aced704f711a59ed903084dc08 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Tue, 2 Sep 2025 14:22:33 +0200
Subject: [PATCH 2/6] format

---
 src/input_affine_form.jl  | 13 +++++-----
 test/input_affine_form.jl | 54 +++++++++++++++++++--------------------
 2 files changed, 34 insertions(+), 33 deletions(-)

diff --git a/src/input_affine_form.jl b/src/input_affine_form.jl
index b0f7ac5062..dde73d6949 100644
--- a/src/input_affine_form.jl
+++ b/src/input_affine_form.jl
@@ -43,12 +43,13 @@ are nonlinear in the inputs, the result may not be meaningful.
 """
 function input_affine_form(eqs, inputs)
     # Extract the input matrix g(x) by taking coefficients of each input
-    g = [Symbolics.coeff(Symbolics.simplify(eq, expand=true), u) for eq in eqs, u in inputs]
-    g = Symbolics.simplify.(g, expand=true)
-    
+    g = [Symbolics.coeff(Symbolics.simplify(eq, expand = true), u)
+         for eq in eqs, u in inputs]
+    g = Symbolics.simplify.(g, expand = true)
+
     # Extract the drift term f(x) by substituting inputs = 0
     f = Symbolics.substitute.(eqs, Ref(Dict(inputs .=> 0)))
-    f = Symbolics.simplify.(f, expand=true)
-    
+    f = Symbolics.simplify.(f, expand = true)
+
     return f, g
-end
\ No newline at end of file
+end
diff --git a/test/input_affine_form.jl b/test/input_affine_form.jl
index fd01e22b85..3234618350 100644
--- a/test/input_affine_form.jl
+++ b/test/input_affine_form.jl
@@ -8,18 +8,18 @@ using StaticArrays
         @variables x1 x2 u1 u2
         state = [x1, x2]
         inputs = [u1, u2]
-        
+
         eqs = [
             -x1 + 2*x2 + u1,
             x1*x2 - x2 + u1 + 2*u2
         ]
-        
+
         f, g = input_affine_form(eqs, inputs)
-        
+
         # Verify reconstruction
         eqs_reconstructed = f + g * inputs
         @test isequal(Symbolics.simplify.(eqs_reconstructed), Symbolics.simplify.(eqs))
-        
+
         # Check dimensions
         @test length(f) == length(eqs)
         @test size(g) == (length(eqs), length(inputs))
@@ -45,19 +45,19 @@ using StaticArrays
         function H(q, q̇)
             return SA[
                 -m * L * sin(q[2]) * q̇[2] + bt * (q̇[1] - R * q̇[2]) / R,
-                -m * grav * L * sin(q[2]) - bt * (q̇[1] - R * q̇[2]),
+                -m * grav * L * sin(q[2]) - bt * (q̇[1] - R * q̇[2])
             ]
         end
         B(q) = SA[Km / R, -Km]
 
         # Convert to control affine form
         function f_seg(x)
-            q, q̇ = x[SA[1,2]], x[SA[3,4]]
+            q, q̇ = x[SA[1, 2]], x[SA[3, 4]]
             return [q̇; -D(q) \ H(q, q̇)]
         end
         function g_seg(x)
-            q, q̇ = x[SA[1,2]], x[SA[3,4]]
-            return [SA[0,0]; D(q) \ B(q)]
+            q, q̇ = x[SA[1, 2]], x[SA[3, 4]]
+            return [SA[0, 0]; D(q) \ B(q)]
         end
 
         # Trace dynamics symbolically
@@ -68,17 +68,17 @@ using StaticArrays
 
         # Extract control-affine form
         fe, ge = input_affine_form(eqs, inputs)
-        
+
         # Test reconstruction
         eqs2 = fe + ge * inputs
-        diff = Symbolics.simplify.(eqs2 - eqs, expand=true)
-        
+        diff = Symbolics.simplify.(eqs2 - eqs, expand = true)
+
         # The difference should be zero or very close to zero symbolically
         # We test numerically since symbolic simplification might not be perfect
-        f2, _ = build_function(fe, x, expression=false)
-        g2, _ = build_function(ge, x, expression=false)
-        
-        for i = 1:10
+        f2, _ = build_function(fe, x, expression = false)
+        g2, _ = build_function(ge, x, expression = false)
+
+        for i in 1:10
             x_val = rand(length(x))
             @test f2(x_val) ≈ f_seg(x_val) rtol=1e-10
             @test g2(x_val) ≈ g_seg(x_val) rtol=1e-10
@@ -90,24 +90,25 @@ using StaticArrays
         @variables x1 x2 x3 u1 u2
         state = [x1, x2, x3]
         inputs = [u1, u2]
-        
+
         eqs = [
             x2,
             x3,
             -x1 - 2*x2 - x3 + u1 + 3*u2
         ]
-        
+
         f, g = input_affine_form(eqs, inputs)
-        
+
         # Expected results
         f_expected = [x2, x3, -x1 - 2*x2 - x3]
         g_expected = [0 0; 0 0; 1 3]
-        
+
         @test isequal(Symbolics.simplify.(f), Symbolics.simplify.(f_expected))
-        
+
         # Test g matrix elements
         for i in 1:size(g, 1), j in 1:size(g, 2)
-            @test isequal(Symbolics.simplify(g[i,j]), g_expected[i,j])
+
+            @test isequal(Symbolics.simplify(g[i, j]), g_expected[i, j])
         end
     end
 
@@ -116,20 +117,19 @@ using StaticArrays
         @variables x1 x2 u
         state = [x1, x2]
         inputs = [u]
-        
+
         eqs = [
             x2,
             -sin(x1) - x2 + u
         ]
-        
+
         f, g = input_affine_form(eqs, inputs)
-        
+
         # Expected results
         f_expected = [x2, -sin(x1) - x2]
         g_expected = reshape([0, 1], 2, 1)
-        
+
         @test isequal(Symbolics.simplify.(f), Symbolics.simplify.(f_expected))
         @test isequal(g, g_expected)
     end
-
-end
\ No newline at end of file
+end

From 457bf96eda13712ea49cde6412cf4e7167916a5c Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Wed, 3 Sep 2025 15:03:43 +0200
Subject: [PATCH 3/6] use `fast_substitute`

---
 src/input_affine_form.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/input_affine_form.jl b/src/input_affine_form.jl
index dde73d6949..bc09c8ef7d 100644
--- a/src/input_affine_form.jl
+++ b/src/input_affine_form.jl
@@ -48,7 +48,7 @@ function input_affine_form(eqs, inputs)
     g = Symbolics.simplify.(g, expand = true)
 
     # Extract the drift term f(x) by substituting inputs = 0
-    f = Symbolics.substitute.(eqs, Ref(Dict(inputs .=> 0)))
+    f = Symbolics.fast_substitute.(eqs, Ref(Dict(inputs .=> 0)))
     f = Symbolics.simplify.(f, expand = true)
 
     return f, g

From f7d06d5ce59231bb34fbf92ad57937aceb6eddad Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Wed, 3 Sep 2025 15:22:52 +0200
Subject: [PATCH 4/6] use linear_expansion

---
 src/input_affine_form.jl | 13 +++----------
 1 file changed, 3 insertions(+), 10 deletions(-)

diff --git a/src/input_affine_form.jl b/src/input_affine_form.jl
index bc09c8ef7d..caadb828e4 100644
--- a/src/input_affine_form.jl
+++ b/src/input_affine_form.jl
@@ -39,17 +39,10 @@ f, g = input_affine_form(eqs, inputs)
 
 # Notes
 The function assumes that the equations are affine in the inputs. If the equations
-are nonlinear in the inputs, the result may not be meaningful.
+are nonlinear in the inputs, an error is thrown.
 """
 function input_affine_form(eqs, inputs)
-    # Extract the input matrix g(x) by taking coefficients of each input
-    g = [Symbolics.coeff(Symbolics.simplify(eq, expand = true), u)
-         for eq in eqs, u in inputs]
-    g = Symbolics.simplify.(g, expand = true)
-
-    # Extract the drift term f(x) by substituting inputs = 0
-    f = Symbolics.fast_substitute.(eqs, Ref(Dict(inputs .=> 0)))
-    f = Symbolics.simplify.(f, expand = true)
-
+    g, f, flag = linear_expansion(eqs, inputs)
+    flag || error("The system is not affine in the inputs.")
     return f, g
 end

From 656aec24e969d4b3a9ebe59180b35b337608f6e2 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Thu, 4 Sep 2025 15:44:15 +0200
Subject: [PATCH 5/6] operate on equations

---
 src/input_affine_form.jl  |  6 +++---
 test/input_affine_form.jl | 20 +++++++++++---------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/input_affine_form.jl b/src/input_affine_form.jl
index caadb828e4..7680b5799f 100644
--- a/src/input_affine_form.jl
+++ b/src/input_affine_form.jl
@@ -29,8 +29,8 @@ inputs = [u1, u2]
 
 # Define system equations: ẋ = F(x, u)
 eqs = [
-    -x1 + 2*x2 + u1,
-    x1*x2 - x2 + u1 + 2*u2
+    D(x1) ~ -x1 + 2*x2 + u1,
+    D(x2) ~ x1*x2 - x2 + u1 + 2*u2
 ]
 
 # Extract control-affine form
@@ -42,7 +42,7 @@ The function assumes that the equations are affine in the inputs. If the equatio
 are nonlinear in the inputs, an error is thrown.
 """
 function input_affine_form(eqs, inputs)
-    g, f, flag = linear_expansion(eqs, inputs)
+    g, f, flag = Symbolics.linear_expansion(getfield.(eqs, :rhs), inputs)
     flag || error("The system is not affine in the inputs.")
     return f, g
 end
diff --git a/test/input_affine_form.jl b/test/input_affine_form.jl
index 3234618350..9533b14ad4 100644
--- a/test/input_affine_form.jl
+++ b/test/input_affine_form.jl
@@ -2,6 +2,8 @@ using ModelingToolkit, Test
 using Symbolics
 using StaticArrays
 
+rhs(eq) = simplify(eq.rhs)
+
 @testset "input_affine_form" begin
     # Test with simple linear system
     @testset "Simple linear system" begin
@@ -9,7 +11,7 @@ using StaticArrays
         state = [x1, x2]
         inputs = [u1, u2]
 
-        eqs = [
+        eqs = D.(state) .~ [
             -x1 + 2*x2 + u1,
             x1*x2 - x2 + u1 + 2*u2
         ]
@@ -18,7 +20,7 @@ using StaticArrays
 
         # Verify reconstruction
         eqs_reconstructed = f + g * inputs
-        @test isequal(Symbolics.simplify.(eqs_reconstructed), Symbolics.simplify.(eqs))
+        @test isequal(Symbolics.simplify.(eqs_reconstructed), rhs.(eqs))
 
         # Check dimensions
         @test length(f) == length(eqs)
@@ -41,7 +43,7 @@ using StaticArrays
         bt = 2 * 1.225
 
         # Dynamics of Segway in Euler-Lagrange form
-        D(q) = [m0 m*L*cos(q[2]); m*L*cos(q[2]) J0]
+        Dq(q) = [m0 m*L*cos(q[2]); m*L*cos(q[2]) J0]
         function H(q, q̇)
             return SA[
                 -m * L * sin(q[2]) * q̇[2] + bt * (q̇[1] - R * q̇[2]) / R,
@@ -53,25 +55,25 @@ using StaticArrays
         # Convert to control affine form
         function f_seg(x)
             q, q̇ = x[SA[1, 2]], x[SA[3, 4]]
-            return [q̇; -D(q) \ H(q, q̇)]
+            return [q̇; -Dq(q) \ H(q, q̇)]
         end
         function g_seg(x)
             q, q̇ = x[SA[1, 2]], x[SA[3, 4]]
-            return [SA[0, 0]; D(q) \ B(q)]
+            return [SA[0, 0]; Dq(q) \ B(q)]
         end
 
         # Trace dynamics symbolically
         @variables q1 q2 qd1 qd2 u
         x = [q1; q2; qd1; qd2]
         inputs = [u]
-        eqs = f_seg(x) + g_seg(x) * u
+        eqs = D.(x) .~ f_seg(x) + g_seg(x) * u
 
         # Extract control-affine form
         fe, ge = input_affine_form(eqs, inputs)
 
         # Test reconstruction
         eqs2 = fe + ge * inputs
-        diff = Symbolics.simplify.(eqs2 - eqs, expand = true)
+        diff = Symbolics.simplify.(eqs2 - rhs.(eqs), expand = true)
 
         # The difference should be zero or very close to zero symbolically
         # We test numerically since symbolic simplification might not be perfect
@@ -91,7 +93,7 @@ using StaticArrays
         state = [x1, x2, x3]
         inputs = [u1, u2]
 
-        eqs = [
+        eqs = D.(state) .~ [
             x2,
             x3,
             -x1 - 2*x2 - x3 + u1 + 3*u2
@@ -118,7 +120,7 @@ using StaticArrays
         state = [x1, x2]
         inputs = [u]
 
-        eqs = [
+        eqs = D.(state) .~ [
             x2,
             -sin(x1) - x2 + u
         ]

From 9cbca5f2bb44d7202532a7d0f8ed494fc7c5f9f9 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Thu, 4 Sep 2025 15:45:21 +0200
Subject: [PATCH 6/6] error on alg equations

---
 src/input_affine_form.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/input_affine_form.jl b/src/input_affine_form.jl
index 7680b5799f..0c446bc37b 100644
--- a/src/input_affine_form.jl
+++ b/src/input_affine_form.jl
@@ -42,6 +42,7 @@ The function assumes that the equations are affine in the inputs. If the equatio
 are nonlinear in the inputs, an error is thrown.
 """
 function input_affine_form(eqs, inputs)
+    any(is_alg_equation, eqs) && error("All equations must be differential equations.")
     g, f, flag = Symbolics.linear_expansion(getfield.(eqs, :rhs), inputs)
     flag || error("The system is not affine in the inputs.")
     return f, g