More code optimization, with a big focus on the smoothing (#70)

* Run even faster - Throughout the code, but especially the smoothing

* Update the version number

* Fix a bug by removing an old `include`

* Fix tests by removing RungeKuttaInit for orders > 5 from testing

* Refactor the diffusion code a bit

* Optimize the FixedDiffusion estimation

* Shorten some `update!` code

* Fix a typo

* Remove an old TODO

* Fix a bug with the FixedDiffusion estimation

* Make the smoothing more compact

Project.toml

@@ -1,7 +1,7 @@
 name = "ProbNumDiffEq"
 uuid = "bf3e78b0-7d74-48a5-b855-9609533b56a5"
 authors = ["Nathanael Bosch"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"

src/ProbNumDiffEq.jl

@@ -54,7 +54,7 @@ _matmul!(C::AbstractMatrix{T}, A::Diagonal{T}, B::AbstractMatrix{T}
          ) where T <: OctavianCompatibleEltypes = (C .= A.diag .* B)
 _matmul!(C::Diagonal{T}, A::AbstractMatrix{T}, B::AbstractMatrix{T}
          ) where T <: OctavianCompatibleEltypes =
-             mul!(C, A, B)
+             @tullio C[i, i] = A[i,j]*B[j,i]
 
 
 # @reexport using PSDMatrices
@@ -105,7 +105,6 @@ include("preconditioning.jl")
 
 # Utils
 include("jacobian.jl")
-include("numerics_tricks.jl")
 
 # Iterated Extended Kalman Smoother
 include("ieks.jl")

src/diffusions.jl

@@ -10,10 +10,12 @@ initial_diffusion(diffusion::AbstractDiffusion, d, q, Eltype) = one(Eltype)
 
 struct FixedDiffusion <: AbstractStaticDiffusion end
 function estimate_diffusion(rule::FixedDiffusion, integ)
-    @unpack d, measurement = integ.cache
+    @unpack d, measurement, m_tmp = integ.cache
     sol_diffusions = integ.sol.diffusions
 
     v, S = measurement.μ, measurement.Σ
+    e, _ = m_tmp.μ, m_tmp.Σ.mat
+    _S = copy!(m_tmp.Σ.mat, S.mat)
 
     if iszero(v)
         return zero(integ.cache.global_diffusion)
@@ -22,7 +24,14 @@ function estimate_diffusion(rule::FixedDiffusion, integ)
         return Inf
     end
 
-    diffusion_t = v' * inv(S) * v / d
+    # diffusion_t = v' * inv(S) * v / d
+    if _S isa Diagonal
+        e .= v ./ _S.diag
+    else
+        S_chol = cholesky!(_S)
+        ldiv!(e, S_chol, v)
+    end
+    diffusion_t = dot(v, e) / d
 
     if integ.success_iter == 0
         @assert length(sol_diffusions) == 0
@@ -73,13 +82,15 @@ function estimate_diffusion(kind::DynamicDiffusion, integ)
     @unpack d, R, H, Qh, measurement, m_tmp = integ.cache
     # @assert all(R .== 0) "The dynamic-diffusion assumes R==0!"
     z = measurement.μ
-    HQH = X_A_Xt!(m_tmp.Σ, Qh, H)
+    e, HQH = m_tmp.μ, m_tmp.Σ
+    X_A_Xt!(HQH, Qh, H)
     if HQH.mat isa Diagonal
-        σ² = dot(z ./ HQH.mat.diag, z) / d
+        e .= z ./ HQH.mat.diag
+        σ² = dot(e, z) / d
         return σ², σ²
     else
         C = cholesky!(HQH.mat)
-        e = ldiv!(m_tmp.μ, C, z)
+        ldiv!(e, C, z)
         σ² = dot(z, e) / d
         return σ², σ²
     end

src/filtering.jl

@@ -39,18 +39,12 @@ function predict_cov!(x_out::SRGaussian, x_curr::SRGaussian, Ah::AbstractMatrix,
     _matmul!(M, L, L')
     chol = cholesky!(Symmetric(M), check=false)
 
-    if issuccess(chol)
-        copy!(x_out.Σ.squareroot, chol.U')
-        mul!(x_out.Σ.mat, chol.U', chol.U)
-    elseif eltype(L) <: Union{Float16, Float32, Float64}
-        Q = lq!(L)
-        copy!(x_out.Σ.squareroot, Q.L)
-        mul!(x_out.Σ.mat, Q.L, Q.L')
-    else
-        Q = qr(L')
-        copy!(x_out.Σ.squareroot, Q.R')
-        mul!(x_out.Σ.mat, Q.R', Q.R)
-    end
+    QL =
+        issuccess(chol) ? Matrix(chol.U)' :
+        eltype(L) <: Union{Float16, Float32, Float64} ? lq!(L).L :
+        qr(L').R'
+    copy!(x_out.Σ.squareroot, QL)
+    _matmul!(x_out.Σ.mat, QL, QL')
     return x_out.Σ
 end
 function predict_cov!(x_out::SRGaussian, x_curr::SRGaussian, Ah::AbstractMatrix, Qh::SRMatrix)

src/smoothing.jl

@@ -22,7 +22,6 @@ function smooth_all!(integ)
         mul!(x_tmp, P, x[i])
         mul!(x_tmp2, P, x[i+1])
         smooth!(x_tmp, x_tmp2, A, Q, integ, diffusions[i])
-        @assert !(any(isnan.(x_tmp.μ)) || any(isnan.(x_tmp.Σ))) "NaNs after smoothing"
         mul!(x[i], PI, x_tmp)
     end
 end
@@ -40,37 +39,31 @@ function smooth!(x_curr, x_next, Ah, Qh, integ, diffusion=1)
     predict_cov!(x_pred, x_curr, Ah, Qh, C1, diffusion)
 
     # Smoothing
-    P_p = x_pred.Σ
-    P_p_s_inv = inv(LowerTriangular(P_p.squareroot))
-    P_p_inv = mul!(x_pred.Σ.mat, P_p_s_inv', P_p_s_inv)
     # G = x_curr.Σ * Ah' * P_p_inv
-    G = _matmul!(G2, mul!(G1, x_curr.Σ, Ah'), P_p_inv)
+    P_p_chol = Cholesky(x_pred.Σ.squareroot, :L, 0)
+    G = rdiv!(_matmul!(G1, x_curr.Σ.mat, Ah'), P_p_chol)
+
     x_curr.μ .+= G * (x_next.μ .- x_pred.μ)
 
     # Joseph-Form:
     M, L = C2.mat, C2.squareroot
     D = length(x_pred.μ)
-    _matmul!(view(L, 1:D, 1:D), (I-G*Ah), x_curr.Σ.squareroot)
-    mul!(view(L, 1:D, D+1:2D), G, sqrt.(diffusion) * Qh.squareroot)
-    _matmul!(view(L, 1:D, 2D+1:3D), G, x_next.Σ.squareroot)
 
-    _matmul!(M, L, L')
-    chol = cholesky!(Symmetric(M), check=false)
+    _matmul!(G2, G, Ah)
+    copy!(view(L, 1:D, 1:D), x_curr.Σ.squareroot)
+    _matmul!(view(L, 1:D, 1:D), G2, x_curr.Σ.squareroot, -1.0, 1.0)
 
-    if issuccess(chol)
-        copy!(x_curr.Σ.squareroot, chol.U')
-        mul!(x_curr.Σ.mat, chol.U', chol.U)
-    elseif eltype(L) <: Union{Float16, Float32, Float64}
-        Q = lq!(L)
-        copy!(x_curr.Σ.squareroot, Q.L)
-        mul!(x_curr.Σ.mat, Q.L, Q.L')
-    else
-        Q = qr(L')
-        copy!(x_curr.Σ.squareroot, Q.R')
-        mul!(x_curr.Σ.mat, Q.R', Q.R)
-    end
+    _matmul!(view(L, 1:D, D+1:2D), _matmul!(G2, G, sqrt.(diffusion)), Qh.squareroot)
+    _matmul!(view(L, 1:D, 2D+1:3D), G, x_next.Σ.squareroot)
+
+    # _matmul!(M, L, L')
+    # chol = cholesky!(Symmetric(M), check=false)
 
-    assert_nonnegative_diagonal(x_curr.Σ)
+    QL = false && issuccess(chol) ? Matrix(chol.U)' :
+        eltype(L) <: Union{Float16, Float32, Float64} ? lq!(L).L :
+        qr(L').R'
+    copy!(x_curr.Σ.squareroot, QL)
+    _matmul!(x_curr.Σ.mat, QL, QL')
 
     return nothing
 end

test/state_init.jl

@@ -60,12 +60,12 @@ end
 
     integ1 = init(prob, EK0(order=8));
     tm_init = integ1.cache.x.μ
-    Proj1 = ProbNumDiffEq.projection(integ1.cache.d, integ1.cache.q)
+    Proj1 = integ1.cache.Proj
 
-    @testset "Order $o" for o in (3, 4, 5, 6, 7)
+    @testset "Order $o" for o in (3, 4, 5)
         integ2 = init(prob, EK0(order=o, initialization=RungeKuttaInit()));
         rk_init = integ2.cache.x.μ
-        Proj2 = ProbNumDiffEq.projection(integ2.cache.d, integ2.cache.q)
+        Proj2 = integ2.cache.Proj
 
         # @info "how do things look?" tm_init rk_init
         # error()
@@ -86,7 +86,9 @@ end
 
         # Test if the covariance covers the true error
         for i in 3:o
-            err = (Proj2(i) * rk_init .- Proj1(i) * tm_init)
+            _rk = Proj2(i) * rk_init
+            _tm = Proj1(i) * tm_init
+            err = _rk .- _tm
             C = ProbNumDiffEq.X_A_Xt(integ2.cache.x.Σ, Proj2(i))
             @assert isdiag(C)
             whitened_err = err ./ sqrt.(diag(C.mat))