From 1d5b4157bfa4edc2dd23e5d246a262104b13a9d0 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 11 May 2022 08:25:45 +0200 Subject: [PATCH] Minimal working example with PSDMatrices.jl (#139) * Minimal working example with PSDMatrices.jl * Get smoothing and plotting to work (but it's _very_ slow) * Fix the wrong index order in the error estimation * Fix some of the tests * Remote squarerootmatrix.jl * Add a `diag(::SRMatrix)` * Add fallbacks for X_A_Xt and X_A_Xt! back in * Fix ALL the tests * Switch to the diag implemented in PSDMatrices.jl * Apply JuliaFormatter.jl * No need to import diag anymore since PSDMats does it * Make many of the preallocations more explicit * General miniature improvements here and there * For now: Don't check for diagonality in `condition_on!` * Moar caching * Improve the callback efficiency * Make the callback a tiny bit nicer * Remove `integ` from the `rk_init_improve` function * Remove `integ` dep to possibly slightly reduce precompilation times * Fix "test/correctness.jl" tests * Remove some old unused argument (thanks to a thrown warning) * Fix the remaining failing tests * Changes from reviewing the PR * Remove the outlier count test * Apply JuliaFormatter * Use LAPACKs geqrf instead of geqrt for QR decompositions * Use an old Cholesky constructor to fix Julia 1.6 compatibility * Replace the oop `triangularize_factor` with an in-place `qr!` * Fix the solution plotting to show stddevs, not variances * Fix outdated docstring * Remove some old commented out lines * Remove old unused lines * Fix a line in `smooth` * Relace a qr! with my custom_qr! * Remove a comment * Replace SRMatrix with PSDMatrix * Fix the plotting standarddev computation * Add the precompilation again --- Project.toml | 2 + src/ProbNumDiffEq.jl | 13 +++-- src/caches.jl | 63 +++++++++++++++---------- src/callbacks.jl | 21 +++++---- src/diffusions.jl | 58 +++++++++++------------ src/filtering/predict.jl | 37 +++++++-------- src/filtering/smooth.jl | 48 ++++++++----------- src/filtering/update.jl | 6 ++- src/gaussians.jl | 6 +-- src/initialization/classicsolverinit.jl | 20 ++++---- src/initialization/common.jl | 5 +- src/initialization/taylormode.jl | 7 +-- src/perform_step.jl | 19 ++++---- src/priors.jl | 2 +- src/solution.jl | 7 +-- src/solution_plotting.jl | 2 +- src/solution_sampling.jl | 6 +-- src/squarerootmatrix.jl | 50 -------------------- test/filtering.jl | 51 ++++++++++---------- test/preconditioning.jl | 8 ++-- test/priors.jl | 8 ++-- test/solution.jl | 14 +----- test/specific_problems.jl | 6 +-- test/state_init.jl | 5 +- 24 files changed, 204 insertions(+), 260 deletions(-) delete mode 100644 src/squarerootmatrix.jl diff --git a/Project.toml b/Project.toml index 66e3970e3..7e92cde25 100644 --- a/Project.toml +++ b/Project.toml @@ -11,6 +11,7 @@ GaussianDistributions = "43dcc890-d446-5863-8d1a-14597580bb8d" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +PSDMatrices = "fe68d972-6fd8-4755-bdf0-97d4c54cefdc" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" @@ -28,6 +29,7 @@ ForwardDiff = "0.10" GaussianDistributions = "0.5" Octavian = "0.3" OrdinaryDiffEq = "6.2 - 6.10" +PSDMatrices = "0.4.2" RecipesBase = "1" RecursiveArrayTools = "2" Reexport = "0.2, 1.0" diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index 0b185ab87..f4385849a 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -22,6 +22,9 @@ stack(x) = copy(reduce(hcat, x)') using LinearAlgebra import LinearAlgebra: mul! +"""LAPACK.geqrf! seems to be faster on small matrices than LAPACK.geqrt!""" +custom_qr!(A) = qr!(A) +custom_qr!(A::StridedMatrix{<:LinearAlgebra.BlasFloat}) = QR(LAPACK.geqrf!(A)...) using TaylorSeries using TaylorIntegration @reexport using StructArrays @@ -60,16 +63,12 @@ _matmul!( B::AbstractVecOrMat{T}, ) where {T<:LinearAlgebra.BlasFloat} = matmul!(C, A, B) -# @reexport using PSDMatrices -# import PSDMatrices: X_A_Xt -include("squarerootmatrix.jl") -const SRMatrix = SquarerootMatrix -export SRMatrix +@reexport using PSDMatrices +import PSDMatrices: X_A_Xt, X_A_Xt! X_A_Xt(A, X) = X * A * X' X_A_Xt!(out, A, X) = (out .= X * A * X') apply_diffusion(Q, diffusion::Diagonal) = X_A_Xt(Q, sqrt.(diffusion)) -apply_diffusion(Q::SRMatrix, diffusion::Number) = - SRMatrix(sqrt.(diffusion) * Q.squareroot, diffusion * Q.mat) +apply_diffusion(Q::PSDMatrix, diffusion::Number) = PSDMatrix(sqrt.(diffusion) * Q.R) # All the Gaussian things @reexport using GaussianDistributions diff --git a/src/caches.jl b/src/caches.jl index 45d7d52b4..1f693d1f4 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -65,6 +65,13 @@ mutable struct GaussianODEFilterCache{ G1::matType G2::matType covmatcache::matType + Smat::matType + C_dxd::matType + C_dxD::matType + C_Dxd::matType + C_DxD::matType + C_2DxD::matType + C_3DxD::matType default_diffusion::diffusionType local_diffusion::diffusionType global_diffusion::diffusionType @@ -94,8 +101,6 @@ function OrdinaryDiffEq.alg_cache( calck, ::Val{IIP}, ) where {IIP,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits} - initialize_derivatives = true - if u isa Number error("We currently don't support scalar-valued problems") end @@ -126,10 +131,7 @@ function OrdinaryDiffEq.alg_cache( A, Q = ibm(d, q, uElType) initial_variance = ones(uElType, D) - x0 = Gaussian( - zeros(uElType, D), - SRMatrix(diagm(sqrt.(initial_variance)), diagm(initial_variance)), - ) + x0 = Gaussian(zeros(uElType, D), PSDMatrix(diagm(sqrt.(initial_variance)))) # Measurement model R = zeros(uElType, d, d) @@ -140,23 +142,29 @@ function OrdinaryDiffEq.alg_cache( du = f isa DynamicalODEFunction ? similar(u[2, :]) : similar(u) ddu = f isa DynamicalODEFunction ? zeros(uElType, d, 2d) : zeros(uElType, d, d) v = similar(h) - S = - # alg isa EK0 ? SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) : - SRMatrix(zeros(uElType, d, D), zeros(uElType, d, d)) + S = PSDMatrix(zeros(uElType, D, d)) measurement = Gaussian(v, S) pu_tmp = f isa DynamicalODEFunction ? - Gaussian(zeros(uElType, 2d), SRMatrix(zeros(uElType, 2d, D))) : similar(measurement) + Gaussian(zeros(uElType, 2d), PSDMatrix(zeros(uElType, D, 2d))) : copy(measurement) K = zeros(uElType, D, d) G = zeros(uElType, D, D) - C1 = SRMatrix(zeros(uElType, D, 2D), zeros(uElType, D, D)) - C2 = SRMatrix(zeros(uElType, D, 3D), zeros(uElType, D, D)) - covmatcache = similar(G) + C1 = PSDMatrix(zeros(uElType, 2D, D)) + C2 = PSDMatrix(zeros(uElType, 3D, D)) + Smat = zeros(uElType, d, d) + covmatcache = copy(G) + + C_dxd = zeros(uElType, d, d) + C_dxD = zeros(uElType, d, D) + C_Dxd = zeros(uElType, D, d) + C_DxD = zeros(uElType, D, D) + C_2DxD = zeros(uElType, 2D, D) + C_3DxD = zeros(uElType, 3D, D) if alg isa EK1FDB H = [E1; E2] v = [v; v] - S = SRMatrix(zeros(uElType, 2d, D), zeros(uElType, 2d, 2d)) + S = PSDMatrix(zeros(uElType, D, 2d)) measurement = Gaussian(v, S) K = zeros(uElType, D, 2d) end @@ -166,17 +174,17 @@ function OrdinaryDiffEq.alg_cache( copy!(x0.Σ, apply_diffusion(x0.Σ, initdiff)) Ah, Qh = copy(A), copy(Q) - u_pred = similar(u) - u_filt = similar(u) - tmp = similar(u) + u_pred = copy(u) + u_filt = copy(u) + tmp = copy(u) x_pred = copy(x0) - x_filt = similar(x0) - x_tmp = similar(x0) - x_tmp2 = similar(x0) - m_tmp = similar(measurement) - K2 = similar(K) - G2 = similar(G) - err_tmp = similar(du) + x_filt = copy(x0) + x_tmp = copy(x0) + x_tmp2 = copy(x0) + m_tmp = copy(measurement) + K2 = copy(K) + G2 = copy(G) + err_tmp = copy(du) # Things for calc_J uf = get_uf(f, t, p, Val(IIP)) @@ -251,6 +259,13 @@ function OrdinaryDiffEq.alg_cache( G, G2, covmatcache, + Smat, + C_dxd, + C_dxD, + C_Dxd, + C_DxD, + C_2DxD, + C_3DxD, initdiff, initdiff * NaN, initdiff * NaN, diff --git a/src/callbacks.jl b/src/callbacks.jl index be4f13b3c..16cb3d05a 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -1,13 +1,16 @@ -function manifoldupdate!(integ, residualf; maxiters=100, ϵ₁=1e-25, ϵ₂=1e-15) - m, C = integ.cache.x +function manifoldupdate!(cache, residualf; maxiters=100, ϵ₁=1e-25, ϵ₂=1e-15) + m, C = cache.x # Create some caches - @unpack SolProj, tmp, H, x_tmp, x_tmp2 = integ.cache + @unpack SolProj, tmp, H, x_tmp, x_tmp2 = cache + D = cache.d * (cache.q + 1) z_tmp = residualf(mul!(tmp, SolProj, m)) result = DiffResults.JacobianResult(z_tmp, tmp) d = length(z_tmp) H = H[1:d, :] - S = SquarerootMatrix(C.squareroot[1:d, :], C.mat[1:d, 1:d]) + K1, K2 = cache.C_DxD[:, 1:d], cache.C_2DxD[1:D, 1:d] + + S = PSDMatrix(C.R[:, 1:d]) m_tmp, C_tmp = x_tmp m_i = copy(m) @@ -23,8 +26,7 @@ function manifoldupdate!(integ, residualf; maxiters=100, ϵ₁=1e-25, ϵ₂=1e-1 X_A_Xt!(S, C, H) # m_i_new, C_i_new = update(x, Gaussian(z .+ (H * (m - m_i)), S), H) - # More efficient update with less allocations: - K = C * (H' / S) + K = _matmul!(K2, C.R', _matmul!(K1, C.R, H' / S)) m_tmp .= m_i .- m mul!(z_tmp, H, m_tmp) z_tmp .-= z @@ -38,7 +40,7 @@ function manifoldupdate!(integ, residualf; maxiters=100, ϵ₁=1e-25, ϵ₂=1e-1 m_i = m_i_new end - copy!(integ.cache.x, Gaussian(m_i_new, C_i_new)) + copy!(cache.x, Gaussian(m_i_new, C_i_new)) return nothing end @@ -70,6 +72,9 @@ function ManifoldUpdate( kwargs..., ) condition(u, t, integ) = true - affect!(integ) = manifoldupdate!(integ, residual; maxiters=maxiters, ϵ₁=ϵ₁, ϵ₂=ϵ₂) + affect!(integ) = begin + manifoldupdate!(integ.cache, residual; maxiters=maxiters, ϵ₁=ϵ₁, ϵ₂=ϵ₂) + mul!(view(integ.u, :), integ.cache.SolProj, integ.cache.x.μ) + end return DiscreteCallback(condition, affect!, args...; kwargs...) end diff --git a/src/diffusions.jl b/src/diffusions.jl index b650aa7d2..ccfd59456 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -17,7 +17,8 @@ A local diffusion parameter is estimated at each step. Works well with adaptive """ struct DynamicDiffusion <: AbstractDynamicDiffusion end initial_diffusion(diffusion::DynamicDiffusion, d, q, Eltype) = one(Eltype) -estimate_local_diffusion(kind::DynamicDiffusion, integ) = local_scalar_diffusion(integ) +estimate_local_diffusion(kind::DynamicDiffusion, integ) = + local_scalar_diffusion(integ.cache) """ DynamicMVDiffusion() @@ -31,7 +32,8 @@ scales of the different dimensions vary a lot. struct DynamicMVDiffusion <: AbstractDynamicDiffusion end initial_diffusion(diffusionmodel::DynamicMVDiffusion, d, q, Eltype) = kron(Diagonal(ones(Eltype, d)), Diagonal(ones(Eltype, q + 1))) -estimate_local_diffusion(kind::DynamicMVDiffusion, integ) = local_diagonal_diffusion(integ) +estimate_local_diffusion(kind::DynamicMVDiffusion, integ) = + local_diagonal_diffusion(integ.cache) """ FixedDiffusion(; initial_diffusion=1.0, calibrate=true) @@ -48,20 +50,16 @@ Base.@kwdef struct FixedDiffusion{T<:Number} <: AbstractStaticDiffusion end initial_diffusion(diffusionmodel::FixedDiffusion, d, q, Eltype) = diffusionmodel.initial_diffusion * one(Eltype) -estimate_local_diffusion(kind::FixedDiffusion, integ) = local_scalar_diffusion(integ) +estimate_local_diffusion(kind::FixedDiffusion, integ) = local_scalar_diffusion(integ.cache) function estimate_global_diffusion(rule::FixedDiffusion, integ) - @unpack d, measurement, m_tmp = integ.cache + @unpack d, measurement, m_tmp, Smat = 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 _S isa Diagonal - e .= v ./ _S.diag - else - S_chol = cholesky!(_S) - ldiv!(e, S_chol, v) - end + e = m_tmp.μ + _S = _matmul!(Smat, S.R', S.R) + S_chol = cholesky!(_S) + ldiv!(e, S_chol, v) diffusion_t = dot(v, e) / d if integ.success_iter == 0 @@ -96,12 +94,15 @@ function initial_diffusion(diffusionmodel::FixedMVDiffusion, d, q, Eltype) @assert initdiff isa Number || length(initdiff) == d return kron(Diagonal(initdiff .* ones(Eltype, d)), Diagonal(ones(Eltype, q + 1))) end -estimate_local_diffusion(kind::FixedMVDiffusion, integ) = local_diagonal_diffusion(integ) +estimate_local_diffusion(kind::FixedMVDiffusion, integ) = + local_diagonal_diffusion(integ.cache) function estimate_global_diffusion(kind::FixedMVDiffusion, integ) @unpack q, measurement = integ.cache v, S = measurement.μ, measurement.Σ - S_11 = diag(S)[1] + # S_11 = diag(S)[1] + c1 = view(S.R, :, 1) + S_11 = dot(c1, c1) Σ_ii = v .^ 2 ./ S_11 Σ = Diagonal(Σ_ii) @@ -132,25 +133,20 @@ where ``z, H, Q`` are taken from the passed integrator. For more background information - N. Bosch, P. Hennig, F. Tronarp: **Calibrated Adaptive Probabilistic ODE Solvers** (2021) """ -function local_scalar_diffusion(integ) - @unpack d, R, H, Qh, measurement, m_tmp = integ.cache +function local_scalar_diffusion(cache) + @unpack d, R, H, Qh, measurement, m_tmp, Smat = cache z = measurement.μ e, HQH = m_tmp.μ, m_tmp.Σ X_A_Xt!(HQH, Qh, H) - if HQH.mat isa Diagonal - e .= z ./ HQH.mat.diag - σ² = dot(e, z) / d - return σ² - else - C = cholesky!(HQH.mat) - ldiv!(e, C, z) - σ² = dot(z, e) / d - return σ² - end + HQHmat = _matmul!(Smat, HQH.R', HQH.R) + C = cholesky!(HQHmat) + ldiv!(e, C, z) + σ² = dot(z, e) / d + return σ² end """ - local_diagonal_diffusion(integ) + local_diagonal_diffusion(cache) Compute the local, scalar diffusion estimate. @@ -164,11 +160,13 @@ where ``z, H, Q`` are taken from the passed integrator. For more background information - N. Bosch, P. Hennig, F. Tronarp: **Calibrated Adaptive Probabilistic ODE Solvers** (2021) """ -function local_diagonal_diffusion(integ) - @unpack q, H, Qh, measurement, m_tmp = integ.cache +function local_diagonal_diffusion(cache) + @unpack q, H, Qh, measurement, m_tmp = cache z = measurement.μ HQH = X_A_Xt!(m_tmp.Σ, Qh, H) - Q0_11 = diag(HQH)[1] + # Q0_11 = diag(HQH)[1] + c1 = view(HQH.R, :, 1) + Q0_11 = dot(c1, c1) Σ_ii = z .^ 2 ./ Q0_11 # Σ_ii .= max.(Σ_ii, eps(eltype(Σ_ii))) diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index f70b062f9..204f4d8ea 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -12,8 +12,8 @@ predict(x::Gaussian, A::AbstractMatrix, Q::AbstractMatrix) = Gaussian(predict_mean(x, A), predict_cov(x, A, Q)) predict_mean(x::Gaussian, A::AbstractMatrix) = A * x.μ predict_cov(x::Gaussian, A::AbstractMatrix, Q::AbstractMatrix) = A * x.Σ * A' + Q -predict_cov(x::SRGaussian, A::AbstractMatrix, Q::SRMatrix) = - SRMatrix(qr([A * x.Σ.squareroot Q.squareroot]').R') +predict_cov(x::SRGaussian, A::AbstractMatrix, Q::PSDMatrix) = + PSDMatrix(qr([x.Σ.R * A'; Q.R]).R) """ predict!(x_out, x_curr, Ah, Qh, cachemat) @@ -21,7 +21,7 @@ predict_cov(x::SRGaussian, A::AbstractMatrix, Q::SRMatrix) = In-place and square-root implementation of [`predict`](@ref) which saves the result into `x_out`. -Only works with `ProbNumDiffEq.SquarerootMatrix` types as `Ah`, `Qh`, and in the +Only works with `PSDMatrices.PSDMatrix` types as `Ah`, `Qh`, and in the covariances of `x_curr` and `x_out` (both of type `Gaussian`). To prevent allocations, a cache matrix `cachemat` of size ``D \\times 2D`` (where ``D \\times D`` is the size of `Ah` and `Qh`) needs to be passed. @@ -32,12 +32,13 @@ function predict!( x_out::SRGaussian, x_curr::SRGaussian, Ah::AbstractMatrix, - Qh::SRMatrix, - cachemat::SRMatrix, + Qh::PSDMatrix, + C_DxD::AbstractMatrix, + C_2DxD::AbstractMatrix, diffusion=1, ) predict_mean!(x_out, x_curr, Ah) - predict_cov!(x_out, x_curr, Ah, Qh, cachemat, diffusion) + predict_cov!(x_out, x_curr, Ah, Qh, C_DxD, C_2DxD, diffusion) return x_out end @@ -50,23 +51,21 @@ function predict_cov!( x_out::SRGaussian, x_curr::SRGaussian, Ah::AbstractMatrix, - Qh::SRMatrix, - cachemat::SRMatrix, + Qh::PSDMatrix, + C_DxD::AbstractMatrix, + C_2DxD::AbstractMatrix, diffusion=1, ) - M, L = cachemat.mat, cachemat.squareroot - D, D = size(Qh.mat) + R, M = C_2DxD, C_DxD + D, D = size(Qh) - _matmul!(view(L, 1:D, 1:D), Ah, x_curr.Σ.squareroot) - _matmul!(view(L, 1:D, D+1:2D), sqrt.(diffusion), Qh.squareroot) - _matmul!(M, L, L') + _matmul!(view(R, 1:D, 1:D), x_curr.Σ.R, Ah') + _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion)) + _matmul!(M, R', R) chol = cholesky!(Symmetric(M), check=false) - QL = - issuccess(chol) ? Matrix(chol.U)' : - eltype(L) <: Union{Float16,Float32,Float64} ? lq!(L).L : qr(L').R' - # QL = eltype(L) <: Union{Float16,Float32,Float64} ? lq!(L).L : qr(L').R' - copy!(x_out.Σ.squareroot, QL) - _matmul!(x_out.Σ.mat, QL, QL') + Q_R = issuccess(chol) ? Matrix(chol.U) : custom_qr!(R).R + copy!(x_out.Σ.R, Q_R) + # _matmul!(x_out.Σ.mat, QL, QL') return x_out.Σ end diff --git a/src/filtering/smooth.jl b/src/filtering/smooth.jl index af9389914..4bc46136c 100644 --- a/src/filtering/smooth.jl +++ b/src/filtering/smooth.jl @@ -46,24 +46,21 @@ function smooth( x_curr::SRGaussian, x_next_smoothed::SRGaussian, Ah::AbstractMatrix, - Qh::SRMatrix, + Qh::PSDMatrix, ) x_pred = predict(x_curr, Ah, Qh) - P_p = x_pred.Σ - P_p_inv = inv(P_p) - - G = x_curr.Σ * Ah' * P_p_inv + G = Matrix(x_curr.Σ) * Ah' / x_pred.Σ smoothed_mean = x_curr.μ + G * (x_next_smoothed.μ - x_pred.μ) _R = [ - x_curr.Σ.squareroot' * (I - G * Ah)' - Qh.squareroot' * G' - x_next_smoothed.Σ.squareroot' * G' + x_curr.Σ.R * (I - G * Ah)' + Qh.R * G' + x_next_smoothed.Σ.R * G' ] - _, P_s_R = qr(_R) - smoothed_cov = SRMatrix(P_s_R') + P_s_R = qr(_R).R + smoothed_cov = PSDMatrix(P_s_R) x_curr_smoothed = Gaussian(smoothed_mean, smoothed_cov) return x_curr_smoothed, G @@ -84,7 +81,7 @@ function smooth!( x_curr::SRGaussian, x_next::SRGaussian, Ah::AbstractMatrix, - Qh::SquarerootMatrix, + Qh::PSDMatrix, cache::GaussianODEFilterCache, diffusion::Union{Number,Diagonal}=1, ) @@ -93,40 +90,33 @@ function smooth!( # x_next is the state at time t_{n+1}, already smoothed, which we use for smoothing @unpack d, q = cache @unpack x_pred = cache - @unpack C1, G1, G2, C2 = cache + @unpack C1, G1, G2, C2, C_DxD, C_2DxD, C_3DxD = cache # Prediction: t -> t+1 predict_mean!(x_pred, x_curr, Ah) - predict_cov!(x_pred, x_curr, Ah, Qh, C1, diffusion) + predict_cov!(x_pred, x_curr, Ah, Qh, C_DxD, C_2DxD, diffusion) # Smoothing # G = 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) + P_p_chol = Cholesky(x_pred.Σ.R, :U, 0) + G = rdiv!(_matmul!(G1, x_curr.Σ.R', _matmul!(G2, x_curr.Σ.R, Ah')), P_p_chol) # x_curr.μ .+= G * (x_next.μ .- x_pred.μ) # less allocations: x_pred.μ .-= x_next.μ _matmul!(x_curr.μ, G, x_pred.μ, -1, 1) # Joseph-Form: - M, L = C2.mat, C2.squareroot + R = C_3DxD _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) - - _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) + copy!(view(R, 1:D, 1:D), x_curr.Σ.R') + _matmul!(view(R, 1:D, 1:D), x_curr.Σ.R, G2', -1.0, 1.0) - # _matmul!(M, L, L') - # chol = cholesky!(Symmetric(M), check=false) - succ = false && issuccess(chol) + _matmul!(view(R, D+1:2D, 1:D), Qh.R, _matmul!(G2, G, sqrt.(diffusion))') + _matmul!(view(R, 2D+1:3D, 1:D), x_next.Σ.R, G') - QL = - succ ? 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') + Q_R = custom_qr!(R).R + copy!(x_curr.Σ.R, Q_R) return nothing end diff --git a/src/filtering/update.jl b/src/filtering/update.jl index 47c4b5b03..13cfcaab8 100644 --- a/src/filtering/update.jl +++ b/src/filtering/update.jl @@ -67,7 +67,11 @@ function update!( D = length(m_p) # K = P_p * H' / S - S_chol = cholesky!(S) + if S isa PSDMatrix + S_chol = Cholesky(custom_qr!(S.R).R, :U, 0) + else + S_chol = cholesky(S) + end K = _matmul!(K_cache, Matrix(P_p), H') rdiv!(K, S_chol) diff --git a/src/gaussians.jl b/src/gaussians.jl index 749510caa..66a66b98c 100644 --- a/src/gaussians.jl +++ b/src/gaussians.jl @@ -1,7 +1,7 @@ # const PSDGaussian{T} = Gaussian{Vector{T}, PSDMatrix{T}} # const PSDGaussianList{T} = StructArray{PSDGaussian{T}} -const SRGaussian{T,S,M} = Gaussian{Vector{T},SRMatrix{T,S,M}} -const SRGaussianList{T,S,M} = StructArray{SRGaussian{T,S,M}} +const SRGaussian{T,S} = Gaussian{Vector{T},PSDMatrix{T,S}} +const SRGaussianList{T,S} = StructArray{SRGaussian{T,S}} copy(P::Gaussian) = Gaussian(copy(P.μ), copy(P.Σ)) similar(P::Gaussian) = Gaussian(similar(P.μ), similar(P.Σ)) @@ -18,7 +18,7 @@ size(g::Gaussian) = size(g.μ) ndims(g::Gaussian) = ndims(g.μ) Base.:*(M, g::SRGaussian) = Gaussian(M * g.μ, X_A_Xt(g.Σ, M)) -# GaussianDistributions.whiten(Σ::SRMatrix, z) = Σ.L\z +# GaussianDistributions.whiten(Σ::PSDMatrix, z) = Σ.L\z function _gaussian_mul!(g_out::SRGaussian, M::AbstractMatrix, g_in::SRGaussian) _matmul!(g_out.μ, M, g_in.μ) diff --git a/src/initialization/classicsolverinit.jl b/src/initialization/classicsolverinit.jl index 1e63c731e..b5369c0a0 100644 --- a/src/initialization/classicsolverinit.jl +++ b/src/initialization/classicsolverinit.jl @@ -1,4 +1,4 @@ -function initial_update!(integ, cache, init::ClassicSolverInit) +function initial_update!(integ, cache, ::ClassicSolverInit) @unpack u, f, p, t = integ @unpack d, x, Proj = cache q = integ.alg.order @@ -13,10 +13,11 @@ function initial_update!(integ, cache, init::ClassicSolverInit) # Initialize on u0; taking special care for DynamicalODEProblems is_secondorder = integ.f isa DynamicalODEFunction _u = is_secondorder ? view(u.x[2], :) : view(u, :) - condition_on!(x, Proj(0), _u, m_tmp.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) + Mcache = cache.C_DxD + condition_on!(x, Proj(0), _u, m_tmp.Σ, K1, x_tmp.Σ, Mcache, cache) is_secondorder ? f.f1(du, u.x[1], u.x[2], p, t) : f(du, u, p, t) integ.destats.nf += 1 - condition_on!(x, Proj(1), view(du, :), m_tmp.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) + condition_on!(x, Proj(1), view(du, :), m_tmp.Σ, K1, x_tmp.Σ, Mcache, cache) if q < 2 return @@ -33,7 +34,7 @@ function initial_update!(integ, cache, init::ClassicSolverInit) ForwardDiff.jacobian!(ddu, (du, u) -> f(du, u, p, t), du, u) end ddfddu = ddu * view(du, :) + view(dfdt, :) - condition_on!(x, Proj(2), ddfddu, m_tmp.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) + condition_on!(x, Proj(2), ddfddu, m_tmp.Σ, K1, x_tmp.Σ, Mcache, cache) if q < 3 return end @@ -79,10 +80,10 @@ function initial_update!(integ, cache, init::ClassicSolverInit) # Filter & smooth to fit these values! us = [u for u in sol.u] - return rk_init_improve(integ, cache, sol.t, us, dt) + return rk_init_improve(cache, sol.t, us, dt) end -function rk_init_improve(integ, cache::GaussianODEFilterCache, ts, us, dt) +function rk_init_improve(cache::GaussianODEFilterCache, ts, us, dt) @unpack A, Q = cache @unpack x, x_pred, x_filt, measurement = cache @@ -100,15 +101,14 @@ function rk_init_improve(integ, cache::GaussianODEFilterCache, ts, us, dt) (u isa RecursiveArrayTools.ArrayPartition) && (u = u.x[2]) # for 2ndOrderODEs u = view(u, :) # just in case the problem is matrix-valued predict_mean!(x_pred, x, A) - predict_cov!(x_pred, x, A, Q, cache.C1, cache.default_diffusion) + predict_cov!(x_pred, x, A, Q, cache.C_DxD, cache.C_2DxD, cache.default_diffusion) push!(preds, copy(x_pred)) H = cache.E0 * PI measurement.μ .= H * x_pred.μ .- u X_A_Xt!(measurement.Σ, x_pred.Σ, H) - M_cache, S_cache = cache.x_tmp2.Σ.mat, cache.m_tmp.Σ - update!(x_filt, x_pred, measurement, H, cache.K1, M_cache, S_cache) + update!(x_filt, x_pred, measurement, H, cache.K1, cache.C_DxD, cache.m_tmp.Σ) push!(filts, copy(x_filt)) x = x_filt @@ -119,7 +119,7 @@ function rk_init_improve(integ, cache::GaussianODEFilterCache, ts, us, dt) xf = filts[i-1] xs = filts[i] xp = preds[i-1] # Since `preds` is one shorter - smooth!(xf, xs, A, Q, integ.cache, 1) + smooth!(xf, xs, A, Q, cache, 1) end _gaussian_mul!(cache.x, PI, filts[1]) diff --git a/src/initialization/common.jl b/src/initialization/common.jl index 46a09052d..172c3441a 100644 --- a/src/initialization/common.jl +++ b/src/initialization/common.jl @@ -55,17 +55,18 @@ function condition_on!( Kcache, covcache, Mcache, + cache, ) S = Scache X_A_Xt!(S, x.Σ, H) - @assert isdiag(S) + # @assert isdiag(Matrix(S)) S_diag = diag(S) if any(iszero.(S_diag)) # could happen with a singular mass-matrix S_diag .+= 1e-20 end - _matmul!(Kcache, x.Σ.mat, H') + _matmul!(Kcache, x.Σ.R', _matmul!(cache.C_Dxd, x.Σ.R, H')) K = Kcache ./= S_diag' # x.μ .+= K*(data - z) diff --git a/src/initialization/taylormode.jl b/src/initialization/taylormode.jl index 30ee046b8..8fa4f7525 100644 --- a/src/initialization/taylormode.jl +++ b/src/initialization/taylormode.jl @@ -12,10 +12,7 @@ function initial_update!(integ, cache, init::TaylorModeInit) f_derivatives = taylormode_get_derivatives(u, f, p, t, q) integ.destats.nf += q @assert length(0:q) == length(f_derivatives) - m_cache = Gaussian( - zeros(eltype(u), d), - SRMatrix(zeros(eltype(u), d, D), zeros(eltype(u), d, d)), - ) + m_cache = Gaussian(zeros(eltype(u), d), PSDMatrix(zeros(eltype(u), D, d))) for (o, df) in zip(0:q, f_derivatives) if f isa DynamicalODEFunction @assert df isa ArrayPartition @@ -27,7 +24,7 @@ function initial_update!(integ, cache, init::TaylorModeInit) df = df[:] end - condition_on!(x, pmat, df, m_cache.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) + condition_on!(x, pmat, df, m_cache.Σ, K1, x_tmp.Σ, cache.C_DxD, cache) end end diff --git a/src/perform_step.jl b/src/perform_step.jl index b276a72b8..8ceebfa4a 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -96,7 +96,7 @@ function OrdinaryDiffEq.perform_step!( # Predict the covariance, using either the local or global diffusion extrapolation_diff = isdynamic(cache.diffusionmodel) ? cache.local_diffusion : cache.default_diffusion - predict_cov!(x_pred, x, Ah, Qh, cache.C1, extrapolation_diff) + predict_cov!(x_pred, x, Ah, Qh, cache.C_DxD, cache.C_2DxD, extrapolation_diff) # Compute measurement covariance only now; likelihood computation is currently broken compute_measurement_covariance!(cache) @@ -300,8 +300,8 @@ compute_measurement_covariance!(cache) = function update!(integ, prediction) @unpack measurement, H, R, x_filt = integ.cache - @unpack K1, K2, x_tmp2, m_tmp = integ.cache - update!(x_filt, prediction, measurement, H, K1, x_tmp2.Σ.mat, m_tmp.Σ) + @unpack K1, K2, x_tmp2, m_tmp, C_DxD = integ.cache + update!(x_filt, prediction, measurement, H, K1, C_DxD, m_tmp.Σ) return x_filt end @@ -352,7 +352,7 @@ Computes a local error estimate, as E_i = ( σ_{loc}^2 ⋅ (H Q(h) H^T)_{ii} )^(1/2) ``` To save allocations, the function modifies the given `cache` and writes into -`cache.m_tmp.Σ.squareroot` during some computations. +`cache.C_Dxd` during some computations. """ function estimate_errors!(cache::GaussianODEFilterCache) @unpack local_diffusion, Qh, H, d = cache @@ -361,17 +361,16 @@ function estimate_errors!(cache::GaussianODEFilterCache) return Inf end - L = cache.m_tmp.Σ.squareroot + R = cache.C_Dxd if local_diffusion isa Diagonal - _matmul!(L, H, sqrt.(local_diffusion) * Qh.squareroot) - error_estimate = sqrt.(diag(L * L')) + _matmul!(R, Qh.R * sqrt.(local_diffusion), H') + error_estimate = sqrt.(diag(PSDMatrix(R))) return view(error_estimate, 1:d) - elseif local_diffusion isa Number - _matmul!(L, H, Qh.squareroot) + _matmul!(R, Qh.R, H') # error_estimate = local_diffusion .* diag(L*L') - @tullio error_estimate[i] := L[i, j] * L[i, j] + error_estimate = diag(PSDMatrix(R)) error_estimate .*= local_diffusion # @info "it's small anyways I guess?" error_estimate cache.measurement.μ .^ 2 diff --git a/src/priors.jl b/src/priors.jl index 5a2bca999..806488f0e 100644 --- a/src/priors.jl +++ b/src/priors.jl @@ -32,7 +32,7 @@ function ibm(d::Integer, q::Integer, ::Type{elType}=typeof(1.0)) where {elType} end end QL_breve = cholesky(Q_breve).L - Q = SRMatrix(kron(I(d), QL_breve), kron(I(d), Q_breve)) + Q = PSDMatrix(kron(I(d), QL_breve')) return A, Q end diff --git a/src/solution.jl b/src/solution.jl index cc634e817..195ee55d2 100644 --- a/src/solution.jl +++ b/src/solution.jl @@ -71,11 +71,8 @@ function DiffEqBase.build_solution( d = length(prob.u0) uElType = eltype(prob.u0) D = d - pu_cov = - # alg isa EK0 && !(prob.f isa DynamicalODEFunction) ? - # SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) : - SRMatrix(zeros(uElType, d, D), zeros(uElType, d, d)) - x_cov = SRMatrix(zeros(uElType, d, d), zeros(uElType, d, d)) + pu_cov = PSDMatrix(zeros(uElType, d, D)) + x_cov = PSDMatrix(zeros(uElType, d, d)) pu = StructArray{Gaussian{Vector{uElType},typeof(pu_cov)}}(undef, 0) x_filt = StructArray{Gaussian{Vector{uElType},typeof(x_cov)}}(undef, 0) x_smooth = copy(x_filt) diff --git a/src/solution_plotting.jl b/src/solution_plotting.jl index a5bf229cf..272adc4df 100644 --- a/src/solution_plotting.jl +++ b/src/solution_plotting.jl @@ -17,7 +17,7 @@ times = times[tstart.<=times.<=tend] end values = stack(mean(sol_rvs)) - stds = stack(std(sol_rvs)) + stds = sqrt.(stack(diag.(sol_rvs.Σ))) if isnothing(vars) ribbon --> ribbon_width * stds diff --git a/src/solution_sampling.jl b/src/solution_sampling.jl index b0a509d39..ff7e3ec06 100644 --- a/src/solution_sampling.jl +++ b/src/solution_sampling.jl @@ -2,14 +2,14 @@ # Sampling from a solution ######################################################################################## """ - _rand(x::Gaussian{<:Vector,<:SRMatrix}, n::Integer=1) + _rand(x::Gaussian{<:Vector,<:PSDMatrix}, n::Integer=1) Sample from a Gaussian with a `ProbNumDiffEq.SquarerootMatrix` covariance. Uses the existing covariance square root to make the sampling more stable. """ function _rand(x::SRGaussian, n::Integer=1) m, C = x.μ, x.Σ - sample = m .+ C.squareroot * randn(length(m), n) + sample = m .+ C.R' * randn(size(C.R, 1), n) return sample end @@ -51,7 +51,7 @@ function sample_states(ts, xs, diffusions, difftimes, posterior, n::Int=1) x_prev_p = P * xs[i] prev_sample_p, _ = - smooth(x_prev_p, Gaussian(sample_p, SRMatrix(zeros(D, D))), A, Qh) + smooth(x_prev_p, Gaussian(sample_p, PSDMatrix(zeros(D, D))), A, Qh) # sample_path[i, :, j] .= PI*prev_sample_p.μ sample_path[i, :, j] .= PI * _rand(prev_sample_p)[:] diff --git a/src/squarerootmatrix.jl b/src/squarerootmatrix.jl deleted file mode 100644 index 6460663af..000000000 --- a/src/squarerootmatrix.jl +++ /dev/null @@ -1,50 +0,0 @@ -""" -A type for positive semi-definite matrices. - -Relates to PSDMatrices.jl, and I might move this to PSDMatrices.jl in the -future, but having this here allowed for easier development. -""" - -struct SquarerootMatrix{T<:Real,S<:AbstractMatrix,M<:AbstractMatrix} <: AbstractMatrix{T} - squareroot::S - mat::M - SquarerootMatrix(S::AbstractMatrix{T}, mat::AbstractMatrix{T}) where {T} = - new{T,typeof(S),typeof(mat)}(S, mat) -end -SquarerootMatrix(S) = SquarerootMatrix(S, S * S') - -Base.Matrix(M::SquarerootMatrix) = M.mat -# TODO Maybe cache the above, into M.mat or something. But do so lazily! - -Base.size(M::SquarerootMatrix) = size(M.mat) -# getindex is expensive, since each call performs a matrix multiplication -Base.getindex(M::SquarerootMatrix, I::Vararg{Int,N}) where {N} = getindex(M.mat, I...) -Base.copy(M::SquarerootMatrix) = SquarerootMatrix( - # I want the copy to look exactly as the original one - M.squareroot isa LinearAlgebra.Adjoint ? copy(M.squareroot')' : copy(M.squareroot), - copy(M.mat), -) -Base.copy!(dst::SquarerootMatrix, src::SquarerootMatrix) = - (Base.copy!(dst.squareroot, src.squareroot); - Base.copy!(dst.mat, src.mat); - dst) -Base.similar(M::SquarerootMatrix) = SquarerootMatrix(similar(M.squareroot), similar(M.mat)) - -X_A_Xt(M::SquarerootMatrix, X::AbstractMatrix) = SquarerootMatrix(X * M.squareroot) -X_A_Xt!(out::SquarerootMatrix, M::SquarerootMatrix, X::AbstractMatrix) = begin - _matmul!(out.squareroot, X, M.squareroot) - _matmul!(out.mat, out.squareroot, out.squareroot') - return out -end -X_A_Xt!(out::SquarerootMatrix, M::SquarerootMatrix, D::Diagonal) = begin - # Basically just to optimize PI*Q*PI - out.squareroot .= D.diag .* M.squareroot - out.mat .= D.diag .* M.mat .* D.diag' - return out -end - -Base.inv(M::SquarerootMatrix) = Base.inv(M.mat) -LinearAlgebra.diag(M::SquarerootMatrix) = LinearAlgebra.diag(M.mat) -LinearAlgebra.cholesky(M::SquarerootMatrix) = LinearAlgebra.cholesky(M.mat) -LinearAlgebra.cholesky!(M::SquarerootMatrix) = LinearAlgebra.cholesky!(M.mat) -LinearAlgebra.logdet(M::SquarerootMatrix) = LinearAlgebra.logdet(M.mat) diff --git a/test/filtering.jl b/test/filtering.jl index 2f878972d..e36115190 100644 --- a/test/filtering.jl +++ b/test/filtering.jl @@ -10,12 +10,12 @@ using LinearAlgebra # Setup d = 5 m = rand(d) - L_p = Matrix(LowerTriangular(rand(d, d))) - P = L_p * L_p' + R_p = Matrix(LowerTriangular(rand(d, d))) + P = R_p'R_p A = rand(d, d) - L_Q = Matrix(LowerTriangular(rand(d, d))) - Q = L_Q * L_Q' + R_Q = Matrix(LowerTriangular(rand(d, d))) + Q = R_Q'R_Q # PREDICT m_p = A * m @@ -30,14 +30,13 @@ using LinearAlgebra @test P_p == x_out.Σ end - @testset "predict! with SRMatrix" begin - x_curr = Gaussian(m, SRMatrix(L_p)) + @testset "predict! with PSDMatrix" begin + x_curr = Gaussian(m, PSDMatrix(R_p)) x_out = copy(x_curr) - Q_SR = SRMatrix(L_Q) - cache = SRMatrix(zeros(d, 2d)) - ProbNumDiffEq.predict!(x_out, x_curr, A, Q_SR, cache) + Q_SR = PSDMatrix(R_Q) + ProbNumDiffEq.predict!(x_out, x_curr, A, Q_SR, zeros(d, d), zeros(2d, d)) @test m_p == x_out.μ - @test P_p ≈ x_out.Σ + @test P_p ≈ Matrix(x_out.Σ) end end @@ -45,14 +44,14 @@ end # Setup d = 5 m_p = rand(d) - L_P_p = Matrix(LowerTriangular(rand(d, d))) - P_p = L_P_p * L_P_p' + R_P_p = Matrix(LowerTriangular(rand(d, d))) + P_p = R_P_p'R_P_p # Measure o = 3 H = rand(o, d) - L_R = 0.0 * rand(o, o) - R = L_R * L_R' + R_R = 0.0 * rand(o, o) + R = R_R'R_R z_data = zeros(o) z = H * m_p @@ -80,7 +79,7 @@ end m_tmp = copy(measurement) ProbNumDiffEq.update!(x_out, x_pred, measurement, H, K_cache, M_cache, m_tmp.Σ) @test m ≈ x_out.μ - @test P ≈ x_out.Σ + @test P ≈ Matrix(x_out.Σ) end end @@ -88,13 +87,13 @@ end # Setup d = 5 m, m_s = rand(d), rand(d) - L_P, L_P_s = Matrix(LowerTriangular(rand(d, d))), Matrix(LowerTriangular(rand(d, d))) - P, P_s = L_P * L_P', L_P_s * L_P_s' + R_P, R_P_s = Matrix(LowerTriangular(rand(d, d))), Matrix(LowerTriangular(rand(d, d))) + P, P_s = R_P'R_P, R_P_s'R_P_s A = rand(d, d) - L_Q = Matrix(LowerTriangular(rand(d, d))) - Q = L_Q * L_Q' - Q_SR = SRMatrix(L_Q) + R_Q = Matrix(LowerTriangular(rand(d, d))) + Q = R_Q'R_Q + Q_SR = PSDMatrix(R_Q) # PREDICT first m_p = A * m @@ -105,19 +104,19 @@ end m_smoothed = m + G * (m_s - m_p) P_smoothed = P + G * (P_s - P_p) * G' - x_curr = Gaussian(m, SRMatrix(L_P)) - x_smoothed = Gaussian(m_s, SRMatrix(L_P_s)) + x_curr = Gaussian(m, P) + x_smoothed = Gaussian(m_s, P_s) @testset "smooth" begin x_out, _ = ProbNumDiffEq.smooth(x_curr, x_smoothed, A, Q) @test m_smoothed ≈ x_out.μ @test P_smoothed ≈ x_out.Σ end - @testset "smooth with SRMatrix" begin - x_curr = Gaussian(m, SRMatrix(L_P)) - x_smoothed = Gaussian(m_s, SRMatrix(L_P_s)) + @testset "smooth with PSDMatrix" begin + x_curr = Gaussian(m, PSDMatrix(R_P)) + x_smoothed = Gaussian(m_s, PSDMatrix(R_P_s)) x_out, _ = ProbNumDiffEq.smooth(x_curr, x_smoothed, A, Q_SR) @test m_smoothed ≈ x_out.μ - @test P_smoothed ≈ x_out.Σ + @test P_smoothed ≈ Matrix(x_out.Σ) end end diff --git a/test/preconditioning.jl b/test/preconditioning.jl index 4d2878b65..a4afc525b 100644 --- a/test/preconditioning.jl +++ b/test/preconditioning.jl @@ -22,18 +22,18 @@ prob = remake(prob, tspan=(0.0, 10.0)) A_p, Q_p = ProbNumDiffEq.ibm(d, q) Ah_p = A_p - Qh_p = Q_p * σ^2 + Qh_p = PSDMatrix(σ * Q_p.R) # First test that they're both equivalent D = d * (q + 1) P, PI = ProbNumDiffEq.init_preconditioner(d, q) ProbNumDiffEq.make_preconditioner!(P, h, d, q) ProbNumDiffEq.make_preconditioner_inv!(PI, h, d, q) - @test Qh_p ≈ P * Qh * P' + @test Matrix(Qh_p) ≈ P * Qh * P' @test Ah_p ≈ P * Ah * PI # Check that the preconditioning actually helps # @info "Condition numbers" cond(Qh) cond(Matrix(Qh_p)) cond(Ah) cond(Ah_p) - @test cond(Qh) > cond(Matrix(Qh_p)) - @test cond(Qh) > cond(Matrix(Qh_p))^2 + @test cond(Matrix(Qh)) > cond(Matrix(Qh_p)) + @test cond(Matrix(Qh)) > cond(Matrix(Qh_p))^2 end diff --git a/test/priors.jl b/test/priors.jl index ba628cffe..2f2e2a39f 100644 --- a/test/priors.jl +++ b/test/priors.jl @@ -44,7 +44,7 @@ end d, q = 2, 2 A, Q = ProbNumDiffEq.ibm(d, q) - Qh = Q * σ^2 + Qh = PSDMatrix(σ * Q.R) AH_22_PRE = [ 1 2 1 0 0 0 @@ -65,15 +65,15 @@ end 0 0 0 1/3 1/2 1/1 ] - @test AH_22_PRE ≈ A - @test QH_22_PRE ≈ Qh + @test AH_22_PRE ≈ Matrix(A) + @test QH_22_PRE ≈ Matrix(Qh) end @testset "Verify correct prior dim" begin prob = prob_ode_lotkavoltera d = length(prob.u0) for q in 1:5 - integ = init(prob, EK0(order=q), initialize_derivatives=false) + integ = init(prob, EK0(order=q)) @test length(integ.cache.x.μ) == d * (q + 1) sol = solve!(integ) @test length(integ.cache.x.μ) == d * (q + 1) diff --git a/test/solution.jl b/test/solution.jl index 2971e6c3a..58a58c550 100644 --- a/test/solution.jl +++ b/test/solution.jl @@ -37,7 +37,7 @@ import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_lotkavoltera @testset "Correct initial values" begin @test sol.pu[1].μ == prob.u0 - @test iszero(sol.pu[1].Σ) + @test iszero(sol.pu[1].Σ.R) end # Interpolation @@ -61,7 +61,6 @@ import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_lotkavoltera n_samples = 2 samples = ProbNumDiffEq.sample(sol, n_samples) - dsamples, dts = ProbNumDiffEq.dense_sample(sol, n_samples) @test samples isa Array @@ -70,11 +69,6 @@ import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_lotkavoltera @test n == length(sol.u[1]) @test o == n_samples - u = ProbNumDiffEq.stack(sol.u) - stds = sqrt.(ProbNumDiffEq.stack(diag.(sol.pu.Σ))) - outlier_count = sum(abs.(u .- samples) .> 3stds) - @test outlier_count < 0.05 * m * n * o - # Dense sampling dense_samples, dense_times = ProbNumDiffEq.dense_sample(sol, n_samples) m, n, o = size(dense_samples) @@ -87,7 +81,6 @@ import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_lotkavoltera n_samples = 2 samples = ProbNumDiffEq.sample_states(sol, n_samples) - dsamples, dts = ProbNumDiffEq.dense_sample_states(sol, n_samples) @test samples isa Array @@ -96,11 +89,6 @@ import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_lotkavoltera @test n == length(sol.u[1]) * (sol.interp.q + 1) @test o == n_samples - x = ProbNumDiffEq.stack(mean(sol.x_smooth)) - stds = sqrt.(ProbNumDiffEq.stack(diag.(sol.x_smooth.Σ))) - outlier_count = sum(abs.(x .- samples) .> 3stds) - @test outlier_count < 0.05 * m * n * o - # Dense sampling dense_samples, dense_times = ProbNumDiffEq.dense_sample_states(sol, n_samples) m, n, o = size(dense_samples) diff --git a/test/specific_problems.jl b/test/specific_problems.jl index 16aae271f..7c57592ad 100644 --- a/test/specific_problems.jl +++ b/test/specific_problems.jl @@ -103,10 +103,10 @@ end e = m0'm0 H = 2m0'E0 - S = H * P * H' + SR = P.R * H' + S = SR'SR - S_inv = inv(S) - K = P * H' * S_inv + K = P.R' * (P.R * (H' / S)) mnew = m + K * (2 .- e) Pnew = ProbNumDiffEq.X_A_Xt(P, (I - K * H)) # + X_A_Xt(R, K) diff --git a/test/state_init.jl b/test/state_init.jl index a03e980ac..58e2ae755 100644 --- a/test/state_init.jl +++ b/test/state_init.jl @@ -85,8 +85,9 @@ end _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)) + Cm = Matrix(C) + @assert isdiag(Cm) + whitened_err = err ./ sqrt.(diag(Cm)) @test all(abs.(whitened_err) .< 4e-1) end end