add root-locus methods for matrix gains (#1012)

* add root-locus methods for matrix gains

* rm comment

* add kwargs...

Author: baggepinnen

src/root_locus.jl

@@ -3,7 +3,7 @@ import ControlSystemsBase.RootLocusResult
 @userplot Rlocusplot
 
 
-function getpoles(G, K::Number)
+function getpoles(G, K::Number; kwargs...)
     issiso(G) || error("root locus only supports SISO systems")
     G isa TransferFunction || (G = tf(G))
     P = numpoly(G)[]
@@ -69,17 +69,115 @@ function getpoles(G, K::AbstractVector{T}) where {T<:Number}
     copy(poleout'), K
 end
 
+"""
+    getpoles(sys::StateSpace, K::AbstractMatrix; tol = 1e-2, initial_stepsize = 1e-3, output=false)
+
+Compute the poles of the closed-loop system defined by `sys` with feedback gains `γ*K` where `γ` is a scalar that ranges from 0 to 1.
+
+If `output = true`, `K` is assumed to be an output feedback matrix of dim `(nu, ny)`
+"""
+function getpoles(sys::StateSpace, K_matrix::AbstractMatrix; tol = 1e-2, initial_stepsize = 1e-3, output=false)
+    (; A, B, C) = sys
+    nx = size(A, 1) # State dimension
+    ny = size(C, 1) # Output dimension
+    tol = tol*nx # Scale tolerance with state dimension
+    # Check for compatibility of K_matrix dimensions with B
+    if size(K_matrix, 2) != (output ? ny : nx)
+        error("The number of columns in K_matrix ($(size(K_matrix, 2))) must match the state dimension ($(nx)) or output dimension ($(ny)) depending on whether output feedback is used.")
+    end
+    if size(K_matrix, 1) != size(B, 2)
+        error("The number of rows in K_matrix ($(size(K_matrix, 1))) must match the number of inputs (columns of B, which is $(size(B, 2))).")
+    end
+
+    if output
+        # We bake C into K here to avoid repeated multiplications below
+        K_matrix = K_matrix * C
+    end
+
+    poleout_list = Vector{Vector{ComplexF64}}() # To store pole sets at each accepted step
+    k_scalars_collected = Float64[] # To store accepted k_scalar values
+
+    prevpoles = ComplexF64[] # Initialize prevpoles for the first iteration
+
+    stepsize = initial_stepsize
+    k_scalar = 0.0
+
+    # Initial poles at k_scalar = 0.0
+    A_cl_initial = A - 0.0 * B * K_matrix
+    initial_poles = eigvals(A_cl_initial)
+    push!(poleout_list, initial_poles)
+    push!(k_scalars_collected, 0.0)
+    prevpoles = initial_poles # Set prevpoles for the first actual step
+
+    while k_scalar < 1.0
+        # Propose a new k_scalar value
+        next_k_scalar = min(1.0, k_scalar + stepsize)
+
+        # Calculate poles for the proposed next_k_scalar
+        A_cl_proposed = A - next_k_scalar * B * K_matrix
+        current_poles_proposed = eigvals(A_cl_proposed)
+
+        # Calculate cost using Hungarian algorithm
+        D = zeros(nx, nx)
+        for r in 1:nx
+            for c in 1:nx
+                D[r, c] = abs(current_poles_proposed[r] - prevpoles[c])
+            end
+        end
+        assignment, cost = Hungarian.hungarian(D)
+
+        # Adaptive step size logic
+        if cost > 2 * tol # Cost is too high, reject step and reduce stepsize
+            stepsize /= 2.0
+            # Ensure stepsize doesn't become too small
+            if stepsize < 100eps()
+                @warn "Step size became extremely small, potentially stuck. Breaking loop."
+                break
+            end
+            # Do not update k_scalar, try again with smaller stepsize
+        else # Step is acceptable or too small
+            # Sort poles using the assignment from Hungarian algorithm
+            temppoles = zeros(ComplexF64, nx)
+            for j = 1:nx
+                temppoles[assignment[j]] = current_poles_proposed[j]
+            end
+            current_poles_sorted = temppoles
+
+            # Accept the step
+            push!(poleout_list, current_poles_sorted)
+            push!(k_scalars_collected, next_k_scalar)
+            prevpoles = current_poles_sorted # Update prevpoles for the next iteration
+            k_scalar = next_k_scalar # Advance k_scalar
+
+            if cost < tol # Cost is too low, increase stepsize
+                stepsize *= 1.1
+            end
+            # Cap stepsize to prevent overshooting 1.0 significantly in a single step
+            stepsize = min(stepsize, 1e-1)
+        end
+
+        # Break if k_scalar has reached or exceeded 1.0
+        if k_scalar >= 1.0
+            break
+        end
+    end
+
+    return hcat(poleout_list...)' |> copy, k_scalars_collected .* Ref(K_matrix) # Return transposed pole matrix and k_values
+end
+
 
 """
     roots, Z, K = rlocus(P::LTISystem, K = 500)
 
 Compute the root locus of the SISO LTISystem `P` with a negative feedback loop and feedback gains between 0 and `K`. `rlocus` will use an adaptive step-size algorithm to determine the values of the feedback gains used to generate the plot.
 
 `roots` is a complex matrix containing the poles trajectories of the closed-loop `1+k⋅G(s)` as a function of `k`, `Z` contains the zeros of the open-loop system `G(s)` and `K` the values of the feedback gain.
+
+If `K` is a matrix and `P` a `StateSpace` system, the poles are computed as `K` ranges from `0*K` to `1*K`. In this case, `K` is assumed to be a state-feedback matrix of dimension `(nu, nx)`. To compute the poles for output feedback, use, pass `output = true` and `K` of dimension `(nu, ny)`.
 """
-function rlocus(P, K)
+function rlocus(P, K; kwargs...)
     Z = tzeros(P)
-    roots, K = getpoles(P,K)
+    roots, K = getpoles(P, K; kwargs...)
     ControlSystemsBase.RootLocusResult(roots, Z, K, P)
 end
 
@@ -89,6 +187,7 @@ rlocus(P; K=500) = rlocus(P, K)
 # This will be called on plot(rlocus(sys, args...))
 @recipe function rootlocusresultplot(r::RootLocusResult)
     roots, Z, K = r
+    array_K = eltype(K) <: AbstractArray
     redata = real.(roots)
     imdata = imag.(roots)
     all_redata = [vec(redata); real.(Z)]
@@ -103,7 +202,9 @@ rlocus(P; K=500) = rlocus(P, K)
     form(k, p) = Printf.@sprintf("%.4f", k) * "  pole=" * Printf.@sprintf("%.3f%+.3fim", real(p), imag(p))
     @series begin
         legend --> false
-        hover := "K=" .* form.(K,roots)
+        if !array_K
+            hover := "K=" .* form.(K,roots)
+        end
         label := ""
         redata, imdata
     end
@@ -121,6 +222,15 @@ rlocus(P; K=500) = rlocus(P, K)
         label --> "Open-loop poles"
         redata[1,:], imdata[1,:]
     end
+    if array_K
+        @series begin
+            seriestype := :scatter
+            markershape --> :diamond
+            markersize --> 10
+            label --> "Closed-loop poles"
+            redata[end,:], imdata[end,:]
+        end
+    end
 end
 
 
@@ -129,5 +239,10 @@ end
 
 Plot the root locus of the SISO LTISystem `P` as computed by `rlocus`.
 """
-@recipe rlocusplot(::Type{Rlocusplot}, p::Rlocusplot; K=500) =
-    rlocus(p.args[1]; K=K)
+@recipe function rlocusplot(::Type{Rlocusplot}, p::Rlocusplot; K=500, output=false)
+    if length(p.args) >= 2
+        rlocus(p.args[1], p.args[2]; output)
+    else
+        rlocus(p.args[1]; K=K)
+    end
+end

test/runtests_controlsystems.jl

@@ -20,5 +20,6 @@ using ControlSystems
     @testset "rootlocus" begin
         @info "Testing rootlocus"
         include("test_rootlocus.jl")
+        include("test_root_locus_matrix.jl")
     end
 end

test/test_root_locus_matrix.jl

@@ -0,0 +1,102 @@
+using ControlSystems
+using ControlSystems: getpoles
+using Test
+using Plots
+
+@testset "Root Locus with Matrix Feedback" begin
+    
+    @testset "State Feedback Matrix" begin
+        # Test with double mass model
+        P = DemoSystems.double_mass_model()
+        
+        # Design a simple state feedback matrix
+        K_state = [1.0 0.5 0.2 0.1]  # nu x nx matrix (1x4)
+        
+        # Test that getpoles works with matrix feedback
+        roots, K_values = getpoles(P, K_state)
+        
+        # Basic sanity checks
+        @test size(roots, 2) == size(P.A, 1)  # Should have nx poles
+        @test length(K_values) == size(roots, 1)  # K_values should match number of steps
+        @test eltype(K_values) <: AbstractMatrix  # K_values should be matrices
+        
+        # Test that initial K_value is zero matrix
+        @test all(K_values[1] .≈ 0.0)
+        
+        # Test that final K_value is the input matrix
+        @test K_values[end] ≈ K_state
+        
+        # Test poles are continuous (no sudden jumps)
+        # The poles should change smoothly
+        pole_diffs = diff(roots, dims=1)
+        max_diff = maximum(abs, pole_diffs)
+        @test max_diff < 1.0  # Poles shouldn't jump too much between steps
+        
+        # Test that rlocus works with matrix K
+        result = rlocus(P, K_state)
+        @test result isa ControlSystemsBase.RootLocusResult
+        @test size(result.roots) == size(roots)
+        
+        # Test plotting (should not error)
+        plt = plot(result)
+        @test plt isa Plots.Plot
+    end
+    
+    @testset "Output Feedback Matrix" begin
+        P = DemoSystems.double_mass_model(outputs=1:2)
+        
+        # Design output feedback matrix
+        # P has 2 outputs, 1 input, so K should be 1x2
+        K_output = [1.0 0.5]  # nu x ny matrix (1x2)
+        
+        # Test output feedback
+        roots, K_values = getpoles(P, K_output; output=true)
+        
+        # Basic checks
+        @test size(roots, 2) == size(P.A, 1)  # Should have nx poles
+        @test length(K_values) == size(roots, 1)
+        @test eltype(K_values) <: AbstractMatrix
+        
+        # Test rlocus with output feedback
+        result = rlocus(P, K_output; output=true)
+        @test result isa ControlSystemsBase.RootLocusResult
+        
+        # Test plotting
+        plt = plot(result)
+        @test plt isa Plots.Plot
+    end
+    
+    @testset "Error Handling" begin
+        P = DemoSystems.double_mass_model()
+        
+        # Test wrong dimensions for state feedback
+        K_wrong = [1.0 0.5 0.2]  # Wrong size (1x3 instead of 1x4)
+        @test_throws ErrorException getpoles(P, K_wrong)
+        
+        # Test wrong dimensions for output feedback
+        K_wrong_output = [1.0 0.5 0.2]  # Wrong size (1x3 instead of 1x2)
+        @test_throws ErrorException getpoles(P, K_wrong_output; output=true)
+        
+        # Test wrong number of rows
+        K_wrong_rows = [1.0 0.5 0.2 0.1; 2.0 1.0 0.4 0.2]  # 2x4 instead of 1x4
+        @test_throws ErrorException getpoles(P, K_wrong_rows)
+    end
+    
+
+    @testset "Adaptive Step Size" begin
+        P = DemoSystems.double_mass_model()
+        K_state = [1.0 0.5 0.2 0.1]
+        
+        # Test with different tolerances
+        roots_loose, K_loose = getpoles(P, K_state; tol=1e-1)
+        roots_tight, K_tight = getpoles(P, K_state; tol=1e-4)
+        
+        # Tighter tolerance should result in more steps
+        @test length(K_tight) >= length(K_loose)
+        
+        # Both should start and end at the same points
+        @test roots_loose[1, :] ≈ roots_tight[1, :]  # Same initial poles
+        @test roots_loose[end, :] ≈ roots_tight[end, :] atol=1e-3  # Similar final poles
+    end
+    
+end