From d3b517a55ed30d4573ccad3004e6ba5d2efe4c83 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <cont-frb@ulund.org>
Date: Wed, 23 Sep 2020 07:48:17 +0200
Subject: [PATCH] Accept controller returning scalar

---
 src/timeresp.jl       | 5 +++--
 test/test_timeresp.jl | 7 ++++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/timeresp.jl b/src/timeresp.jl
index d3db3f2c6..b9fb2aaee 100644
--- a/src/timeresp.jl
+++ b/src/timeresp.jl
@@ -143,9 +143,10 @@ function lsim(sys::StateSpace, u::Function, t::AbstractVector;
         x0::VecOrMat=zeros(sys.nx), method::Symbol=:cont)
     ny, nu = size(sys)
     nx = sys.nx
+    u0 = u(x0,1)
     if length(x0) != nx
         error("size(x0) must match the number of states of sys")
-    elseif size(u(x0,1)) != (nu,) && size(u(x0,1)) != (nu,1)
+    elseif !(u0 isa Number && nu == 1) && (size(u0) != (nu,) && size(u0) != (nu,1))
         error("return value of u must be of size nu")
     end
     T = promote_type(Float64, eltype(x0))
@@ -207,7 +208,7 @@ function ltitr(A::AbstractMatrix{T}, B::AbstractMatrix{T}, u::Function, t,
 
     for i=1:iters
         x[:,i] = x0
-        uout[:,i] = u(x0,t[i])
+        uout[:,i] .= u(x0,t[i])
         x0 = A * x0 + B * uout[:,i]
     end
     return transpose(x), transpose(uout)
diff --git a/test/test_timeresp.jl b/test/test_timeresp.jl
index 3d0037235..63334bea8 100644
--- a/test/test_timeresp.jl
+++ b/test/test_timeresp.jl
@@ -33,10 +33,15 @@ sysd = c2d(sys, 0.1)[1]
 sysdfb = ss(sysd.A-sysd.B*L, sysd.B, sysd.C, sysd.D, 0.1)
 #Simulate without input
 yd, td, xd = lsim(sysdfb, zeros(501), t, x0=x0)
-
 @test y ≈ yd
 @test x ≈ xd
 
+# Test that the discrete lsim accepts u function that returns scalar
+L = lqr(sysd,Q,R)
+u(x,i) = -L*x
+yd, td, xd = lsim(sysd, u, t, x0=x0)
+@test norm(y - yd)/norm(y) < 0.05 # Since the cost matrices are not discretized, these will differ a bit
+@test norm(x - xd)/norm(x) < 0.05
 
 #Test step and impulse Continuous
 t0 = 0:0.05:2