add JET analysis (#141)

* add JET analysis

to check for dynamic dispatch. Also improve allocation performance to allow for completely allocation free operation of KF and UKF

* further allocation improvements

* ?

* add test_jet file

* tweak allocs

Project.toml

@@ -59,10 +59,11 @@ julia = "1.7"
 
 [extras]
 ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ControlSystemsBase", "MonteCarloMeasurements", "Optim", "Test", "Plots"]
+test = ["ControlSystemsBase", "JET", "MonteCarloMeasurements", "Optim", "Test", "Plots"]

docs/src/parameter_estimation.md

@@ -21,7 +21,7 @@ N = 2000 # Number of particles
 
 const dg = MvNormal(ny,1.0)          # Measurement noise Distribution
 const df = MvNormal(nx,1.0)          # Dynamics noise Distribution
-const d0 = MvNormal(randn(nx),2.0)   # Initial state Distribution
+const d0 = MvNormal(@SVector(randn(nx)),2.0)   # Initial state Distribution
 
 const A = SA[1 0.1; 0 1]
 const B = @SMatrix [0.0 0.1; 1 0.1]
@@ -56,7 +56,7 @@ the correct value for the simulated data is 1 (the simulated system is the same
 We can do the same with a Kalman filter
 
 ```@example ml_map
-eye(n) = Matrix{Float64}(I,n,n)
+eye(n) = SMatrix{n,n}(1.0I(n))
 llskf = map(svec) do s
     kfs = KalmanFilter(A, B, C, 0, s^2*eye(nx), eye(ny), d0)
     loglik(kfs, u, y, p)
@@ -230,11 +230,11 @@ We simulate the system using the `rollout` function and add some noise to the me
 Tperiod = 200
 t = 0:Ts:1000
 u = vcat.(0.25 .* sign.(sin.(2pi/Tperiod .* t)) .+ 0.25)
-u = vcat.(u,u)
+u = SVector{nu}.(vcat.(u,u))
 x0 = Float64[2,2,3,3]
 x = LowLevelParticleFilters.rollout(discrete_dynamics, x0, u)[1:end-1]
 y = measurement.(x, u, 0, 0)
-y = [y .+ 0.01randn(ny) for y in y]
+y = [y .+ 0.01.*randn.() for y in y]
 
 plot(
     plot(reduce(hcat, x)', title="State"),
@@ -271,9 +271,9 @@ nothing # hide
 We then define a nonlinear state estimator, we will use the [`UnscentedKalmanFilter`](@ref), and solve the filtering problem. We start by an initial state estimate ``x_0`` that is slightly off for the parameter ``a_1``
 ```@example paramest
 nx = 5
-R1 = Diagonal([0.1, 0.1, 0.1, 0.1, 0.0001])
-R2 = Diagonal((1e-2)^2 * ones(ny))
-x0 = [2, 2, 3, 3, 0.02]
+R1 = SMatrix{nx,nx}(Diagonal([0.1, 0.1, 0.1, 0.1, 0.0001]))
+R2 = SMatrix{ny,ny}(Diagonal((1e-2)^2 * ones(ny)))
+x0 = SA[2, 2, 3, 3, 0.02]
 
 kf = UnscentedKalmanFilter(discrete_dynamics_params, measurement, R1, R2, MvNormal(x0, R1); ny, nu)
 
@@ -322,11 +322,11 @@ Tperiod = 200
 t = 0:Ts:1000
 u1 = vcat.(0.25 .* sign.(sin.(2pi/Tperiod .* (t ./ 40).^2)) .+ 0.25)
 u2 = vcat.(0.25 .* sign.(sin.(2pi/Tperiod .* (t ./ 40).^2 .+ pi/2)) .+ 0.25)
-u  = vcat.(u1,u2)
-x0 = Float64[2,2,3,3]
+u  = SVector{nu}.(vcat.(u1,u2))
+x0 = SA[2.0,2,3,3]
 x = LowLevelParticleFilters.rollout(discrete_dynamics, x0, u, p_true)[1:end-1]
 y = measurement.(x, u, 0, 0)
-y = [y .+ 0.01randn(ny) for y in y]
+y = [y .+ 0.01 .* randn.() for y in y]
 
 plot(
     plot(reduce(hcat, x)', title="State"),
@@ -335,15 +335,15 @@ plot(
 ```
 
 
-This time, we define a cost function for the optimizer to optimize, we'll use the sum of squared errors (`sse`). It's important to define the UKF with an initial state distribution with the same element type as the parameter vector so that automatic differentiation through the state estimator works, hence the explicit casting `T.(x0)`.
+This time, we define a cost function for the optimizer to optimize, we'll use the sum of squared errors (`sse`). It's important to define the UKF with an initial state distribution with the same element type as the parameter vector so that automatic differentiation through the state estimator works, hence the explicit casting `T.(x0)` and `T.(R1)`.
 ```@example paramest
 nx = 4
-R1 = Diagonal([0.1, 0.1, 0.1, 0.1])
-R2 = Diagonal((1e-2)^2 * ones(ny))
-x0 = [2, 2, 3, 3]
+R1 = SMatrix{nx,nx}(Diagonal([0.1, 0.1, 0.1, 0.1]))
+R2 = SMatrix{ny,ny}(Diagonal((1e-2)^2 * ones(ny)))
+x0 = SA[2.0, 2, 3, 3]
 
 function cost(p::Vector{T}) where T
-    kf = UnscentedKalmanFilter(discrete_dynamics, measurement, R1, R2, MvNormal(T.(x0), R1); ny, nu)
+    kf = UnscentedKalmanFilter(discrete_dynamics, measurement, R1, R2, MvNormal(T.(x0), T.(R1)); ny, nu)
     LowLevelParticleFilters.sse(kf, u, y, p) # Sum of squared prediction errors
 end
 nothing # hide
@@ -390,7 +390,7 @@ Below, we optimize the sum of squared residuals again, but this time we do it us
 using LeastSquaresOptim
 
 function residuals!(res, p::Vector{T}) where T
-    kf = UnscentedKalmanFilter(discrete_dynamics, measurement, R1, R2, MvNormal(T.(x0), R1); ny, nu)
+    kf = UnscentedKalmanFilter(discrete_dynamics, measurement, R1, R2, MvNormal(T.(x0), T.(R1)); ny, nu)
     LowLevelParticleFilters.prediction_errors!(res, kf, u, y, p) 
 end
 

src/ekf.jl

@@ -45,6 +45,11 @@ function Base.getproperty(ekf::EKF, s::Symbol) where EKF <: AbstractExtendedKalm
     return getproperty(getfield(ekf, :kf), s)
 end
 
+function Base.setproperty!(ekf::ExtendedKalmanFilter, s::Symbol, val)
+    s ∈ fieldnames(typeof(ekf)) && return setproperty!(ekf, s, val)
+    setproperty!(getfield(ekf, :kf), s, val) # Forward to inner filter
+end
+
 function Base.propertynames(ekf::EKF, private::Bool=false) where EKF <: AbstractExtendedKalmanFilter
     return (fieldnames(EKF)..., propertynames(ekf.kf, private)...)
 end

src/filtering.jl

@@ -45,11 +45,13 @@ function predict!(kf::AbstractKalmanFilter, u, p=parameters(kf), t::Real = index
     @unpack A,B,x,R = kf
     At = get_mat(A, x, u, p, t)
     Bt = get_mat(B, x, u, p, t)
-    x .= At*x .+ Bt*u |> vec
+    kf.x = At*x .+ Bt*u |> vec
     if kf.α == 1
-        R .= symmetrize(At*R*At') + R1
+        Ru = symmetrize(At*R*At')
+        kf.R = Ru + R1
     else
-        R .= symmetrize(kf.α*At*R*At') + R1
+        Ru = symmetrize(kf.α*At*R*At')
+        kf.R = Ru + R1
     end
     kf.t[] += 1
 end
@@ -81,13 +83,13 @@ function correct!(kf::AbstractKalmanFilter, u, y, p=parameters(kf), t::Real = in
     Dt = get_mat(D, x, u, p, t)
     e   = y .- Ct*x
     if !iszero(D)
-        e .-= Dt*u
+        e -= Dt*u
     end
     S   = symmetrize(Ct*R*Ct') + R2
     Sᵪ  = cholesky(S)
     K   = (R*Ct')/Sᵪ
-    x .+= K*e
-    R  .= symmetrize((I - K*Ct)*R) # WARNING against I .- A
+    kf.x += K*e
+    kf.R  = symmetrize((I - K*Ct)*R) # WARNING against I .- A
     ll = logpdf(MvNormal(PDMat(S, Sᵪ)), e)# - 1/2*logdet(S) # logdet is included in logpdf
     (; ll, e, S, Sᵪ, K)
 end

src/kalman.jl

@@ -1,6 +1,26 @@
 abstract type AbstractKalmanFilter <: AbstractFilter end
 
-@with_kw struct KalmanFilter{AT,BT,CT,DT,R1T,R2T,R2DT,D0T,XT,RT,P,αT} <: AbstractKalmanFilter
+function convert_cov_type(R1, R)
+    if R isa SMatrix || R isa Matrix
+        return R
+    elseif R1 isa SMatrix
+        return SMatrix{size(R1,1),size(R1,2)}(R)
+    elseif R1 isa Matrix
+        return Matrix(R)
+    else
+        return Matrix(R)
+    end
+end
+
+function convert_x0_type(μ)
+    if μ isa Vector || μ isa SVector
+        return μ
+    else
+        return Vector(μ)
+    end
+end
+
+@with_kw mutable struct KalmanFilter{AT,BT,CT,DT,R1T,R2T,R2DT,D0T,XT,RT,P,αT} <: AbstractKalmanFilter
     A::AT
     B::BT
     C::CT
@@ -47,7 +67,9 @@ function KalmanFilter(A,B,C,D,R1,R2,d0=MvNormal(Matrix(R1)); p = SciMLBase.NullP
     if check
         maximum(abs, eigvals(A isa SMatrix ? Matrix(A) : A)) ≥ 2 && @warn "The dynamics matrix A has eigenvalues with absolute value ≥ 2. This is either a highly unstable system, or you have forgotten to discretize a continuous-time model. If you are sure that the system is provided in discrete time, you can disable this warning by setting check=false." maxlog=1
     end
-    KalmanFilter(A,B,C,D,R1,R2,MvNormal(Matrix(R2)), d0, Vector(d0.μ), Matrix(d0.Σ), Ref(1), p, α)
+    R = convert_cov_type(R1, d0.Σ)
+    x0 = convert_x0_type(d0.μ)
+    KalmanFilter(A,B,C,D,R1,R2,MvNormal(Matrix(R2)), d0, x0, R, Ref(1), p, α)
 end
 
 function Base.propertynames(kf::KF, private::Bool=false) where KF <: AbstractKalmanFilter
@@ -95,7 +117,7 @@ end
 Reset the initial distribution of the state. Optionally, a new mean vector `x0` can be provided.
 """
 function reset!(kf::AbstractKalmanFilter; x0 = kf.d0.μ)
-    kf.x .= Vector(x0)
-    kf.R .= copy(Matrix(kf.d0.Σ))
+    kf.x = convert_x0_type(x0)
+    kf.R = convert_cov_type(kf.R1, kf.d0.Σ)# typeof(kf.R1)(kf.d0.Σ)
     kf.t[] = 1
 end

src/ukf.jl

@@ -32,13 +32,14 @@ end
 
 abstract type AbstractUnscentedKalmanFilter <: AbstractKalmanFilter end
 
-@with_kw struct UnscentedKalmanFilter{DT,MT,R1T,R2T,D0T,VT,XT,RT,P} <: AbstractUnscentedKalmanFilter
+@with_kw mutable struct UnscentedKalmanFilter{DT,MT,R1T,R2T,D0T,VT,YVT,XT,RT,P} <: AbstractUnscentedKalmanFilter
     dynamics::DT
     measurement::MT
     R1::R1T
     R2::R2T
     d0::D0T
     xs::Vector{VT}
+    ys::Vector{YVT}
     x::XT
     R::RT
     t::Base.RefValue{Int} = Ref(1)
@@ -77,7 +78,14 @@ The one exception where this will not work is when calling `simulate`, which ass
 """
 function UnscentedKalmanFilter(dynamics,measurement,R1,R2,d0=MvNormal(Matrix(R1)); p = SciMLBase.NullParameters(), nu::Int, ny::Int)
     xs = sigmapoints(mean(d0), cov(d0), static = !has_ip(dynamics))
-    UnscentedKalmanFilter(dynamics,measurement,R1,R2, d0, xs, Vector(d0.μ), Matrix(d0.Σ), Ref(1), ny, nu, p)
+    if has_ip(measurement)
+        ys = [zeros(promote_type(eltype(xs[1]), eltype(d0)), ny) for _ in 1:length(xs)]
+    else
+        ys = [@SVector zeros(promote_type(eltype(xs[1]), eltype(d0)), ny) for _ in 1:length(xs)]
+    end
+    R = convert_cov_type(R1, d0.Σ)
+    x0 = convert_x0_type(d0.μ)
+    UnscentedKalmanFilter(dynamics,measurement,R1,R2, d0, xs, ys, x0, R, Ref(1), ny, nu, p)
 end
 
 sample_state(kf::AbstractUnscentedKalmanFilter, p=parameters(kf); noise=true) = noise ? rand(kf.d0) : mean(kf.d0)
@@ -86,10 +94,10 @@ sample_measurement(kf::AbstractUnscentedKalmanFilter, x, u, p=parameters(kf), t=
 measurement(kf::AbstractUnscentedKalmanFilter) = kf.measurement
 dynamics(kf::AbstractUnscentedKalmanFilter) = kf.dynamics
 
-#                                x(k+1)          x            u             p           t
-@inline has_ip(fun) = hasmethod(fun, (AbstractArray,AbstractArray,AbstractArray,AbstractArray,Real))
+#                                        x(k+1)          x            u             p           t
+@inline has_ip(fun) = hasmethod(fun, Tuple{AbstractArray,AbstractArray,AbstractArray,AbstractArray,Real})
 
-function predict!(ukf::UnscentedKalmanFilter, u, p = parameters(ukf), t::Real = index(ukf); R1 = get_mat(ukf.R1, ukf.x, u, p, t))
+function predict!(ukf::UnscentedKalmanFilter{DT}, u, p = parameters(ukf), t::Real = index(ukf); R1 = get_mat(ukf.R1, ukf.x, u, p, t)) where DT
     @unpack dynamics,measurement,x,xs,R = ukf
     ns = length(xs)
     sigmapoints!(xs,eltype(xs)(x),R)
@@ -105,41 +113,66 @@ function predict!(ukf::UnscentedKalmanFilter, u, p = parameters(ukf), t::Real =
             xs[i] = dynamics(xs[i], u, p, t)
         end
     end
-    x .= mean(xs)
-    R .= symmetrize(cov(xs)) .+ R1
+    ukf.x = safe_mean(xs)
+    ukf.R = symmetrize(safe_cov(xs)) .+ R1
     ukf.t[] += 1
 end
 
+# The functions below are JET-safe from dynamic dispatch if called with static arrays
+safe_mean(xs) = mean(xs)
+function safe_mean(xs::Vector{<:SVector})
+    m = xs[1]
+    for i = 2:length(xs)
+        m += xs[i]
+    end
+    m ./ length(xs)
+end
+
+safe_cov(xs) = cov(xs)
+function safe_cov(xs::Vector{<:SVector})
+    m = safe_mean(xs)
+    P = 0 .* m*m'
+    for i in eachindex(xs)
+        e = xs[i] .- m
+        P += e*e'
+    end
+    c = P ./ (length(xs) - 1)
+    c
+end
+
 
-function correct!(ukf::UnscentedKalmanFilter, u, y, p=parameters(ukf), t::Real = index(ukf); R2 = get_mat(ukf.R2, ukf.x, u, p, t))
-    @unpack measurement,x,xs,R,R1 = ukf
+
+function correct!(ukf::UnscentedKalmanFilter{<: Any, MT}, u, y, p=parameters(ukf), t::Real = index(ukf); R2 = get_mat(ukf.R2, ukf.x, u, p, t)) where MT
+    @unpack measurement,x,xs,ys,R,R1 = ukf
     n = length(xs[1])
     m = length(y)
     ns = length(xs)
     sigmapoints!(xs,eltype(xs)(x),R) # Update sigmapoints here since untransformed points required
     C = @SMatrix zeros(n,m)
-    ys = map(xs) do x
+    for i = eachindex(xs,ys)
         if has_ip(measurement)
-            yi = zeros(promote_type(eltype(x), eltype(u)), m)
-            measurement(yi, x, u, p, t)
-            yi
+            measurement(ys[i], xs[i], u, p, t)
         else
-            measurement(x, u, p, t)
+            ys[i] = measurement(xs[i], u, p, t)
         end
     end
-    ym = mean(ys)
+    ym = safe_mean(ys)
     @inbounds for i in eachindex(ys) # Cross cov between x and y
         d   = ys[i]-ym
         ca  = (xs[i]-x)*d'
         C  += ca
     end
     e   = y .- ym
-    S   = symmetrize(cov(ys)) + R2 # cov of y
+    S   = symmetrize(safe_cov(ys)) + R2 # cov of y
     Sᵪ  = cholesky(S)
     K   = (C./ns)/Sᵪ # ns normalization to make it a covariance matrix
-    x .+= K*e
+    ukf.x += K*e
     # mul!(x, K, e, 1, 1) # K and e will be SVectors if ukf correctly initialized
-    RmKSKT!(R, K, S)
+    if R isa SMatrix
+        ukf.R = symmetrize(R - K*S*K')
+    else
+        RmKSKT!(R, K, S)
+    end
     ll = logpdf(MvNormal(PDMat(S,Sᵪ)), e) #- 1/2*logdet(S) # logdet is included in logpdf
     (; ll, e, S, Sᵪ, K)
 end
@@ -274,6 +307,11 @@ function Base.getproperty(ukf::DAEUnscentedKalmanFilter, s::Symbol)
     getproperty(getfield(ukf, :ukf), s) # Forward to inner filter
 end
 
+function Base.setproperty!(ukf::DAEUnscentedKalmanFilter, s::Symbol, val)
+    s ∈ fieldnames(typeof(ukf)) && return setproperty!(ukf, s, val)
+    setproperty!(getfield(ukf, :ukf), s, val) # Forward to inner filter
+end
+
 Base.propertynames(ukf::DAEUnscentedKalmanFilter) = (fieldnames(typeof(ukf))..., propertynames(getfield(ukf, :ukf))...)
 
 state(ukf::DAEUnscentedKalmanFilter) = ukf.xz
@@ -334,9 +372,9 @@ function predict!(ukf::DAEUnscentedKalmanFilter, u, p = parameters(ukf), t::Inte
             xs[i],_ = get_x_z(xzs[i])
         end
     end
-    x .= mean(xs) # xz or xs here? Answer: Covariance is associated only with x
+    ukf.x = mean(xs) # xz or xs here? Answer: Covariance is associated only with x
     xz .= calc_xz(ukf, xz, u, p, t, x)
-    R .= symmetrize(cov(xs)) + get_mat(R1, x, u, p, t)
+    ukf.R = symmetrize(cov(xs)) + get_mat(R1, x, u, p, t)
     ukf.t[] += 1
 end
 
@@ -364,10 +402,14 @@ function correct!(ukf::DAEUnscentedKalmanFilter, u, y, p = parameters(ukf), t::I
     S   = symmetrize(cov(ys)) + R2 # cov of y
     Sᵪ = cholesky(S)
     K   = (C./ns)/Sᵪ # ns normalization to make it a covariance matrix
-    x .+= K*e
+    ukf.x += K*e
     xz .= calc_xz(ukf, xz, u, p, t, x)
     # mul!(x, K, e, 1, 1) # K and e will be SVectors if ukf correctly initialized
-    RmKSKT!(R, K, S)
+    if R isa SMatrix
+        ukf.R = symmetrize(R - K*S*K')
+    else
+        RmKSKT!(R, K, S)
+    end
     ll = logpdf(MvNormal(PDMat(S,Sᵪ)), e) #- 1/2*logdet(S) # logdet is included in logpdf
     (; ll, e, S, Sᵪ, K)
 end

test/runtests.jl

@@ -262,6 +262,10 @@ mvnormal(μ::AbstractVector{<:Real}, σ::Real) = MvNormal(μ, float(σ) ^ 2 * I)
 end
 
 
+@testset "jet" begin
+    @info "Testing jet"
+    include("test_jet.jl")
+end
 
 
 @testset "Advanced filters" begin

test/test_ekf.jl

@@ -1,6 +1,7 @@
 using LowLevelParticleFilters, ForwardDiff, Distributions
 const LLPF = LowLevelParticleFilters
 using ControlSystemsBase
+using StaticArrays, LinearAlgebra, Test
 
 n = 2   # Dimension of state
 m = 1   # Dimension of input