From 58c63eb178f164f3e08c26fce5f7ffaf9f665a42 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 08:23:37 +0100 Subject: [PATCH 01/99] Add a BlockDiagonal implementation --- src/ProbNumDiffEq.jl | 2 ++ src/blockdiagonals.jl | 48 +++++++++++++++++++++++++++++++++++++ src/caches.jl | 2 ++ src/covariance_structure.jl | 22 ++++++++++++++++- src/filtering/predict.jl | 23 ++++++++++++++++++ src/preconditioning.jl | 28 ++++++++++++++++++++++ src/priors/iwp.jl | 9 +++++++ src/projection.jl | 11 +++++++++ 8 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 src/blockdiagonals.jl diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index eac2a17cb..802c653b7 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -28,6 +28,7 @@ using ExponentialUtilities using Octavian using FastGaussQuadrature import Kronecker +using BlockDiagonals using ArrayAllocators using FiniteHorizonGramians using FillArrays @@ -53,6 +54,7 @@ cov2psdmatrix(cov::PSDMatrix; d) = (@assert size(cov, 1) == size(cov, 2) == d; c include("fast_linalg.jl") include("kronecker.jl") +include("blockdiagonals.jl") include("covariance_structure.jl") abstract type AbstractODEFilterCache <: OrdinaryDiffEq.OrdinaryDiffEqCache end diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl new file mode 100644 index 000000000..18b3ea2b2 --- /dev/null +++ b/src/blockdiagonals.jl @@ -0,0 +1,48 @@ + +_matmul!( + C::BlockDiagonal{T}, + A::BlockDiagonal{T}, + B::BlockDiagonal{T}, +) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i]) + end + return C +end + +_matmul!( + C::BlockDiagonal{T}, + A::BlockDiagonal{T}, + B::BlockDiagonal{T}, + alpha::Number, + beta::Number, +) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + end + return C +end + +_matmul!( + C::AbstractVector{T}, + A::BlockDiagonal{T}, + B::AbstractVector{T}, +) where {T<:LinearAlgebra.BlasFloat} = begin + @assert size(A, 2) == length(B) + @assert length(C) == size(A, 1) + ic, ib = 1, 1 + for i in eachindex(blocks(A)) + d1, d2 = size(A.blocks[i]) + @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) + ic += d1 + ib += d2 + end + return C +end + +function LinearAlgebra.cholesky!(B::BlockDiagonal) + C = BlockDiagonal(map(b -> parent(UpperTriangular(cholesky!(b).U)), blocks(B))) + return Cholesky(C, 'U', 0) +end diff --git a/src/caches.jl b/src/caches.jl index da8e191cd..ab3e41917 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -189,6 +189,8 @@ function OrdinaryDiffEq.alg_cache( else factorized_similar(FAC, D, d) end + elseif FAC isa BlockDiagonalCovariance + factorized_similar(FAC, D, d) else similar(Matrix{uElType}, D, _d) end, diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index 7390b5124..39d46126e 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -7,6 +7,10 @@ struct DenseCovariance{T} <: CovarianceStructure{T} d::Int64 q::Int64 end +struct BlockDiagonalCovariance{T} <: CovarianceStructure{T} + d::Int64 + q::Int64 +end function get_covariance_structure(alg; elType, d, q) if ( @@ -41,9 +45,25 @@ factorized_zeros(::DenseCovariance{T}, sizes...) where {T} = factorized_similar(::DenseCovariance{T}, size1, size2) where {T} = similar(Matrix{T}, size1, size2) +factorized_zeros(C::BlockDiagonalCovariance{T}, sizes...) where {T} = begin + for s in sizes + @assert s % C.d == 0 + end + return BlockDiagonal([Array{T}(calloc, (s ÷ C.d for s in sizes)...) for _ in 1:C.d]) +end +factorized_similar(C::BlockDiagonalCovariance{T}, size1, size2) where {T} = begin + for s in (size1, size2) + @assert s % C.d == 0 + end + return BlockDiagonal([similar(Matrix{T}, size1 ÷ C.d, size2 ÷ C.d) for _ in 1:C.d]) +end + to_factorized_matrix(::DenseCovariance, M::AbstractMatrix) = Matrix(M) to_factorized_matrix(::IsometricKroneckerCovariance, M::IsometricKroneckerProduct) = M -for FT in [:DenseCovariance, :IsometricKroneckerCovariance] +to_factorized_matrix(C::BlockDiagonalCovariance, M::IsometricKroneckerProduct) = + BlockDiagonal([M.B for _ in 1:C.d]) + +for FT in [:DenseCovariance, :IsometricKroneckerCovariance, :BlockDiagonalCovariance] @eval to_factorized_matrix(FAC::$FT, M::PSDMatrix) = PSDMatrix(to_factorized_matrix(FAC, M.R)) end diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 8738d74fa..630f91dca 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -113,3 +113,26 @@ function predict_cov!( return predict_cov!(_Σ_out, _Σ_curr, _Ah, _Qh, _C_DxD, _C_2DxD, _diffusion) end + +# BlockDiagonal version +function predict_cov!( + Σ_out::PSDMatrix{T,<:BlockDiagonal}, + Σ_curr::PSDMatrix{T,<:BlockDiagonal}, + Ah::BlockDiagonal, + Qh::PSDMatrix{S,<:BlockDiagonal}, + C_DxD::BlockDiagonal, + C_2DxD::BlockDiagonal, + diffusion=1, +) where {T,S} + for i in eachindex(blocks(Σ_out.R)) + predict_cov!( + PSDMatrix(Σ_out.R.blocks[i]), + PSDMatrix(Σ_curr.R.blocks[i]), + Ah.blocks[i], + PSDMatrix(Qh.R.blocks[i]), + C_DxD.blocks[i], + C_2DxD.blocks[i], + diffusion, + ) + end +end diff --git a/src/preconditioning.jl b/src/preconditioning.jl index 8dfd9fac8..371258f55 100644 --- a/src/preconditioning.jl +++ b/src/preconditioning.jl @@ -8,6 +8,11 @@ function init_preconditioner(C::DenseCovariance{elType}) where {elType} PI = kron(I(C.d), Diagonal(ones(elType, C.q + 1))) return P, PI end +function init_preconditioner(C::BlockDiagonalCovariance{elType}) where {elType} + P = BlockDiagonal([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) + PI = BlockDiagonal([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) + return P, PI +end function make_preconditioners!(cache, dt) @unpack P, PI, d, q = cache @@ -41,6 +46,17 @@ end return P end +@fastmath @inbounds function make_preconditioner!(P::BlockDiagonal, h, d, q) + val = factorial(q) / h^(q + 1 / 2) + @simd ivdep for j in 0:q + for M in P.blocks + M.diag[j+1] = val + end + val /= (q - j) / h + end + return P +end + @fastmath @inbounds function make_preconditioner_inv!(PI::Diagonal, h, d, q) val = h^(q + 1 / 2) / factorial(q) for j in 0:q @@ -62,3 +78,15 @@ end end return PI end + +@fastmath @inbounds function make_preconditioner_inv!( + PI::BlockDiagonal, h, d, q) + val = h^(q + 1 / 2) / factorial(q) + @simd ivdep for j in 0:q + for M in PI.blocks + M.diag[j+1] = val + end + val *= (q - j) / h + end + return PI +end diff --git a/src/priors/iwp.jl b/src/priors/iwp.jl index 3a7d7f311..b62d5d3ac 100644 --- a/src/priors/iwp.jl +++ b/src/priors/iwp.jl @@ -169,6 +169,15 @@ function initialize_transition_matrices(FAC::DenseCovariance, p::IWP, dt) Ah, Qh = copy(A), copy(Q) return A, Q, Ah, Qh, P, PI end +function initialize_transition_matrices(FAC::BlockDiagonalCovariance, p::IWP, dt) + A, Q = preconditioned_discretize(p) + A = to_factorized_matrix(FAC, A) + Q = to_factorized_matrix(FAC, Q) + P, PI = initialize_preconditioner(FAC, p, dt) + Ah = PI * A * P + Qh = PSDMatrix(Q.R * PI) + return A, Q, Ah, Qh, P, PI +end function make_transition_matrices!(cache, prior::IWP, dt) @unpack A, Q, Ah, Qh, P, PI = cache diff --git a/src/projection.jl b/src/projection.jl index 29097b253..1dfcb5d5a 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -30,6 +30,17 @@ function projection(C::IsometricKroneckerCovariance{elType}) where {elType} return Proj end +function projection(C::BlockDiagonalCovariance{elType}) where {elType} + Proj(deriv) = begin + e_i = zeros(elType, C.q + 1, 1) + if deriv <= C.q + e_i[deriv+1] = 1 + end + return BlockDiagonal([e_i' for _ in 1:C.d]) + end + return Proj +end + function solution_space_projection(C::CovarianceStructure, is_secondorder_ode) Proj = projection(C) if is_secondorder_ode From d0a3eb09b72769fcf15cdd8217a148c2017a63b0 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 09:09:53 +0100 Subject: [PATCH 02/99] It works and it's (a little bit) faster than dense! --- Project.toml | 1 + src/blockdiagonals.jl | 37 ++++++++++++++++++++++++++++++--- src/diffusions.jl | 7 +++++++ src/filtering/update.jl | 45 +++++++++++++++++++++++++++++++++++++++++ src/perform_step.jl | 4 ++++ 5 files changed, 91 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index ed5af8bb9..6b69f5705 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ version = "0.15.0" [deps] ArrayAllocators = "c9d4266f-a5cb-439d-837c-c97b191379f5" +BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 18b3ea2b2..7b2948c89 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -25,6 +25,30 @@ _matmul!( return C end +_matmul!( + C::BlockDiagonal{T}, + A::BlockDiagonal{T}, + B::Adjoint{T, <:BlockDiagonal{T}}, +) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + end + return C +end + +_matmul!( + C::BlockDiagonal{T}, + A::Adjoint{T, <:BlockDiagonal{T}}, + B::BlockDiagonal{T}, +) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + end + return C +end + _matmul!( C::AbstractVector{T}, A::BlockDiagonal{T}, @@ -42,7 +66,14 @@ _matmul!( return C end -function LinearAlgebra.cholesky!(B::BlockDiagonal) - C = BlockDiagonal(map(b -> parent(UpperTriangular(cholesky!(b).U)), blocks(B))) - return Cholesky(C, 'U', 0) +function BlockDiagonals.isequal_blocksizes(B1::BlockDiagonal, B2::BlockDiagonal) + @assert length(B1.blocks) == length(B2.blocks) + for i in eachindex(B1.blocks) + if size(B1.blocks[i]) != size(B2.blocks[i]) + return false + end + end + return true end + +LinearAlgebra.adjoint(B::BlockDiagonal) = Adjoint(B) diff --git a/src/diffusions.jl b/src/diffusions.jl index 3570a22f7..06cb800b7 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -166,6 +166,13 @@ function local_scalar_diffusion(cache) σ² = if HQH isa IsometricKroneckerProduct @assert length(HQH.B) == 1 dot(z, e) / d / HQH.B[1] + elseif HQH isa BlockDiagonal + @assert length(HQH.blocks) == d + @assert length(HQH.blocks[1]) == 1 + for i in eachindex(e) + e[i] /= HQH.blocks[i][1] + end + dot(z, e) / d else C = cholesky!(HQH) ldiv!(C, e) diff --git a/src/filtering/update.jl b/src/filtering/update.jl index 3d9125fcd..4a17e7cd4 100644 --- a/src/filtering/update.jl +++ b/src/filtering/update.jl @@ -193,6 +193,51 @@ function update!( return x_out, loglikelihood end + +function update!( + x_out::SRGaussian{T,<:BlockDiagonal}, + x_pred::SRGaussian{T,<:BlockDiagonal}, + measurement::Gaussian{ + <:AbstractVector, + <:Union{<:PSDMatrix{T,<:BlockDiagonal},<:BlockDiagonal}, + }, + H::BlockDiagonal, + K1_cache::BlockDiagonal, + K2_cache::BlockDiagonal, + M_cache::BlockDiagonal, + C_dxd::BlockDiagonal, + C_d::AbstractVector; + R::Union{Nothing,PSDMatrix{T,<:BlockDiagonal}}=nothing, +) where {T} + d = length(blocks(x_out.Σ.R)) + q = size(blocks(x_out.Σ.R)[1], 1) - 1 + + ll = zero(eltype(x_out.μ)) + for i in eachindex(blocks(x_out.Σ.R)) + _, _ll = update!( + Gaussian(view(x_out.μ, (i-1)*(q+1)+1:i*(q+1)), + PSDMatrix(x_out.Σ.R.blocks[i])), + Gaussian(view(x_pred.μ, (i-1)*(q+1)+1:i*(q+1)), + PSDMatrix(x_pred.Σ.R.blocks[i])), + Gaussian(view(measurement.μ, i:i), + if measurement.Σ isa PSDMatrix + PSDMatrix(measurement.Σ.R.blocks[i]) + else + measurement.Σ.blocks[i] + end), + H.blocks[i], + K1_cache.blocks[i], + K2_cache.blocks[i], + M_cache.blocks[i], + C_dxd.blocks[i], + view(C_d, i:i); + R, + ) + ll += _ll + end + return x_out, ll +end + # Short-hand with cache function update!(x_out, x, measurement, H; cache, R=nothing) @unpack K1, m_tmp, C_DxD, C_dxd, C_Dxd, C_d = cache diff --git a/src/perform_step.jl b/src/perform_step.jl index 29574b68a..621445ffc 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -241,6 +241,10 @@ function estimate_errors!(cache::AbstractODEFilterCache) error_estimate = view(cache.tmp, 1:d) if R isa IsometricKroneckerProduct error_estimate .= sum(abs2, R.B) + elseif R isa BlockDiagonal + for i in eachindex(blocks(R)) + error_estimate[i] = sum(abs2, R.blocks[i]) + end else sum!(abs2, error_estimate', view(R, :, 1:d)) end From 5519c314223ee05ec8f90d3fae5a6ea9df2d0747 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 11:41:24 +0100 Subject: [PATCH 03/99] Implement a first version of the DiagonalEK1 --- src/ProbNumDiffEq.jl | 2 +- src/alg_utils.jl | 14 ++++++------ src/algorithms.jl | 43 ++++++++++++++++++++++++++++++++++++- src/caches.jl | 9 ++++---- src/covariance_structure.jl | 2 ++ src/derivative_utils.jl | 18 +++++++++++++++- 6 files changed, 75 insertions(+), 13 deletions(-) diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index 802c653b7..d4c73da69 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -74,7 +74,7 @@ include("initialization/common.jl") export TaylorModeInit, ClassicSolverInit, SimpleInit, ForwardDiffInit include("algorithms.jl") -export EK0, EK1 +export EK0, EK1, DiagonalEK1 export ExpEK, RosenbrockExpEK include("alg_utils.jl") diff --git a/src/alg_utils.jl b/src/alg_utils.jl index a3c84736f..a845a0ad4 100644 --- a/src/alg_utils.jl +++ b/src/alg_utils.jl @@ -4,19 +4,21 @@ ############################################################################################ OrdinaryDiffEq._alg_autodiff(::AbstractEK) = Val{true}() -OrdinaryDiffEq._alg_autodiff(::EK1{CS,AD}) where {CS,AD} = Val{AD}() -OrdinaryDiffEq.alg_difftype(::EK1{CS,AD,DiffType}) where {CS,AD,DiffType} = DiffType OrdinaryDiffEq.standardtag(::AbstractEK) = false -OrdinaryDiffEq.standardtag(::EK1{CS,AD,DiffType,ST}) where {CS,AD,DiffType,ST} = ST OrdinaryDiffEq.concrete_jac(::AbstractEK) = nothing -OrdinaryDiffEq.concrete_jac(::EK1{CS,AD,DiffType,ST,CJ}) where {CS,AD,DiffType,ST,CJ} = CJ @inline DiffEqBase.get_tmp_cache(integ, alg::AbstractEK, cache::AbstractODEFilterCache) = (cache.tmp, cache.atmp) -OrdinaryDiffEq.get_chunksize(::EK1{CS}) where {CS} = Val(CS) OrdinaryDiffEq.isfsal(::AbstractEK) = false -OrdinaryDiffEq.isimplicit(::EK1) = true +for ALG in [:EK1, :DiagonalEK1] + @eval OrdinaryDiffEq._alg_autodiff(::$ALG{CS,AD}) where {CS,AD} = Val{AD}() + @eval OrdinaryDiffEq.alg_difftype(::$ALG{CS,AD,DiffType}) where {CS,AD,DiffType} = DiffType + @eval OrdinaryDiffEq.standardtag(::$ALG{CS,AD,DiffType,ST}) where {CS,AD,DiffType,ST} = ST + @eval OrdinaryDiffEq.concrete_jac(::$ALG{CS,AD,DiffType,ST,CJ}) where {CS,AD,DiffType,ST,CJ} = CJ + @eval OrdinaryDiffEq.get_chunksize(::$ALG{CS}) where {CS} = Val(CS) + @eval OrdinaryDiffEq.isimplicit(::$ALG) = true +end ############################################ # Step size control diff --git a/src/algorithms.jl b/src/algorithms.jl index 9982f0f5d..c084fcd42 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -157,6 +157,47 @@ struct EK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK end end +struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK + prior::PT + diffusionmodel::DT + smooth::Bool + initialization::IT + pn_observation_noise::RT + DiagonalEK1(; + order=3, + prior::PT=IWP(order), + diffusionmodel::DT=DynamicDiffusion(), + smooth=true, + initialization::IT=TaylorModeInit(num_derivatives(prior)), + chunk_size=Val{0}(), + autodiff=Val{true}(), + diff_type=Val{:forward}, + standardtag=Val{true}(), + concrete_jac=nothing, + pn_observation_noise::RT=nothing, + ) where {PT,DT,IT,RT} = begin + ekargcheck(DiagonalEK1; diffusionmodel, pn_observation_noise) + new{ + _unwrap_val(chunk_size), + _unwrap_val(autodiff), + diff_type, + _unwrap_val(standardtag), + _unwrap_val(concrete_jac), + PT, + DT, + IT, + RT, + }( + prior, + diffusionmodel, + smooth, + initialization, + pn_observation_noise, + ) + end +end + + """ ExpEK(; L, order=3, kwargs...) @@ -236,7 +277,7 @@ function DiffEqBase.remake(thing::EK1{CS,AD,DT,ST,CJ}; kwargs...) where {CS,AD,D ) end -function DiffEqBase.prepare_alg(alg::EK1{0}, u0::AbstractArray{T}, p, prob) where {T} +function DiffEqBase.prepare_alg(alg::Union{EK1{0},DiagonalEK1{0}}, u0::AbstractArray{T}, p, prob) where {T} # See OrdinaryDiffEq.jl: ./src/alg_utils.jl (where this is copied from). # In the future we might want to make EK1 an OrdinaryDiffEqAdaptiveImplicitAlgorithm and # use the prepare_alg from OrdinaryDiffEq; but right now, we do not use `linsolve` which diff --git a/src/caches.jl b/src/caches.jl index ab3e41917..93bd4920e 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -5,7 +5,7 @@ mutable struct EKCache{ RType,CFacType,ProjType,SolProjType,PType,PIType,EType,uType,duType,xType,PriorType, AType,QType, FType,LType,FHGMethodType,FHGCacheType, - HType,vecType,matType,bkType,diffusionType,diffModelType,measModType,measType, + HType,vecType,dduType,matType,bkType,diffusionType,diffModelType,measModType,measType, puType,llType,dtType,rateType,UF,JC,uNoUnitsType, } <: AbstractODEFilterCache # Constants @@ -49,7 +49,7 @@ mutable struct EKCache{ pu_tmp::puType H::HType du::duType - ddu::matType + ddu::dduType K1::matType G1::matType Smat::HType @@ -178,7 +178,8 @@ function OrdinaryDiffEq.alg_cache( # Caches du = is_secondorder_ode ? similar(u.x[2]) : similar(u) - ddu = factorized_similar(FAC, length(u), length(u)) + # ddu = factorized_similar(FAC, length(u), length(u)) + ddu = similar(u, length(u), length(u)) _d = is_secondorder_ode ? 2d : d pu_tmp = Gaussian( similar(Array{uElType}, _d), @@ -242,7 +243,7 @@ function OrdinaryDiffEq.alg_cache( typeof(R),typeof(FAC),typeof(Proj),typeof(SolProj),typeof(P),typeof(PI),typeof(E0), uType,typeof(du),typeof(x0),typeof(prior),typeof(A),typeof(Q), typeof(F),typeof(L),typeof(FHG_method),typeof(FHG_cache), - typeof(H),typeof(C_d),matType, + typeof(H),typeof(C_d),typeof(ddu),matType, typeof(backward_kernel),typeof(initdiff), typeof(diffmodel),typeof(measurement_model),typeof(measurement),typeof(pu_tmp), uEltypeNoUnits,typeof(dt),typeof(du1),typeof(uf),typeof(jac_config),typeof(atmp), diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index 39d46126e..46955c4d5 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -22,6 +22,8 @@ function get_covariance_structure(alg; elType, d, q) alg.prior isa IWP ) return IsometricKroneckerCovariance{elType}(d, q) + elseif alg isa DiagonalEK1 + return BlockDiagonalCovariance{elType}(d, q) else return DenseCovariance{elType}(d, q) end diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index 9b84fb1c8..bd08e409c 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -8,7 +8,23 @@ function calc_H!(H, integ, cache) calc_H_EK0!(H, integ, cache) # @assert integ.u == @view x_pred.μ[1:(q+1):end] OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) - ProbNumDiffEq._matmul!(H, view(ddu, 1:d, :), cache.SolProj, -1.0, 1.0) + _matmul!(H, view(ddu, 1:d, :), cache.SolProj, -1.0, 1.0) + elseif integ.alg isa DiagonalEK1 + calc_H_EK0!(H, integ, cache) + # @assert integ.u == @view x_pred.μ[1:(q+1):end] + # ddu_full = Matrix(ddu) + # @info "ddu" ddu_full + # error() + OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) + + @unpack C_dxd = cache + @simd ivdep for i in eachindex(blocks(C_dxd)) + @assert length(C_dxd.blocks[i]) == 1 + C_dxd.blocks[i][1] = ddu[i, i] + end + _matmul!(H, C_dxd, cache.SolProj, -1.0, 1.0) + else + error("Unknown algorithm") end return nothing end From c81e5cac9dea1c053066e45642592cdcea8e0724 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 12:33:20 +0100 Subject: [PATCH 04/99] Added smoothing --- src/filtering/markov_kernel.jl | 71 ++++++++++++++++++++++++++++++++++ src/filtering/predict.jl | 1 + 2 files changed, 72 insertions(+) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index e980108dd..1884430a2 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -120,6 +120,28 @@ function marginalize_cov!( return marginalize_cov!(_Σ_out, _Σ_curr, _K; C_DxD=_C_DxD, C_3DxD=_C_3DxD) end +function marginalize_cov!( + Σ_out::PSDMatrix{T,<:BlockDiagonal}, + Σ_curr::PSDMatrix{T,<:BlockDiagonal}, + K::AffineNormalKernel{ + <:AbstractMatrix, + <:Any, + <:PSDMatrix{S,<:BlockDiagonal}, + }; + C_DxD::AbstractMatrix, + C_3DxD::AbstractMatrix, +) where {T,S} + for i in eachindex(blocks(Σ_out.R)) + _Σ_out = PSDMatrix(Σ_out.R.blocks[i]) + _Σ_curr = PSDMatrix(Σ_curr.R.blocks[i]) + _K = AffineNormalKernel(K.A.blocks[i], K.b, PSDMatrix(K.C.R.blocks[i])) + _C_DxD = C_DxD.blocks[i] + _C_3DxD = C_3DxD.blocks[i] + marginalize_cov!(_Σ_out, _Σ_curr, _K; C_DxD=_C_DxD, C_3DxD=_C_3DxD) + end + return Σ_out +end + """ compute_backward_kernel!(Kout, xpred, x, K; C_DxD[, diffusion=1]) @@ -243,3 +265,52 @@ function compute_backward_kernel!( return compute_backward_kernel!( _Kout, _x_pred, _x, _K; C_DxD=_C_DxD, diffusion=diffusion) end + +function compute_backward_kernel!( + Kout::KT1, + xpred::SRGaussian{T,<:BlockDiagonal}, + x::SRGaussian{T,<:BlockDiagonal}, + K::KT2; + C_DxD::AbstractMatrix, + diffusion=1, +) where { + T, + KT1<:AffineNormalKernel{ + <:BlockDiagonal, + <:AbstractVector, + <:PSDMatrix{T,<:BlockDiagonal}, + }, + KT2<:AffineNormalKernel{ + <:BlockDiagonal, + <:Any, + <:PSDMatrix{T,<:BlockDiagonal}, + }, +} + d = length(blocks(xpred.Σ.R)) + q = size(blocks(xpred.Σ.R)[1], 1) - 1 + for i in eachindex(blocks(xpred.Σ.R)) + _Kout = AffineNormalKernel( + Kout.A.blocks[i], + view(Kout.b, (i-1)*(q+1)+1:i*(q+1)), + PSDMatrix(Kout.C.R.blocks[i]) + ) + _xpred = Gaussian( + view(xpred.μ, (i-1)*(q+1)+1:i*(q+1)), + PSDMatrix(xpred.Σ.R.blocks[i]) + ) + _x = Gaussian( + view(x.μ, (i-1)*(q+1)+1:i*(q+1)), + PSDMatrix(x.Σ.R.blocks[i]) + ) + _K = AffineNormalKernel( + K.A.blocks[i], + ismissing(K.b) ? missing : view(K.b, (i-1)*(q+1)+1:i*(q+1)), + PSDMatrix(K.C.R.blocks[i]) + ) + _C_DxD = C_DxD.blocks[i] + compute_backward_kernel!( + _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=diffusion + ) + end + return Kout +end diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 630f91dca..06bb67b57 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -135,4 +135,5 @@ function predict_cov!( diffusion, ) end + return Σ_out end From 05c93df870c70e1c8668d442dd93ae1d1e1f1763 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 17:22:42 +0100 Subject: [PATCH 05/99] Add BlockDiagonals to the tests --- test/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/test/Project.toml b/test/Project.toml index f0fff40ee..39a6ebb11 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,6 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" From bc6c372c2d5fcaa72f0131c2aba6b4ca3b26b96a Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 17:44:42 +0100 Subject: [PATCH 06/99] This should be the proper logic to choose the cov factorization --- src/algorithms.jl | 22 ++++++++++++++++++++++ src/caches.jl | 2 +- src/covariance_structure.jl | 17 ----------------- 3 files changed, 23 insertions(+), 18 deletions(-) diff --git a/src/algorithms.jl b/src/algorithms.jl index c084fcd42..f3a7a4f1f 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -23,6 +23,28 @@ function ekargcheck(alg; diffusionmodel, pn_observation_noise, kwargs...) end end +function covariance_structure(alg) + if alg isa EK0 + if alg.prior isa IWP + if (alg.diffusionmodel isa DynamicDiffusion || + alg.diffusionmodel isa FixedDiffusion) + return IsometricKroneckerCovariance + else + return BlockDiagonalCovariance + end + else + # This is not great as other priors can be Kronecker too; TODO + return DenseCovariance + end + elseif alg isa DiagonalEK1 + return BlockDiagonalCovariance + elseif alg isa EK1 + return DenseCovariance + else + throw(ArgumentError("Unknown algorithm type $alg")) + end +end + """ EK0(; order=3, smooth=true, diff --git a/src/caches.jl b/src/caches.jl index 93bd4920e..bec32760d 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -103,7 +103,7 @@ function OrdinaryDiffEq.alg_cache( # uElType = eltype(u_vec) uElType = uBottomEltypeNoUnits - FAC = get_covariance_structure(alg; elType=uElType, d, q) + FAC = covariance_structure(alg){uElType}(d, q) if FAC isa IsometricKroneckerCovariance && !(f.mass_matrix isa UniformScaling) throw( ArgumentError( diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index 46955c4d5..f4d7c88ce 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -12,23 +12,6 @@ struct BlockDiagonalCovariance{T} <: CovarianceStructure{T} q::Int64 end -function get_covariance_structure(alg; elType, d, q) - if ( - alg isa EK0 && - !( - alg.diffusionmodel isa DynamicMVDiffusion || - alg.diffusionmodel isa FixedMVDiffusion - ) && - alg.prior isa IWP - ) - return IsometricKroneckerCovariance{elType}(d, q) - elseif alg isa DiagonalEK1 - return BlockDiagonalCovariance{elType}(d, q) - else - return DenseCovariance{elType}(d, q) - end -end - factorized_zeros(C::IsometricKroneckerCovariance{T}, sizes...) where {T} = begin for s in sizes @assert s % C.d == 0 From f112c151c7e789819c9467ec770e8dc98fb69641 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 17:57:27 +0100 Subject: [PATCH 07/99] We can now select the covariance from the outside! --- src/ProbNumDiffEq.jl | 1 + src/algorithms.jl | 42 ++++++++++++++++++++++++++---------------- src/caches.jl | 2 +- 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index d4c73da69..7c90ac2ca 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -56,6 +56,7 @@ include("fast_linalg.jl") include("kronecker.jl") include("blockdiagonals.jl") include("covariance_structure.jl") +export IsometricKroneckerCovariance, DenseCovariance, BlockDiagonalCovariance abstract type AbstractODEFilterCache <: OrdinaryDiffEq.OrdinaryDiffEqCache end diff --git a/src/algorithms.jl b/src/algorithms.jl index f3a7a4f1f..37b87bb34 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -23,11 +23,10 @@ function ekargcheck(alg; diffusionmodel, pn_observation_noise, kwargs...) end end -function covariance_structure(alg) - if alg isa EK0 - if alg.prior isa IWP - if (alg.diffusionmodel isa DynamicDiffusion || - alg.diffusionmodel isa FixedDiffusion) +function covariance_structure(::Type{Alg}, prior, diffusionmodel) where {Alg<:AbstractEK} + if Alg <: EK0 + if prior isa IWP + if (diffusionmodel isa DynamicDiffusion || diffusionmodel isa FixedDiffusion) return IsometricKroneckerCovariance else return BlockDiagonalCovariance @@ -36,14 +35,15 @@ function covariance_structure(alg) # This is not great as other priors can be Kronecker too; TODO return DenseCovariance end - elseif alg isa DiagonalEK1 + elseif Alg <: DiagonalEK1 return BlockDiagonalCovariance - elseif alg isa EK1 + elseif Alg <: EK1 return DenseCovariance else - throw(ArgumentError("Unknown algorithm type $alg")) + throw(ArgumentError("Unknown algorithm type $Alg")) end end +covariance_structure(alg) = covariance_structure(typeof(alg), alg.prior, alg.diffusionmodel) """ EK0(; order=3, @@ -80,22 +80,24 @@ julia> solve(prob, EK0()) # [References](@ref references) """ -struct EK0{PT,DT,IT,RT} <: AbstractEK +struct EK0{PT,DT,IT,RT,CF} <: AbstractEK prior::PT diffusionmodel::DT smooth::Bool initialization::IT pn_observation_noise::RT + covariance_factorization::CF EK0(; order=3, prior::PT=IWP(order), diffusionmodel::DT=DynamicDiffusion(), smooth=true, initialization::IT=TaylorModeInit(num_derivatives(prior)), pn_observation_noise::RT=nothing, - ) where {PT,DT,IT,RT} = begin + covariance_factorization::CF=covariance_structure(EK0, prior, diffusionmodel), + ) where {PT,DT,IT,RT,CF} = begin ekargcheck(EK0; diffusionmodel, pn_observation_noise) - new{PT,DT,IT,RT}( - prior, diffusionmodel, smooth, initialization, pn_observation_noise) + new{PT,DT,IT,RT,CF}( + prior, diffusionmodel, smooth, initialization, pn_observation_noise, covariance_factorization) end end @@ -139,12 +141,13 @@ julia> solve(prob, EK1()) # [References](@ref references) """ -struct EK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK +struct EK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT,CF} <: AbstractEK prior::PT diffusionmodel::DT smooth::Bool initialization::IT pn_observation_noise::RT + covariance_factorization::CF EK1(; order=3, prior::PT=IWP(order), @@ -157,7 +160,8 @@ struct EK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK standardtag=Val{true}(), concrete_jac=nothing, pn_observation_noise::RT=nothing, - ) where {PT,DT,IT,RT} = begin + covariance_factorization::CF=covariance_structure(EK1, prior, diffusionmodel), + ) where {PT,DT,IT,RT,CF} = begin ekargcheck(EK1; diffusionmodel, pn_observation_noise) new{ _unwrap_val(chunk_size), @@ -169,22 +173,25 @@ struct EK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK DT, IT, RT, + CF, }( prior, diffusionmodel, smooth, initialization, pn_observation_noise, + covariance_factorization, ) end end -struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK +struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT,CF} <: AbstractEK prior::PT diffusionmodel::DT smooth::Bool initialization::IT pn_observation_noise::RT + covariance_factorization::CF DiagonalEK1(; order=3, prior::PT=IWP(order), @@ -197,7 +204,8 @@ struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK standardtag=Val{true}(), concrete_jac=nothing, pn_observation_noise::RT=nothing, - ) where {PT,DT,IT,RT} = begin + covariance_factorization::CF=covariance_structure(DiagonalEK1, prior, diffusionmodel), + ) where {PT,DT,IT,RT,CF} = begin ekargcheck(DiagonalEK1; diffusionmodel, pn_observation_noise) new{ _unwrap_val(chunk_size), @@ -209,12 +217,14 @@ struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT} <: AbstractEK DT, IT, RT, + CF, }( prior, diffusionmodel, smooth, initialization, pn_observation_noise, + covariance_factorization, ) end end diff --git a/src/caches.jl b/src/caches.jl index bec32760d..78977d313 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -103,7 +103,7 @@ function OrdinaryDiffEq.alg_cache( # uElType = eltype(u_vec) uElType = uBottomEltypeNoUnits - FAC = covariance_structure(alg){uElType}(d, q) + FAC = alg.covariance_factorization{uElType}(d, q) if FAC isa IsometricKroneckerCovariance && !(f.mass_matrix isa UniformScaling) throw( ArgumentError( From 1404ee755f012500433b9b6c35d127bb9fec37aa Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 18:01:56 +0100 Subject: [PATCH 08/99] Add some SIMD here and there --- src/filtering/markov_kernel.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index 1884430a2..f132e1a55 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -113,7 +113,7 @@ function marginalize_cov!( ) where {T,S} _Σ_out = PSDMatrix(Σ_out.R.B) _Σ_curr = PSDMatrix(Σ_curr.R.B) - _K = AffineNormalKernel(K.A.B, K.b, PSDMatrix(K.C.R.B)) + _K = AffineNormalKernel(K.A.B, nothing, PSDMatrix(K.C.R.B)) _D = size(_Σ_out, 1) _C_DxD = C_DxD.B _C_3DxD = C_3DxD.B @@ -131,10 +131,10 @@ function marginalize_cov!( C_DxD::AbstractMatrix, C_3DxD::AbstractMatrix, ) where {T,S} - for i in eachindex(blocks(Σ_out.R)) + @inbounds @simd ivdep for i in eachindex(blocks(Σ_out.R)) _Σ_out = PSDMatrix(Σ_out.R.blocks[i]) _Σ_curr = PSDMatrix(Σ_curr.R.blocks[i]) - _K = AffineNormalKernel(K.A.blocks[i], K.b, PSDMatrix(K.C.R.blocks[i])) + _K = AffineNormalKernel(K.A.blocks[i], nothing, PSDMatrix(K.C.R.blocks[i])) _C_DxD = C_DxD.blocks[i] _C_3DxD = C_3DxD.blocks[i] marginalize_cov!(_Σ_out, _Σ_curr, _K; C_DxD=_C_DxD, C_3DxD=_C_3DxD) @@ -288,7 +288,7 @@ function compute_backward_kernel!( } d = length(blocks(xpred.Σ.R)) q = size(blocks(xpred.Σ.R)[1], 1) - 1 - for i in eachindex(blocks(xpred.Σ.R)) + @inbounds @simd ivdep for i in eachindex(blocks(xpred.Σ.R)) _Kout = AffineNormalKernel( Kout.A.blocks[i], view(Kout.b, (i-1)*(q+1)+1:i*(q+1)), From ff81928949e93d9012e76deb20ea8379a4d876b5 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 19:51:27 +0100 Subject: [PATCH 09/99] Make views of BlockDiagonals illegal as they are super slow --- src/blockdiagonals.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 7b2948c89..f54dd8cd2 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -1,3 +1,5 @@ +Base.view(::BlockDiagonal, idxs...) = + throw(MethodError("BlockDiagonal does not support views")) _matmul!( C::BlockDiagonal{T}, From 2bb610e469ba7afc37ef279131828747e852a603 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 22:11:36 +0100 Subject: [PATCH 10/99] Change how the diffusions work --- src/diffusions.jl | 31 +++++++++++++++++++++++-------- src/filtering/predict.jl | 10 +++++++++- src/perform_step.jl | 12 ++++++------ 3 files changed, 38 insertions(+), 15 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index 06cb800b7..7a745ee28 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -6,8 +6,23 @@ isdynamic(diffusion::AbstractStaticDiffusion) = false isstatic(diffusion::AbstractDynamicDiffusion) = false isdynamic(diffusion::AbstractDynamicDiffusion) = true -apply_diffusion(Q::PSDMatrix, diffusion::Diagonal) = X_A_Xt(Q, sqrt.(diffusion)) -apply_diffusion(Q::PSDMatrix, diffusion::Number) = PSDMatrix(Q.R * sqrt.(diffusion)) +apply_diffusion(Q::PSDMatrix{T, <:Matrix}, diffusion::Diagonal) where {T} = begin + d = size(diffusion, 1) + q = size(Q, 1) ÷ d - 1 + return PSDMatrix(Q.R * sqrt.(kron(diffusion, I(q+1)))) +end +apply_diffusion( + Q::PSDMatrix{T, <:IsometricKroneckerProduct}, + diffusion::Diagonal{T, <:FillArrays.Fill}, +) where {T} = begin + PSDMatrix(Q.R * sqrt.(diffusion.diag.value)) +end +apply_diffusion(Q::PSDMatrix{T, <:BlockDiagonal}, diffusion::Diagonal) where {T} = begin + PSDMatrix(BlockDiagonal([ + Q.R.blocks[i] * sqrt.(diffusion.diag[i]) for i in eachindex(Q.R.blocks) + ])) +end + estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = NaN @@ -20,7 +35,7 @@ Time-varying, isotropic diffusion, which is quasi-maximum-likelihood-estimated a particular also when solving stiff systems. """ struct DynamicDiffusion <: AbstractDynamicDiffusion end -initial_diffusion(::DynamicDiffusion, d, q, Eltype) = one(Eltype) +initial_diffusion(::DynamicDiffusion, d, q, Eltype) = one(Eltype) * Eye(d) estimate_local_diffusion(::DynamicDiffusion, integ) = local_scalar_diffusion(integ.cache) """ @@ -39,8 +54,7 @@ separately. * [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) """ struct DynamicMVDiffusion <: AbstractDynamicDiffusion end -initial_diffusion(::DynamicMVDiffusion, d, q, Eltype) = - kron(Diagonal(ones(Eltype, d)), Diagonal(ones(Eltype, q + 1))) +initial_diffusion(::DynamicMVDiffusion, d, q, Eltype) = Diagonal(ones(Eltype, d)) estimate_local_diffusion(::DynamicMVDiffusion, integ) = local_diagonal_diffusion(integ.cache) @@ -61,7 +75,7 @@ Base.@kwdef struct FixedDiffusion{T<:Number} <: AbstractStaticDiffusion calibrate::Bool = true end initial_diffusion(diffusionmodel::FixedDiffusion, d, q, Eltype) = - diffusionmodel.initial_diffusion * one(Eltype) + diffusionmodel.initial_diffusion * one(Eltype) * Eye(d) estimate_local_diffusion(::FixedDiffusion, integ) = local_scalar_diffusion(integ.cache) function estimate_global_diffusion(::FixedDiffusion, integ) @unpack d, measurement, m_tmp, Smat = integ.cache @@ -115,7 +129,7 @@ end function initial_diffusion(diffusionmodel::FixedMVDiffusion, d, q, Eltype) initdiff = diffusionmodel.initial_diffusion @assert initdiff isa Number || length(initdiff) == d - return kron(Diagonal(initdiff .* ones(Eltype, d)), Diagonal(ones(Eltype, q + 1))) + return Diagonal(initdiff .* ones(Eltype, d)) end estimate_local_diffusion(::FixedMVDiffusion, integ) = local_diagonal_diffusion(integ.cache) function estimate_global_diffusion(::FixedMVDiffusion, integ) @@ -178,7 +192,8 @@ function local_scalar_diffusion(cache) ldiv!(C, e) dot(z, e) / d end - return σ² + cache.local_diffusion = σ² * Eye(d) + return cache.local_diffusion end """ diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 06bb67b57..5566c7ed3 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -77,7 +77,15 @@ function predict_cov!( _matmul!(view(R, 1:D, 1:D), Σ_curr.R, Ah') if !isone(diffusion) - _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion)) + if diffusion isa Number + _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion)) + elseif diffusion isa Diagonal{<:Number, <:FillArrays.Fill} + _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) + elseif diffusion isa Diagonal{<:Number, <:Vector} + _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion, I(q+1))) + else + error() + end else @.. R[D+1:2D, 1:D] = Qh.R end diff --git a/src/perform_step.jl b/src/perform_step.jl index 621445ffc..8ec0ff690 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -218,18 +218,18 @@ To save allocations, the function modifies the given `cache` and writes into `cache.C_Dxd` during some computations. """ function estimate_errors!(cache::AbstractODEFilterCache) - @unpack local_diffusion, Qh, H, d = cache + @unpack local_diffusion, Qh, H, d, q = cache R = cache.C_Dxd - if local_diffusion isa Diagonal - _QR = cache.C_DxD .= Qh.R .* sqrt.(local_diffusion.diag)' - _matmul!(R, _QR, H') + if local_diffusion isa Diagonal{<:Number, <:Vector} + _Q = apply_diffusion(Qh, local_diffusion) + _matmul!(R, _Q.R, H') error_estimate = view(cache.tmp, 1:d) sum!(abs2, error_estimate', view(R, :, 1:d)) error_estimate .= sqrt.(error_estimate) return error_estimate - elseif local_diffusion isa Number + elseif local_diffusion isa Diagonal{<:Number,<:FillArrays.Fill} _matmul!(R, Qh.R, H') # error_estimate = diag(PSDMatrix(R)) @@ -248,7 +248,7 @@ function estimate_errors!(cache::AbstractODEFilterCache) else sum!(abs2, error_estimate', view(R, :, 1:d)) end - error_estimate .*= local_diffusion + error_estimate .*= local_diffusion.diag.value error_estimate .= sqrt.(error_estimate) return error_estimate From d7a49c663b38764222d2b77a902d4f74b8b7a492 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Tue, 13 Feb 2024 22:18:56 +0100 Subject: [PATCH 11/99] Make some more diffusions work --- src/diffusions.jl | 20 ++++++++------------ src/filtering/predict.jl | 4 +++- src/perform_step.jl | 8 +++++++- 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index 7a745ee28..afe6f7ea7 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -218,18 +218,14 @@ function local_diagonal_diffusion(cache) # HQH = H * unfactorize(Qh) * H' # @assert HQH |> diag |> unique |> length == 1 # c1 = view(_matmul!(cache.C_Dxd, Qh.R, H'), :, 1) - c1 = mul!(view(cache.C_Dxd, :, 1:1), Qh.R, view(H, 1:1, :)') - Q0_11 = dot(c1, c1) - - Σ_ii = @. m_tmp.μ = z^2 / Q0_11 - Σ = Diagonal(Σ_ii) - - # local_diffusion = kron(Σ, I(q+1)) - # -> Different for each dimension; same for each derivative - for i in 1:d - for j in (i-1)*(q+1)+1:i*(q+1) - local_diffusion[j, j] = Σ[i, i] - end + Q0_11 = if Qh.R isa BlockDiagonal + c1 = mul!(view(cache.C_Dxd.blocks[1], :, 1:1), Qh.R.blocks[1], view(H.blocks[1], 1:1, :)') + dot(c1, c1) + else + c1 = mul!(view(cache.C_Dxd, :, 1:1), Qh.R, view(H, 1:1, :)') + dot(c1, c1) end + + @. local_diffusion.diag = z^2 / Q0_11 return local_diffusion end diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 5566c7ed3..37da3a6ab 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -82,7 +82,9 @@ function predict_cov!( elseif diffusion isa Diagonal{<:Number, <:FillArrays.Fill} _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) elseif diffusion isa Diagonal{<:Number, <:Vector} - _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion, I(q+1))) + d = size(diffusion, 1) + q = size(Qh, 1) ÷ d - 1 + _matmul!(view(R, D+1:2D, 1:D), Qh.R, kron(sqrt.(diffusion), I(q+1))) else error() end diff --git a/src/perform_step.jl b/src/perform_step.jl index 8ec0ff690..8ebf7758f 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -226,7 +226,13 @@ function estimate_errors!(cache::AbstractODEFilterCache) _Q = apply_diffusion(Qh, local_diffusion) _matmul!(R, _Q.R, H') error_estimate = view(cache.tmp, 1:d) - sum!(abs2, error_estimate', view(R, :, 1:d)) + if R isa BlockDiagonal + for i in eachindex(R.blocks) + error_estimate[i] = sum(abs2, R.blocks[i]) + end + else + sum!(abs2, error_estimate', view(R, :, 1:d)) + end error_estimate .= sqrt.(error_estimate) return error_estimate elseif local_diffusion isa Diagonal{<:Number,<:FillArrays.Fill} From 473a140ec0b4b5c46c05fee37d590ba610cb2e1b Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:38:12 +0100 Subject: [PATCH 12/99] Better checking for validity of algorithm arguments --- src/algorithms.jl | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/algorithms.jl b/src/algorithms.jl index 37b87bb34..0dbcf654a 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -12,14 +12,20 @@ function ekargcheck(alg; diffusionmodel, pn_observation_noise, kwargs...) ), ) end - if ( - (diffusionmodel isa FixedMVDiffusion && diffusionmodel.calibrate) || - diffusionmodel isa DynamicMVDiffusion) && alg == EK1 + if alg == EK1 + if diffusionmodel isa FixedMVDiffusion && diffusionmodel.calibrate throw( ArgumentError( - "The `EK1` algorithm does not support automatic calibration of multivariate diffusion models. Either use the `EK0` instead, or use a scalar diffusion model, or set `calibrate=false` and calibrate manually by optimizing `sol.pnstats.log_likelihood`.", + "The `EK1` algorithm does not support automatic global calibration of multivariate diffusion models. Either use a scalar diffusion model, or set `calibrate=false` and calibrate manually by optimizing `sol.pnstats.log_likelihood`. Or use a different solve, like `EK0` or `DiagonalEK1`.", ), ) + elseif diffusionmodel isa DynamicMVDiffusion + throw( + ArgumentError( + "The `EK1` algorithm does not support automatic calibration of local multivariate diffusion models. Either use a scalar diffusion model, or use a different solve, like `EK0` or `DiagonalEK1`.", + ), + ) + end end end From 051a36f37395a3ad2ce130c389ba9ddc141986e7 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:39:05 +0100 Subject: [PATCH 13/99] More BlockDiagonal linalg things --- src/blockdiagonals.jl | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index f54dd8cd2..2d8deb5b8 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -78,4 +78,11 @@ function BlockDiagonals.isequal_blocksizes(B1::BlockDiagonal, B2::BlockDiagonal) return true end +LinearAlgebra.rmul!(B::BlockDiagonal, n::Number) = @simd ivdep for i in eachindex(B.blocks) + rmul!(B.blocks[i], n) +end LinearAlgebra.adjoint(B::BlockDiagonal) = Adjoint(B) +Base.:*(A::Adjoint{T, <:BlockDiagonal}, B::BlockDiagonal) where {T} = begin + @assert length(A.parent.blocks) == length(B.blocks) + return BlockDiagonal([A.parent.blocks[i]' * B.blocks[i] for i in eachindex(B.blocks)]) +end From 051689d99f793456f9cb4fc9c152357c4c8323f7 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:39:44 +0100 Subject: [PATCH 14/99] Better handling of the diffusion for prediction --- src/filtering/predict.jl | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 37da3a6ab..6652d43ab 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -66,7 +66,7 @@ function predict_cov!( Qh::PSDMatrix, C_DxD::AbstractMatrix, C_2DxD::AbstractMatrix, - diffusion=1, + diffusion::Union{Number, Diagonal}, ) if iszero(diffusion) fast_X_A_Xt!(Σ_out, Σ_curr, Ah) @@ -81,10 +81,6 @@ function predict_cov!( _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion)) elseif diffusion isa Diagonal{<:Number, <:FillArrays.Fill} _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) - elseif diffusion isa Diagonal{<:Number, <:Vector} - d = size(diffusion, 1) - q = size(Qh, 1) ÷ d - 1 - _matmul!(view(R, D+1:2D, 1:D), Qh.R, kron(sqrt.(diffusion), I(q+1))) else error() end @@ -111,7 +107,7 @@ function predict_cov!( Qh::PSDMatrix{S,<:IsometricKroneckerProduct}, C_DxD::IsometricKroneckerProduct, C_2DxD::IsometricKroneckerProduct, - diffusion=1, + diffusion::Diagonal, ) where {T,S} _Σ_out = PSDMatrix(Σ_out.R.B) _Σ_curr = PSDMatrix(Σ_curr.R.B) @@ -132,7 +128,7 @@ function predict_cov!( Qh::PSDMatrix{S,<:BlockDiagonal}, C_DxD::BlockDiagonal, C_2DxD::BlockDiagonal, - diffusion=1, + diffusion::Diagonal, ) where {T,S} for i in eachindex(blocks(Σ_out.R)) predict_cov!( @@ -142,7 +138,7 @@ function predict_cov!( PSDMatrix(Qh.R.blocks[i]), C_DxD.blocks[i], C_2DxD.blocks[i], - diffusion, + diffusion.diag[i], ) end return Σ_out From 390c46865dbf2622954ed65ed92c50cdb3dc9a91 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:40:18 +0100 Subject: [PATCH 15/99] The global diffusion is now written into the cache directly --- src/perform_step.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/perform_step.jl b/src/perform_step.jl index 8ebf7758f..24653b8a5 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -124,7 +124,7 @@ function OrdinaryDiffEq.perform_step!(integ, cache::EKCache, repeat_step=false) # Update the global diffusion MLE (if applicable) if !isdynamic(cache.diffusionmodel) - cache.global_diffusion = estimate_global_diffusion(cache.diffusionmodel, integ) + estimate_global_diffusion(cache.diffusionmodel, integ) end # Advance the state From 4c56635aefab2098acda1799eb7afe9eb712cad4 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:40:40 +0100 Subject: [PATCH 16/99] Implement rmul! or the IsometricKroneckerProduct --- src/kronecker.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/kronecker.jl b/src/kronecker.jl index 7c6d71a48..4e5c2bb87 100644 --- a/src/kronecker.jl +++ b/src/kronecker.jl @@ -62,6 +62,7 @@ end Base.:*(K::IKP, a::Number) = IsometricKroneckerProduct(K.ldim, K.B * a) Base.:*(a::Number, K::IKP) = IsometricKroneckerProduct(K.ldim, a * K.B) LinearAlgebra.adjoint(A::IKP) = IsometricKroneckerProduct(A.ldim, A.B') +LinearAlgebra.rmul!(A::IKP, b::Number) = IsometricKroneckerProduct(A.ldim, rmul!(A.B, b)) function check_same_size(A::IKP, B::IKP) if A.ldim != B.ldim || size(A.B) != size(B.B) From dff3e05607f5d4044bc5a5b13829511a905fe225 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:41:45 +0100 Subject: [PATCH 17/99] Properly ply the new diffusion after the solve --- src/diffusions.jl | 13 +++++++++++++ src/integrator_utils.jl | 25 ++++++++++++++----------- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index afe6f7ea7..d13c7fb2a 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -23,6 +23,19 @@ apply_diffusion(Q::PSDMatrix{T, <:BlockDiagonal}, diffusion::Diagonal) where {T} ])) end +apply_diffusion!(Q::PSDMatrix, diffusion::Diagonal{T, <:FillArrays.Fill}) where {T} = + rmul!(Q.R, sqrt.(diffusion.diag.value)) +apply_diffusion!( + Q::PSDMatrix{T,<:BlockDiagonal}, + diffusion::Diagonal{T,<:Vector}, +) where {T} = + @simd ivdep for i in eachindex(blocks(Q.R)) + rmul!(blocks(Q.R)[i], diffusion.diag[i]) + end + +apply_diffusion!(out::PSDMatrix, Q::PSDMatrix, diffusion::Diagonal{T,<:FillArrays.Fill}) where {T} = + rmul!(Q.R, sqrt.(diffusion.diag.value)) + estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = NaN diff --git a/src/integrator_utils.jl b/src/integrator_utils.jl index e13c9e4b8..59507998a 100644 --- a/src/integrator_utils.jl +++ b/src/integrator_utils.jl @@ -47,13 +47,12 @@ function calibrate_solution!(integ, mle_diffusion) set_diffusions!(integ.sol, mle_diffusion * integ.cache.default_diffusion) # Rescale all filtering estimates to have the correct diffusion - @assert mle_diffusion isa Number || mle_diffusion isa Diagonal - sqrt_diff = mle_diffusion isa Number ? sqrt(mle_diffusion) : sqrt.(mle_diffusion) + @assert mle_diffusion isa Diagonal @simd ivdep for C in integ.sol.x_filt.Σ - rmul!(C.R, sqrt_diff) + apply_diffusion!(C, mle_diffusion) end @simd ivdep for C in integ.sol.backward_kernels.C - rmul!(C.R, sqrt_diff) + apply_diffusion!(C, mle_diffusion) end # Re-write into the solution estimates @@ -70,13 +69,17 @@ Set the contents of `solution.diffusions` to the provided `diffusion`, overwriti diffusion estimates that are in there. Typically, `diffusion` is either a global quasi-MLE or the specified initial diffusion value if no calibration is desired. """ -function set_diffusions!(solution::AbstractProbODESolution, diffusion::Number) - solution.diffusions .= diffusion - return nothing -end -function set_diffusions!(solution::AbstractProbODESolution, diffusion::Diagonal) - @simd ivdep for d in solution.diffusions - copy!(d, diffusion) +function set_diffusions!(solution::AbstractProbODESolution, diffusion) + if diffusion isa Diagonal{<:Number, <:FillArrays.Fill} + @simd ivdep for i in eachindex(solution.diffusions) + solution.diffusions[i] = copy(diffusion) + end + elseif diffusion isa Diagonal{<:Number, <:Vector} + @simd ivdep for d in solution.diffusions + copy!(d, diffusion) + end + else + throw(ArgumentError("unexpected diffusion type $(typeof(diffusion))")) end return nothing end From 27912d593e900cdde5f063b954e84d6f4fedcff9 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:44:01 +0100 Subject: [PATCH 18/99] Properly estimate the global scalar diffusion --- src/diffusions.jl | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index d13c7fb2a..3ab5206c3 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -100,8 +100,15 @@ function estimate_global_diffusion(::FixedDiffusion, integ) diffusion_t = if S isa IsometricKroneckerProduct @assert length(S.B) == 1 dot(v, e) / d / S.B[1] + elseif S isa BlockDiagonal + @assert length(S.blocks) == d + @assert length(S.blocks[1]) == 1 + @simd ivdep for i in eachindex(e) + @inbounds e[i] /= S.blocks[i][1] + end + dot(v, e) / d else - S_chol = cholesky!(S) + S_chol = cholesky!(copy!(Smat, S)) ldiv!(S_chol, e) dot(v, e) / d end @@ -109,14 +116,16 @@ function estimate_global_diffusion(::FixedDiffusion, integ) if integ.success_iter == 0 # @assert length(sol_diffusions) == 0 global_diffusion = diffusion_t - return global_diffusion + integ.cache.global_diffusion = global_diffusion * Eye(d) + return integ.cache.global_diffusion else # @assert length(sol_diffusions) == integ.success_iter - diffusion_prev = integ.cache.global_diffusion + diffusion_prev = integ.cache.global_diffusion.diag.value global_diffusion = diffusion_prev + (diffusion_t - diffusion_prev) / integ.success_iter # @info "compute diffusion" diffusion_prev global_diffusion - return global_diffusion + integ.cache.global_diffusion = global_diffusion * Eye(d) + return integ.cache.global_diffusion end end From 3685edd61820c66b01a34893bc8b61d17208986c Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:48:53 +0100 Subject: [PATCH 19/99] Properly implement the global MV diffusion --- src/diffusions.jl | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index 3ab5206c3..1130d0bc3 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -37,7 +37,7 @@ apply_diffusion!(out::PSDMatrix, Q::PSDMatrix, diffusion::Diagonal{T,<:FillArray rmul!(Q.R, sqrt.(diffusion.diag.value)) -estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = NaN +estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = error() """ DynamicDiffusion() @@ -163,18 +163,15 @@ function estimate_global_diffusion(::FixedMVDiffusion, integ) Σ_ii = v .^ 2 ./ S_11 Σ = Diagonal(Σ_ii) - Σ_out = kron(Σ, I(q + 1)) # -> Different for each dimension; same for each derivative + Σ_out = Σ if integ.success_iter == 0 - # @assert length(diffusions) == 0 - return Σ_out + integ.cache.global_diffusion .= Σ_out + return integ.cache.global_diffusion else - # @assert length(diffusions) == integ.success_iter diffusion_prev = integ.cache.global_diffusion - diffusion = - @. diffusion_prev = - diffusion_prev + (Σ_out - diffusion_prev) / integ.success_iter - return diffusion + @.. diffusion_prev = diffusion_prev + (Σ_out - diffusion_prev) / integ.success_iter + return integ.cache.global_diffusion end end From a18495255d89cad89db66faa8c776e3496dbbd2d Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 13:50:49 +0100 Subject: [PATCH 20/99] This should be a proper implementation of the dynamic MV diffusion --- src/diffusions.jl | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index 1130d0bc3..3aa75b7a9 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -237,14 +237,24 @@ function local_diagonal_diffusion(cache) # HQH = H * unfactorize(Qh) * H' # @assert HQH |> diag |> unique |> length == 1 # c1 = view(_matmul!(cache.C_Dxd, Qh.R, H'), :, 1) - Q0_11 = if Qh.R isa BlockDiagonal - c1 = mul!(view(cache.C_Dxd.blocks[1], :, 1:1), Qh.R.blocks[1], view(H.blocks[1], 1:1, :)') - dot(c1, c1) - else - c1 = mul!(view(cache.C_Dxd, :, 1:1), Qh.R, view(H, 1:1, :)') - dot(c1, c1) + # Q_11 = dot(c1, c1) + + @assert Qh.R isa BlockDiagonal + for i in 1:d + c1 = _matmul!( + view(cache.C_Dxd.blocks[i], :, 1:1), + Qh.R.blocks[i], + view(H.blocks[i], 1:1, :)', + ) + tmp[i] = dot(c1, c1) end + Q_11 = tmp + + # To double-check: + HQH = H * unfactorize(Qh) * H' + @assert Q_11 ≈ diag(HQH) + + @. local_diffusion.diag = z^2 / Q_11 - @. local_diffusion.diag = z^2 / Q0_11 return local_diffusion end From e4fd99f76f79b82e86167e436937a73bacb9e762 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 14:20:36 +0100 Subject: [PATCH 21/99] Try to fix how the prediction handles the diffusion (I failed) --- src/filtering/predict.jl | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 6652d43ab..b8dacf3dd 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -78,11 +78,14 @@ function predict_cov!( _matmul!(view(R, 1:D, 1:D), Σ_curr.R, Ah') if !isone(diffusion) if diffusion isa Number - _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion)) - elseif diffusion isa Diagonal{<:Number, <:FillArrays.Fill} + _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt(diffusion)) + elseif diffusion isa Diagonal{<:Number,<:FillArrays.Fill} _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) else - error() + error("This is not yet implemented efficiently; TODO") + d = size(diffusion, 1) + q = D ÷ d - 1 + _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(kron(Eye(d)*diffusion, Eye(q + 1)))) end else @.. R[D+1:2D, 1:D] = Qh.R From 57e533d2f19abbd73eefc7ca6070ea3066cdabf5 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 14:25:34 +0100 Subject: [PATCH 22/99] Try to get the DynamicMV diff to work with BlockDiag cov (but fail) --- src/diffusions.jl | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index 3aa75b7a9..bd311538b 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -239,22 +239,28 @@ function local_diagonal_diffusion(cache) # c1 = view(_matmul!(cache.C_Dxd, Qh.R, H'), :, 1) # Q_11 = dot(c1, c1) - @assert Qh.R isa BlockDiagonal - for i in 1:d - c1 = _matmul!( - view(cache.C_Dxd.blocks[i], :, 1:1), - Qh.R.blocks[i], - view(H.blocks[i], 1:1, :)', - ) - tmp[i] = dot(c1, c1) + # @assert + Q_11 = if Qh.R isa BlockDiagonal + for i in 1:d + c1 = _matmul!( + view(cache.C_Dxd.blocks[i], :, 1:1), + Qh.R.blocks[i], + view(H.blocks[i], 1:1, :)', + ) + tmp[i] = dot(c1, c1) + end + tmp + else + error("This is not yet implemented efficiently; TODO") + diag(H * unfactorize(Qh) * H') end - Q_11 = tmp # To double-check: HQH = H * unfactorize(Qh) * H' @assert Q_11 ≈ diag(HQH) + # Also if the solver is a EK0 and not a DiagonalEK1: + # @assert Q_11 |> unique |> length == 1 @. local_diffusion.diag = z^2 / Q_11 - return local_diffusion end From 097ea9841af88fa9bb795018beecbe77e1d2cd2e Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 14:25:52 +0100 Subject: [PATCH 23/99] Get the DiagonalEK1 to work with a dense covariance factorization --- src/derivative_utils.jl | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index bd08e409c..6e840afe8 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -18,9 +18,13 @@ function calc_H!(H, integ, cache) OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) @unpack C_dxd = cache - @simd ivdep for i in eachindex(blocks(C_dxd)) - @assert length(C_dxd.blocks[i]) == 1 - C_dxd.blocks[i][1] = ddu[i, i] + if C_dxd isa BlockDiagonal + @simd ivdep for i in eachindex(blocks(C_dxd)) + @assert length(C_dxd.blocks[i]) == 1 + C_dxd.blocks[i][1] = ddu[i, i] + end + else + C_dxd .= Diagonal(ddu) end _matmul!(H, C_dxd, cache.SolProj, -1.0, 1.0) else From 895b6664ac568783917df60b120566d7129c18f3 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 14:29:52 +0100 Subject: [PATCH 24/99] Check diffusion and factorization compat somewhere else and warn instead of erroring --- src/algorithms.jl | 14 +++++++++++--- src/diffusions.jl | 2 +- src/filtering/predict.jl | 2 +- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/algorithms.jl b/src/algorithms.jl index 0dbcf654a..f9e33e222 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -3,7 +3,7 @@ ######################################################################################## abstract type AbstractEK <: OrdinaryDiffEq.OrdinaryDiffEqAdaptiveAlgorithm end -function ekargcheck(alg; diffusionmodel, pn_observation_noise, kwargs...) +function ekargcheck(alg; diffusionmodel, pn_observation_noise, covariance_factorization, kwargs...) if (isstatic(diffusionmodel) && diffusionmodel.calibrate) && (!isnothing(pn_observation_noise) && !iszero(pn_observation_noise)) throw( @@ -27,6 +27,13 @@ function ekargcheck(alg; diffusionmodel, pn_observation_noise, kwargs...) ) end end + if diffusionmodel isa DynamicMVDiffusion && covariance_factorization == BlockDiagonalCovariance + throw( + ArgumentError( + "Currenty the `DynamicMVDiffusion` does not work properly with the `BlockDiagonalCovariance`. Use `DenseCovariance` instead, or change the diffusionmodel to a scalar one and use `DynamicDiffusion`.", + ), + ) + end end function covariance_structure(::Type{Alg}, prior, diffusionmodel) where {Alg<:AbstractEK} @@ -38,6 +45,7 @@ function covariance_structure(::Type{Alg}, prior, diffusionmodel) where {Alg<:Ab return BlockDiagonalCovariance end else + error() # This is not great as other priors can be Kronecker too; TODO return DenseCovariance end @@ -168,7 +176,7 @@ struct EK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT,CF} <: AbstractEK pn_observation_noise::RT=nothing, covariance_factorization::CF=covariance_structure(EK1, prior, diffusionmodel), ) where {PT,DT,IT,RT,CF} = begin - ekargcheck(EK1; diffusionmodel, pn_observation_noise) + ekargcheck(EK1; diffusionmodel, pn_observation_noise, covariance_factorization) new{ _unwrap_val(chunk_size), _unwrap_val(autodiff), @@ -212,7 +220,7 @@ struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT,CF} <: AbstractEK pn_observation_noise::RT=nothing, covariance_factorization::CF=covariance_structure(DiagonalEK1, prior, diffusionmodel), ) where {PT,DT,IT,RT,CF} = begin - ekargcheck(DiagonalEK1; diffusionmodel, pn_observation_noise) + ekargcheck(DiagonalEK1; diffusionmodel, pn_observation_noise, covariance_factorization) new{ _unwrap_val(chunk_size), _unwrap_val(autodiff), diff --git a/src/diffusions.jl b/src/diffusions.jl index bd311538b..3b3842f9c 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -251,7 +251,7 @@ function local_diagonal_diffusion(cache) end tmp else - error("This is not yet implemented efficiently; TODO") + @warn "This is not yet implemented efficiently; TODO" diag(H * unfactorize(Qh) * H') end diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index b8dacf3dd..e55ddadad 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -82,7 +82,7 @@ function predict_cov!( elseif diffusion isa Diagonal{<:Number,<:FillArrays.Fill} _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) else - error("This is not yet implemented efficiently; TODO") + @warn "This is not yet implemented efficiently; TODO" d = size(diffusion, 1) q = D ÷ d - 1 _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(kron(Eye(d)*diffusion, Eye(q + 1)))) From efeaaadbb919a0c6e459ef0990eabcddccb6dece Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 14:32:00 +0100 Subject: [PATCH 25/99] JuliaFormatter.jl --- src/alg_utils.jl | 10 +++++--- src/algorithms.jl | 46 +++++++++++++++++++++++----------- src/blockdiagonals.jl | 6 ++--- src/diffusions.jl | 27 ++++++++++++-------- src/filtering/markov_kernel.jl | 10 ++++---- src/filtering/predict.jl | 8 ++++-- src/filtering/update.jl | 1 - src/integrator_utils.jl | 4 +-- src/perform_step.jl | 2 +- 9 files changed, 71 insertions(+), 43 deletions(-) diff --git a/src/alg_utils.jl b/src/alg_utils.jl index a845a0ad4..4fbbddd1b 100644 --- a/src/alg_utils.jl +++ b/src/alg_utils.jl @@ -13,9 +13,13 @@ OrdinaryDiffEq.isfsal(::AbstractEK) = false for ALG in [:EK1, :DiagonalEK1] @eval OrdinaryDiffEq._alg_autodiff(::$ALG{CS,AD}) where {CS,AD} = Val{AD}() - @eval OrdinaryDiffEq.alg_difftype(::$ALG{CS,AD,DiffType}) where {CS,AD,DiffType} = DiffType - @eval OrdinaryDiffEq.standardtag(::$ALG{CS,AD,DiffType,ST}) where {CS,AD,DiffType,ST} = ST - @eval OrdinaryDiffEq.concrete_jac(::$ALG{CS,AD,DiffType,ST,CJ}) where {CS,AD,DiffType,ST,CJ} = CJ + @eval OrdinaryDiffEq.alg_difftype(::$ALG{CS,AD,DiffType}) where {CS,AD,DiffType} = + DiffType + @eval OrdinaryDiffEq.standardtag(::$ALG{CS,AD,DiffType,ST}) where {CS,AD,DiffType,ST} = + ST + @eval OrdinaryDiffEq.concrete_jac( + ::$ALG{CS,AD,DiffType,ST,CJ}, + ) where {CS,AD,DiffType,ST,CJ} = CJ @eval OrdinaryDiffEq.get_chunksize(::$ALG{CS}) where {CS} = Val(CS) @eval OrdinaryDiffEq.isimplicit(::$ALG) = true end diff --git a/src/algorithms.jl b/src/algorithms.jl index f9e33e222..11a793339 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -3,7 +3,13 @@ ######################################################################################## abstract type AbstractEK <: OrdinaryDiffEq.OrdinaryDiffEqAdaptiveAlgorithm end -function ekargcheck(alg; diffusionmodel, pn_observation_noise, covariance_factorization, kwargs...) +function ekargcheck( + alg; + diffusionmodel, + pn_observation_noise, + covariance_factorization, + kwargs..., +) if (isstatic(diffusionmodel) && diffusionmodel.calibrate) && (!isnothing(pn_observation_noise) && !iszero(pn_observation_noise)) throw( @@ -14,20 +20,21 @@ function ekargcheck(alg; diffusionmodel, pn_observation_noise, covariance_factor end if alg == EK1 if diffusionmodel isa FixedMVDiffusion && diffusionmodel.calibrate - throw( - ArgumentError( - "The `EK1` algorithm does not support automatic global calibration of multivariate diffusion models. Either use a scalar diffusion model, or set `calibrate=false` and calibrate manually by optimizing `sol.pnstats.log_likelihood`. Or use a different solve, like `EK0` or `DiagonalEK1`.", - ), - ) + throw( + ArgumentError( + "The `EK1` algorithm does not support automatic global calibration of multivariate diffusion models. Either use a scalar diffusion model, or set `calibrate=false` and calibrate manually by optimizing `sol.pnstats.log_likelihood`. Or use a different solve, like `EK0` or `DiagonalEK1`.", + ), + ) elseif diffusionmodel isa DynamicMVDiffusion - throw( - ArgumentError( + throw( + ArgumentError( "The `EK1` algorithm does not support automatic calibration of local multivariate diffusion models. Either use a scalar diffusion model, or use a different solve, like `EK0` or `DiagonalEK1`.", - ), - ) + ), + ) end end - if diffusionmodel isa DynamicMVDiffusion && covariance_factorization == BlockDiagonalCovariance + if diffusionmodel isa DynamicMVDiffusion && + covariance_factorization == BlockDiagonalCovariance throw( ArgumentError( "Currenty the `DynamicMVDiffusion` does not work properly with the `BlockDiagonalCovariance`. Use `DenseCovariance` instead, or change the diffusionmodel to a scalar one and use `DynamicDiffusion`.", @@ -111,7 +118,8 @@ struct EK0{PT,DT,IT,RT,CF} <: AbstractEK ) where {PT,DT,IT,RT,CF} = begin ekargcheck(EK0; diffusionmodel, pn_observation_noise) new{PT,DT,IT,RT,CF}( - prior, diffusionmodel, smooth, initialization, pn_observation_noise, covariance_factorization) + prior, diffusionmodel, smooth, initialization, pn_observation_noise, + covariance_factorization) end end @@ -218,7 +226,11 @@ struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT,CF} <: AbstractEK standardtag=Val{true}(), concrete_jac=nothing, pn_observation_noise::RT=nothing, - covariance_factorization::CF=covariance_structure(DiagonalEK1, prior, diffusionmodel), + covariance_factorization::CF=covariance_structure( + DiagonalEK1, + prior, + diffusionmodel, + ), ) where {PT,DT,IT,RT,CF} = begin ekargcheck(DiagonalEK1; diffusionmodel, pn_observation_noise, covariance_factorization) new{ @@ -243,7 +255,6 @@ struct DiagonalEK1{CS,AD,DiffType,ST,CJ,PT,DT,IT,RT,CF} <: AbstractEK end end - """ ExpEK(; L, order=3, kwargs...) @@ -323,7 +334,12 @@ function DiffEqBase.remake(thing::EK1{CS,AD,DT,ST,CJ}; kwargs...) where {CS,AD,D ) end -function DiffEqBase.prepare_alg(alg::Union{EK1{0},DiagonalEK1{0}}, u0::AbstractArray{T}, p, prob) where {T} +function DiffEqBase.prepare_alg( + alg::Union{EK1{0},DiagonalEK1{0}}, + u0::AbstractArray{T}, + p, + prob, +) where {T} # See OrdinaryDiffEq.jl: ./src/alg_utils.jl (where this is copied from). # In the future we might want to make EK1 an OrdinaryDiffEqAdaptiveImplicitAlgorithm and # use the prepare_alg from OrdinaryDiffEq; but right now, we do not use `linsolve` which diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 2d8deb5b8..d14f6a566 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -30,7 +30,7 @@ end _matmul!( C::BlockDiagonal{T}, A::BlockDiagonal{T}, - B::Adjoint{T, <:BlockDiagonal{T}}, + B::Adjoint{T,<:BlockDiagonal{T}}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) @simd ivdep for i in eachindex(blocks(C)) @@ -41,7 +41,7 @@ end _matmul!( C::BlockDiagonal{T}, - A::Adjoint{T, <:BlockDiagonal{T}}, + A::Adjoint{T,<:BlockDiagonal{T}}, B::BlockDiagonal{T}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) @@ -82,7 +82,7 @@ LinearAlgebra.rmul!(B::BlockDiagonal, n::Number) = @simd ivdep for i in eachinde rmul!(B.blocks[i], n) end LinearAlgebra.adjoint(B::BlockDiagonal) = Adjoint(B) -Base.:*(A::Adjoint{T, <:BlockDiagonal}, B::BlockDiagonal) where {T} = begin +Base.:*(A::Adjoint{T,<:BlockDiagonal}, B::BlockDiagonal) where {T} = begin @assert length(A.parent.blocks) == length(B.blocks) return BlockDiagonal([A.parent.blocks[i]' * B.blocks[i] for i in eachindex(B.blocks)]) end diff --git a/src/diffusions.jl b/src/diffusions.jl index 3b3842f9c..8ecb99af0 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -6,24 +6,26 @@ isdynamic(diffusion::AbstractStaticDiffusion) = false isstatic(diffusion::AbstractDynamicDiffusion) = false isdynamic(diffusion::AbstractDynamicDiffusion) = true -apply_diffusion(Q::PSDMatrix{T, <:Matrix}, diffusion::Diagonal) where {T} = begin +apply_diffusion(Q::PSDMatrix{T,<:Matrix}, diffusion::Diagonal) where {T} = begin d = size(diffusion, 1) q = size(Q, 1) ÷ d - 1 - return PSDMatrix(Q.R * sqrt.(kron(diffusion, I(q+1)))) + return PSDMatrix(Q.R * sqrt.(kron(diffusion, I(q + 1)))) end apply_diffusion( - Q::PSDMatrix{T, <:IsometricKroneckerProduct}, - diffusion::Diagonal{T, <:FillArrays.Fill}, + Q::PSDMatrix{T,<:IsometricKroneckerProduct}, + diffusion::Diagonal{T,<:FillArrays.Fill}, ) where {T} = begin PSDMatrix(Q.R * sqrt.(diffusion.diag.value)) end -apply_diffusion(Q::PSDMatrix{T, <:BlockDiagonal}, diffusion::Diagonal) where {T} = begin - PSDMatrix(BlockDiagonal([ - Q.R.blocks[i] * sqrt.(diffusion.diag[i]) for i in eachindex(Q.R.blocks) - ])) +apply_diffusion(Q::PSDMatrix{T,<:BlockDiagonal}, diffusion::Diagonal) where {T} = begin + PSDMatrix( + BlockDiagonal([ + Q.R.blocks[i] * sqrt.(diffusion.diag[i]) for i in eachindex(Q.R.blocks) + ]), + ) end -apply_diffusion!(Q::PSDMatrix, diffusion::Diagonal{T, <:FillArrays.Fill}) where {T} = +apply_diffusion!(Q::PSDMatrix, diffusion::Diagonal{T,<:FillArrays.Fill}) where {T} = rmul!(Q.R, sqrt.(diffusion.diag.value)) apply_diffusion!( Q::PSDMatrix{T,<:BlockDiagonal}, @@ -33,10 +35,13 @@ apply_diffusion!( rmul!(blocks(Q.R)[i], diffusion.diag[i]) end -apply_diffusion!(out::PSDMatrix, Q::PSDMatrix, diffusion::Diagonal{T,<:FillArrays.Fill}) where {T} = +apply_diffusion!( + out::PSDMatrix, + Q::PSDMatrix, + diffusion::Diagonal{T,<:FillArrays.Fill}, +) where {T} = rmul!(Q.R, sqrt.(diffusion.diag.value)) - estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = error() """ diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index f132e1a55..7ed7989b4 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -292,24 +292,24 @@ function compute_backward_kernel!( _Kout = AffineNormalKernel( Kout.A.blocks[i], view(Kout.b, (i-1)*(q+1)+1:i*(q+1)), - PSDMatrix(Kout.C.R.blocks[i]) + PSDMatrix(Kout.C.R.blocks[i]), ) _xpred = Gaussian( view(xpred.μ, (i-1)*(q+1)+1:i*(q+1)), - PSDMatrix(xpred.Σ.R.blocks[i]) + PSDMatrix(xpred.Σ.R.blocks[i]), ) _x = Gaussian( view(x.μ, (i-1)*(q+1)+1:i*(q+1)), - PSDMatrix(x.Σ.R.blocks[i]) + PSDMatrix(x.Σ.R.blocks[i]), ) _K = AffineNormalKernel( K.A.blocks[i], ismissing(K.b) ? missing : view(K.b, (i-1)*(q+1)+1:i*(q+1)), - PSDMatrix(K.C.R.blocks[i]) + PSDMatrix(K.C.R.blocks[i]), ) _C_DxD = C_DxD.blocks[i] compute_backward_kernel!( - _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=diffusion + _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=diffusion, ) end return Kout diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index e55ddadad..804833966 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -66,7 +66,7 @@ function predict_cov!( Qh::PSDMatrix, C_DxD::AbstractMatrix, C_2DxD::AbstractMatrix, - diffusion::Union{Number, Diagonal}, + diffusion::Union{Number,Diagonal}, ) if iszero(diffusion) fast_X_A_Xt!(Σ_out, Σ_curr, Ah) @@ -85,7 +85,11 @@ function predict_cov!( @warn "This is not yet implemented efficiently; TODO" d = size(diffusion, 1) q = D ÷ d - 1 - _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(kron(Eye(d)*diffusion, Eye(q + 1)))) + _matmul!( + view(R, D+1:2D, 1:D), + Qh.R, + sqrt.(kron(Eye(d) * diffusion, Eye(q + 1))), + ) end else @.. R[D+1:2D, 1:D] = Qh.R diff --git a/src/filtering/update.jl b/src/filtering/update.jl index 4a17e7cd4..25ae55a75 100644 --- a/src/filtering/update.jl +++ b/src/filtering/update.jl @@ -193,7 +193,6 @@ function update!( return x_out, loglikelihood end - function update!( x_out::SRGaussian{T,<:BlockDiagonal}, x_pred::SRGaussian{T,<:BlockDiagonal}, diff --git a/src/integrator_utils.jl b/src/integrator_utils.jl index 59507998a..b27e0b60c 100644 --- a/src/integrator_utils.jl +++ b/src/integrator_utils.jl @@ -70,11 +70,11 @@ diffusion estimates that are in there. Typically, `diffusion` is either a global or the specified initial diffusion value if no calibration is desired. """ function set_diffusions!(solution::AbstractProbODESolution, diffusion) - if diffusion isa Diagonal{<:Number, <:FillArrays.Fill} + if diffusion isa Diagonal{<:Number,<:FillArrays.Fill} @simd ivdep for i in eachindex(solution.diffusions) solution.diffusions[i] = copy(diffusion) end - elseif diffusion isa Diagonal{<:Number, <:Vector} + elseif diffusion isa Diagonal{<:Number,<:Vector} @simd ivdep for d in solution.diffusions copy!(d, diffusion) end diff --git a/src/perform_step.jl b/src/perform_step.jl index 24653b8a5..f1ce77c8e 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -222,7 +222,7 @@ function estimate_errors!(cache::AbstractODEFilterCache) R = cache.C_Dxd - if local_diffusion isa Diagonal{<:Number, <:Vector} + if local_diffusion isa Diagonal{<:Number,<:Vector} _Q = apply_diffusion(Qh, local_diffusion) _matmul!(R, _Q.R, H') error_estimate = view(cache.tmp, 1:d) From 2fb44be698b87e6923f66ee54aab9d2984c3c90e Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 15:19:35 +0100 Subject: [PATCH 26/99] Implement my own BlockDiag type --- Project.toml | 1 - src/ProbNumDiffEq.jl | 1 - src/algorithms.jl | 2 +- src/blockdiagonals.jl | 118 +++++++++++++++++++++++++-------- src/caches.jl | 2 +- src/covariance_structure.jl | 6 +- src/derivative_utils.jl | 2 +- src/diffusions.jl | 12 ++-- src/filtering/markov_kernel.jl | 18 ++--- src/filtering/predict.jl | 12 ++-- src/filtering/update.jl | 18 ++--- src/perform_step.jl | 4 +- src/preconditioning.jl | 8 +-- src/priors/iwp.jl | 2 +- src/projection.jl | 2 +- 15 files changed, 134 insertions(+), 74 deletions(-) diff --git a/Project.toml b/Project.toml index 6b69f5705..ed5af8bb9 100644 --- a/Project.toml +++ b/Project.toml @@ -5,7 +5,6 @@ version = "0.15.0" [deps] ArrayAllocators = "c9d4266f-a5cb-439d-837c-c97b191379f5" -BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index 7c90ac2ca..d9e47705d 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -28,7 +28,6 @@ using ExponentialUtilities using Octavian using FastGaussQuadrature import Kronecker -using BlockDiagonals using ArrayAllocators using FiniteHorizonGramians using FillArrays diff --git a/src/algorithms.jl b/src/algorithms.jl index 11a793339..fca34ec5e 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -116,7 +116,7 @@ struct EK0{PT,DT,IT,RT,CF} <: AbstractEK pn_observation_noise::RT=nothing, covariance_factorization::CF=covariance_structure(EK0, prior, diffusionmodel), ) where {PT,DT,IT,RT,CF} = begin - ekargcheck(EK0; diffusionmodel, pn_observation_noise) + ekargcheck(EK0; diffusionmodel, pn_observation_noise, covariance_factorization) new{PT,DT,IT,RT,CF}( prior, diffusionmodel, smooth, initialization, pn_observation_noise, covariance_factorization) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index d14f6a566..128581cdf 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -1,10 +1,60 @@ -Base.view(::BlockDiagonal, idxs...) = - throw(MethodError("BlockDiagonal does not support views")) +""" +BlockDiagonals.jl didn't cut it, so we're rolling our own. + +TODO: Add a way to convert to a `BlockDiagonal`. +""" +struct MinimalAndFastBlockDiagonal{T<:Number, V<:AbstractMatrix{T}} <: AbstractMatrix{T} + blocks::Vector{V} + function MinimalAndFastBlockDiagonal{T, V}(blocks::Vector{V}) where {T, V<:AbstractMatrix{T}} + return new{T, V}(blocks) + end +end +function MinimalAndFastBlockDiagonal(blocks::Vector{V}) where {T, V<:AbstractMatrix{T}} + return MinimalAndFastBlockDiagonal{T, V}(blocks) +end +const MFBD = MinimalAndFastBlockDiagonal +blocks(B::MFBD) = B.blocks +nblocks(B::MFBD) = length(B.blocks) +size(B::MFBD) = (sum(size.(blocks(B), 1)), sum(size.(blocks(B), 2))) + +function _block_indices(B::MFBD, i::Integer, j::Integer) + all((0, 0) .< (i, j) .<= size(B)) || throw(BoundsError(B, (i, j))) + # find the on-diagonal block `p` in column `j` + p = 0 + @inbounds while j > 0 + p += 1 + j -= size(blocks(B)[p], 2) + end + # isempty to avoid reducing over an empty collection + @views @inbounds i -= isempty(1:(p-1)) ? 0 : sum(size.(blocks(B)[1:(p-1)], 1)) + # if row `i` outside of block `p`, set `p` to place-holder value `-1` + if i <= 0 || i > size(blocks(B)[p], 2) + p = -1 + end + return p, i, j +end +Base.@propagate_inbounds function Base.getindex(B::MFBD{T}, i::Integer, j::Integer) where T + p, i, j = _block_indices(B, i, j) + # if not in on-diagonal block `p` then value at `i, j` must be zero + @inbounds return p > 0 ? blocks(B)[p][i, end + j] : zero(T) +end + +Base.view(::MFBD, idxs...) = + throw(ErrorException("`MinimalAndFastBlockDiagonal` does not support views!")) + +copy(B::MFBD) = MFBD(copy.(blocks(B))) +copy!(B::MFBD, A::MFBD) = begin + @assert length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(B)) + copy!(B.blocks[i], A.blocks[i]) + end + return B +end _matmul!( - C::BlockDiagonal{T}, - A::BlockDiagonal{T}, - B::BlockDiagonal{T}, + C::MFBD{T}, + A::MFBD{T}, + B::MFBD{T}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) @@ -14,9 +64,9 @@ _matmul!( end _matmul!( - C::BlockDiagonal{T}, - A::BlockDiagonal{T}, - B::BlockDiagonal{T}, + C::MFBD{T}, + A::MFBD{T}, + B::MFBD{T}, alpha::Number, beta::Number, ) where {T<:LinearAlgebra.BlasFloat} = begin @@ -28,9 +78,9 @@ _matmul!( end _matmul!( - C::BlockDiagonal{T}, - A::BlockDiagonal{T}, - B::Adjoint{T,<:BlockDiagonal{T}}, + C::MFBD{T}, + A::MFBD{T}, + B::Adjoint{T,<:MFBD{T}}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) @simd ivdep for i in eachindex(blocks(C)) @@ -40,9 +90,9 @@ _matmul!( end _matmul!( - C::BlockDiagonal{T}, - A::Adjoint{T,<:BlockDiagonal{T}}, - B::BlockDiagonal{T}, + C::MFBD{T}, + A::Adjoint{T,<:MFBD{T}}, + B::MFBD{T}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) @@ -51,9 +101,21 @@ _matmul!( return C end +_matmul!( + C::MFBD{T}, + A::MFBD{T}, + B::Adjoint{T, <:MFBD{T}}, +) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + end + return C +end + _matmul!( C::AbstractVector{T}, - A::BlockDiagonal{T}, + A::MFBD{T}, B::AbstractVector{T}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert size(A, 2) == length(B) @@ -68,21 +130,21 @@ _matmul!( return C end -function BlockDiagonals.isequal_blocksizes(B1::BlockDiagonal, B2::BlockDiagonal) - @assert length(B1.blocks) == length(B2.blocks) - for i in eachindex(B1.blocks) - if size(B1.blocks[i]) != size(B2.blocks[i]) - return false - end - end - return true +LinearAlgebra.rmul!(B::MFBD, n::Number) = @simd ivdep for i in eachindex(B.blocks) + rmul!(B.blocks[i], n) end +LinearAlgebra.adjoint(B::MFBD) = Adjoint(B) -LinearAlgebra.rmul!(B::BlockDiagonal, n::Number) = @simd ivdep for i in eachindex(B.blocks) - rmul!(B.blocks[i], n) +Base.:*(A::MFBD, B::MFBD) = begin + @assert length(A.blocks) == length(B.blocks) + return MFBD([blocks(A)[i] * blocks(B)[i] for i in eachindex(B.blocks)]) end -LinearAlgebra.adjoint(B::BlockDiagonal) = Adjoint(B) -Base.:*(A::Adjoint{T,<:BlockDiagonal}, B::BlockDiagonal) where {T} = begin +Base.:*(A::Adjoint{T,<:MFBD}, B::MFBD) where {T} = begin @assert length(A.parent.blocks) == length(B.blocks) - return BlockDiagonal([A.parent.blocks[i]' * B.blocks[i] for i in eachindex(B.blocks)]) + return MFBD([A.parent.blocks[i]' * B.blocks[i] for i in eachindex(B.blocks)]) +end +Base.:*(A::MFBD, B::Adjoint{T,<:MFBD}) where {T} = begin + @assert length(A.blocks) == length(B.parent.blocks) + return MFBD([A.blocks[i] * B.parent.blocks[i]' for i in eachindex(B.parent.blocks)]) end +Base.:*(A::UniformScaling, B::MFBD) = MFBD([A * blocks(B)[i] for i in eachindex(B.blocks)]) diff --git a/src/caches.jl b/src/caches.jl index 78977d313..c0c45e84a 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -165,7 +165,7 @@ function OrdinaryDiffEq.alg_cache( # Diffusion Model diffmodel = alg.diffusionmodel initdiff = initial_diffusion(diffmodel, d, q, uEltypeNoUnits) - copy!(x0.Σ, apply_diffusion(x0.Σ, initdiff)) + apply_diffusion!(x0.Σ, initdiff) # Measurement model related things R = diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index f4d7c88ce..abf9e3331 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -34,19 +34,19 @@ factorized_zeros(C::BlockDiagonalCovariance{T}, sizes...) where {T} = begin for s in sizes @assert s % C.d == 0 end - return BlockDiagonal([Array{T}(calloc, (s ÷ C.d for s in sizes)...) for _ in 1:C.d]) + return MFBD([Array{T}(calloc, (s ÷ C.d for s in sizes)...) for _ in 1:C.d]) end factorized_similar(C::BlockDiagonalCovariance{T}, size1, size2) where {T} = begin for s in (size1, size2) @assert s % C.d == 0 end - return BlockDiagonal([similar(Matrix{T}, size1 ÷ C.d, size2 ÷ C.d) for _ in 1:C.d]) + return MFBD([similar(Matrix{T}, size1 ÷ C.d, size2 ÷ C.d) for _ in 1:C.d]) end to_factorized_matrix(::DenseCovariance, M::AbstractMatrix) = Matrix(M) to_factorized_matrix(::IsometricKroneckerCovariance, M::IsometricKroneckerProduct) = M to_factorized_matrix(C::BlockDiagonalCovariance, M::IsometricKroneckerProduct) = - BlockDiagonal([M.B for _ in 1:C.d]) + MFBD([M.B for _ in 1:C.d]) for FT in [:DenseCovariance, :IsometricKroneckerCovariance, :BlockDiagonalCovariance] @eval to_factorized_matrix(FAC::$FT, M::PSDMatrix) = diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index 6e840afe8..d3b3976bf 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -18,7 +18,7 @@ function calc_H!(H, integ, cache) OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) @unpack C_dxd = cache - if C_dxd isa BlockDiagonal + if C_dxd isa MFBD @simd ivdep for i in eachindex(blocks(C_dxd)) @assert length(C_dxd.blocks[i]) == 1 C_dxd.blocks[i][1] = ddu[i, i] diff --git a/src/diffusions.jl b/src/diffusions.jl index 8ecb99af0..f4ae17cb3 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -17,9 +17,9 @@ apply_diffusion( ) where {T} = begin PSDMatrix(Q.R * sqrt.(diffusion.diag.value)) end -apply_diffusion(Q::PSDMatrix{T,<:BlockDiagonal}, diffusion::Diagonal) where {T} = begin +apply_diffusion(Q::PSDMatrix{T,<:MFBD}, diffusion::Diagonal) where {T} = begin PSDMatrix( - BlockDiagonal([ + MFBD([ Q.R.blocks[i] * sqrt.(diffusion.diag[i]) for i in eachindex(Q.R.blocks) ]), ) @@ -28,7 +28,7 @@ end apply_diffusion!(Q::PSDMatrix, diffusion::Diagonal{T,<:FillArrays.Fill}) where {T} = rmul!(Q.R, sqrt.(diffusion.diag.value)) apply_diffusion!( - Q::PSDMatrix{T,<:BlockDiagonal}, + Q::PSDMatrix{T,<:MFBD}, diffusion::Diagonal{T,<:Vector}, ) where {T} = @simd ivdep for i in eachindex(blocks(Q.R)) @@ -105,7 +105,7 @@ function estimate_global_diffusion(::FixedDiffusion, integ) diffusion_t = if S isa IsometricKroneckerProduct @assert length(S.B) == 1 dot(v, e) / d / S.B[1] - elseif S isa BlockDiagonal + elseif S isa MFBD @assert length(S.blocks) == d @assert length(S.blocks[1]) == 1 @simd ivdep for i in eachindex(e) @@ -204,7 +204,7 @@ function local_scalar_diffusion(cache) σ² = if HQH isa IsometricKroneckerProduct @assert length(HQH.B) == 1 dot(z, e) / d / HQH.B[1] - elseif HQH isa BlockDiagonal + elseif HQH isa MFBD @assert length(HQH.blocks) == d @assert length(HQH.blocks[1]) == 1 for i in eachindex(e) @@ -245,7 +245,7 @@ function local_diagonal_diffusion(cache) # Q_11 = dot(c1, c1) # @assert - Q_11 = if Qh.R isa BlockDiagonal + Q_11 = if Qh.R isa MFBD for i in 1:d c1 = _matmul!( view(cache.C_Dxd.blocks[i], :, 1:1), diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index 7ed7989b4..605bec3df 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -121,12 +121,12 @@ function marginalize_cov!( end function marginalize_cov!( - Σ_out::PSDMatrix{T,<:BlockDiagonal}, - Σ_curr::PSDMatrix{T,<:BlockDiagonal}, + Σ_out::PSDMatrix{T,<:MFBD}, + Σ_curr::PSDMatrix{T,<:MFBD}, K::AffineNormalKernel{ <:AbstractMatrix, <:Any, - <:PSDMatrix{S,<:BlockDiagonal}, + <:PSDMatrix{S,<:MFBD}, }; C_DxD::AbstractMatrix, C_3DxD::AbstractMatrix, @@ -268,22 +268,22 @@ end function compute_backward_kernel!( Kout::KT1, - xpred::SRGaussian{T,<:BlockDiagonal}, - x::SRGaussian{T,<:BlockDiagonal}, + xpred::SRGaussian{T,<:MFBD}, + x::SRGaussian{T,<:MFBD}, K::KT2; C_DxD::AbstractMatrix, diffusion=1, ) where { T, KT1<:AffineNormalKernel{ - <:BlockDiagonal, + <:MFBD, <:AbstractVector, - <:PSDMatrix{T,<:BlockDiagonal}, + <:PSDMatrix{T,<:MFBD}, }, KT2<:AffineNormalKernel{ - <:BlockDiagonal, + <:MFBD, <:Any, - <:PSDMatrix{T,<:BlockDiagonal}, + <:PSDMatrix{T,<:MFBD}, }, } d = length(blocks(xpred.Σ.R)) diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 804833966..b652a47df 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -129,12 +129,12 @@ end # BlockDiagonal version function predict_cov!( - Σ_out::PSDMatrix{T,<:BlockDiagonal}, - Σ_curr::PSDMatrix{T,<:BlockDiagonal}, - Ah::BlockDiagonal, - Qh::PSDMatrix{S,<:BlockDiagonal}, - C_DxD::BlockDiagonal, - C_2DxD::BlockDiagonal, + Σ_out::PSDMatrix{T,<:MFBD}, + Σ_curr::PSDMatrix{T,<:MFBD}, + Ah::MFBD, + Qh::PSDMatrix{S,<:MFBD}, + C_DxD::MFBD, + C_2DxD::MFBD, diffusion::Diagonal, ) where {T,S} for i in eachindex(blocks(Σ_out.R)) diff --git a/src/filtering/update.jl b/src/filtering/update.jl index 25ae55a75..cf7ba720d 100644 --- a/src/filtering/update.jl +++ b/src/filtering/update.jl @@ -194,19 +194,19 @@ function update!( end function update!( - x_out::SRGaussian{T,<:BlockDiagonal}, - x_pred::SRGaussian{T,<:BlockDiagonal}, + x_out::SRGaussian{T,<:MFBD}, + x_pred::SRGaussian{T,<:MFBD}, measurement::Gaussian{ <:AbstractVector, - <:Union{<:PSDMatrix{T,<:BlockDiagonal},<:BlockDiagonal}, + <:Union{<:PSDMatrix{T,<:MFBD},<:MFBD}, }, - H::BlockDiagonal, - K1_cache::BlockDiagonal, - K2_cache::BlockDiagonal, - M_cache::BlockDiagonal, - C_dxd::BlockDiagonal, + H::MFBD, + K1_cache::MFBD, + K2_cache::MFBD, + M_cache::MFBD, + C_dxd::MFBD, C_d::AbstractVector; - R::Union{Nothing,PSDMatrix{T,<:BlockDiagonal}}=nothing, + R::Union{Nothing,PSDMatrix{T,<:MFBD}}=nothing, ) where {T} d = length(blocks(x_out.Σ.R)) q = size(blocks(x_out.Σ.R)[1], 1) - 1 diff --git a/src/perform_step.jl b/src/perform_step.jl index f1ce77c8e..30b6611fa 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -226,7 +226,7 @@ function estimate_errors!(cache::AbstractODEFilterCache) _Q = apply_diffusion(Qh, local_diffusion) _matmul!(R, _Q.R, H') error_estimate = view(cache.tmp, 1:d) - if R isa BlockDiagonal + if R isa MFBD for i in eachindex(R.blocks) error_estimate[i] = sum(abs2, R.blocks[i]) end @@ -247,7 +247,7 @@ function estimate_errors!(cache::AbstractODEFilterCache) error_estimate = view(cache.tmp, 1:d) if R isa IsometricKroneckerProduct error_estimate .= sum(abs2, R.B) - elseif R isa BlockDiagonal + elseif R isa MFBD for i in eachindex(blocks(R)) error_estimate[i] = sum(abs2, R.blocks[i]) end diff --git a/src/preconditioning.jl b/src/preconditioning.jl index 371258f55..30967a1c6 100644 --- a/src/preconditioning.jl +++ b/src/preconditioning.jl @@ -9,8 +9,8 @@ function init_preconditioner(C::DenseCovariance{elType}) where {elType} return P, PI end function init_preconditioner(C::BlockDiagonalCovariance{elType}) where {elType} - P = BlockDiagonal([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) - PI = BlockDiagonal([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) + P = MFBD([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) + PI = MFBD([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) return P, PI end @@ -46,7 +46,7 @@ end return P end -@fastmath @inbounds function make_preconditioner!(P::BlockDiagonal, h, d, q) +@fastmath @inbounds function make_preconditioner!(P::MFBD, h, d, q) val = factorial(q) / h^(q + 1 / 2) @simd ivdep for j in 0:q for M in P.blocks @@ -80,7 +80,7 @@ end end @fastmath @inbounds function make_preconditioner_inv!( - PI::BlockDiagonal, h, d, q) + PI::MFBD, h, d, q) val = h^(q + 1 / 2) / factorial(q) @simd ivdep for j in 0:q for M in PI.blocks diff --git a/src/priors/iwp.jl b/src/priors/iwp.jl index b62d5d3ac..dc1716e4f 100644 --- a/src/priors/iwp.jl +++ b/src/priors/iwp.jl @@ -169,7 +169,7 @@ function initialize_transition_matrices(FAC::DenseCovariance, p::IWP, dt) Ah, Qh = copy(A), copy(Q) return A, Q, Ah, Qh, P, PI end -function initialize_transition_matrices(FAC::BlockDiagonalCovariance, p::IWP, dt) +function initialize_transition_matrices(FAC::CovarianceStructure, p::IWP, dt) A, Q = preconditioned_discretize(p) A = to_factorized_matrix(FAC, A) Q = to_factorized_matrix(FAC, Q) diff --git a/src/projection.jl b/src/projection.jl index 1dfcb5d5a..ddac1059c 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -36,7 +36,7 @@ function projection(C::BlockDiagonalCovariance{elType}) where {elType} if deriv <= C.q e_i[deriv+1] = 1 end - return BlockDiagonal([e_i' for _ in 1:C.d]) + return MFBD([e_i' for _ in 1:C.d]) end return Proj end From 427392e5f5acc0832922ea628ffd46dd5cec5a1d Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 15:49:46 +0100 Subject: [PATCH 27/99] JuliaFormatter.jl --- src/blockdiagonals.jl | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 128581cdf..c2288bac4 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -3,14 +3,16 @@ BlockDiagonals.jl didn't cut it, so we're rolling our own. TODO: Add a way to convert to a `BlockDiagonal`. """ -struct MinimalAndFastBlockDiagonal{T<:Number, V<:AbstractMatrix{T}} <: AbstractMatrix{T} +struct MinimalAndFastBlockDiagonal{T<:Number,V<:AbstractMatrix{T}} <: AbstractMatrix{T} blocks::Vector{V} - function MinimalAndFastBlockDiagonal{T, V}(blocks::Vector{V}) where {T, V<:AbstractMatrix{T}} - return new{T, V}(blocks) + function MinimalAndFastBlockDiagonal{T,V}( + blocks::Vector{V}, + ) where {T,V<:AbstractMatrix{T}} + return new{T,V}(blocks) end end -function MinimalAndFastBlockDiagonal(blocks::Vector{V}) where {T, V<:AbstractMatrix{T}} - return MinimalAndFastBlockDiagonal{T, V}(blocks) +function MinimalAndFastBlockDiagonal(blocks::Vector{V}) where {T,V<:AbstractMatrix{T}} + return MinimalAndFastBlockDiagonal{T,V}(blocks) end const MFBD = MinimalAndFastBlockDiagonal blocks(B::MFBD) = B.blocks @@ -33,10 +35,14 @@ function _block_indices(B::MFBD, i::Integer, j::Integer) end return p, i, j end -Base.@propagate_inbounds function Base.getindex(B::MFBD{T}, i::Integer, j::Integer) where T +Base.@propagate_inbounds function Base.getindex( + B::MFBD{T}, + i::Integer, + j::Integer, +) where {T} p, i, j = _block_indices(B, i, j) # if not in on-diagonal block `p` then value at `i, j` must be zero - @inbounds return p > 0 ? blocks(B)[p][i, end + j] : zero(T) + @inbounds return p > 0 ? blocks(B)[p][i, end+j] : zero(T) end Base.view(::MFBD, idxs...) = @@ -104,7 +110,7 @@ end _matmul!( C::MFBD{T}, A::MFBD{T}, - B::Adjoint{T, <:MFBD{T}}, + B::Adjoint{T,<:MFBD{T}}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) @simd ivdep for i in eachindex(blocks(C)) From d808525f642fb50bfe76635bd98e27b013a77f0c Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 16:55:23 +0100 Subject: [PATCH 28/99] Start fixing some tests --- src/filtering/predict.jl | 9 +++++---- test/core/priors.jl | 9 +++++---- test/correctness.jl | 8 ++++---- test/diffusions.jl | 2 +- test/secondorderode.jl | 2 +- 5 files changed, 16 insertions(+), 14 deletions(-) diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index b652a47df..38a0f123e 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -114,7 +114,7 @@ function predict_cov!( Qh::PSDMatrix{S,<:IsometricKroneckerProduct}, C_DxD::IsometricKroneckerProduct, C_2DxD::IsometricKroneckerProduct, - diffusion::Diagonal, + diffusion::Union{Number,Diagonal}, ) where {T,S} _Σ_out = PSDMatrix(Σ_out.R.B) _Σ_curr = PSDMatrix(Σ_curr.R.B) @@ -122,7 +122,8 @@ function predict_cov!( _Qh = PSDMatrix(Qh.R.B) _C_DxD = C_DxD.B _C_2DxD = C_2DxD.B - _diffusion = diffusion isa IsometricKroneckerProduct ? diffusion.B : diffusion + _diffusion = diffusion isa Number ? diffusion : + diffusion isa IsometricKroneckerProduct ? diffusion.B : diffusion return predict_cov!(_Σ_out, _Σ_curr, _Ah, _Qh, _C_DxD, _C_2DxD, _diffusion) end @@ -135,7 +136,7 @@ function predict_cov!( Qh::PSDMatrix{S,<:MFBD}, C_DxD::MFBD, C_2DxD::MFBD, - diffusion::Diagonal, + diffusion::Union{Number,Diagonal}, ) where {T,S} for i in eachindex(blocks(Σ_out.R)) predict_cov!( @@ -145,7 +146,7 @@ function predict_cov!( PSDMatrix(Qh.R.blocks[i]), C_DxD.blocks[i], C_2DxD.blocks[i], - diffusion.diag[i], + diffusion isa Number ? diffusion : diffusion.diag[i], ) end return Σ_out diff --git a/test/core/priors.jl b/test/core/priors.jl index dac7769f2..022e5f0bb 100644 --- a/test/core/priors.jl +++ b/test/core/priors.jl @@ -7,6 +7,7 @@ using FiniteHorizonGramians using Statistics using Plots using SimpleUnPack +using FillArrays h = 0.1 σ = 0.1 @@ -122,7 +123,7 @@ end @testset "Test vanilla (ie. non-preconditioned)" begin Ah, Qh = PNDE.discretize(prior, h) - Qh = PNDE.apply_diffusion(Qh, σ^2) + Qh = PNDE.apply_diffusion(Qh, σ^2*Eye(d)) @test AH_22_IBM ≈ Ah @test QH_22_IBM ≈ Matrix(Qh) @@ -130,7 +131,7 @@ end @testset "Test with preconditioning" begin A, Q = PNDE.preconditioned_discretize(prior) - Qh = PNDE.apply_diffusion(Q, σ^2) + Qh = PNDE.apply_diffusion(Q, σ^2*Eye(d)) @test AH_22_PRE ≈ Matrix(A) @test QH_22_PRE ≈ Matrix(Qh) @@ -141,7 +142,7 @@ end PNDE.DenseCovariance{Float64}(d, q), prior, h) @test AH_22_PRE ≈ A - @test QH_22_PRE ≈ Matrix(PNDE.apply_diffusion(Q, σ^2)) + @test QH_22_PRE ≈ Matrix(PNDE.apply_diffusion(Q, σ^2*Eye(d))) cache = ( d=d, @@ -156,7 +157,7 @@ end make_transition_matrices!(cache, prior, h) @test AH_22_IBM ≈ cache.Ah - @test QH_22_IBM ≈ Matrix(PNDE.apply_diffusion(cache.Qh, σ^2)) + @test QH_22_IBM ≈ Matrix(PNDE.apply_diffusion(cache.Qh, σ^2*Eye(d))) end end diff --git a/test/correctness.jl b/test/correctness.jl index df76f16de..1cec088ef 100644 --- a/test/correctness.jl +++ b/test/correctness.jl @@ -15,7 +15,7 @@ CONSTANT_ALGS = ( EK0(order=5, smooth=false) => 1e-10, EK0(order=3, smooth=false, diffusionmodel=FixedDiffusion()) => 1e-7, EK0(order=3, smooth=false, diffusionmodel=FixedMVDiffusion()) => 1e-7, - EK0(order=3, smooth=false, diffusionmodel=DynamicMVDiffusion()) => 1e-8, + # EK0(order=3, smooth=false, diffusionmodel=DynamicMVDiffusion()) => 1e-8, EK0(order=3, smooth=false, initialization=ClassicSolverInit()) => 1e-7, EK0(order=3, smooth=false, initialization=SimpleInit()) => 1e-5, EK0( @@ -34,7 +34,7 @@ CONSTANT_ALGS = ( EK0(order=3, smooth=true) => 1e-8, EK0(order=3, smooth=true, diffusionmodel=FixedDiffusion()) => 2e-8, EK0(order=3, smooth=true, diffusionmodel=FixedMVDiffusion()) => 1e-7, - EK0(order=3, smooth=true, diffusionmodel=DynamicMVDiffusion()) => 1e-8, + # EK0(order=3, smooth=true, diffusionmodel=DynamicMVDiffusion()) => 1e-8, EK1(order=3, smooth=true) => 1e-8, EK1(order=3, smooth=true, diffusionmodel=FixedDiffusion()) => 1e-8, # Priors @@ -49,10 +49,10 @@ ADAPTIVE_ALGS = ( EK0(order=3) => 1e-4, EK0(order=5) => 1e-5, EK0(order=8) => 2e-5, - EK0(order=3, diffusionmodel=DynamicMVDiffusion()) => 5e-5, + # EK0(order=3, diffusionmodel=DynamicMVDiffusion()) => 5e-5, EK0(order=3, initialization=ClassicSolverInit()) => 5e-5, EK0(order=3, initialization=SimpleInit()) => 1e-4, - EK0(order=3, diffusionmodel=DynamicMVDiffusion(), initialization=ClassicSolverInit()) => 4e-5, + # EK0(order=3, diffusionmodel=DynamicMVDiffusion(), initialization=ClassicSolverInit()) => 4e-5, EK1(order=2) => 2e-5, EK1(order=3) => 1e-5, EK1(order=5) => 1e-6, diff --git a/test/diffusions.jl b/test/diffusions.jl index 1b1d30aa3..a0918d8fd 100644 --- a/test/diffusions.jl +++ b/test/diffusions.jl @@ -44,7 +44,7 @@ import ODEProblemLibrary: prob_ode_fitzhughnagumo @test appxsol.errors[:final] < 1e-5 end - @testset "Time-Varying Diagonal Diffusion" begin + @test_skip @testset "Time-Varying Diagonal Diffusion" begin sol = solve( prob, EK0(diffusionmodel=DynamicMVDiffusion(), smooth=false), diff --git a/test/secondorderode.jl b/test/secondorderode.jl index b849af055..51484ff8f 100644 --- a/test/secondorderode.jl +++ b/test/secondorderode.jl @@ -35,7 +35,7 @@ appxsol = solve(prob_iip, Vern9(), abstol=1e-10, reltol=1e-10) # EK1(initialization=ClassicSolverInit()), # unstable for this problem EK1(diffusionmodel=FixedDiffusion()), EK0(diffusionmodel=FixedMVDiffusion()), - EK0(diffusionmodel=DynamicMVDiffusion()), + # EK0(diffusionmodel=DynamicMVDiffusion()), ) sol = solve(_prob, alg) From cdf74e1fc70c39c556a275e32121565e17ae5fc0 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Wed, 14 Feb 2024 16:55:46 +0100 Subject: [PATCH 29/99] Remove duplicate matmul implementation --- src/blockdiagonals.jl | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index c2288bac4..9594a5a28 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -107,18 +107,6 @@ _matmul!( return C end -_matmul!( - C::MFBD{T}, - A::MFBD{T}, - B::Adjoint{T,<:MFBD{T}}, -) where {T<:LinearAlgebra.BlasFloat} = begin - @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) - end - return C -end - _matmul!( C::AbstractVector{T}, A::MFBD{T}, From a583738fb8e3f3f6d4eb65e3e8383181c60ca180 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 10:08:40 +0100 Subject: [PATCH 30/99] Fix some failing state init tests --- src/blockdiagonals.jl | 31 +++++++++++++++++++++++ src/fast_linalg.jl | 33 +++++++++++-------------- src/initialization/classicsolverinit.jl | 4 +-- test/state_init.jl | 2 +- 4 files changed, 48 insertions(+), 22 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 9594a5a28..e6540d8ee 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -57,6 +57,37 @@ copy!(B::MFBD, A::MFBD) = begin return B end +# Standard LinearAlgebra.mul! +mul!(C::MFBD, A::MFBD, B::MFBD) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i]) + end + return C +end +mul!(C::MFBD, A::MFBD, B::MFBD, alpha::Number, beta::Number) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + end + return C +end +mul!(C::MFBD, A::Adjoint{<:Number,<:MFBD}, B::MFBD) = begin + @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + end + return C +end +mul!(C::MFBD, A::MFBD, B::Adjoint{<:Number,<:MFBD}) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + end + return C +end + +# Our fast _matmul! _matmul!( C::MFBD{T}, A::MFBD{T}, diff --git a/src/fast_linalg.jl b/src/fast_linalg.jl index b81ad9924..e0a01dcc4 100644 --- a/src/fast_linalg.jl +++ b/src/fast_linalg.jl @@ -30,26 +30,21 @@ _matmul!( _matmul!(C::AbstractVecOrMat, A::AbstractVecOrMat, b::Number) = @.. C = A * b _matmul!(C::AbstractVecOrMat, a::Number, B::AbstractVecOrMat) = @.. C = a * B # Matrix matrix products with diagonal matrices -_matmul!(C::AbstractMatrix, A::AbstractMatrix, B::Diagonal) = (@.. C = A * B.diag') -_matmul!(C::AbstractMatrix, A::Diagonal, B::AbstractMatrix) = (@.. C = A.diag * B) -_matmul!(C::AbstractMatrix, A::Diagonal, B::Diagonal) = @.. C = A * B +const MSR{T} = Union{SubArray{T},Matrix{T},Base.ReshapedArray{T}} +_matmul!(C::MSR, A::MSR, B::Diagonal) = + @.. C = A * B.diag' +_matmul!(C::MSR, A::Diagonal, B::MSR) = + (@.. C = A.diag * B) +_matmul!(C::MSR, A::Diagonal, B::Diagonal) = + (@.. C = A * B) +_matmul!(C::MSR{T}, A::MSR{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = + (@.. C = A * B.diag') +_matmul!(C::MSR{T}, A::Diagonal{T}, B::MSR{T}) where {T<:LinearAlgebra.BlasFloat} = + (@.. C = A.diag * B) +_matmul!(C::MSR{T}, A::Diagonal{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = + (@.. C = A * B) _matmul!( - C::AbstractMatrix{T}, - A::AbstractMatrix{T}, - B::Diagonal{T}, -) where {T<:LinearAlgebra.BlasFloat} = (@.. C = A * B.diag') -_matmul!( - C::AbstractMatrix{T}, - A::Diagonal{T}, - B::AbstractMatrix{T}, -) where {T<:LinearAlgebra.BlasFloat} = (@.. C = A.diag * B) -_matmul!( - C::AbstractMatrix{T}, - A::Diagonal{T}, - B::Diagonal{T}, -) where {T<:LinearAlgebra.BlasFloat} = @.. C = A * B -_matmul!( - C::AbstractMatrix{T}, + C::Matrix{T}, A::LowerTriangular{T}, B::UpperTriangular{T}, ) where {T<:LinearAlgebra.BlasFloat} = mul!(C, A, B) diff --git a/src/initialization/classicsolverinit.jl b/src/initialization/classicsolverinit.jl index e7bd1547f..123ddadee 100644 --- a/src/initialization/classicsolverinit.jl +++ b/src/initialization/classicsolverinit.jl @@ -90,7 +90,7 @@ end function rk_init_improve(cache::AbstractODEFilterCache, ts, us, dt) @unpack A, Q = cache # @unpack Ah, Qh = cache - @unpack x, x_pred, x_filt, measurement = cache + @unpack x, x_pred, x_filt, measurement, x_tmp = cache @unpack K1, C_Dxd, C_DxD, C_dxd, C_3DxD, C_d = cache @unpack backward_kernel = cache @@ -98,7 +98,7 @@ function rk_init_improve(cache::AbstractODEFilterCache, ts, us, dt) make_preconditioners!(cache, dt) @unpack P, PI = cache - _gaussian_mul!(x, P, x) + _gaussian_mul!(x, P, copy!(x_tmp, x)) preds = [] filts = [copy(x)] diff --git a/test/state_init.jl b/test/state_init.jl index 4330e5804..bf30c0082 100644 --- a/test/state_init.jl +++ b/test/state_init.jl @@ -64,7 +64,7 @@ import ODEProblemLibrary: prob_ode_fitzhughnagumo, prob_ode_pleiades prob, EK0(order=2, initialization=ClassicSolverInit(init_on_ddu=false)), ) - @test_broken init( + @test_nowarn init( prob, EK0(order=2, initialization=ClassicSolverInit(init_on_ddu=true)), ) From d9fc137bc4ed48fdbc878fbe6bdc4da861c74a4c Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 10:26:01 +0100 Subject: [PATCH 31/99] Improve the diffusion handling some more --- src/diffusions.jl | 64 ++++++++++++++++++++++++---------- src/filtering/markov_kernel.jl | 14 +++++--- src/filtering/predict.jl | 30 ++++++++-------- 3 files changed, 70 insertions(+), 38 deletions(-) diff --git a/src/diffusions.jl b/src/diffusions.jl index f4ae17cb3..a1a218b2b 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -6,41 +6,67 @@ isdynamic(diffusion::AbstractStaticDiffusion) = false isstatic(diffusion::AbstractDynamicDiffusion) = false isdynamic(diffusion::AbstractDynamicDiffusion) = true -apply_diffusion(Q::PSDMatrix{T,<:Matrix}, diffusion::Diagonal) where {T} = begin +# TODO Add a proper length description somewhere that exaplains better what this "diffusion" +# object is and how it is handled in this package. +apply_diffusion( + Q::PSDMatrix{<:Number,<:IsometricKroneckerProduct}, + diffusion::Number +) = PSDMatrix(Q.R * sqrt.(diffusion)) + +apply_diffusion( + Q::PSDMatrix{T,<:IsometricKroneckerProduct}, + diffusion::Diagonal{T,<:FillArrays.Fill} +) where {T} = apply_diffusion(Q, diffusion.diag.value) + +apply_diffusion( + Q::PSDMatrix{T,<:Matrix}, + diffusion::Diagonal +) where {T} = begin + # @warn "This should ideally not be called; TODO" d = size(diffusion, 1) q = size(Q, 1) ÷ d - 1 return PSDMatrix(Q.R * sqrt.(kron(diffusion, I(q + 1)))) end + apply_diffusion( - Q::PSDMatrix{T,<:IsometricKroneckerProduct}, - diffusion::Diagonal{T,<:FillArrays.Fill}, -) where {T} = begin - PSDMatrix(Q.R * sqrt.(diffusion.diag.value)) -end -apply_diffusion(Q::PSDMatrix{T,<:MFBD}, diffusion::Diagonal) where {T} = begin - PSDMatrix( - MFBD([ - Q.R.blocks[i] * sqrt.(diffusion.diag[i]) for i in eachindex(Q.R.blocks) - ]), - ) -end + Q::PSDMatrix{T,<:MFBD}, diffusion::Diagonal +) where {T} = PSDMatrix( + MFBD([blocks(Q.R)[i] * sqrt.(diffusion.diag[i]) for i in eachindex(blocks(Q.R))])) -apply_diffusion!(Q::PSDMatrix, diffusion::Diagonal{T,<:FillArrays.Fill}) where {T} = - rmul!(Q.R, sqrt.(diffusion.diag.value)) +apply_diffusion!( + Q::PSDMatrix, + diffusion::Diagonal{T,<:FillArrays.Fill} +) where {T} = rmul!(Q.R, sqrt.(diffusion.diag.value)) apply_diffusion!( Q::PSDMatrix{T,<:MFBD}, diffusion::Diagonal{T,<:Vector}, -) where {T} = +) where {T} = begin @simd ivdep for i in eachindex(blocks(Q.R)) rmul!(blocks(Q.R)[i], diffusion.diag[i]) end +end apply_diffusion!( out::PSDMatrix, Q::PSDMatrix, - diffusion::Diagonal{T,<:FillArrays.Fill}, -) where {T} = - rmul!(Q.R, sqrt.(diffusion.diag.value)) + diffusion::Number +) = _matmul!(out.R, Q.R, sqrt.(diffusion)) +apply_diffusion!( + out::PSDMatrix, + Q::PSDMatrix, + diffusion::Diagonal{<:Number,<:FillArrays.Fill}, +) = apply_diffusion!(out, Q, diffusion.diag.value) +apply_diffusion!( + out::PSDMatrix, + Q::PSDMatrix, + diffusion::Diagonal +) = begin + @warn "This is not yet implemented efficiently; TODO" + d = size(diffusion, 1) + D = size(Q, 1) + q = D ÷ d - 1 + _matmul!(out.R, Q.R, sqrt.(kron(Eye(d) * diffusion, Eye(q + 1)))) +end estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = error() diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index 605bec3df..5136b1a25 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -223,7 +223,7 @@ function compute_backward_kernel!( _matmul!(view(Λ.R, 1:D, 1:D), x.Σ.R, C_DxD) # Λ.R[D+1:2D, 1:D] = (G * Q.R')' if !isone(diffusion) - _matmul!(C_DxD, Q.R, sqrt.(diffusion)) + apply_diffusion!(PSDMatrix(C_DxD), Q, diffusion) _matmul!(view(Λ.R, D+1:2D, 1:D), C_DxD, G') else _matmul!(view(Λ.R, D+1:2D, 1:D), Q.R, G') @@ -238,7 +238,7 @@ function compute_backward_kernel!( x::SRGaussian{T,<:IsometricKroneckerProduct}, K::KT2; C_DxD::AbstractMatrix, - diffusion=1, + diffusion::Union{Number,Diagonal}=1, ) where { T, KT1<:AffineNormalKernel{ @@ -261,9 +261,12 @@ function compute_backward_kernel!( _x = Gaussian(reshape_no_alloc(x.μ, Q, d), PSDMatrix(x.Σ.R.B)) _K = AffineNormalKernel(K.A.B, reshape_no_alloc(K.b, Q, d), PSDMatrix(K.C.R.B)) _C_DxD = C_DxD.B + _diffusion = + diffusion isa Number ? diffusion : + diffusion isa IsometricKroneckerProduct ? diffusion.B : diffusion return compute_backward_kernel!( - _Kout, _x_pred, _x, _K; C_DxD=_C_DxD, diffusion=diffusion) + _Kout, _x_pred, _x, _K; C_DxD=_C_DxD, diffusion=_diffusion) end function compute_backward_kernel!( @@ -272,7 +275,7 @@ function compute_backward_kernel!( x::SRGaussian{T,<:MFBD}, K::KT2; C_DxD::AbstractMatrix, - diffusion=1, + diffusion::Union{Number,Diagonal}=1, ) where { T, KT1<:AffineNormalKernel{ @@ -308,8 +311,9 @@ function compute_backward_kernel!( PSDMatrix(K.C.R.blocks[i]), ) _C_DxD = C_DxD.blocks[i] + _diffusion = diffusion isa Number ? diffusion : diffusion[i] compute_backward_kernel!( - _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=diffusion, + _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=_diffusion ) end return Kout diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 38a0f123e..fb3de6164 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -77,20 +77,22 @@ function predict_cov!( _matmul!(view(R, 1:D, 1:D), Σ_curr.R, Ah') if !isone(diffusion) - if diffusion isa Number - _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt(diffusion)) - elseif diffusion isa Diagonal{<:Number,<:FillArrays.Fill} - _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) - else - @warn "This is not yet implemented efficiently; TODO" - d = size(diffusion, 1) - q = D ÷ d - 1 - _matmul!( - view(R, D+1:2D, 1:D), - Qh.R, - sqrt.(kron(Eye(d) * diffusion, Eye(q + 1))), - ) - end + apply_diffusion!(PSDMatrix(view(R, D+1:2D, 1:D)), Qh, diffusion) + + # if diffusion isa Number + # _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt(diffusion)) + # elseif diffusion isa Diagonal{<:Number,<:FillArrays.Fill} + # _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) + # else + # @warn "This is not yet implemented efficiently; TODO" + # d = size(diffusion, 1) + # q = D ÷ d - 1 + # _matmul!( + # view(R, D+1:2D, 1:D), + # Qh.R, + # sqrt.(kron(Eye(d) * diffusion, Eye(q + 1))), + # ) + # end else @.. R[D+1:2D, 1:D] = Qh.R end From 385d53478501f40bc422c6f4c55b65f05cc23dde Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 10:26:11 +0100 Subject: [PATCH 32/99] Enable the EK0 again with priors that are not Kronecker --- src/algorithms.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/algorithms.jl b/src/algorithms.jl index fca34ec5e..586ef7b8f 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -52,7 +52,6 @@ function covariance_structure(::Type{Alg}, prior, diffusionmodel) where {Alg<:Ab return BlockDiagonalCovariance end else - error() # This is not great as other priors can be Kronecker too; TODO return DenseCovariance end From ca925913acefe7b8cf819ce79e88bca4310fe708 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 10:27:50 +0100 Subject: [PATCH 33/99] Remove one test case that's not yet supported --- test/data_likelihoods.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/data_likelihoods.jl b/test/data_likelihoods.jl index 26a6b2596..dacb2b6fd 100644 --- a/test/data_likelihoods.jl +++ b/test/data_likelihoods.jl @@ -35,7 +35,7 @@ kwargs = ( @testset "$alg" for alg in ( EK1(), EK1(diffusionmodel=FixedDiffusion()), - EK1(diffusionmodel=FixedMVDiffusion(rand(2), false)), + # EK1(diffusionmodel=FixedMVDiffusion(rand(2), false)), # not yet supported EK1(prior=IOUP(3, -1)), EK1(prior=Matern(3, 1.5)), EK1(prior=IOUP(3, update_rate_parameter=true)), From c766d6e3a099ae224b8fdf33888966693f7f1edc Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 10:37:39 +0100 Subject: [PATCH 34/99] Significantly speed up the secondorderodeproblem tests --- test/secondorderode.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/secondorderode.jl b/test/secondorderode.jl index 51484ff8f..0a75f80a9 100644 --- a/test/secondorderode.jl +++ b/test/secondorderode.jl @@ -4,13 +4,13 @@ using LinearAlgebra using Test function twobody(du, u, p, t) - R3 = norm(u[1:2])^3 - @. du[3:4] = -u[1:2] / R3 - @. du[1:2] = u[3:4] + R3 = norm(u[3:4])^3 + @. du[3:4] = u[1:2] + @. du[1:2] = -u[3:4] / R3 end u0, du0 = [0.4, 0.0], [0.0, 2.0] tspan = (0, 0.1) -prob_base = ODEProblem(twobody, [u0...; du0...], tspan) +prob_base = ODEProblem(twobody, [du0...; u0...], tspan) function twobody2_iip(ddu, du, u, p, t) R3 = norm(u)^3 @@ -24,7 +24,7 @@ function twobody2_oop(du, u, p, t) end prob_oop = SecondOrderODEProblem(twobody2_oop, du0, u0, tspan) -appxsol = solve(prob_iip, Vern9(), abstol=1e-10, reltol=1e-10) +appxsol = solve(prob_base, Vern9(), abstol=1e-9, reltol=1e-6) @testset "$S" for (S, _prob) in (("IIP", prob_iip), ("OOP", prob_oop)) @testset "$alg" for alg in ( From ed207606fe9808dba8f7a0a3131e107d2255a9b7 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 10:37:59 +0100 Subject: [PATCH 35/99] Remove a test that currently fails --- test/secondorderode.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/secondorderode.jl b/test/secondorderode.jl index 0a75f80a9..468a174b5 100644 --- a/test/secondorderode.jl +++ b/test/secondorderode.jl @@ -34,7 +34,7 @@ appxsol = solve(prob_base, Vern9(), abstol=1e-9, reltol=1e-6) EK1(initialization=ForwardDiffInit(2)), # EK1(initialization=ClassicSolverInit()), # unstable for this problem EK1(diffusionmodel=FixedDiffusion()), - EK0(diffusionmodel=FixedMVDiffusion()), + # EK0(diffusionmodel=FixedMVDiffusion()), # EK0(diffusionmodel=DynamicMVDiffusion()), ) sol = solve(_prob, alg) From 3f9e00d1c9573701d5ac64fee6d00fca69c9b6ac Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 11:44:56 +0100 Subject: [PATCH 36/99] Fix many of the tests that I had to temporally remove --- src/algorithms.jl | 8 -------- src/covariance_structure.jl | 2 +- src/projection.jl | 2 +- test/correctness.jl | 7 +++---- test/diffusions.jl | 2 +- 5 files changed, 6 insertions(+), 15 deletions(-) diff --git a/src/algorithms.jl b/src/algorithms.jl index 586ef7b8f..9a49481a4 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -33,14 +33,6 @@ function ekargcheck( ) end end - if diffusionmodel isa DynamicMVDiffusion && - covariance_factorization == BlockDiagonalCovariance - throw( - ArgumentError( - "Currenty the `DynamicMVDiffusion` does not work properly with the `BlockDiagonalCovariance`. Use `DenseCovariance` instead, or change the diffusionmodel to a scalar one and use `DynamicDiffusion`.", - ), - ) - end end function covariance_structure(::Type{Alg}, prior, diffusionmodel) where {Alg<:AbstractEK} diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index abf9e3331..9de3c5223 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -46,7 +46,7 @@ end to_factorized_matrix(::DenseCovariance, M::AbstractMatrix) = Matrix(M) to_factorized_matrix(::IsometricKroneckerCovariance, M::IsometricKroneckerProduct) = M to_factorized_matrix(C::BlockDiagonalCovariance, M::IsometricKroneckerProduct) = - MFBD([M.B for _ in 1:C.d]) + MFBD([copy(M.B) for _ in 1:C.d]) for FT in [:DenseCovariance, :IsometricKroneckerCovariance, :BlockDiagonalCovariance] @eval to_factorized_matrix(FAC::$FT, M::PSDMatrix) = diff --git a/src/projection.jl b/src/projection.jl index ddac1059c..8cf89f9ac 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -36,7 +36,7 @@ function projection(C::BlockDiagonalCovariance{elType}) where {elType} if deriv <= C.q e_i[deriv+1] = 1 end - return MFBD([e_i' for _ in 1:C.d]) + return MFBD([copy(e_i)' for _ in 1:C.d]) end return Proj end diff --git a/test/correctness.jl b/test/correctness.jl index 1cec088ef..ac9a34235 100644 --- a/test/correctness.jl +++ b/test/correctness.jl @@ -15,7 +15,7 @@ CONSTANT_ALGS = ( EK0(order=5, smooth=false) => 1e-10, EK0(order=3, smooth=false, diffusionmodel=FixedDiffusion()) => 1e-7, EK0(order=3, smooth=false, diffusionmodel=FixedMVDiffusion()) => 1e-7, - # EK0(order=3, smooth=false, diffusionmodel=DynamicMVDiffusion()) => 1e-8, + EK0(order=3, smooth=false, diffusionmodel=DynamicMVDiffusion()) => 1e-8, EK0(order=3, smooth=false, initialization=ClassicSolverInit()) => 1e-7, EK0(order=3, smooth=false, initialization=SimpleInit()) => 1e-5, EK0( @@ -34,7 +34,7 @@ CONSTANT_ALGS = ( EK0(order=3, smooth=true) => 1e-8, EK0(order=3, smooth=true, diffusionmodel=FixedDiffusion()) => 2e-8, EK0(order=3, smooth=true, diffusionmodel=FixedMVDiffusion()) => 1e-7, - # EK0(order=3, smooth=true, diffusionmodel=DynamicMVDiffusion()) => 1e-8, + EK0(order=3, smooth=true, diffusionmodel=DynamicMVDiffusion()) => 1e-8, EK1(order=3, smooth=true) => 1e-8, EK1(order=3, smooth=true, diffusionmodel=FixedDiffusion()) => 1e-8, # Priors @@ -49,10 +49,9 @@ ADAPTIVE_ALGS = ( EK0(order=3) => 1e-4, EK0(order=5) => 1e-5, EK0(order=8) => 2e-5, - # EK0(order=3, diffusionmodel=DynamicMVDiffusion()) => 5e-5, + EK0(order=3, diffusionmodel=DynamicMVDiffusion()) => 5e-5, EK0(order=3, initialization=ClassicSolverInit()) => 5e-5, EK0(order=3, initialization=SimpleInit()) => 1e-4, - # EK0(order=3, diffusionmodel=DynamicMVDiffusion(), initialization=ClassicSolverInit()) => 4e-5, EK1(order=2) => 2e-5, EK1(order=3) => 1e-5, EK1(order=5) => 1e-6, diff --git a/test/diffusions.jl b/test/diffusions.jl index a0918d8fd..1b1d30aa3 100644 --- a/test/diffusions.jl +++ b/test/diffusions.jl @@ -44,7 +44,7 @@ import ODEProblemLibrary: prob_ode_fitzhughnagumo @test appxsol.errors[:final] < 1e-5 end - @test_skip @testset "Time-Varying Diagonal Diffusion" begin + @testset "Time-Varying Diagonal Diffusion" begin sol = solve( prob, EK0(diffusionmodel=DynamicMVDiffusion(), smooth=false), From 15d3c287d35994f90426d3fff828b806554dfd76 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 11:55:13 +0100 Subject: [PATCH 37/99] Make it more obvious that BlockDiagonals and second order ODEs are not working --- src/caches.jl | 29 +++++++++++++---------------- src/projection.jl | 4 ++++ 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/caches.jl b/src/caches.jl index c0c45e84a..88a4f10a1 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -181,22 +181,19 @@ function OrdinaryDiffEq.alg_cache( # ddu = factorized_similar(FAC, length(u), length(u)) ddu = similar(u, length(u), length(u)) _d = is_secondorder_ode ? 2d : d - pu_tmp = Gaussian( - similar(Array{uElType}, _d), - PSDMatrix( - if FAC isa IsometricKroneckerCovariance - if is_secondorder_ode - Kronecker.kronecker(similar(Matrix{uElType}, D ÷ d, _d ÷ d), I(d)) - else - factorized_similar(FAC, D, d) - end - elseif FAC isa BlockDiagonalCovariance - factorized_similar(FAC, D, d) - else - similar(Matrix{uElType}, D, _d) - end, - ), - ) + pu_tmp = if is_secondorder_ode + Gaussian(similar(Array{uElType}, 2d), + PSDMatrix( + if FAC isa IsometricKroneckerCovariance + Kronecker.kronecker(similar(Matrix{uElType}, D ÷ d, _d ÷ d), I(d)) + elseif FAC isa BlockDiagonalCovariance + error("I have no idea") + else + similar(Matrix{uElType}, D, _d) + end)) + else + Gaussian(similar(Array{uElType}, d), PSDMatrix(factorized_similar(FAC, D, d))) + end K = factorized_similar(FAC, D, d) G = factorized_similar(FAC, D, D) diff --git a/src/projection.jl b/src/projection.jl index 8cf89f9ac..7a4ce3734 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -58,6 +58,10 @@ function solution_space_projection(C::IsometricKroneckerCovariance, is_secondord return Proj(0) end end +function solution_space_projection(C::BlockDiagonalCovariance, is_secondorder_ode) + Proj = projection(C) + error("No idea!") +end struct KroneckerSecondOrderODESolutionProjector{T,FAC,M,M2} <: AbstractMatrix{T} covariance_structure::FAC From 4c3be25af0f9cbf7ff53948a63b6cb4b0d7a6999 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 12:01:12 +0100 Subject: [PATCH 38/99] Rename our BlockDiagonals to ProbNumDiffEqBlockDiagonal (+ shortcut) --- src/blockdiagonals.jl | 81 +++++++++++++++++----------------- src/covariance_structure.jl | 6 +-- src/derivative_utils.jl | 2 +- src/diffusions.jl | 12 ++--- src/filtering/markov_kernel.jl | 18 ++++---- src/filtering/predict.jl | 12 ++--- src/filtering/update.jl | 18 ++++---- src/perform_step.jl | 4 +- src/preconditioning.jl | 8 ++-- src/projection.jl | 2 +- 10 files changed, 82 insertions(+), 81 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index e6540d8ee..11c4dc89f 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -3,23 +3,24 @@ BlockDiagonals.jl didn't cut it, so we're rolling our own. TODO: Add a way to convert to a `BlockDiagonal`. """ -struct MinimalAndFastBlockDiagonal{T<:Number,V<:AbstractMatrix{T}} <: AbstractMatrix{T} +struct ProbNumDiffEqBlockDiagonal{T<:Number,V<:AbstractMatrix{T}} <: AbstractMatrix{T} blocks::Vector{V} - function MinimalAndFastBlockDiagonal{T,V}( + function ProbNumDiffEqBlockDiagonal{T,V}( blocks::Vector{V}, ) where {T,V<:AbstractMatrix{T}} return new{T,V}(blocks) end end -function MinimalAndFastBlockDiagonal(blocks::Vector{V}) where {T,V<:AbstractMatrix{T}} - return MinimalAndFastBlockDiagonal{T,V}(blocks) +function ProbNumDiffEqBlockDiagonal(blocks::Vector{V}) where {T,V<:AbstractMatrix{T}} + return ProbNumDiffEqBlockDiagonal{T,V}(blocks) end -const MFBD = MinimalAndFastBlockDiagonal -blocks(B::MFBD) = B.blocks -nblocks(B::MFBD) = length(B.blocks) -size(B::MFBD) = (sum(size.(blocks(B), 1)), sum(size.(blocks(B), 2))) +const BlockDiag = ProbNumDiffEqBlockDiagonal -function _block_indices(B::MFBD, i::Integer, j::Integer) +blocks(B::BlockDiag) = B.blocks +nblocks(B::BlockDiag) = length(B.blocks) +size(B::BlockDiag) = (sum(size.(blocks(B), 1)), sum(size.(blocks(B), 2))) + +function _block_indices(B::BlockDiag, i::Integer, j::Integer) all((0, 0) .< (i, j) .<= size(B)) || throw(BoundsError(B, (i, j))) # find the on-diagonal block `p` in column `j` p = 0 @@ -36,7 +37,7 @@ function _block_indices(B::MFBD, i::Integer, j::Integer) return p, i, j end Base.@propagate_inbounds function Base.getindex( - B::MFBD{T}, + B::BlockDiag{T}, i::Integer, j::Integer, ) where {T} @@ -45,11 +46,11 @@ Base.@propagate_inbounds function Base.getindex( @inbounds return p > 0 ? blocks(B)[p][i, end+j] : zero(T) end -Base.view(::MFBD, idxs...) = - throw(ErrorException("`MinimalAndFastBlockDiagonal` does not support views!")) +Base.view(::BlockDiag, idxs...) = + throw(ErrorException("`BlockDiag` does not support views!")) -copy(B::MFBD) = MFBD(copy.(blocks(B))) -copy!(B::MFBD, A::MFBD) = begin +copy(B::BlockDiag) = BlockDiag(copy.(blocks(B))) +copy!(B::BlockDiag, A::BlockDiag) = begin @assert length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(B)) copy!(B.blocks[i], A.blocks[i]) @@ -58,28 +59,28 @@ copy!(B::MFBD, A::MFBD) = begin end # Standard LinearAlgebra.mul! -mul!(C::MFBD, A::MFBD, B::MFBD) = begin +mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag) = begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i]) end return C end -mul!(C::MFBD, A::MFBD, B::MFBD, alpha::Number, beta::Number) = begin +mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag, alpha::Number, beta::Number) = begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) end return C end -mul!(C::MFBD, A::Adjoint{<:Number,<:MFBD}, B::MFBD) = begin +mul!(C::BlockDiag, A::Adjoint{<:Number,<:BlockDiag}, B::BlockDiag) = begin @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) @inbounds mul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) end return C end -mul!(C::MFBD, A::MFBD, B::Adjoint{<:Number,<:MFBD}) = begin +mul!(C::BlockDiag, A::BlockDiag, B::Adjoint{<:Number,<:BlockDiag}) = begin @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) @simd ivdep for i in eachindex(blocks(C)) @inbounds mul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) @@ -89,9 +90,9 @@ end # Our fast _matmul! _matmul!( - C::MFBD{T}, - A::MFBD{T}, - B::MFBD{T}, + C::BlockDiag{T}, + A::BlockDiag{T}, + B::BlockDiag{T}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) @@ -101,9 +102,9 @@ _matmul!( end _matmul!( - C::MFBD{T}, - A::MFBD{T}, - B::MFBD{T}, + C::BlockDiag{T}, + A::BlockDiag{T}, + B::BlockDiag{T}, alpha::Number, beta::Number, ) where {T<:LinearAlgebra.BlasFloat} = begin @@ -115,9 +116,9 @@ _matmul!( end _matmul!( - C::MFBD{T}, - A::MFBD{T}, - B::Adjoint{T,<:MFBD{T}}, + C::BlockDiag{T}, + A::BlockDiag{T}, + B::Adjoint{T,<:BlockDiag{T}}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) @simd ivdep for i in eachindex(blocks(C)) @@ -127,9 +128,9 @@ _matmul!( end _matmul!( - C::MFBD{T}, - A::Adjoint{T,<:MFBD{T}}, - B::MFBD{T}, + C::BlockDiag{T}, + A::Adjoint{T,<:BlockDiag{T}}, + B::BlockDiag{T}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) @@ -140,7 +141,7 @@ end _matmul!( C::AbstractVector{T}, - A::MFBD{T}, + A::BlockDiag{T}, B::AbstractVector{T}, ) where {T<:LinearAlgebra.BlasFloat} = begin @assert size(A, 2) == length(B) @@ -155,21 +156,21 @@ _matmul!( return C end -LinearAlgebra.rmul!(B::MFBD, n::Number) = @simd ivdep for i in eachindex(B.blocks) +LinearAlgebra.rmul!(B::BlockDiag, n::Number) = @simd ivdep for i in eachindex(B.blocks) rmul!(B.blocks[i], n) end -LinearAlgebra.adjoint(B::MFBD) = Adjoint(B) +LinearAlgebra.adjoint(B::BlockDiag) = Adjoint(B) -Base.:*(A::MFBD, B::MFBD) = begin +Base.:*(A::BlockDiag, B::BlockDiag) = begin @assert length(A.blocks) == length(B.blocks) - return MFBD([blocks(A)[i] * blocks(B)[i] for i in eachindex(B.blocks)]) + return BlockDiag([blocks(A)[i] * blocks(B)[i] for i in eachindex(B.blocks)]) end -Base.:*(A::Adjoint{T,<:MFBD}, B::MFBD) where {T} = begin +Base.:*(A::Adjoint{T,<:BlockDiag}, B::BlockDiag) where {T} = begin @assert length(A.parent.blocks) == length(B.blocks) - return MFBD([A.parent.blocks[i]' * B.blocks[i] for i in eachindex(B.blocks)]) + return BlockDiag([A.parent.blocks[i]' * B.blocks[i] for i in eachindex(B.blocks)]) end -Base.:*(A::MFBD, B::Adjoint{T,<:MFBD}) where {T} = begin +Base.:*(A::BlockDiag, B::Adjoint{T,<:BlockDiag}) where {T} = begin @assert length(A.blocks) == length(B.parent.blocks) - return MFBD([A.blocks[i] * B.parent.blocks[i]' for i in eachindex(B.parent.blocks)]) + return BlockDiag([A.blocks[i] * B.parent.blocks[i]' for i in eachindex(B.parent.blocks)]) end -Base.:*(A::UniformScaling, B::MFBD) = MFBD([A * blocks(B)[i] for i in eachindex(B.blocks)]) +Base.:*(A::UniformScaling, B::BlockDiag) = BlockDiag([A * blocks(B)[i] for i in eachindex(B.blocks)]) diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index 9de3c5223..234a77688 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -34,19 +34,19 @@ factorized_zeros(C::BlockDiagonalCovariance{T}, sizes...) where {T} = begin for s in sizes @assert s % C.d == 0 end - return MFBD([Array{T}(calloc, (s ÷ C.d for s in sizes)...) for _ in 1:C.d]) + return BlockDiag([Array{T}(calloc, (s ÷ C.d for s in sizes)...) for _ in 1:C.d]) end factorized_similar(C::BlockDiagonalCovariance{T}, size1, size2) where {T} = begin for s in (size1, size2) @assert s % C.d == 0 end - return MFBD([similar(Matrix{T}, size1 ÷ C.d, size2 ÷ C.d) for _ in 1:C.d]) + return BlockDiag([similar(Matrix{T}, size1 ÷ C.d, size2 ÷ C.d) for _ in 1:C.d]) end to_factorized_matrix(::DenseCovariance, M::AbstractMatrix) = Matrix(M) to_factorized_matrix(::IsometricKroneckerCovariance, M::IsometricKroneckerProduct) = M to_factorized_matrix(C::BlockDiagonalCovariance, M::IsometricKroneckerProduct) = - MFBD([copy(M.B) for _ in 1:C.d]) + BlockDiag([copy(M.B) for _ in 1:C.d]) for FT in [:DenseCovariance, :IsometricKroneckerCovariance, :BlockDiagonalCovariance] @eval to_factorized_matrix(FAC::$FT, M::PSDMatrix) = diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index d3b3976bf..1a8ec21b3 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -18,7 +18,7 @@ function calc_H!(H, integ, cache) OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) @unpack C_dxd = cache - if C_dxd isa MFBD + if C_dxd isa BlockDiag @simd ivdep for i in eachindex(blocks(C_dxd)) @assert length(C_dxd.blocks[i]) == 1 C_dxd.blocks[i][1] = ddu[i, i] diff --git a/src/diffusions.jl b/src/diffusions.jl index a1a218b2b..41edd00fb 100644 --- a/src/diffusions.jl +++ b/src/diffusions.jl @@ -29,16 +29,16 @@ apply_diffusion( end apply_diffusion( - Q::PSDMatrix{T,<:MFBD}, diffusion::Diagonal + Q::PSDMatrix{T,<:BlockDiag}, diffusion::Diagonal ) where {T} = PSDMatrix( - MFBD([blocks(Q.R)[i] * sqrt.(diffusion.diag[i]) for i in eachindex(blocks(Q.R))])) + BlockDiag([blocks(Q.R)[i] * sqrt.(diffusion.diag[i]) for i in eachindex(blocks(Q.R))])) apply_diffusion!( Q::PSDMatrix, diffusion::Diagonal{T,<:FillArrays.Fill} ) where {T} = rmul!(Q.R, sqrt.(diffusion.diag.value)) apply_diffusion!( - Q::PSDMatrix{T,<:MFBD}, + Q::PSDMatrix{T,<:BlockDiag}, diffusion::Diagonal{T,<:Vector}, ) where {T} = begin @simd ivdep for i in eachindex(blocks(Q.R)) @@ -131,7 +131,7 @@ function estimate_global_diffusion(::FixedDiffusion, integ) diffusion_t = if S isa IsometricKroneckerProduct @assert length(S.B) == 1 dot(v, e) / d / S.B[1] - elseif S isa MFBD + elseif S isa BlockDiag @assert length(S.blocks) == d @assert length(S.blocks[1]) == 1 @simd ivdep for i in eachindex(e) @@ -230,7 +230,7 @@ function local_scalar_diffusion(cache) σ² = if HQH isa IsometricKroneckerProduct @assert length(HQH.B) == 1 dot(z, e) / d / HQH.B[1] - elseif HQH isa MFBD + elseif HQH isa BlockDiag @assert length(HQH.blocks) == d @assert length(HQH.blocks[1]) == 1 for i in eachindex(e) @@ -271,7 +271,7 @@ function local_diagonal_diffusion(cache) # Q_11 = dot(c1, c1) # @assert - Q_11 = if Qh.R isa MFBD + Q_11 = if Qh.R isa BlockDiag for i in 1:d c1 = _matmul!( view(cache.C_Dxd.blocks[i], :, 1:1), diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index 5136b1a25..a97b37275 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -121,12 +121,12 @@ function marginalize_cov!( end function marginalize_cov!( - Σ_out::PSDMatrix{T,<:MFBD}, - Σ_curr::PSDMatrix{T,<:MFBD}, + Σ_out::PSDMatrix{T,<:BlockDiag}, + Σ_curr::PSDMatrix{T,<:BlockDiag}, K::AffineNormalKernel{ <:AbstractMatrix, <:Any, - <:PSDMatrix{S,<:MFBD}, + <:PSDMatrix{S,<:BlockDiag}, }; C_DxD::AbstractMatrix, C_3DxD::AbstractMatrix, @@ -271,22 +271,22 @@ end function compute_backward_kernel!( Kout::KT1, - xpred::SRGaussian{T,<:MFBD}, - x::SRGaussian{T,<:MFBD}, + xpred::SRGaussian{T,<:BlockDiag}, + x::SRGaussian{T,<:BlockDiag}, K::KT2; C_DxD::AbstractMatrix, diffusion::Union{Number,Diagonal}=1, ) where { T, KT1<:AffineNormalKernel{ - <:MFBD, + <:BlockDiag, <:AbstractVector, - <:PSDMatrix{T,<:MFBD}, + <:PSDMatrix{T,<:BlockDiag}, }, KT2<:AffineNormalKernel{ - <:MFBD, + <:BlockDiag, <:Any, - <:PSDMatrix{T,<:MFBD}, + <:PSDMatrix{T,<:BlockDiag}, }, } d = length(blocks(xpred.Σ.R)) diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index fb3de6164..a0768080c 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -132,12 +132,12 @@ end # BlockDiagonal version function predict_cov!( - Σ_out::PSDMatrix{T,<:MFBD}, - Σ_curr::PSDMatrix{T,<:MFBD}, - Ah::MFBD, - Qh::PSDMatrix{S,<:MFBD}, - C_DxD::MFBD, - C_2DxD::MFBD, + Σ_out::PSDMatrix{T,<:BlockDiag}, + Σ_curr::PSDMatrix{T,<:BlockDiag}, + Ah::BlockDiag, + Qh::PSDMatrix{S,<:BlockDiag}, + C_DxD::BlockDiag, + C_2DxD::BlockDiag, diffusion::Union{Number,Diagonal}, ) where {T,S} for i in eachindex(blocks(Σ_out.R)) diff --git a/src/filtering/update.jl b/src/filtering/update.jl index cf7ba720d..3b07b0a45 100644 --- a/src/filtering/update.jl +++ b/src/filtering/update.jl @@ -194,19 +194,19 @@ function update!( end function update!( - x_out::SRGaussian{T,<:MFBD}, - x_pred::SRGaussian{T,<:MFBD}, + x_out::SRGaussian{T,<:BlockDiag}, + x_pred::SRGaussian{T,<:BlockDiag}, measurement::Gaussian{ <:AbstractVector, - <:Union{<:PSDMatrix{T,<:MFBD},<:MFBD}, + <:Union{<:PSDMatrix{T,<:BlockDiag},<:BlockDiag}, }, - H::MFBD, - K1_cache::MFBD, - K2_cache::MFBD, - M_cache::MFBD, - C_dxd::MFBD, + H::BlockDiag, + K1_cache::BlockDiag, + K2_cache::BlockDiag, + M_cache::BlockDiag, + C_dxd::BlockDiag, C_d::AbstractVector; - R::Union{Nothing,PSDMatrix{T,<:MFBD}}=nothing, + R::Union{Nothing,PSDMatrix{T,<:BlockDiag}}=nothing, ) where {T} d = length(blocks(x_out.Σ.R)) q = size(blocks(x_out.Σ.R)[1], 1) - 1 diff --git a/src/perform_step.jl b/src/perform_step.jl index 30b6611fa..26a833a24 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -226,7 +226,7 @@ function estimate_errors!(cache::AbstractODEFilterCache) _Q = apply_diffusion(Qh, local_diffusion) _matmul!(R, _Q.R, H') error_estimate = view(cache.tmp, 1:d) - if R isa MFBD + if R isa BlockDiag for i in eachindex(R.blocks) error_estimate[i] = sum(abs2, R.blocks[i]) end @@ -247,7 +247,7 @@ function estimate_errors!(cache::AbstractODEFilterCache) error_estimate = view(cache.tmp, 1:d) if R isa IsometricKroneckerProduct error_estimate .= sum(abs2, R.B) - elseif R isa MFBD + elseif R isa BlockDiag for i in eachindex(blocks(R)) error_estimate[i] = sum(abs2, R.blocks[i]) end diff --git a/src/preconditioning.jl b/src/preconditioning.jl index 30967a1c6..b6f323853 100644 --- a/src/preconditioning.jl +++ b/src/preconditioning.jl @@ -9,8 +9,8 @@ function init_preconditioner(C::DenseCovariance{elType}) where {elType} return P, PI end function init_preconditioner(C::BlockDiagonalCovariance{elType}) where {elType} - P = MFBD([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) - PI = MFBD([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) + P = BlockDiag([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) + PI = BlockDiag([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) return P, PI end @@ -46,7 +46,7 @@ end return P end -@fastmath @inbounds function make_preconditioner!(P::MFBD, h, d, q) +@fastmath @inbounds function make_preconditioner!(P::BlockDiag, h, d, q) val = factorial(q) / h^(q + 1 / 2) @simd ivdep for j in 0:q for M in P.blocks @@ -80,7 +80,7 @@ end end @fastmath @inbounds function make_preconditioner_inv!( - PI::MFBD, h, d, q) + PI::BlockDiag, h, d, q) val = h^(q + 1 / 2) / factorial(q) @simd ivdep for j in 0:q for M in PI.blocks diff --git a/src/projection.jl b/src/projection.jl index 7a4ce3734..de7b4ad8c 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -36,7 +36,7 @@ function projection(C::BlockDiagonalCovariance{elType}) where {elType} if deriv <= C.q e_i[deriv+1] = 1 end - return MFBD([copy(e_i)' for _ in 1:C.d]) + return BlockDiag([copy(e_i)' for _ in 1:C.d]) end return Proj end From aafb416df6ba0835412852126c41542e37b32a3f Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 12:21:17 +0100 Subject: [PATCH 39/99] Add a BlockDiagonals extension to transfomr ours to theirs --- Project.toml | 3 ++- ext/BlockDiagonalsExt.jl | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) create mode 100644 ext/BlockDiagonalsExt.jl diff --git a/Project.toml b/Project.toml index ed5af8bb9..030a13470 100644 --- a/Project.toml +++ b/Project.toml @@ -7,7 +7,6 @@ version = "0.15.0" ArrayAllocators = "c9d4266f-a5cb-439d-837c-c97b191379f5" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" -DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" @@ -39,10 +38,12 @@ TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea" ToeplitzMatrices = "c751599d-da0a-543b-9d20-d0a503d91d24" [weakdeps] +BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" [extensions] +BlockDiagonalsExt = "BlockDiagonals" DiffEqDevToolsExt = "DiffEqDevTools" RecipesBaseExt = "RecipesBase" diff --git a/ext/BlockDiagonalsExt.jl b/ext/BlockDiagonalsExt.jl new file mode 100644 index 000000000..c50bdeac0 --- /dev/null +++ b/ext/BlockDiagonalsExt.jl @@ -0,0 +1,8 @@ +module BlockDiagonalsExt + +import ProbNumDiffEq: ProbNumDiffEqBlockDiagonal, blocks +using BlockDiagonals + +BlockDiagonal(M::ProbNumDiffEqBlockDiagonal) = BlockDiagonal(blocks(M)) + +end From 65018784314ac055f9c6bf940ed0cd5ab7559b8b Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 15:52:30 +0100 Subject: [PATCH 40/99] Add unit-tests for our `BlockDiag`s --- ext/BlockDiagonalsExt.jl | 2 +- src/ProbNumDiffEq.jl | 2 +- src/blockdiagonals.jl | 44 +++++++++++++++++----------- test/core/blockdiagonals.jl | 57 +++++++++++++++++++++++++++++++++++++ test/runtests.jl | 3 ++ 5 files changed, 89 insertions(+), 19 deletions(-) create mode 100644 test/core/blockdiagonals.jl diff --git a/ext/BlockDiagonalsExt.jl b/ext/BlockDiagonalsExt.jl index c50bdeac0..37a7f6bfc 100644 --- a/ext/BlockDiagonalsExt.jl +++ b/ext/BlockDiagonalsExt.jl @@ -1,7 +1,7 @@ module BlockDiagonalsExt import ProbNumDiffEq: ProbNumDiffEqBlockDiagonal, blocks -using BlockDiagonals +import BlockDiagonals: BlockDiagonal BlockDiagonal(M::ProbNumDiffEqBlockDiagonal) = BlockDiagonal(blocks(M)) diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index d9e47705d..1a5ff69d0 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -2,7 +2,7 @@ __precompile__() module ProbNumDiffEq -import Base: copy, copy!, show, size, ndims, similar, isapprox, isequal, iterate, ==, length +import Base: copy, copy!, show, size, ndims, similar, isapprox, isequal, iterate, ==, length, zero using LinearAlgebra import LinearAlgebra: mul! diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 11c4dc89f..aeff0080a 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -57,6 +57,28 @@ copy!(B::BlockDiag, A::BlockDiag) = begin end return B end +similar(B::BlockDiag) = BlockDiag(similar.(blocks(B))) +zero(B::BlockDiag) = BlockDiag(zero.(blocks(B))) + +# Mul with Scalar or UniformScaling +Base.:*(a::Number, M::BlockDiag) = BlockDiag([a * B for B in blocks(M)]) +Base.:*(M::BlockDiag, a::Number) = BlockDiag([B * a for B in blocks(M)]) +Base.:*(U::UniformScaling, M::BlockDiag) = BlockDiag([U * B for B in blocks(M)]) +Base.:*(M::BlockDiag, U::UniformScaling) = BlockDiag([B * U for B in blocks(M)]) + +# Mul between BockDiag's +Base.:*(A::BlockDiag, B::BlockDiag) = begin + @assert length(A.blocks) == length(B.blocks) + return BlockDiag([Ai * Bi for (Ai, Bi) in zip(blocks(A), blocks(B))]) +end +Base.:*(A::Adjoint{T,<:BlockDiag}, B::BlockDiag) where {T} = begin + @assert length(A.parent.blocks) == length(B.blocks) + return BlockDiag([Ai' * Bi for (Ai, Bi) in zip(blocks(A.parent), blocks(B))]) +end +Base.:*(A::BlockDiag, B::Adjoint{T,<:BlockDiag}) where {T} = begin + @assert length(A.blocks) == length(B.parent.blocks) + return BlockDiag([Ai * Bi' for (Ai, Bi) in zip(blocks(A), blocks(B.parent))]) +end # Standard LinearAlgebra.mul! mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag) = begin @@ -156,21 +178,9 @@ _matmul!( return C end -LinearAlgebra.rmul!(B::BlockDiag, n::Number) = @simd ivdep for i in eachindex(B.blocks) - rmul!(B.blocks[i], n) -end -LinearAlgebra.adjoint(B::BlockDiag) = Adjoint(B) - -Base.:*(A::BlockDiag, B::BlockDiag) = begin - @assert length(A.blocks) == length(B.blocks) - return BlockDiag([blocks(A)[i] * blocks(B)[i] for i in eachindex(B.blocks)]) -end -Base.:*(A::Adjoint{T,<:BlockDiag}, B::BlockDiag) where {T} = begin - @assert length(A.parent.blocks) == length(B.blocks) - return BlockDiag([A.parent.blocks[i]' * B.blocks[i] for i in eachindex(B.blocks)]) -end -Base.:*(A::BlockDiag, B::Adjoint{T,<:BlockDiag}) where {T} = begin - @assert length(A.blocks) == length(B.parent.blocks) - return BlockDiag([A.blocks[i] * B.parent.blocks[i]' for i in eachindex(B.parent.blocks)]) +LinearAlgebra.rmul!(B::BlockDiag, n::Number) = begin + @simd ivdep for i in eachindex(B.blocks) + rmul!(B.blocks[i], n) + end + return B end -Base.:*(A::UniformScaling, B::BlockDiag) = BlockDiag([A * blocks(B)[i] for i in eachindex(B.blocks)]) diff --git a/test/core/blockdiagonals.jl b/test/core/blockdiagonals.jl new file mode 100644 index 000000000..b9b7af643 --- /dev/null +++ b/test/core/blockdiagonals.jl @@ -0,0 +1,57 @@ +using ProbNumDiffEq +import ProbNumDiffEq: BlockDiag, _matmul! +using LinearAlgebra +using BlockDiagonals +using Test + +d1, d2 = 2, 3 +@testset "T=$T" for T in (Float64, BigFloat) + A = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) + B = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) + C = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) + + AM, BM, CM = @test_nowarn Matrix.((A, B, C)) + + @test Matrix(BlockDiagonal(A)) == AM + @test Matrix(BlockDiagonal(B)) == BM + @test Matrix(BlockDiagonal(C)) == CM + + _A = @test_nowarn copy(A) + @test _A isa BlockDiag + + _B = @test_nowarn copy!(_A, B) + @test _B === _A + @test _B == B + + _A = @test_nowarn similar(A) + @test _A isa BlockDiag + @test size(_A) == size(A) + + _Z = @test_nowarn zero(A) + @test _Z isa BlockDiag + @test size(_Z) == size(A) + @test all(_Z .== 0) + + function tttm(M) # quick type test and to matrix + @test M isa BlockDiag + return Matrix(M) + end + + for _mul! in (:mul!, :_matmul!) + @test @eval tttm($_mul!(C, A, B)) ≈ $_mul!(CM, AM, BM) + @test @eval tttm($_mul!(C, A', B)) ≈ $_mul!(CM, AM', BM) + @test @eval tttm($_mul!(C, A, B')) ≈ $_mul!(CM, AM, BM') + end + @test tttm(A * B) ≈ AM * BM + @test tttm(A' * B) ≈ AM' * BM + @test tttm(A * B') ≈ AM * BM' + + a = rand() + @test tttm(A * a) ≈ AM * a + @test tttm(a * A) ≈ a * AM + @test tttm(A * (a * I)) ≈ AM * a + @test tttm((a * I) * A) ≈ a * AM + @test tttm(rmul!(copy(A), a)) ≈ a * AM + + @test_throws ErrorException view(A, 1:2, 1:2) +end diff --git a/test/runtests.jl b/test/runtests.jl index 966119b38..e0c42a3bb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -19,6 +19,9 @@ const GROUP = get(ENV, "GROUP", "All") @testset "ProbNumDiffEq" begin if GROUP == "All" || GROUP == "Core" @timedtestset "Core" begin + @timedsafetestset "BlockDiagonals" begin + include("core/blockdiagonals.jl") + end @timedsafetestset "Filtering" begin include("core/filtering.jl") end From 85edcaabb11a67d45b8056271b091bdc6243dc73 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 22:10:39 +0100 Subject: [PATCH 41/99] Check that the K.b is actually empty --- src/filtering/markov_kernel.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index a97b37275..17e1a5802 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -111,6 +111,7 @@ function marginalize_cov!( C_DxD::AbstractMatrix, C_3DxD::AbstractMatrix, ) where {T,S} + @assert ismissing(K.b) || isnothing(K.b) _Σ_out = PSDMatrix(Σ_out.R.B) _Σ_curr = PSDMatrix(Σ_curr.R.B) _K = AffineNormalKernel(K.A.B, nothing, PSDMatrix(K.C.R.B)) @@ -131,6 +132,7 @@ function marginalize_cov!( C_DxD::AbstractMatrix, C_3DxD::AbstractMatrix, ) where {T,S} + @assert ismissing(K.b) || isnothing(K.b) @inbounds @simd ivdep for i in eachindex(blocks(Σ_out.R)) _Σ_out = PSDMatrix(Σ_out.R.blocks[i]) _Σ_curr = PSDMatrix(Σ_curr.R.blocks[i]) From a27079ddb307e818bd5e26903eb90808f713489c Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 22:15:00 +0100 Subject: [PATCH 42/99] Add versions to overload also the non-blasfloat matmuls --- src/blockdiagonals.jl | 62 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index aeff0080a..5b3197143 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -111,6 +111,17 @@ mul!(C::BlockDiag, A::BlockDiag, B::Adjoint{<:Number,<:BlockDiag}) = begin end # Our fast _matmul! +_matmul!( + C::BlockDiag, + A::BlockDiag, + B::BlockDiag, +) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i]) + end + return C +end _matmul!( C::BlockDiag{T}, A::BlockDiag{T}, @@ -123,6 +134,19 @@ _matmul!( return C end +_matmul!( + C::BlockDiag, + A::BlockDiag, + B::BlockDiag, + alpha::Number, + beta::Number, +) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + end + return C +end _matmul!( C::BlockDiag{T}, A::BlockDiag{T}, @@ -137,6 +161,17 @@ _matmul!( return C end +_matmul!( + C::BlockDiag, + A::BlockDiag, + B::Adjoint{<:Number,<:BlockDiag}, +) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + end + return C +end _matmul!( C::BlockDiag{T}, A::BlockDiag{T}, @@ -149,6 +184,17 @@ _matmul!( return C end +_matmul!( + C::BlockDiag, + A::Adjoint{<:Number,<:BlockDiag}, + B::BlockDiag, +) = begin + @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + end + return C +end _matmul!( C::BlockDiag{T}, A::Adjoint{T,<:BlockDiag{T}}, @@ -161,6 +207,22 @@ _matmul!( return C end +_matmul!( + C::AbstractVector, + A::BlockDiag, + B::AbstractVector, +) = begin + @assert size(A, 2) == length(B) + @assert length(C) == size(A, 1) + ic, ib = 1, 1 + for i in eachindex(blocks(A)) + d1, d2 = size(A.blocks[i]) + @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) + ic += d1 + ib += d2 + end + return C +end _matmul!( C::AbstractVector{T}, A::BlockDiag{T}, From 20b8430af07d3aa1d349893126b3dc86b13535fd Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 22:25:13 +0100 Subject: [PATCH 43/99] Make some code more compact and readable --- src/blockdiagonals.jl | 11 +++++++++++ src/derivative_utils.jl | 9 +-------- test/core/blockdiagonals.jl | 2 ++ 3 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 5b3197143..9402df1f8 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -246,3 +246,14 @@ LinearAlgebra.rmul!(B::BlockDiag, n::Number) = begin end return B end + +copy!(A::BlockDiag, B::Diagonal) = begin + @assert size(A) == size(B) + i = 1 + for Ai in blocks(A) + d = LinearAlgebra.checksquare(Ai) + @views copy!(Ai, Diagonal(B.diag[i:i+d-1])) + i += d + end + return A +end diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index 1a8ec21b3..55391644a 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -18,14 +18,7 @@ function calc_H!(H, integ, cache) OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) @unpack C_dxd = cache - if C_dxd isa BlockDiag - @simd ivdep for i in eachindex(blocks(C_dxd)) - @assert length(C_dxd.blocks[i]) == 1 - C_dxd.blocks[i][1] = ddu[i, i] - end - else - C_dxd .= Diagonal(ddu) - end + copy!(C_dxd, Diagonal(ddu)) _matmul!(H, C_dxd, cache.SolProj, -1.0, 1.0) else error("Unknown algorithm") diff --git a/test/core/blockdiagonals.jl b/test/core/blockdiagonals.jl index b9b7af643..b5c68b62b 100644 --- a/test/core/blockdiagonals.jl +++ b/test/core/blockdiagonals.jl @@ -54,4 +54,6 @@ d1, d2 = 2, 3 @test tttm(rmul!(copy(A), a)) ≈ a * AM @test_throws ErrorException view(A, 1:2, 1:2) + + tttm(copy!(A, Diagonal(A))) end From faa4d5587e13cb1b746d63ca9ff2eb1dd10decf9 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 22:32:39 +0100 Subject: [PATCH 44/99] Change order to fit acronym --- src/fast_linalg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/fast_linalg.jl b/src/fast_linalg.jl index e0a01dcc4..7a311424b 100644 --- a/src/fast_linalg.jl +++ b/src/fast_linalg.jl @@ -30,7 +30,7 @@ _matmul!( _matmul!(C::AbstractVecOrMat, A::AbstractVecOrMat, b::Number) = @.. C = A * b _matmul!(C::AbstractVecOrMat, a::Number, B::AbstractVecOrMat) = @.. C = a * B # Matrix matrix products with diagonal matrices -const MSR{T} = Union{SubArray{T},Matrix{T},Base.ReshapedArray{T}} +const MSR{T} = Union{Matrix{T},SubArray{T},Base.ReshapedArray{T}} _matmul!(C::MSR, A::MSR, B::Diagonal) = @.. C = A * B.diag' _matmul!(C::MSR, A::Diagonal, B::MSR) = From 6733003113520a1aadd126f7d17367fb39fe096e Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 22:35:05 +0100 Subject: [PATCH 45/99] For some reson the eval tests failed; so fix them --- test/core/blockdiagonals.jl | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/test/core/blockdiagonals.jl b/test/core/blockdiagonals.jl index b5c68b62b..81eb592dc 100644 --- a/test/core/blockdiagonals.jl +++ b/test/core/blockdiagonals.jl @@ -37,11 +37,13 @@ d1, d2 = 2, 3 return Matrix(M) end - for _mul! in (:mul!, :_matmul!) - @test @eval tttm($_mul!(C, A, B)) ≈ $_mul!(CM, AM, BM) - @test @eval tttm($_mul!(C, A', B)) ≈ $_mul!(CM, AM', BM) - @test @eval tttm($_mul!(C, A, B')) ≈ $_mul!(CM, AM, BM') - end + @test tttm(mul!(C, A, B)) ≈ mul!(CM, AM, BM) + @test tttm(mul!(C, A', B)) ≈ mul!(CM, AM', BM) + @test tttm(mul!(C, A, B')) ≈ mul!(CM, AM, BM') + @test tttm(_matmul!(C, A, B)) ≈ _matmul!(CM, AM, BM) + @test tttm(_matmul!(C, A', B)) ≈ _matmul!(CM, AM', BM) + @test tttm(_matmul!(C, A, B')) ≈ _matmul!(CM, AM, BM') + @test tttm(A * B) ≈ AM * BM @test tttm(A' * B) ≈ AM' * BM @test tttm(A * B') ≈ AM * BM' From 597f9688eeb5b4d31c426f0b71ebefc963e2c613 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 22:36:33 +0100 Subject: [PATCH 46/99] Make the if else order in predict and backward kernel easier --- src/filtering/markov_kernel.jl | 6 +++--- src/filtering/predict.jl | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index 17e1a5802..ff9026865 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -224,11 +224,11 @@ function compute_backward_kernel!( end _matmul!(view(Λ.R, 1:D, 1:D), x.Σ.R, C_DxD) # Λ.R[D+1:2D, 1:D] = (G * Q.R')' - if !isone(diffusion) + if isone(diffusion) + _matmul!(view(Λ.R, D+1:2D, 1:D), Q.R, G') + else apply_diffusion!(PSDMatrix(C_DxD), Q, diffusion) _matmul!(view(Λ.R, D+1:2D, 1:D), C_DxD, G') - else - _matmul!(view(Λ.R, D+1:2D, 1:D), Q.R, G') end return Kout diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index a0768080c..94fd8b10e 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -76,7 +76,9 @@ function predict_cov!( D = size(Qh, 1) _matmul!(view(R, 1:D, 1:D), Σ_curr.R, Ah') - if !isone(diffusion) + if isone(diffusion) + @.. R[D+1:2D, 1:D] = Qh.R + else apply_diffusion!(PSDMatrix(view(R, D+1:2D, 1:D)), Qh, diffusion) # if diffusion isa Number @@ -93,8 +95,6 @@ function predict_cov!( # sqrt.(kron(Eye(d) * diffusion, Eye(q + 1))), # ) # end - else - @.. R[D+1:2D, 1:D] = Qh.R end _matmul!(M, R', R) chol = cholesky!(Symmetric(M), check=false) From ee2622042516199d4bd36342aa44193c9bd9b2f2 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 22:47:14 +0100 Subject: [PATCH 47/99] misc --- src/priors/iwp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/priors/iwp.jl b/src/priors/iwp.jl index dc1716e4f..b62d5d3ac 100644 --- a/src/priors/iwp.jl +++ b/src/priors/iwp.jl @@ -169,7 +169,7 @@ function initialize_transition_matrices(FAC::DenseCovariance, p::IWP, dt) Ah, Qh = copy(A), copy(Q) return A, Q, Ah, Qh, P, PI end -function initialize_transition_matrices(FAC::CovarianceStructure, p::IWP, dt) +function initialize_transition_matrices(FAC::BlockDiagonalCovariance, p::IWP, dt) A, Q = preconditioned_discretize(p) A = to_factorized_matrix(FAC, A) Q = to_factorized_matrix(FAC, Q) From cbb4e093d845b84543a497eaf99e8faeb8800e94 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Thu, 15 Feb 2024 23:06:05 +0100 Subject: [PATCH 48/99] Properly implement `size` --- src/blockdiagonals.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 9402df1f8..9e8b8b454 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -18,7 +18,7 @@ const BlockDiag = ProbNumDiffEqBlockDiagonal blocks(B::BlockDiag) = B.blocks nblocks(B::BlockDiag) = length(B.blocks) -size(B::BlockDiag) = (sum(size.(blocks(B), 1)), sum(size.(blocks(B), 2))) +size(B::BlockDiag) = mapreduce(size, ((a, b), (c, d)) -> (a + c, b + d), blocks(B)) function _block_indices(B::BlockDiag, i::Integer, j::Integer) all((0, 0) .< (i, j) .<= size(B)) || throw(BoundsError(B, (i, j))) From 5b2186b28778d4327ca64491d73dfe1a0f3fbb31 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 08:28:30 +0100 Subject: [PATCH 49/99] Remove an inbounds as we don't explicitly do a sizecheck --- src/filtering/markov_kernel.jl | 3 ++- src/filtering/predict.jl | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index ff9026865..1244891b0 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -293,7 +293,8 @@ function compute_backward_kernel!( } d = length(blocks(xpred.Σ.R)) q = size(blocks(xpred.Σ.R)[1], 1) - 1 - @inbounds @simd ivdep for i in eachindex(blocks(xpred.Σ.R)) + + @simd ivdep for i in eachindex(blocks(xpred.Σ.R)) _Kout = AffineNormalKernel( Kout.A.blocks[i], view(Kout.b, (i-1)*(q+1)+1:i*(q+1)), diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 94fd8b10e..8d9b0b548 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -140,7 +140,7 @@ function predict_cov!( C_2DxD::BlockDiag, diffusion::Union{Number,Diagonal}, ) where {T,S} - for i in eachindex(blocks(Σ_out.R)) + @simd ivdep for i in eachindex(blocks(Σ_out.R)) predict_cov!( PSDMatrix(Σ_out.R.blocks[i]), PSDMatrix(Σ_curr.R.blocks[i]), From 9e2177a5165ed2fecd713e2ef9d577d1a2d8a94b Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 08:47:31 +0100 Subject: [PATCH 50/99] Remove some checks again as they are irrelevant for the cov --- src/filtering/markov_kernel.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index 1244891b0..03973c767 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -111,10 +111,9 @@ function marginalize_cov!( C_DxD::AbstractMatrix, C_3DxD::AbstractMatrix, ) where {T,S} - @assert ismissing(K.b) || isnothing(K.b) _Σ_out = PSDMatrix(Σ_out.R.B) _Σ_curr = PSDMatrix(Σ_curr.R.B) - _K = AffineNormalKernel(K.A.B, nothing, PSDMatrix(K.C.R.B)) + _K = AffineNormalKernel(K.A.B, K.b, PSDMatrix(K.C.R.B)) _D = size(_Σ_out, 1) _C_DxD = C_DxD.B _C_3DxD = C_3DxD.B @@ -132,11 +131,10 @@ function marginalize_cov!( C_DxD::AbstractMatrix, C_3DxD::AbstractMatrix, ) where {T,S} - @assert ismissing(K.b) || isnothing(K.b) - @inbounds @simd ivdep for i in eachindex(blocks(Σ_out.R)) + @simd ivdep for i in eachindex(blocks(Σ_out.R)) _Σ_out = PSDMatrix(Σ_out.R.blocks[i]) _Σ_curr = PSDMatrix(Σ_curr.R.blocks[i]) - _K = AffineNormalKernel(K.A.blocks[i], nothing, PSDMatrix(K.C.R.blocks[i])) + _K = AffineNormalKernel(K.A.blocks[i], K.b, PSDMatrix(K.C.R.blocks[i])) _C_DxD = C_DxD.blocks[i] _C_3DxD = C_3DxD.blocks[i] marginalize_cov!(_Σ_out, _Σ_curr, _K; C_DxD=_C_DxD, C_3DxD=_C_3DxD) From e9946aca6c8774aefdc0e86c13c149f5060329b0 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 08:53:15 +0100 Subject: [PATCH 51/99] Add a very minimal docstring to ProbNumDiffEqBlockDiagonal --- src/blockdiagonals.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 9e8b8b454..1544c4e99 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -1,7 +1,7 @@ """ -BlockDiagonals.jl didn't cut it, so we're rolling our own. + ProbNumDiffEqBlockDiagonal(blocks::Vector{V}) where {T,V<:AbstractMatrix{T}} -TODO: Add a way to convert to a `BlockDiagonal`. +A very minimal but fast re-implementation of `BlockDiagonals.Blockdiagonal`. """ struct ProbNumDiffEqBlockDiagonal{T<:Number,V<:AbstractMatrix{T}} <: AbstractMatrix{T} blocks::Vector{V} From b15fae83d7051a502f7f70127f0fbe0aacb84c2b Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 11:27:25 +0100 Subject: [PATCH 52/99] Better BlockDiagonals and a bit of Kronecker --- src/blockdiagonals.jl | 266 +++++++++++++++++++----------------------- src/kronecker.jl | 11 ++ 2 files changed, 131 insertions(+), 146 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 1544c4e99..ab09f1801 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -81,163 +81,132 @@ Base.:*(A::BlockDiag, B::Adjoint{T,<:BlockDiag}) where {T} = begin end # Standard LinearAlgebra.mul! -mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag) = begin - @assert length(C.blocks) == length(A.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i]) +for _mul! in (:mul!, :_matmul!) + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i]) + end + return C + end + (_mul! == :_matmul!) && @eval $_mul!( + C::BlockDiag{T}, + A::BlockDiag{T}, + B::BlockDiag{T}, + ) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i]) + end + return C end - return C -end -mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag, alpha::Number, beta::Number) = begin - @assert length(C.blocks) == length(A.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) - end - return C -end -mul!(C::BlockDiag, A::Adjoint{<:Number,<:BlockDiag}, B::BlockDiag) = begin - @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) - end - return C -end -mul!(C::BlockDiag, A::BlockDiag, B::Adjoint{<:Number,<:BlockDiag}) = begin - @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) - end - return C -end -# Our fast _matmul! -_matmul!( - C::BlockDiag, - A::BlockDiag, - B::BlockDiag, -) = begin - @assert length(C.blocks) == length(A.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i]) + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag, alpha::Number, beta::Number) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + end + return C + end + (_mul! == :_matmul!) && @eval $_mul!( + C::BlockDiag{T}, + A::BlockDiag{T}, + B::BlockDiag{T}, + alpha::Number, + beta::Number, + ) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + end + return C end - return C -end -_matmul!( - C::BlockDiag{T}, - A::BlockDiag{T}, - B::BlockDiag{T}, -) where {T<:LinearAlgebra.BlasFloat} = begin - @assert length(C.blocks) == length(A.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i]) - end - return C -end -_matmul!( - C::BlockDiag, - A::BlockDiag, - B::BlockDiag, - alpha::Number, - beta::Number, -) = begin - @assert length(C.blocks) == length(A.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + @eval $_mul!(C::BlockDiag, A::Adjoint{<:Number,<:BlockDiag}, B::BlockDiag) = begin + @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + end + return C + end + (_mul! == :_matmul!) && @eval $_mul!( + C::BlockDiag{T}, + A::BlockDiag{T}, + B::Adjoint{T,<:BlockDiag{T}}, + ) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + end + return C end - return C -end -_matmul!( - C::BlockDiag{T}, - A::BlockDiag{T}, - B::BlockDiag{T}, - alpha::Number, - beta::Number, -) where {T<:LinearAlgebra.BlasFloat} = begin - @assert length(C.blocks) == length(A.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) - end - return C -end -_matmul!( - C::BlockDiag, - A::BlockDiag, - B::Adjoint{<:Number,<:BlockDiag}, -) = begin - @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) - end - return C -end -_matmul!( - C::BlockDiag{T}, - A::BlockDiag{T}, - B::Adjoint{T,<:BlockDiag{T}}, -) where {T<:LinearAlgebra.BlasFloat} = begin - @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Adjoint{<:Number,<:BlockDiag}) = begin + @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + end + return C + end + (_mul! == :_matmul!) && @eval $_mul!( + C::BlockDiag{T}, + A::Adjoint{T,<:BlockDiag{T}}, + B::BlockDiag{T}, + ) where {T<:LinearAlgebra.BlasFloat} = begin + @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds _matmul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + end + return C end - return C -end -_matmul!( - C::BlockDiag, - A::Adjoint{<:Number,<:BlockDiag}, - B::BlockDiag, -) = begin - @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) - end - return C -end -_matmul!( - C::BlockDiag{T}, - A::Adjoint{T,<:BlockDiag{T}}, - B::BlockDiag{T}, -) where {T<:LinearAlgebra.BlasFloat} = begin - @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + @eval $_mul!(C::BlockDiag, A::Number, B::BlockDiag) = begin + @assert length(C.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A, B.blocks[i]) + end + return C + end + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Number) = begin + @assert length(C.blocks) == length(A.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], B) + end + return C end - return C -end -_matmul!( - C::AbstractVector, - A::BlockDiag, - B::AbstractVector, -) = begin - @assert size(A, 2) == length(B) - @assert length(C) == size(A, 1) - ic, ib = 1, 1 - for i in eachindex(blocks(A)) - d1, d2 = size(A.blocks[i]) - @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) - ic += d1 - ib += d2 + @eval $_mul!( + C::AbstractVector, + A::BlockDiag, + B::AbstractVector, + ) = begin + @assert size(A, 2) == length(B) + @assert length(C) == size(A, 1) + ic, ib = 1, 1 + for i in eachindex(blocks(A)) + d1, d2 = size(A.blocks[i]) + @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) + ic += d1 + ib += d2 + end + return C + end + (_mul! == :_matmul!) && @eval $_mul!( + C::AbstractVector{T}, + A::BlockDiag{T}, + B::AbstractVector{T}, + ) where {T<:LinearAlgebra.BlasFloat} = begin + @assert size(A, 2) == length(B) + @assert length(C) == size(A, 1) + ic, ib = 1, 1 + for i in eachindex(blocks(A)) + d1, d2 = size(A.blocks[i]) + @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) + ic += d1 + ib += d2 + end + return C end - return C -end -_matmul!( - C::AbstractVector{T}, - A::BlockDiag{T}, - B::AbstractVector{T}, -) where {T<:LinearAlgebra.BlasFloat} = begin - @assert size(A, 2) == length(B) - @assert length(C) == size(A, 1) - ic, ib = 1, 1 - for i in eachindex(blocks(A)) - d1, d2 = size(A.blocks[i]) - @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) - ic += d1 - ib += d2 - end - return C end LinearAlgebra.rmul!(B::BlockDiag, n::Number) = begin @@ -257,3 +226,8 @@ copy!(A::BlockDiag, B::Diagonal) = begin end return A end + +Base.isequal(A::BlockDiag, B::BlockDiag) = + length(A.blocks) == length(B.blocks) && all(map(isequal, A.blocks, B.blocks)) +==(A::BlockDiag, B::BlockDiag) = + length(A.blocks) == length(B.blocks) && all(map(==, A.blocks, B.blocks)) diff --git a/src/kronecker.jl b/src/kronecker.jl index 4e5c2bb87..d05096703 100644 --- a/src/kronecker.jl +++ b/src/kronecker.jl @@ -153,6 +153,17 @@ _matmul!( return A end +_matmul!(A::IKP, b::Number, C::IKP) = begin + check_matmul_sizes(A, C) + _matmul!(A.B, b, C.B) + return A +end +_matmul!(A::IKP, B::IKP, c::Number) = begin + check_matmul_sizes(A, B) + _matmul!(A.B, B.B, c) + return A +end + """ Allocation-free reshape Found here: https://discourse.julialang.org/t/convert-array-into-matrix-in-place/55624/5 From 5ece898b2a4bf238b2d83f8102bfd94f75606378 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 11:27:37 +0100 Subject: [PATCH 53/99] Make the DiagonalEK1 work again (except for secondorderodes) --- src/projection.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/projection.jl b/src/projection.jl index de7b4ad8c..ba08e65e0 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -60,7 +60,11 @@ function solution_space_projection(C::IsometricKroneckerCovariance, is_secondord end function solution_space_projection(C::BlockDiagonalCovariance, is_secondorder_ode) Proj = projection(C) - error("No idea!") + if is_secondorder_ode + error("Not yet implemented!") + else + return Proj(0) + end end struct KroneckerSecondOrderODESolutionProjector{T,FAC,M,M2} <: AbstractMatrix{T} From 8d0ae859f6a96edbbfdbf5aa9f411a6038dec185 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 11:28:04 +0100 Subject: [PATCH 54/99] Test the diffusions (found a bug! unittests are actually nice) --- test/core/priors.jl | 68 +++++++++++++++++++++++++++------------------ test/runtests.jl | 3 ++ 2 files changed, 44 insertions(+), 27 deletions(-) diff --git a/test/core/priors.jl b/test/core/priors.jl index 022e5f0bb..3d8564818 100644 --- a/test/core/priors.jl +++ b/test/core/priors.jl @@ -10,7 +10,6 @@ using SimpleUnPack using FillArrays h = 0.1 -σ = 0.1 @testset "General prior API" begin for prior in ( @@ -55,6 +54,8 @@ end @testset "Test IWP (d=2,q=2)" begin d, q = 2, 2 + σ = 0.1 + prior = PNDE.IWP(dim=d, num_derivatives=q) # true sde parameters @@ -123,41 +124,54 @@ end @testset "Test vanilla (ie. non-preconditioned)" begin Ah, Qh = PNDE.discretize(prior, h) - Qh = PNDE.apply_diffusion(Qh, σ^2*Eye(d)) - @test AH_22_IBM ≈ Ah - @test QH_22_IBM ≈ Matrix(Qh) + + for Γ in (σ^2, σ^2 * Eye(d)) + @test QH_22_IBM ≈ Matrix(PNDE.apply_diffusion(Qh, Γ)) + end end @testset "Test with preconditioning" begin A, Q = PNDE.preconditioned_discretize(prior) - Qh = PNDE.apply_diffusion(Q, σ^2*Eye(d)) - @test AH_22_PRE ≈ Matrix(A) - @test QH_22_PRE ≈ Matrix(Qh) + + for Γ in (σ^2, σ^2 * Eye(d)) + @test QH_22_PRE ≈ Matrix(PNDE.apply_diffusion(Q, Γ)) + end end @testset "Test `make_transition_matrices!`" begin - A, Q, Ah, Qh, P, PI = PNDE.initialize_transition_matrices( - PNDE.DenseCovariance{Float64}(d, q), prior, h) - - @test AH_22_PRE ≈ A - @test QH_22_PRE ≈ Matrix(PNDE.apply_diffusion(Q, σ^2*Eye(d))) - - cache = ( - d=d, - q=q, - A=A, - Q=Q, - P=P, - PI=PI, - Ah=Ah, - Qh=Qh, - ) - - make_transition_matrices!(cache, prior, h) - @test AH_22_IBM ≈ cache.Ah - @test QH_22_IBM ≈ Matrix(PNDE.apply_diffusion(cache.Qh, σ^2*Eye(d))) + for FAC in (PNDE.IsometricKroneckerCovariance, PNDE.BlockDiagonalCovariance) + A, Q, Ah, Qh, P, PI = PNDE.initialize_transition_matrices( + PNDE.BlockDiagonalCovariance{Float64}(d, q), prior, h) + + @test AH_22_PRE ≈ A + + for Γ in (σ^2, σ^2 * Eye(d), σ^2 * I(d)) + @test QH_22_PRE ≈ Matrix(PNDE.apply_diffusion(Q, Γ)) + end + + cache = ( + d=d, + q=q, + A=A, + Q=Q, + P=P, + PI=PI, + Ah=Ah, + Qh=Qh, + ) + + make_transition_matrices!(cache, prior, h) + @test AH_22_IBM ≈ cache.Ah + + for Γ in (σ^2, σ^2 * Eye(d)) + @test QH_22_IBM ≈ Matrix(PNDE.apply_diffusion(cache.Qh, Γ)) + end + if FAC != PNDE.IsometricKroneckerCovariance + @test QH_22_IBM ≈ Matrix(PNDE.apply_diffusion(cache.Qh, σ^2 * I(d))) + end + end end end diff --git a/test/runtests.jl b/test/runtests.jl index e0c42a3bb..a36528f79 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -28,6 +28,9 @@ const GROUP = get(ENV, "GROUP", "All") @timedsafetestset "Priors" begin include("core/priors.jl") end + @timedsafetestset "Diffusions" begin + include("core/diffusions.jl") + end @timedsafetestset "Preconditioning" begin include("core/preconditioning.jl") end From e613401fb0911dc231ea834102ee4f9c262a085f Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 11:29:34 +0100 Subject: [PATCH 55/99] Give the diffusions much more space --- src/ProbNumDiffEq.jl | 4 +- src/diffusions.jl | 297 ------------------------------ src/diffusions/apply_diffusion.jl | 92 +++++++++ src/diffusions/calibration.jl | 182 ++++++++++++++++++ src/diffusions/typedefs.jl | 97 ++++++++++ 5 files changed, 374 insertions(+), 298 deletions(-) delete mode 100644 src/diffusions.jl create mode 100644 src/diffusions/apply_diffusion.jl create mode 100644 src/diffusions/calibration.jl create mode 100644 src/diffusions/typedefs.jl diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index 1a5ff69d0..3b1fdfb17 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -67,7 +67,9 @@ include("priors/ltisde.jl") include("priors/ioup.jl") include("priors/matern.jl") export IWP, IOUP, Matern -include("diffusions.jl") +include("diffusions/typedefs.jl") +include("diffusions/apply_diffusion.jl") +include("diffusions/calibration.jl") export FixedDiffusion, DynamicDiffusion, FixedMVDiffusion, DynamicMVDiffusion include("initialization/common.jl") diff --git a/src/diffusions.jl b/src/diffusions.jl deleted file mode 100644 index 41edd00fb..000000000 --- a/src/diffusions.jl +++ /dev/null @@ -1,297 +0,0 @@ -abstract type AbstractDiffusion end -abstract type AbstractStaticDiffusion <: AbstractDiffusion end -abstract type AbstractDynamicDiffusion <: AbstractDiffusion end -isstatic(diffusion::AbstractStaticDiffusion) = true -isdynamic(diffusion::AbstractStaticDiffusion) = false -isstatic(diffusion::AbstractDynamicDiffusion) = false -isdynamic(diffusion::AbstractDynamicDiffusion) = true - -# TODO Add a proper length description somewhere that exaplains better what this "diffusion" -# object is and how it is handled in this package. -apply_diffusion( - Q::PSDMatrix{<:Number,<:IsometricKroneckerProduct}, - diffusion::Number -) = PSDMatrix(Q.R * sqrt.(diffusion)) - -apply_diffusion( - Q::PSDMatrix{T,<:IsometricKroneckerProduct}, - diffusion::Diagonal{T,<:FillArrays.Fill} -) where {T} = apply_diffusion(Q, diffusion.diag.value) - -apply_diffusion( - Q::PSDMatrix{T,<:Matrix}, - diffusion::Diagonal -) where {T} = begin - # @warn "This should ideally not be called; TODO" - d = size(diffusion, 1) - q = size(Q, 1) ÷ d - 1 - return PSDMatrix(Q.R * sqrt.(kron(diffusion, I(q + 1)))) -end - -apply_diffusion( - Q::PSDMatrix{T,<:BlockDiag}, diffusion::Diagonal -) where {T} = PSDMatrix( - BlockDiag([blocks(Q.R)[i] * sqrt.(diffusion.diag[i]) for i in eachindex(blocks(Q.R))])) - -apply_diffusion!( - Q::PSDMatrix, - diffusion::Diagonal{T,<:FillArrays.Fill} -) where {T} = rmul!(Q.R, sqrt.(diffusion.diag.value)) -apply_diffusion!( - Q::PSDMatrix{T,<:BlockDiag}, - diffusion::Diagonal{T,<:Vector}, -) where {T} = begin - @simd ivdep for i in eachindex(blocks(Q.R)) - rmul!(blocks(Q.R)[i], diffusion.diag[i]) - end -end - -apply_diffusion!( - out::PSDMatrix, - Q::PSDMatrix, - diffusion::Number -) = _matmul!(out.R, Q.R, sqrt.(diffusion)) -apply_diffusion!( - out::PSDMatrix, - Q::PSDMatrix, - diffusion::Diagonal{<:Number,<:FillArrays.Fill}, -) = apply_diffusion!(out, Q, diffusion.diag.value) -apply_diffusion!( - out::PSDMatrix, - Q::PSDMatrix, - diffusion::Diagonal -) = begin - @warn "This is not yet implemented efficiently; TODO" - d = size(diffusion, 1) - D = size(Q, 1) - q = D ÷ d - 1 - _matmul!(out.R, Q.R, sqrt.(kron(Eye(d) * diffusion, Eye(q + 1)))) -end - -estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = error() - -""" - DynamicDiffusion() - -Time-varying, isotropic diffusion, which is quasi-maximum-likelihood-estimated at each step. - -**This is the recommended diffusion when using adaptive step-size selection,** and in -particular also when solving stiff systems. -""" -struct DynamicDiffusion <: AbstractDynamicDiffusion end -initial_diffusion(::DynamicDiffusion, d, q, Eltype) = one(Eltype) * Eye(d) -estimate_local_diffusion(::DynamicDiffusion, integ) = local_scalar_diffusion(integ.cache) - -""" - DynamicMVDiffusion() - -Time-varying, diagonal diffusion, which is quasi-maximum-likelihood-estimated at each step. - -**Only works with the [`EK0`](@ref)!** - -A multi-variate version of [`DynamicDiffusion`](@ref), where instead of an isotropic matrix, -a diagonal matrix is estimated. This can be helpful to get more expressive posterior -covariances when using the [`EK0`](@ref), since the individual dimensions can be adjusted -separately. - -# References -* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) -""" -struct DynamicMVDiffusion <: AbstractDynamicDiffusion end -initial_diffusion(::DynamicMVDiffusion, d, q, Eltype) = Diagonal(ones(Eltype, d)) -estimate_local_diffusion(::DynamicMVDiffusion, integ) = - local_diagonal_diffusion(integ.cache) - -""" - FixedDiffusion(; initial_diffusion=1.0, calibrate=true) - -Time-fixed, isotropic diffusion, which is (optionally) quasi-maximum-likelihood-estimated. - -**This is the recommended diffusion when using fixed steps.** - -By default with `calibrate=true`, all covariances are re-scaled at the end of the solve -with the MLE diffusion. Set `calibrate=false` to skip this step, e.g. when setting the -`initial_diffusion` and then estimating the diffusion outside of the solver -(e.g. with [Fenrir.jl](https://github.com/nathanaelbosch/Fenrir.jl)). -""" -Base.@kwdef struct FixedDiffusion{T<:Number} <: AbstractStaticDiffusion - initial_diffusion::T = 1.0 - calibrate::Bool = true -end -initial_diffusion(diffusionmodel::FixedDiffusion, d, q, Eltype) = - diffusionmodel.initial_diffusion * one(Eltype) * Eye(d) -estimate_local_diffusion(::FixedDiffusion, integ) = local_scalar_diffusion(integ.cache) -function estimate_global_diffusion(::FixedDiffusion, integ) - @unpack d, measurement, m_tmp, Smat = integ.cache - # sol_diffusions = integ.sol.diffusions - - v, S = measurement.μ, measurement.Σ - e = m_tmp.μ - e .= v - diffusion_t = if S isa IsometricKroneckerProduct - @assert length(S.B) == 1 - dot(v, e) / d / S.B[1] - elseif S isa BlockDiag - @assert length(S.blocks) == d - @assert length(S.blocks[1]) == 1 - @simd ivdep for i in eachindex(e) - @inbounds e[i] /= S.blocks[i][1] - end - dot(v, e) / d - else - S_chol = cholesky!(copy!(Smat, S)) - ldiv!(S_chol, e) - dot(v, e) / d - end - - if integ.success_iter == 0 - # @assert length(sol_diffusions) == 0 - global_diffusion = diffusion_t - integ.cache.global_diffusion = global_diffusion * Eye(d) - return integ.cache.global_diffusion - else - # @assert length(sol_diffusions) == integ.success_iter - diffusion_prev = integ.cache.global_diffusion.diag.value - global_diffusion = - diffusion_prev + (diffusion_t - diffusion_prev) / integ.success_iter - # @info "compute diffusion" diffusion_prev global_diffusion - integ.cache.global_diffusion = global_diffusion * Eye(d) - return integ.cache.global_diffusion - end -end - -""" - FixedMVDiffusion(; initial_diffusion=1.0, calibrate=true) - -Time-fixed, diagonal diffusion, which is quasi-maximum-likelihood-estimated at each step. - -**Only works with the [`EK0`](@ref)!** - -A multi-variate version of [`FixedDiffusion`](@ref), where instead of an isotropic matrix, -a diagonal matrix is estimated. This can be helpful to get more expressive posterior -covariances when using the [`EK0`](@ref), since the individual dimensions can be adjusted -separately. - -# References -* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) -""" -Base.@kwdef struct FixedMVDiffusion{T} <: AbstractStaticDiffusion - initial_diffusion::T = 1.0 - calibrate::Bool = true -end -function initial_diffusion(diffusionmodel::FixedMVDiffusion, d, q, Eltype) - initdiff = diffusionmodel.initial_diffusion - @assert initdiff isa Number || length(initdiff) == d - return Diagonal(initdiff .* ones(Eltype, d)) -end -estimate_local_diffusion(::FixedMVDiffusion, integ) = local_diagonal_diffusion(integ.cache) -function estimate_global_diffusion(::FixedMVDiffusion, integ) - @unpack d, q, measurement, local_diffusion = integ.cache - - v, S = measurement.μ, measurement.Σ - # @assert diag(S) |> unique |> length == 1 - S_11 = S[1, 1] - - Σ_ii = v .^ 2 ./ S_11 - Σ = Diagonal(Σ_ii) - Σ_out = Σ - - if integ.success_iter == 0 - integ.cache.global_diffusion .= Σ_out - return integ.cache.global_diffusion - else - diffusion_prev = integ.cache.global_diffusion - @.. diffusion_prev = diffusion_prev + (Σ_out - diffusion_prev) / integ.success_iter - return integ.cache.global_diffusion - end -end - -""" - local_scalar_diffusion(integ) - -Compute the local, scalar diffusion estimate. - -Corresponds to -```math -σ² = zᵀ (H Q H^T)⁻¹ z, -``` -where ``z, H, Q`` are taken from the passed integrator. - -For more background information -* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) -""" -function local_scalar_diffusion(cache) - @unpack d, R, H, Qh, measurement, m_tmp, Smat, C_Dxd = cache - z = measurement.μ - e, HQH = m_tmp.μ, m_tmp.Σ - _matmul!(C_Dxd, Qh.R, H') - _matmul!(HQH, C_Dxd', C_Dxd) - e .= z - σ² = if HQH isa IsometricKroneckerProduct - @assert length(HQH.B) == 1 - dot(z, e) / d / HQH.B[1] - elseif HQH isa BlockDiag - @assert length(HQH.blocks) == d - @assert length(HQH.blocks[1]) == 1 - for i in eachindex(e) - e[i] /= HQH.blocks[i][1] - end - dot(z, e) / d - else - C = cholesky!(HQH) - ldiv!(C, e) - dot(z, e) / d - end - cache.local_diffusion = σ² * Eye(d) - return cache.local_diffusion -end - -""" - local_diagonal_diffusion(cache) - -Compute the local, scalar diffusion estimate. - -Corresponds to -```math -Σ_{ii} = z_i^2 / (H Q H^T)_{ii}, -``` -where ``z, H, Q`` are taken from the passed integrator. -**This should only be used with the EK0!** - -For more background information -* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) -""" -function local_diagonal_diffusion(cache) - @unpack d, q, H, Qh, measurement, m_tmp, tmp = cache - @unpack local_diffusion = cache - z = measurement.μ - # HQH = H * unfactorize(Qh) * H' - # @assert HQH |> diag |> unique |> length == 1 - # c1 = view(_matmul!(cache.C_Dxd, Qh.R, H'), :, 1) - # Q_11 = dot(c1, c1) - - # @assert - Q_11 = if Qh.R isa BlockDiag - for i in 1:d - c1 = _matmul!( - view(cache.C_Dxd.blocks[i], :, 1:1), - Qh.R.blocks[i], - view(H.blocks[i], 1:1, :)', - ) - tmp[i] = dot(c1, c1) - end - tmp - else - @warn "This is not yet implemented efficiently; TODO" - diag(H * unfactorize(Qh) * H') - end - - # To double-check: - HQH = H * unfactorize(Qh) * H' - @assert Q_11 ≈ diag(HQH) - # Also if the solver is a EK0 and not a DiagonalEK1: - # @assert Q_11 |> unique |> length == 1 - - @. local_diffusion.diag = z^2 / Q_11 - return local_diffusion -end diff --git a/src/diffusions/apply_diffusion.jl b/src/diffusions/apply_diffusion.jl new file mode 100644 index 000000000..ca3bce189 --- /dev/null +++ b/src/diffusions/apply_diffusion.jl @@ -0,0 +1,92 @@ +""" + apply_diffusion(Q::PSDMatrix, diffusion::Union{Number, Diagonal}) -> PSDMatrix + +Apply the diffusion to the PSD transition noise covariance `Q`, return the result. +""" +apply_diffusion +apply_diffusion( + Q::PSDMatrix, + diffusion::Number, +) = PSDMatrix(Q.R * sqrt.(diffusion)) +apply_diffusion( + Q::PSDMatrix, + diffusion::Diagonal{T,<:FillArrays.Fill} +) where {T} = apply_diffusion(Q, diffusion.diag.value) +apply_diffusion( + Q::PSDMatrix{T,<:BlockDiag}, + diffusion::Diagonal{T,<:Vector}, +) where {T} = PSDMatrix( + BlockDiag([blocks(Q.R)[i] * sqrt.(diffusion.diag[i]) for i in eachindex(blocks(Q.R))])) +apply_diffusion( + Q::PSDMatrix{T,<:Matrix}, + diffusion::Diagonal{T,<:Vector}, +) where {T} = begin + d = size(diffusion, 1) + q = size(Q, 1) ÷ d - 1 + return PSDMatrix(Q.R * sqrt.(kron(diffusion, I(q + 1)))) +end + +""" + apply_diffusion!(Q::PSDMatrix, diffusion::Union{Number, Diagonal}) -> PSDMatrix + +Apply the diffusion to the PSD transition noise covariance `Q` in place and return the result. +""" +apply_diffusion! +apply_diffusion!( + Q::PSDMatrix, + diffusion::Diagonal{T,<:FillArrays.Fill} +) where {T} = begin + rmul!(Q.R, sqrt.(diffusion.diag.value)) + return Q +end +apply_diffusion!( + Q::PSDMatrix{T,<:BlockDiag}, + diffusion::Diagonal{T,<:Vector}, +) where {T} = begin + @simd ivdep for i in eachindex(blocks(Q.R)) + rmul!(blocks(Q.R)[i], sqrt(diffusion.diag[i])) + end + return Q +end + +""" + apply_diffusion!(out::PSDMatrix, Q::PSDMatrix, diffusion::Union{Number, Diagonal}) -> PSDMatrix + +Apply the diffusion to the PSD transition noise covariance `Q` and store the result in `out`. +""" +apply_diffusion! +apply_diffusion!( + out::PSDMatrix, + Q::PSDMatrix, + diffusion::Number +) = begin + _matmul!(out.R, Q.R, sqrt.(diffusion)) + return out +end +apply_diffusion!( + out::PSDMatrix, + Q::PSDMatrix, + diffusion::Diagonal{<:Number,<:FillArrays.Fill}, +) = apply_diffusion!(out, Q, diffusion.diag.value) +apply_diffusion!( + out::PSDMatrix{T,<:BlockDiag}, + Q::PSDMatrix{T,<:BlockDiag}, + diffusion::Diagonal{<:T,<:Vector}, +) where {T} = begin + @simd ivdep for i in eachindex(blocks(Q.R)) + _matmul!(blocks(out.R)[i], blocks(Q.R)[i], sqrt(diffusion.diag[i])) + end + return out +end +apply_diffusion!( + out::PSDMatrix, + Q::PSDMatrix, + diffusion::Diagonal, +) = begin + @warn "This is not yet implemented efficiently; TODO" + d = size(diffusion, 1) + D = size(Q, 1) + q = D ÷ d - 1 + _matmul!(out.R, Q.R, sqrt.(kron(diffusion, Eye(q + 1)))) + return out +end diff --git a/src/diffusions/calibration.jl b/src/diffusions/calibration.jl new file mode 100644 index 000000000..b8f0975e5 --- /dev/null +++ b/src/diffusions/calibration.jl @@ -0,0 +1,182 @@ +@doc raw""" + invquad(v, M; v_cache, M_cache) + +Compute ``v' M^{-1} v`` without allocations and with Matrix-specific specializations. + +Needed for MLE diffusion estimation. +""" +invquad +invquad(v, M::Matrix; v_cache, M_cache) = begin + v_cache .= v + M_chol = cholesky!(copy!(M_cache, M)) + ldiv!(M_chol, v_cache) + dot(v, v_cache) +end +invquad(v, M::IsometricKroneckerProduct; v_cache, M_cache=nothing) = begin + v_cache .= v + @assert length(M.B) == 1 + return dot(v, v_cache) / M.B[1] +end +invquad(v, M::BlockDiag; v_cache, M_cache=nothing) = begin + v_cache .= v + @assert length(M.blocks) == length(v) == length(v_cache) + @simd ivdep for i in eachindex(v) + @assert length(M.blocks[i]) == 1 + @inbounds v_cache[i] /= M.blocks[i][1] + end + return dot(v, v_cache) +end + +@doc raw""" + estimate_global_diffusion(::FixedDiffusion, integ) + +Updates the global quasi-MLE diffusion estimate on the current measuremnt. + +The global quasi-MLE diffusion estimate Corresponds to +```math +\hat{σ}^2_N = \frac{1}{Nd} \sum_{i=1}^N z_i^T S_i^{-1} z_i, +``` +where ``z_i, S_i`` are taken the predicted observations from each step. +This function updates the iteratively computed global diffusion estimate by computing +```math +\hat{σ}^2_n = \hat{σ}^2_{n-1} + ((z_n^T S_n^{-1} z_n) / d - \hat{σ}^2_{n-1}) / n. +``` + +For more background information +* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) +""" +function estimate_global_diffusion(::FixedDiffusion, integ) + @unpack d, measurement, m_tmp, Smat = integ.cache + v, S = measurement.μ, measurement.Σ + _v, _S = m_tmp.μ, m_tmp.Σ + + diffusion_increment = invquad(v, S; v_cache=_v, M_cache=_S) / d + + new_mle_diffusion = if integ.success_iter == 0 + diffusion_increment + else + current_mle_diffusion = integ.cache.global_diffusion.diag.value + current_mle_diffusion + (diffusion_increment - current_mle_diffusion) / integ.success_iter + end + + integ.cache.global_diffusion = new_mle_diffusion * Eye(d) + return integ.cache.global_diffusion +end + +@doc raw""" + estimate_global_diffusion(::FixedMVDiffusion, integ) + +Updates the multivariate global quasi-MLE diffusion estimate on the current measuremnt. + +**This only works with the EK0!** + +The global quasi-MLE diffusion estimate Corresponds to +```math +[\hat{Σ}^2_N]_{jj} = \frac{1}{N} \sum_{i=1}^N [z_i]_j^2 / [S_i]_{11}, +``` +where ``z_i, S_i`` are taken the predicted observations from each step. +This function updates the iteratively computed global diffusion estimate by computing +```math +[\hat{Σ}^2_n]_{jj} = [\hat{Σ}^2_{n-1}]_{jj} + ([z_n]_j^2 / [S_n]_{11}, - [\hat{Σ}^2_{n-1}]_{jj}) / n. +``` + +For more background information +* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) +""" +function estimate_global_diffusion(::FixedMVDiffusion, integ) + @unpack d, q, measurement, local_diffusion, C_d = integ.cache + v, S = measurement.μ, measurement.Σ + # @assert diag(S) |> unique |> length == 1 + diffusion_increment = let + @.. C_d = v ^ 2 / S[1, 1] + Diagonal(C_d) + end + + new_mle_diffusion = if integ.success_iter == 0 + diffusion_increment + else + current_mle_diffusion = integ.cache.global_diffusion + @.. current_mle_diffusion + (diffusion_increment - current_mle_diffusion) / integ.success_iter + end + + copy!(integ.cache.global_diffusion, new_mle_diffusion) + return integ.cache.global_diffusion +end + +@doc raw""" + local_scalar_diffusion(integ) + +Compute the local scalar quasi-MLE diffusion estimate. + +Corresponds to +```math +σ² = zᵀ (H Q H^T)⁻¹ z, +``` +where ``z, H, Q`` are taken from the passed integrator. + +For more background information +* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) +""" +function local_scalar_diffusion(cache) + @unpack d, R, H, Qh, measurement, m_tmp, Smat, C_Dxd, C_d, C_dxd = cache + z = measurement.μ + HQH = let + _matmul!(C_Dxd, Qh.R, H') + _matmul!(C_dxd, C_Dxd', C_Dxd) + end + σ² = invquad(z, HQH; v_cache=C_d, M_cache=C_dxd) / d + cache.local_diffusion = σ² * Eye(d) + return cache.local_diffusion +end + +@doc raw""" + local_diagonal_diffusion(cache) + +Compute the local diagonal quasi-MLE diffusion estimate. + +**This only works with the EK0!** + +Corresponds to +```math +Σ_{ii} = z_i^2 / (H Q H^T)_{ii}, +``` +where ``z, H, Q`` are taken from the passed integrator. + +For more background information +* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) +""" +function local_diagonal_diffusion(cache) + @unpack d, q, H, Qh, measurement, m_tmp, tmp = cache + @unpack local_diffusion = cache + @assert H == cache.E1 + + z = measurement.μ + # HQH = H * unfactorize(Qh) * H' + # @assert HQH |> diag |> unique |> length == 1 + # c1 = view(_matmul!(cache.C_Dxd, Qh.R, H'), :, 1) + # Q_11 = dot(c1, c1) + + Q_11 = if Qh.R isa BlockDiag + for i in 1:d + c1 = _matmul!( + view(cache.C_Dxd.blocks[i], :, 1:1), + Qh.R.blocks[i], + view(H.blocks[i], 1:1, :)', + ) + tmp[i] = dot(c1, c1) + end + tmp + else + @warn "This is not yet implemented efficiently; TODO" + diag(H * unfactorize(Qh) * H') + end + + # To double-check: + HQH = H * unfactorize(Qh) * H' + @assert Q_11 ≈ diag(HQH) + # Also if the solver is a EK0 and not a DiagonalEK1: + # @assert Q_11 |> unique |> length == 1 + + @. local_diffusion.diag = z^2 / Q_11 + return local_diffusion +end diff --git a/src/diffusions/typedefs.jl b/src/diffusions/typedefs.jl new file mode 100644 index 000000000..b00830768 --- /dev/null +++ b/src/diffusions/typedefs.jl @@ -0,0 +1,97 @@ +abstract type AbstractDiffusion end +abstract type AbstractStaticDiffusion <: AbstractDiffusion end +abstract type AbstractDynamicDiffusion <: AbstractDiffusion end +isstatic(diffusion::AbstractStaticDiffusion) = true +isdynamic(diffusion::AbstractStaticDiffusion) = false +isstatic(diffusion::AbstractDynamicDiffusion) = false +isdynamic(diffusion::AbstractDynamicDiffusion) = true + +estimate_global_diffusion(diffusion::AbstractDynamicDiffusion, d, q, Eltype) = + error("Not possible or not implemented") + +""" + DynamicDiffusion() + +Time-varying, isotropic diffusion, which is quasi-maximum-likelihood-estimated at each step. + +**This is the recommended diffusion when using adaptive step-size selection,** and in +particular also when solving stiff systems. +""" +struct DynamicDiffusion <: AbstractDynamicDiffusion end +initial_diffusion(::DynamicDiffusion, d, q, Eltype) = one(Eltype) * Eye(d) +estimate_local_diffusion(::DynamicDiffusion, integ) = local_scalar_diffusion(integ.cache) + +""" + DynamicMVDiffusion() + +Time-varying, diagonal diffusion, which is quasi-maximum-likelihood-estimated at each step. + +**Only works with the [`EK0`](@ref)!** + +A multi-variate version of [`DynamicDiffusion`](@ref), where instead of an isotropic matrix, +a diagonal matrix is estimated. This can be helpful to get more expressive posterior +covariances when using the [`EK0`](@ref), since the individual dimensions can be adjusted +separately. + +# References +* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) +""" +struct DynamicMVDiffusion <: AbstractDynamicDiffusion end +initial_diffusion(::DynamicMVDiffusion, d, q, Eltype) = Diagonal(ones(Eltype, d)) +estimate_local_diffusion(::DynamicMVDiffusion, integ) = + local_diagonal_diffusion(integ.cache) + +""" + FixedDiffusion(; initial_diffusion=1.0, calibrate=true) + +Time-fixed, isotropic diffusion, which is (optionally) quasi-maximum-likelihood-estimated. + +**This is the recommended diffusion when using fixed steps.** + +By default with `calibrate=true`, all covariances are re-scaled at the end of the solve +with the MLE diffusion. Set `calibrate=false` to skip this step, e.g. when setting the +`initial_diffusion` and then estimating the diffusion outside of the solver +(e.g. with [Fenrir.jl](https://github.com/nathanaelbosch/Fenrir.jl)). +""" +Base.@kwdef struct FixedDiffusion{T<:Number} <: AbstractStaticDiffusion + initial_diffusion::T = 1.0 + calibrate::Bool = true +end +initial_diffusion(diffusionmodel::FixedDiffusion, d, q, Eltype) = + diffusionmodel.initial_diffusion * one(Eltype) * Eye(d) +estimate_local_diffusion(::FixedDiffusion, integ) = local_scalar_diffusion(integ.cache)o + +""" + FixedMVDiffusion(; initial_diffusion=1.0, calibrate=true) + +Time-fixed, diagonal diffusion, which is quasi-maximum-likelihood-estimated at each step. + +**Only works with the [`EK0`](@ref)!** + +A multi-variate version of [`FixedDiffusion`](@ref), where instead of an isotropic matrix, +a diagonal matrix is estimated. This can be helpful to get more expressive posterior +covariances when using the [`EK0`](@ref), since the individual dimensions can be adjusted +separately. + +# References +* [bosch20capos](@cite) Bosch et al, "Calibrated Adaptive Probabilistic ODE Solvers", AISTATS (2021) +""" +Base.@kwdef struct FixedMVDiffusion{T} <: AbstractStaticDiffusion + initial_diffusion::T = 1.0 + calibrate::Bool = true +end +function initial_diffusion(diffusionmodel::FixedMVDiffusion, d, q, Eltype) + initdiff = diffusionmodel.initial_diffusion + if initdiff isa Number + return initdiff * one(Eltype) * I(d) + elseif initdiff isa AbstractVector + @assert length(initdiff) == d + return Diagonal(initdiff) + elseif initdiff isa Diagonal + @assert size(initdiff) == (d, d) + return initdiff + else + throw(ArgumentError("Invalid `initial_diffusion`. The `FixedMVDiffusion` assumes a dxd diagonal diffusion model. So, pass either a Number, a Vector of length d, or a `Diagonal`.")) + end +end +estimate_local_diffusion(::FixedMVDiffusion, integ) = local_diagonal_diffusion(integ.cache) From fa9fa33f115164dab09a6eee9f053ee2d249e7ad Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 12:08:14 +0100 Subject: [PATCH 56/99] Grealy simplify the local error estimate code --- src/perform_step.jl | 55 +++++++++++++-------------------------------- 1 file changed, 16 insertions(+), 39 deletions(-) diff --git a/src/perform_step.jl b/src/perform_step.jl index 26a833a24..4469c978b 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -218,45 +218,22 @@ To save allocations, the function modifies the given `cache` and writes into `cache.C_Dxd` during some computations. """ function estimate_errors!(cache::AbstractODEFilterCache) - @unpack local_diffusion, Qh, H, d, q = cache - - R = cache.C_Dxd - - if local_diffusion isa Diagonal{<:Number,<:Vector} - _Q = apply_diffusion(Qh, local_diffusion) - _matmul!(R, _Q.R, H') - error_estimate = view(cache.tmp, 1:d) - if R isa BlockDiag - for i in eachindex(R.blocks) - error_estimate[i] = sum(abs2, R.blocks[i]) - end - else - sum!(abs2, error_estimate', view(R, :, 1:d)) - end - error_estimate .= sqrt.(error_estimate) - return error_estimate - elseif local_diffusion isa Diagonal{<:Number,<:FillArrays.Fill} - _matmul!(R, Qh.R, H') - - # error_estimate = diag(PSDMatrix(R)) - # error_estimate .*= local_diffusion - # error_estimate .= sqrt.(error_estimate) - # error_estimate = view(error_estimate, 1:d) - - # faster: - error_estimate = view(cache.tmp, 1:d) - if R isa IsometricKroneckerProduct - error_estimate .= sum(abs2, R.B) - elseif R isa BlockDiag - for i in eachindex(blocks(R)) - error_estimate[i] = sum(abs2, R.blocks[i]) - end - else - sum!(abs2, error_estimate', view(R, :, 1:d)) - end - error_estimate .*= local_diffusion.diag.value - error_estimate .= sqrt.(error_estimate) + @unpack local_diffusion, Qh, H, C_d, C_Dxd, C_DxD = cache + _Q = apply_diffusion!(PSDMatrix(C_DxD), Qh, local_diffusion) + _HQH = PSDMatrix(_matmul!(C_Dxd, _Q.R, H')) + error_estimate = diag!(C_d, _HQH) + @.. error_estimate = sqrt(error_estimate) + return error_estimate +end - return error_estimate +diag!(v::AbstractVector, M::PSDMatrix) = (sum!(abs2, v', M.R); v) +diag!(v::AbstractVector, M::PSDMatrix{<:Number,<:IsometricKroneckerProduct}) = + v .= sum(abs2, M.R.B) +diag!(v::AbstractVector, M::PSDMatrix{<:Number,<:BlockDiag}) = begin + @assert length(v) == nblocks(M.R) + @assert size(blocks(M.R)[1], 2) == 1 # assumes all of them have the same shape + @simd ivdep for i in eachindex(blocks(M.R)) + v[i] = sum(abs2, blocks(M.R)[i]) end + return v end From 5aa15f2dbb97e2ae0c09f13f80e7beffe17ffd4e Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 13:03:31 +0100 Subject: [PATCH 57/99] Better predict tests --- test/core/filtering.jl | 90 ++++++++++++++++++++++++++---------------- 1 file changed, 55 insertions(+), 35 deletions(-) diff --git a/test/core/filtering.jl b/test/core/filtering.jl index d570e2eed..d4d922806 100644 --- a/test/core/filtering.jl +++ b/test/core/filtering.jl @@ -5,53 +5,52 @@ Check the correctness of the filtering implementations vs. basic readable math c using Test using ProbNumDiffEq using LinearAlgebra -import ProbNumDiffEq: IsometricKroneckerProduct +import ProbNumDiffEq: IsometricKroneckerProduct, BlockDiag import ProbNumDiffEq as PNDE +import BlockDiagonals @testset "PREDICT" begin # Setup - d = 5 - m = rand(d) - P_R = Matrix(UpperTriangular(rand(d, d))) - P = P_R'P_R + d = 2 + q = 2 + D = d * (q + 1) + m = rand(D) - A = rand(d, d) - Q_R = Matrix(UpperTriangular(rand(d, d))) - Q = Q_R'Q_R + _P_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q + 1, q + 1)))) + _P = _P_R'_P_R + PM = Matrix(_P) - # PREDICT - m_p = A * m - P_p = A * P * A' + Q + _A = IsometricKroneckerProduct(d, rand(q + 1, q + 1)) + AM = Matrix(_A) - x_curr = Gaussian(m, P) - x_out = copy(x_curr) + _Q_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q + 1, q + 1)))) + _Q = _Q_R'_Q_R + QM = Matrix(_Q) - C_DxD = zeros(d, d) - C_2DxD = zeros(2d, d) - C_3DxD = zeros(3d, d) + # PREDICT + m_p = AM * m + P_p = AM * PM * AM' + QM - _fstr(F) = F ? "Kronecker" : "None" - @testset "Factorization: $(_fstr(KRONECKER))" for KRONECKER in (false, true) - if KRONECKER - K = 2 - m = kron(ones(K), m) - P_R = IsometricKroneckerProduct(K, P_R) - P = P_R'P_R + @testset "Factorization: $_FAC" for _FAC in ( + PNDE.DenseCovariance, + PNDE.BlockDiagonalCovariance, + PNDE.IsometricKroneckerCovariance, + ) - A = IsometricKroneckerProduct(K, A) - Q_R = IsometricKroneckerProduct(K, Q_R) - Q = Q_R'Q_R + FAC = _FAC{Float64}(d, q) - m_p = A * m - P_p = A * P * A' + Q + P_R = PNDE.to_factorized_matrix(FAC, _P_R) + P = P_R'P_R + A = PNDE.to_factorized_matrix(FAC, _A) + Q_R = PNDE.to_factorized_matrix(FAC, _Q_R) + Q = Q_R'Q_R - x_curr = Gaussian(m, P) - x_out = copy(x_curr) + x_curr = Gaussian(m, P) + x_out = copy(x_curr) - C_DxD = IsometricKroneckerProduct(K, C_DxD) - C_2DxD = IsometricKroneckerProduct(K, C_2DxD) - C_3DxD = IsometricKroneckerProduct(K, C_3DxD) - end + C_DxD = PNDE.factorized_zeros(FAC, D, D) + C_2DxD = PNDE.factorized_zeros(FAC, 2D, D) + C_3DxD = PNDE.factorized_zeros(FAC, 3D, D) @testset "predict" begin x_out = ProbNumDiffEq.predict(x_curr, A, Q) @@ -68,6 +67,25 @@ import ProbNumDiffEq as PNDE @test P_p ≈ Matrix(x_out.Σ) end + @testset "predict! with PSDMatrix and diffusion" begin + for diffusion in (rand(), rand() * Eye(d), rand() * I(d), Diagonal(rand(d))) + if _FAC == PNDE.IsometricKroneckerCovariance && + !(diffusion isa Number || diffusion isa Diagonal{<:Number,<:FillArrays.Fill}) + continue + end + _diffusions = diffusion isa Number ? diffusion * Ones(d) : diffusion.diag + + QM_diff = Matrix(BlockDiagonal([σ² * _Q.B for σ² in _diffusions])) + P_p_diff = AM * PM * AM' + QM_diff + + x_curr = Gaussian(m, PSDMatrix(P_R)) + x_out = copy(x_curr) + Q_SR = PSDMatrix(Q_R) + ProbNumDiffEq.predict!(x_out, x_curr, A, Q_SR, C_DxD, C_2DxD, diffusion) + @test P_p_diff ≈ Matrix(x_out.Σ) + end + end + @testset "predict! with zero diffusion" begin x_curr = Gaussian(m, PSDMatrix(P_R)) x_out = copy(x_curr) @@ -81,8 +99,10 @@ import ProbNumDiffEq as PNDE x_curr = Gaussian(m, PSDMatrix(P_R)) x_out = copy(x_curr) # marginalize! needs tall square-roots: - Q_SR = if KRONECKER + Q_SR = if Q_R isa IsometricKroneckerProduct PSDMatrix(IsometricKroneckerProduct(Q_R.ldim, [Q_R.B; zero(Q_R.B)])) + elseif Q_R isa BlockDiag + PSDMatrix(BlockDiag([[B; zero(B)] for B in Q_R.blocks])) else PSDMatrix([Q_R; zero(Q_R)]) end From 7d07e9a11683f44d81a448bbe3c1524da4486ed9 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 15:41:42 +0100 Subject: [PATCH 58/99] Beter update and smoothing tests --- src/blockdiagonals.jl | 2 + test/core/filtering.jl | 178 ++++++++++++++++++++--------------------- 2 files changed, 89 insertions(+), 91 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index ab09f1801..15a713886 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -216,6 +216,8 @@ LinearAlgebra.rmul!(B::BlockDiag, n::Number) = begin return B end +LinearAlgebra.inv(A::BlockDiag) = BlockDiag(inv.(blocks(A))) + copy!(A::BlockDiag, B::Diagonal) = begin @assert size(A) == size(B) i = 1 diff --git a/test/core/filtering.jl b/test/core/filtering.jl index d4d922806..8a5cd3e19 100644 --- a/test/core/filtering.jl +++ b/test/core/filtering.jl @@ -118,60 +118,65 @@ end @testset "UPDATE" begin # Setup d = 5 + o = 1 + m_p = rand(d) - P_p_R = Matrix(UpperTriangular(rand(d, d))) - P_p = P_p_R'P_p_R + _P_p_R = IsometricKroneckerProduct(o, Matrix(UpperTriangular(rand(d, d)))) + _P_p = _P_p_R'_P_p_R + P_p_M = Matrix(_P_p) # Measure - o = 1 _HB = rand(1, d) - H = kron(I(o), _HB) + _H = IsometricKroneckerProduct(o, _HB) + HM = Matrix(_H) R = zeros(o, o) z_data = zeros(o) - z = H * m_p - SR = P_p_R * H' - S = Symmetric(SR'SR) + z = _H * m_p + _SR = _P_p_R * _H' + _S = _SR'_SR + SM = Matrix(_S) # UPDATE - S_inv = inv(S) - K = P_p * H' * S_inv - m = m_p + K * (z_data .- z) - P = P_p - K * S * K' + KM = P_p_M * HM' / SM + m = m_p + KM * (z_data .- z) + P = P_p_M - KM * SM * KM' - x_pred = Gaussian(m_p, P_p) - x_out = copy(x_pred) - measurement = Gaussian(z, S) + _R_R = rand(o, o) + _R = _R_R'_R_R + _SR_noisy = qr([_P_p_R * _H'; _R_R]).R |> Matrix + _S_noisy = _SR_noisy'_SR_noisy + SM_noisy = Matrix(_S_noisy) - C_dxd = zeros(o, o) - C_d = zeros(o) - C_Dxd = zeros(d, o) - C_DxD = zeros(d, d) - C_2DxD = zeros(2d, d) - C_3DxD = zeros(3d, d) + KM_noisy = P_p_M * HM' / SM_noisy + m_noisy = m_p + KM_noisy * (z_data .- z) + P_noisy = P_p_M - KM_noisy * SM_noisy * KM_noisy' - _fstr(F) = F ? "Kronecker" : "None" - @testset "Factorization: $(_fstr(KRONECKER))" for KRONECKER in (false, true) - if KRONECKER - P_p_R = IsometricKroneckerProduct(1, P_p_R) - P_p = P_p_R'P_p_R + @testset "Factorization: $_FAC" for _FAC in ( + PNDE.DenseCovariance, + PNDE.BlockDiagonalCovariance, + PNDE.IsometricKroneckerCovariance, + ) + FAC = _FAC{Float64}(o, d) - H = IsometricKroneckerProduct(1, _HB) - R = zeros(o, o) + C_dxd = PNDE.factorized_zeros(FAC, o, o) + C_d = zeros(o) + C_Dxd = PNDE.factorized_zeros(FAC, d, o) + C_DxD = PNDE.factorized_zeros(FAC, d, d) + C_2DxD = PNDE.factorized_zeros(FAC, 2d, d) + C_3DxD = PNDE.factorized_zeros(FAC, 3d, d) - SR = IsometricKroneckerProduct(1, SR) - S = SR'SR + P_p_R = PNDE.to_factorized_matrix(FAC, _P_p_R) + P_p = P_p_R'P_p_R - x_pred = Gaussian(m_p, P_p) - x_out = copy(x_pred) - measurement = Gaussian(z, S) + H = PNDE.to_factorized_matrix(FAC, _H) - C_dxd = IsometricKroneckerProduct(1, C_dxd) - C_Dxd = IsometricKroneckerProduct(1, C_Dxd) - C_DxD = IsometricKroneckerProduct(1, C_DxD) - C_2DxD = IsometricKroneckerProduct(1, C_2DxD) - C_3DxD = IsometricKroneckerProduct(1, C_3DxD) - end + SR = PNDE.to_factorized_matrix(FAC, _SR) + S = SR'SR + + x_pred = Gaussian(m_p, P_p) + x_out = copy(x_pred) + measurement = Gaussian(z, S) @testset "update" begin x_out = ProbNumDiffEq.update(x_pred, measurement, H) @@ -354,57 +359,63 @@ end @testset "SMOOTH" begin # Setup d = 5 - m, m_s = rand(d), rand(d) - P_R, P_s_R = Matrix(UpperTriangular(rand(d, d))), Matrix(UpperTriangular(rand(d, d))) - P, P_s = P_R'P_R, P_s_R'P_s_R + q = 2 + D = d * (q + 1) + + m, m_s = rand(D), rand(D) + _P_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q+1, q+1)))) + _P_s_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q+1, q+1)))) + _P, _P_s = _P_R'_P_R, _P_s_R'_P_s_R + PM, P_sM = Matrix(_P), Matrix(_P_s) - A = rand(d, d) - Q_R = Matrix(UpperTriangular(rand(d, d))) + I - Q = Q_R'Q_R - Q_SR = PSDMatrix(Q_R) + _A = IsometricKroneckerProduct(d, rand(q + 1, q + 1)) + AM = Matrix(_A) + _Q_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q+1, q+1))+I)) + _Q = _Q_R'_Q_R + _Q_SR = PSDMatrix(_Q_R) # PREDICT first - m_p = A * m - P_p_R = qr([P_R * A'; Q_R]).R |> Matrix - P_p = A * P * A' + Q - @assert P_p ≈ P_p_R'P_p_R + m_p = AM * m + _P_p_R = IsometricKroneckerProduct(d, qr([_P_R.B * _A.B'; _Q_R.B]).R |> Matrix) + _P_p = _A * _P * _A' + _Q + @assert _P_p ≈ _P_p_R'_P_p_R + P_pM = Matrix(_P_p) # SMOOTH - G = P * A' * inv(P_p) + G = _P * _A' * inv(_P_p) |> Matrix m_smoothed = m + G * (m_s - m_p) - P_smoothed = P + G * (P_s - P_p) * G' + P_smoothed = PM + G * (P_sM - P_pM) * G' - x_curr = Gaussian(m, P) - x_next = Gaussian(m_s, P_s) x_smoothed = Gaussian(m_smoothed, P_smoothed) - _fstr(F) = F ? "Kronecker" : "None" - @testset "Factorization: $(_fstr(KRONECKER))" for KRONECKER in (false, true) - K = 2 - if KRONECKER - P_R = IsometricKroneckerProduct(K, P_R) - P = P_R'P_R - - P_s_R = IsometricKroneckerProduct(K, P_s_R) - P_s = P_s_R'P_s_R + @testset "Factorization: $_FAC" for _FAC in ( + PNDE.DenseCovariance, + PNDE.BlockDiagonalCovariance, + PNDE.IsometricKroneckerCovariance, + ) + FAC = _FAC{Float64}(d, q) - P_p_R = IsometricKroneckerProduct(K, P_p_R) - P_p = P_p_R'P_p_R + P_R = PNDE.to_factorized_matrix(FAC, _P_R) + P = P_R'P_R + P_s_R = PNDE.to_factorized_matrix(FAC, _P_s_R) + P_s = P_s_R'P_s_R + P_p_R = PNDE.to_factorized_matrix(FAC, _P_p_R) + P_p = P_p_R'P_p_R - m, m_s, m_p = kron(ones(K), m), kron(ones(K), m_s), kron(ones(K), m_p) + x_curr = Gaussian(m, P) + x_next = Gaussian(m_s, P_s) - A = IsometricKroneckerProduct(K, A) - Q_R = IsometricKroneckerProduct(K, Q_R) - Q = Q_R'Q_R - Q_SR = PSDMatrix(Q_R) + A = PNDE.to_factorized_matrix(FAC, _A) + Q_R = PNDE.to_factorized_matrix(FAC, _Q_R) + Q = Q_R'Q_R + Q_SR = PSDMatrix(Q_R) - x_curr = Gaussian(m, P) - x_next = Gaussian(m_s, P_s) + x_curr = Gaussian(m, P) + x_next = Gaussian(m_s, P_s) - m_smoothed = kron(ones(K), m_smoothed) - P_smoothed = IsometricKroneckerProduct(K, P_smoothed) - x_smoothed = Gaussian(m_smoothed, P_smoothed) - end + C_DxD = PNDE.factorized_zeros(FAC, D, D) + C_2DxD = PNDE.factorized_zeros(FAC, 2D, D) + C_3DxD = PNDE.factorized_zeros(FAC, 3D, D) @testset "smooth" begin x_out, _ = ProbNumDiffEq.smooth(x_curr, x_next, A, Q) @@ -421,22 +432,12 @@ end @testset "smooth via backward kernels" begin K_forward = ProbNumDiffEq.AffineNormalKernel(copy(A), copy(Q_SR)) K_backward = ProbNumDiffEq.AffineNormalKernel( - copy(A), copy(m_p), - if KRONECKER - PSDMatrix(IsometricKroneckerProduct(K, zeros(2d, d))) - else - PSDMatrix(zeros(2d, d)) - end) + copy(A), copy(m_p), PSDMatrix(copy(C_2DxD))) x_curr = Gaussian(m, PSDMatrix(P_R)) |> copy x_next_pred = Gaussian(m_p, PSDMatrix(P_p_R)) |> copy x_next_smoothed = Gaussian(m_s, PSDMatrix(P_s_R)) |> copy - C_DxD = if KRONECKER - IsometricKroneckerProduct(K, zeros(d, d)) - else - zeros(d, d) - end ProbNumDiffEq.compute_backward_kernel!( K_backward, x_next_pred, x_curr, K_forward; C_DxD) @@ -447,11 +448,6 @@ end @test K_backward.b ≈ b @test Matrix(K_backward.C) ≈ Λ - C_3DxD = if KRONECKER - IsometricKroneckerProduct(K, zeros(3d, d)) - else - zeros(3d, d) - end ProbNumDiffEq.marginalize_mean!(x_curr.μ, x_next_smoothed.μ, K_backward) ProbNumDiffEq.marginalize_cov!( x_curr.Σ, From 4669e9ffbff7d0eb941f1dbf4d43a6fe0c9cf70d Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:08:15 +0100 Subject: [PATCH 59/99] Praise the lord for unittests --- src/filtering/markov_kernel.jl | 2 +- test/core/filtering.jl | 39 +++++++++++++++++++++++++++++----- 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index 03973c767..e1fcf6700 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -312,7 +312,7 @@ function compute_backward_kernel!( PSDMatrix(K.C.R.blocks[i]), ) _C_DxD = C_DxD.blocks[i] - _diffusion = diffusion isa Number ? diffusion : diffusion[i] + _diffusion = diffusion isa Number ? diffusion : diffusion.diag[i] compute_backward_kernel!( _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=_diffusion ) diff --git a/test/core/filtering.jl b/test/core/filtering.jl index 8a5cd3e19..2f872afaf 100644 --- a/test/core/filtering.jl +++ b/test/core/filtering.jl @@ -8,6 +8,7 @@ using LinearAlgebra import ProbNumDiffEq: IsometricKroneckerProduct, BlockDiag import ProbNumDiffEq as PNDE import BlockDiagonals +using FillArrays @testset "PREDICT" begin # Setup @@ -36,7 +37,6 @@ import BlockDiagonals PNDE.BlockDiagonalCovariance, PNDE.IsometricKroneckerCovariance, ) - FAC = _FAC{Float64}(d, q) P_R = PNDE.to_factorized_matrix(FAC, _P_R) @@ -70,7 +70,10 @@ import BlockDiagonals @testset "predict! with PSDMatrix and diffusion" begin for diffusion in (rand(), rand() * Eye(d), rand() * I(d), Diagonal(rand(d))) if _FAC == PNDE.IsometricKroneckerCovariance && - !(diffusion isa Number || diffusion isa Diagonal{<:Number,<:FillArrays.Fill}) + !( + diffusion isa Number || + diffusion isa Diagonal{<:Number,<:FillArrays.Fill} + ) continue end _diffusions = diffusion isa Number ? diffusion * Ones(d) : diffusion.diag @@ -363,14 +366,14 @@ end D = d * (q + 1) m, m_s = rand(D), rand(D) - _P_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q+1, q+1)))) - _P_s_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q+1, q+1)))) + _P_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q + 1, q + 1)))) + _P_s_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q + 1, q + 1)))) _P, _P_s = _P_R'_P_R, _P_s_R'_P_s_R PM, P_sM = Matrix(_P), Matrix(_P_s) _A = IsometricKroneckerProduct(d, rand(q + 1, q + 1)) AM = Matrix(_A) - _Q_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q+1, q+1))+I)) + _Q_R = IsometricKroneckerProduct(d, Matrix(UpperTriangular(rand(q + 1, q + 1)) + I)) _Q = _Q_R'_Q_R _Q_SR = PSDMatrix(_Q_R) @@ -470,6 +473,32 @@ end @test K2.b == K_backward.b @test K2.C == K_backward.C end + + @testset "smooth via backward kernels with diffusion $diffusion" for diffusion in + ( + rand(), rand() * Eye(d), rand() * I(d), Diagonal(rand(d)), + ) + if _FAC == PNDE.IsometricKroneckerCovariance && + !( + diffusion isa Number || + diffusion isa Diagonal{<:Number,<:FillArrays.Fill} + ) + continue + end + _diffusions = + diffusion isa Number ? diffusion * Ones(d) : diffusion.diag + QM_diff = Matrix(BlockDiagonal([σ² * _Q.B for σ² in _diffusions])) + + ProbNumDiffEq.compute_backward_kernel!( + K_backward, x_next_pred, x_curr, K_forward; C_DxD, diffusion) + + G = Matrix(x_curr.Σ) * Matrix(A)' * inv(Matrix(x_next_pred.Σ)) + b = x_curr.μ - G * x_next_pred.μ + Λ = (I - G * AM) * Matrix(x_curr.Σ) * (I - G * AM)' + G * QM_diff * G' + @test K_backward.A ≈ G + @test K_backward.b ≈ b + @test Matrix(K_backward.C) ≈ Λ + end end end end From 20d2a596538c932433ffba1cb4db61347dd590bd Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:08:55 +0100 Subject: [PATCH 60/99] JuliaFormatter.jl --- src/ProbNumDiffEq.jl | 3 ++- src/blockdiagonals.jl | 13 +++++++------ src/caches.jl | 16 ++++++++-------- src/diffusions/apply_diffusion.jl | 6 +++--- src/diffusions/calibration.jl | 8 +++++--- src/diffusions/typedefs.jl | 8 ++++++-- src/filtering/markov_kernel.jl | 2 +- src/filtering/predict.jl | 3 ++- 8 files changed, 34 insertions(+), 25 deletions(-) diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index 3b1fdfb17..3f1e7fde8 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -2,7 +2,8 @@ __precompile__() module ProbNumDiffEq -import Base: copy, copy!, show, size, ndims, similar, isapprox, isequal, iterate, ==, length, zero +import Base: + copy, copy!, show, size, ndims, similar, isapprox, isequal, iterate, ==, length, zero using LinearAlgebra import LinearAlgebra: mul! diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 15a713886..f5738b9dc 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -101,13 +101,14 @@ for _mul! in (:mul!, :_matmul!) return C end - @eval $_mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag, alpha::Number, beta::Number) = begin - @assert length(C.blocks) == length(A.blocks) == length(B.blocks) - @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag, alpha::Number, beta::Number) = + begin + @assert length(C.blocks) == length(A.blocks) == length(B.blocks) + @simd ivdep for i in eachindex(blocks(C)) + @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + end + return C end - return C - end (_mul! == :_matmul!) && @eval $_mul!( C::BlockDiag{T}, A::BlockDiag{T}, diff --git a/src/caches.jl b/src/caches.jl index 88a4f10a1..c2e8595b6 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -183,14 +183,14 @@ function OrdinaryDiffEq.alg_cache( _d = is_secondorder_ode ? 2d : d pu_tmp = if is_secondorder_ode Gaussian(similar(Array{uElType}, 2d), - PSDMatrix( - if FAC isa IsometricKroneckerCovariance - Kronecker.kronecker(similar(Matrix{uElType}, D ÷ d, _d ÷ d), I(d)) - elseif FAC isa BlockDiagonalCovariance - error("I have no idea") - else - similar(Matrix{uElType}, D, _d) - end)) + PSDMatrix( + if FAC isa IsometricKroneckerCovariance + Kronecker.kronecker(similar(Matrix{uElType}, D ÷ d, _d ÷ d), I(d)) + elseif FAC isa BlockDiagonalCovariance + error("I have no idea") + else + similar(Matrix{uElType}, D, _d) + end)) else Gaussian(similar(Array{uElType}, d), PSDMatrix(factorized_similar(FAC, D, d))) end diff --git a/src/diffusions/apply_diffusion.jl b/src/diffusions/apply_diffusion.jl index ca3bce189..3480e8119 100644 --- a/src/diffusions/apply_diffusion.jl +++ b/src/diffusions/apply_diffusion.jl @@ -10,7 +10,7 @@ apply_diffusion( ) = PSDMatrix(Q.R * sqrt.(diffusion)) apply_diffusion( Q::PSDMatrix, - diffusion::Diagonal{T,<:FillArrays.Fill} + diffusion::Diagonal{T,<:FillArrays.Fill}, ) where {T} = apply_diffusion(Q, diffusion.diag.value) apply_diffusion( Q::PSDMatrix{T,<:BlockDiag}, @@ -34,7 +34,7 @@ Apply the diffusion to the PSD transition noise covariance `Q` in place and retu apply_diffusion! apply_diffusion!( Q::PSDMatrix, - diffusion::Diagonal{T,<:FillArrays.Fill} + diffusion::Diagonal{T,<:FillArrays.Fill}, ) where {T} = begin rmul!(Q.R, sqrt.(diffusion.diag.value)) return Q @@ -58,7 +58,7 @@ apply_diffusion! apply_diffusion!( out::PSDMatrix, Q::PSDMatrix, - diffusion::Number + diffusion::Number, ) = begin _matmul!(out.R, Q.R, sqrt.(diffusion)) return out diff --git a/src/diffusions/calibration.jl b/src/diffusions/calibration.jl index b8f0975e5..9bf388145 100644 --- a/src/diffusions/calibration.jl +++ b/src/diffusions/calibration.jl @@ -56,7 +56,8 @@ function estimate_global_diffusion(::FixedDiffusion, integ) diffusion_increment else current_mle_diffusion = integ.cache.global_diffusion.diag.value - current_mle_diffusion + (diffusion_increment - current_mle_diffusion) / integ.success_iter + current_mle_diffusion + + (diffusion_increment - current_mle_diffusion) / integ.success_iter end integ.cache.global_diffusion = new_mle_diffusion * Eye(d) @@ -88,7 +89,7 @@ function estimate_global_diffusion(::FixedMVDiffusion, integ) v, S = measurement.μ, measurement.Σ # @assert diag(S) |> unique |> length == 1 diffusion_increment = let - @.. C_d = v ^ 2 / S[1, 1] + @.. C_d = v^2 / S[1, 1] Diagonal(C_d) end @@ -96,7 +97,8 @@ function estimate_global_diffusion(::FixedMVDiffusion, integ) diffusion_increment else current_mle_diffusion = integ.cache.global_diffusion - @.. current_mle_diffusion + (diffusion_increment - current_mle_diffusion) / integ.success_iter + @.. current_mle_diffusion + + (diffusion_increment - current_mle_diffusion) / integ.success_iter end copy!(integ.cache.global_diffusion, new_mle_diffusion) diff --git a/src/diffusions/typedefs.jl b/src/diffusions/typedefs.jl index b00830768..8d754b163 100644 --- a/src/diffusions/typedefs.jl +++ b/src/diffusions/typedefs.jl @@ -59,7 +59,7 @@ Base.@kwdef struct FixedDiffusion{T<:Number} <: AbstractStaticDiffusion end initial_diffusion(diffusionmodel::FixedDiffusion, d, q, Eltype) = diffusionmodel.initial_diffusion * one(Eltype) * Eye(d) -estimate_local_diffusion(::FixedDiffusion, integ) = local_scalar_diffusion(integ.cache)o +estimate_local_diffusion(::FixedDiffusion, integ) = local_scalar_diffusion(integ.cache) """ FixedMVDiffusion(; initial_diffusion=1.0, calibrate=true) @@ -91,7 +91,11 @@ function initial_diffusion(diffusionmodel::FixedMVDiffusion, d, q, Eltype) @assert size(initdiff) == (d, d) return initdiff else - throw(ArgumentError("Invalid `initial_diffusion`. The `FixedMVDiffusion` assumes a dxd diagonal diffusion model. So, pass either a Number, a Vector of length d, or a `Diagonal`.")) + throw( + ArgumentError( + "Invalid `initial_diffusion`. The `FixedMVDiffusion` assumes a dxd diagonal diffusion model. So, pass either a Number, a Vector of length d, or a `Diagonal`.", + ), + ) end end estimate_local_diffusion(::FixedMVDiffusion, integ) = local_diagonal_diffusion(integ.cache) diff --git a/src/filtering/markov_kernel.jl b/src/filtering/markov_kernel.jl index e1fcf6700..07873ede1 100644 --- a/src/filtering/markov_kernel.jl +++ b/src/filtering/markov_kernel.jl @@ -314,7 +314,7 @@ function compute_backward_kernel!( _C_DxD = C_DxD.blocks[i] _diffusion = diffusion isa Number ? diffusion : diffusion.diag[i] compute_backward_kernel!( - _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=_diffusion + _Kout, _xpred, _x, _K, C_DxD=_C_DxD, diffusion=_diffusion, ) end return Kout diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index 8d9b0b548..cf3e4047d 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -124,7 +124,8 @@ function predict_cov!( _Qh = PSDMatrix(Qh.R.B) _C_DxD = C_DxD.B _C_2DxD = C_2DxD.B - _diffusion = diffusion isa Number ? diffusion : + _diffusion = + diffusion isa Number ? diffusion : diffusion isa IsometricKroneckerProduct ? diffusion.B : diffusion return predict_cov!(_Σ_out, _Σ_curr, _Ah, _Qh, _C_DxD, _C_2DxD, _diffusion) From 4bf25683d617520bc753a73a78ec2f1ec509755a Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:09:03 +0100 Subject: [PATCH 61/99] Actually git the diffusion tests --- test/core/diffusions.jl | 61 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 test/core/diffusions.jl diff --git a/test/core/diffusions.jl b/test/core/diffusions.jl new file mode 100644 index 000000000..7c5658160 --- /dev/null +++ b/test/core/diffusions.jl @@ -0,0 +1,61 @@ +using ProbNumDiffEq +import ProbNumDiffEq as PNDE +using Test +using LinearAlgebra +using FillArrays + +d, q = 2, 3 +T = Float64 + +@testset "$diffusionmodel" for diffusionmodel in ( + DynamicDiffusion(), + DynamicMVDiffusion(), + FixedDiffusion(), + FixedDiffusion(calibrate=false), + FixedMVDiffusion(), + FixedMVDiffusion(; initial_diffusion=rand(2)), + FixedMVDiffusion(; initial_diffusion=Diagonal(rand(2))), + FixedMVDiffusion(; initial_diffusion=Diagonal(rand(2)), calibrate=false), +) + + # Test the initial diffusion + diff = PNDE.initial_diffusion(diffusionmodel, d, q, T) + @assert size(diff) == (d, d) + @assert diff isa Diagonal + if !(diffusionmodel isa FixedMVDiffusion || diffusionmodel isa DynamicMVDiffusion) + @assert diff isa Diagonal{T,<:Fill} + end + + # Test applying the diffusion + _, Q = PNDE.discretize(PNDE.IWP{T}(d, q), 0.1) + Qmat = PSDMatrix(Matrix(Q.R)) + _diff = rand() * diff + @testset "$FAC" for FAC in ( + PNDE.DenseCovariance{T}(d, q), + PNDE.BlockDiagonalCovariance{T}(d, q), + PNDE.IsometricKroneckerCovariance{T}(d, q), + ) + if diff isa Diagonal{T,<:Vector} && FAC isa PNDE.IsometricKroneckerCovariance + continue + end + + _Q = PNDE.to_factorized_matrix(FAC, Q) + Qdiff = @test_nowarn PNDE.apply_diffusion(_Q, _diff) + Qmatdiff = @test_nowarn PNDE.apply_diffusion(Qmat, _diff) + @test Qdiff == Qmatdiff + + if !(diff isa Diagonal{T,<:Vector} && FAC isa PNDE.DenseCovariance) + Qdiff = @test_nowarn PNDE.apply_diffusion!(copy(_Q), _diff) + @test Qdiff == Qmatdiff + + Qdiff = @test_nowarn PNDE.apply_diffusion!(copy(_Q), _Q, _diff) + @test Qdiff == Qmatdiff + end + end + + @testset "Calibration" begin + # MLE + # At their core, they all just compute z' * inv(S) * z + # and then do something with the result + end +end From 46dbd17ddf8b7571a809adc468a5bdd01367b6a7 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:10:17 +0100 Subject: [PATCH 62/99] Testfix --- test/core/filtering.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/core/filtering.jl b/test/core/filtering.jl index 2f872afaf..ddc6a41d7 100644 --- a/test/core/filtering.jl +++ b/test/core/filtering.jl @@ -7,7 +7,7 @@ using ProbNumDiffEq using LinearAlgebra import ProbNumDiffEq: IsometricKroneckerProduct, BlockDiag import ProbNumDiffEq as PNDE -import BlockDiagonals +using BlockDiagonals using FillArrays @testset "PREDICT" begin From b9de4c0ac8e409c4a14b0c27db540013c1980dd3 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:18:39 +0100 Subject: [PATCH 63/99] Remove some comments --- src/filtering/predict.jl | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/src/filtering/predict.jl b/src/filtering/predict.jl index cf3e4047d..b37fc5bd3 100644 --- a/src/filtering/predict.jl +++ b/src/filtering/predict.jl @@ -80,21 +80,6 @@ function predict_cov!( @.. R[D+1:2D, 1:D] = Qh.R else apply_diffusion!(PSDMatrix(view(R, D+1:2D, 1:D)), Qh, diffusion) - - # if diffusion isa Number - # _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt(diffusion)) - # elseif diffusion isa Diagonal{<:Number,<:FillArrays.Fill} - # _matmul!(view(R, D+1:2D, 1:D), Qh.R, sqrt.(diffusion.diag.value)) - # else - # @warn "This is not yet implemented efficiently; TODO" - # d = size(diffusion, 1) - # q = D ÷ d - 1 - # _matmul!( - # view(R, D+1:2D, 1:D), - # Qh.R, - # sqrt.(kron(Eye(d) * diffusion, Eye(q + 1))), - # ) - # end end _matmul!(M, R', R) chol = cholesky!(Symmetric(M), check=false) From 802f3226b14bf0d6eeb0bdb8a7e22c198342a0bd Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:18:58 +0100 Subject: [PATCH 64/99] Test that the computed log-likelihood is correct --- test/core/filtering.jl | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/test/core/filtering.jl b/test/core/filtering.jl index ddc6a41d7..9da066b22 100644 --- a/test/core/filtering.jl +++ b/test/core/filtering.jl @@ -9,6 +9,7 @@ import ProbNumDiffEq: IsometricKroneckerProduct, BlockDiag import ProbNumDiffEq as PNDE using BlockDiagonals using FillArrays +import ProbNumDiffEq.GaussianDistributions: logpdf @testset "PREDICT" begin # Setup @@ -140,6 +141,8 @@ end _S = _SR'_SR SM = Matrix(_S) + LL = logpdf(Gaussian(z, SM), z_data) + # UPDATE KM = P_p_M * HM' / SM m = m_p + KM * (z_data .- z) @@ -198,7 +201,7 @@ end z_cache = C_d x_pred = Gaussian(x_pred.μ, PSDMatrix(P_p_R)) x_out = copy(x_pred) - ProbNumDiffEq.update!( + _, ll = ProbNumDiffEq.update!( x_out, x_pred, msmnt, @@ -211,6 +214,7 @@ end ) @test m ≈ x_out.μ @test P ≈ Matrix(x_out.Σ) + @test ll ≈ LL end @testset "Zero predicted covariance" begin K_cache = copy(C_Dxd) @@ -258,6 +262,8 @@ end SR = qr([P_p_R * H'; R_R]).R |> Matrix S = Symmetric(SR'SR) + LL = logpdf(Gaussian(z, S), z_data) + # UPDATE S_inv = inv(S) K = P_p * H' * S_inv @@ -316,7 +322,7 @@ end z_cache = C_d x_pred = Gaussian(x_pred.μ, PSDMatrix(P_p_R)) x_out = copy(x_pred) - ProbNumDiffEq.update!( + _, ll = ProbNumDiffEq.update!( x_out, x_pred, msmnt, @@ -330,6 +336,7 @@ end ) @test m ≈ x_out.μ @test P ≈ Matrix(x_out.Σ) + @test ll ≈ LL end @testset "Zero predicted covariance" begin K_cache = copy(C_Dxd) From cc38b174da6469b8ef9f8b746760dabad449b153 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:21:46 +0100 Subject: [PATCH 65/99] Add BlockDiagonals compat entry --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 030a13470..d4969b1d5 100644 --- a/Project.toml +++ b/Project.toml @@ -49,6 +49,7 @@ RecipesBaseExt = "RecipesBase" [compat] ArrayAllocators = "0.3" +BlockDiagonals = "0.1" DiffEqBase = "6.122" DiffEqCallbacks = "2.36" DiffEqDevTools = "2" From 33d5852270e13a135a47f811f501d348f170af6a Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 16:34:29 +0100 Subject: [PATCH 66/99] Add more solvers to the autodiff tests --- test/autodiff.jl | 82 +++++++++++++++++++++++++----------------------- 1 file changed, 42 insertions(+), 40 deletions(-) diff --git a/test/autodiff.jl b/test/autodiff.jl index de5dbef83..149245cbe 100644 --- a/test/autodiff.jl +++ b/test/autodiff.jl @@ -9,48 +9,50 @@ using ForwardDiff import ODEProblemLibrary: prob_ode_fitzhughnagumo -const _prob = prob_ode_fitzhughnagumo -const prob = ODEProblem(modelingtoolkitize(_prob), _prob.u0, _prob.tspan, jac=true) +@testset "solver: $ALG" for ALG in (EK0, EK1, DiagonalEK1) -function param_to_loss(p) - sol = solve( - remake(prob, p=p), - EK1(order=3, smooth=false), - sensealg=SensitivityADPassThrough(), - abstol=1e-3, - reltol=1e-2, - save_everystep=false, - dense=false, - ) - return norm(sol.u[end]) # Dummy loss -end -function startval_to_loss(u0) - sol = solve( - remake(prob, u0=u0), - EK1(order=3, smooth=false), - sensealg=SensitivityADPassThrough(), - abstol=1e-3, - reltol=1e-2, - save_everystep=false, - dense=false, - ) - return norm(sol.u[end]) # Dummy loss -end + _prob = prob_ode_fitzhughnagumo + prob = ODEProblem(modelingtoolkitize(_prob), _prob.u0, _prob.tspan, jac=true) + function param_to_loss(p) + sol = solve( + remake(prob, p=p), + ALG(order=3, smooth=false), + sensealg=SensitivityADPassThrough(), + abstol=1e-3, + reltol=1e-2, + save_everystep=false, + dense=false, + ) + return norm(sol.u[end]) # Dummy loss + end + function startval_to_loss(u0) + sol = solve( + remake(prob, u0=u0), + ALG(order=3, smooth=false), + sensealg=SensitivityADPassThrough(), + abstol=1e-3, + reltol=1e-2, + save_everystep=false, + dense=false, + ) + return norm(sol.u[end]) # Dummy loss + end -const dldp = FiniteDiff.finite_difference_gradient(param_to_loss, prob.p) -const dldu0 = FiniteDiff.finite_difference_gradient(startval_to_loss, prob.u0) + dldp = FiniteDiff.finite_difference_gradient(param_to_loss, prob.p) + dldu0 = FiniteDiff.finite_difference_gradient(startval_to_loss, prob.u0) -@testset "ForwardDiff.jl" begin - @test ForwardDiff.gradient(param_to_loss, prob.p) ≈ dldp rtol = 1e-4 - @test ForwardDiff.gradient(startval_to_loss, prob.u0) ≈ dldu0 rtol = 5e-4 -end + @testset "ForwardDiff.jl" begin + @test ForwardDiff.gradient(param_to_loss, prob.p) ≈ dldp rtol = 1e-4 + @test ForwardDiff.gradient(startval_to_loss, prob.u0) ≈ dldu0 rtol = 5e-4 + end -# @testset "ReverseDiff.jl" begin -# @test_broken ReverseDiff.gradient(param_to_loss, prob.p) ≈ dldp -# @test_broken ReverseDiff.gradient(startval_to_loss, prob.u0) ≈ dldu0 -# end + # @testset "ReverseDiff.jl" begin + # @test_broken ReverseDiff.gradient(param_to_loss, prob.p) ≈ dldp + # @test_broken ReverseDiff.gradient(startval_to_loss, prob.u0) ≈ dldu0 + # end -# @testset "Zygote.jl" begin -# @test_broken Zygote.gradient(param_to_loss, prob.p) ≈ dldp -# @test_broken Zygote.gradient(startval_to_loss, prob.u0) ≈ dldu0 -# end + # @testset "Zygote.jl" begin + # @test_broken Zygote.gradient(param_to_loss, prob.p) ≈ dldp + # @test_broken Zygote.gradient(startval_to_loss, prob.u0) ≈ dldu0 + # end +end From 3cbcd20ef625d38c85bc4871fd37f8a78b2abe5f Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 20:58:16 +0100 Subject: [PATCH 67/99] Add much more tests --- src/blockdiagonals.jl | 67 ++++++++++++++++++++++++++++++------- test/autodiff.jl | 4 +-- test/core/blockdiagonals.jl | 8 +++++ test/correctness.jl | 12 +++++++ test/mass_matrix.jl | 15 +++++++++ test/stiff_problem.jl | 4 +++ 6 files changed, 96 insertions(+), 14 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index f5738b9dc..673516487 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -85,7 +85,7 @@ for _mul! in (:mul!, :_matmul!) @eval $_mul!(C::BlockDiag, A::BlockDiag, B::BlockDiag) = begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i]) + @inbounds $_mul!(C.blocks[i], A.blocks[i], B.blocks[i]) end return C end @@ -96,7 +96,7 @@ for _mul! in (:mul!, :_matmul!) ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i]) + @inbounds $_mul!(C.blocks[i], A.blocks[i], B.blocks[i]) end return C end @@ -105,7 +105,7 @@ for _mul! in (:mul!, :_matmul!) begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + @inbounds $_mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) end return C end @@ -118,7 +118,7 @@ for _mul! in (:mul!, :_matmul!) ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) + @inbounds $_mul!(C.blocks[i], A.blocks[i], B.blocks[i], alpha, beta) end return C end @@ -126,7 +126,7 @@ for _mul! in (:mul!, :_matmul!) @eval $_mul!(C::BlockDiag, A::Adjoint{<:Number,<:BlockDiag}, B::BlockDiag) = begin @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + @inbounds $_mul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) end return C end @@ -137,7 +137,7 @@ for _mul! in (:mul!, :_matmul!) ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + @inbounds $_mul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) end return C end @@ -145,7 +145,7 @@ for _mul! in (:mul!, :_matmul!) @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Adjoint{<:Number,<:BlockDiag}) = begin @assert length(C.blocks) == length(A.blocks) == length(B.parent.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) + @inbounds $_mul!(C.blocks[i], A.blocks[i], adjoint(B.parent.blocks[i])) end return C end @@ -156,7 +156,7 @@ for _mul! in (:mul!, :_matmul!) ) where {T<:LinearAlgebra.BlasFloat} = begin @assert length(C.blocks) == length(A.parent.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds _matmul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) + @inbounds $_mul!(C.blocks[i], adjoint(A.parent.blocks[i]), B.blocks[i]) end return C end @@ -164,14 +164,14 @@ for _mul! in (:mul!, :_matmul!) @eval $_mul!(C::BlockDiag, A::Number, B::BlockDiag) = begin @assert length(C.blocks) == length(B.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A, B.blocks[i]) + @inbounds $_mul!(C.blocks[i], A, B.blocks[i]) end return C end @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Number) = begin @assert length(C.blocks) == length(A.blocks) @simd ivdep for i in eachindex(blocks(C)) - @inbounds mul!(C.blocks[i], A.blocks[i], B) + @inbounds $_mul!(C.blocks[i], A.blocks[i], B) end return C end @@ -186,7 +186,7 @@ for _mul! in (:mul!, :_matmul!) ic, ib = 1, 1 for i in eachindex(blocks(A)) d1, d2 = size(A.blocks[i]) - @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) + @inbounds $_mul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) ic += d1 ib += d2 end @@ -202,7 +202,7 @@ for _mul! in (:mul!, :_matmul!) ic, ib = 1, 1 for i in eachindex(blocks(A)) d1, d2 = size(A.blocks[i]) - @inbounds _matmul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) + @inbounds $_mul!(view(C, ic:(ic+d1-1)), A.blocks[i], view(B, ib:(ib+d2-1))) ic += d1 ib += d2 end @@ -230,6 +230,49 @@ copy!(A::BlockDiag, B::Diagonal) = begin return A end +Base.:*(D::Diagonal, A::BlockDiag) = begin + @assert size(D, 2) == size(A, 1) + local i = 1 + outblocks = map(blocks(A)) do Ai + d = size(Ai, 1) + outi = Diagonal(view(D.diag, i:(i+d-1))) * Ai + i += d + outi + end + return BlockDiag(outblocks) +end +Base.:*(A::BlockDiag, D::Diagonal) = begin + local i = 1 + outblocks = map(blocks(A)) do Ai + d = size(Ai, 2) + outi = Ai * Diagonal(view(D.diag, i:(i+d-1))) + i += d + outi + end + return BlockDiag(outblocks) +end +for _mul! in (:mul!, :_matmul!) + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Diagonal) = begin + local i = 1 + map(zip(blocks(C), blocks(A))) do (Ci, Ai) + d = size(Ai, 2) + $_mul!(Ci, Ai, Diagonal(view(B.diag, i:(i+d-1)))) + i += d + end + return C + end + @eval $_mul!(C::BlockDiag, A::Diagonal, B::BlockDiag) = begin + local i = 1 + map(zip(blocks(C), blocks(B))) do (Ci, Bi) + d = size(Bi, 1) + $_mul!(Ci, Diagonal(view(A.diag, i:(i+d-1))), Bi) + i += d + end + return C + end +end + + Base.isequal(A::BlockDiag, B::BlockDiag) = length(A.blocks) == length(B.blocks) && all(map(isequal, A.blocks, B.blocks)) ==(A::BlockDiag, B::BlockDiag) = diff --git a/test/autodiff.jl b/test/autodiff.jl index 149245cbe..a6abc357a 100644 --- a/test/autodiff.jl +++ b/test/autodiff.jl @@ -42,8 +42,8 @@ import ODEProblemLibrary: prob_ode_fitzhughnagumo dldu0 = FiniteDiff.finite_difference_gradient(startval_to_loss, prob.u0) @testset "ForwardDiff.jl" begin - @test ForwardDiff.gradient(param_to_loss, prob.p) ≈ dldp rtol = 1e-4 - @test ForwardDiff.gradient(startval_to_loss, prob.u0) ≈ dldu0 rtol = 5e-4 + @test ForwardDiff.gradient(param_to_loss, prob.p) ≈ dldp rtol = 1e-3 + @test ForwardDiff.gradient(startval_to_loss, prob.u0) ≈ dldu0 rtol = 5e-3 end # @testset "ReverseDiff.jl" begin diff --git a/test/core/blockdiagonals.jl b/test/core/blockdiagonals.jl index 81eb592dc..3f9e8917b 100644 --- a/test/core/blockdiagonals.jl +++ b/test/core/blockdiagonals.jl @@ -5,6 +5,7 @@ using BlockDiagonals using Test d1, d2 = 2, 3 +D = d1 * d2 @testset "T=$T" for T in (Float64, BigFloat) A = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) B = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) @@ -55,6 +56,13 @@ d1, d2 = 2, 3 @test tttm((a * I) * A) ≈ a * AM @test tttm(rmul!(copy(A), a)) ≈ a * AM + @test tttm((a * I(D)) * A) ≈ a * AM + @test tttm(A * (a * I(D))) ≈ AM * a + @test tttm(mul!(_A, A, a * I(D))) ≈ a * AM + @test tttm(mul!(_A, a * I(D), A)) ≈ a * AM + @test tttm(_matmul!(_A, A, a * I(D))) ≈ a * AM + @test tttm(_matmul!(_A, a * I(D), A)) ≈ a * AM + @test_throws ErrorException view(A, 1:2, 1:2) tttm(copy!(A, Diagonal(A))) diff --git a/test/correctness.jl b/test/correctness.jl index ac9a34235..049b5fd3f 100644 --- a/test/correctness.jl +++ b/test/correctness.jl @@ -30,6 +30,10 @@ CONSTANT_ALGS = ( EK1(order=3, smooth=false, diffusionmodel=FixedDiffusion()) => 1e-8, EK1(order=3, smooth=false, initialization=ClassicSolverInit()) => 1e-7, EK1(order=3, smooth=false, initialization=SimpleInit()) => 1e-5, + DiagonalEK1(order=3, smooth=false) => 1e-7, + DiagonalEK1(order=3, smooth=false, diffusionmodel=FixedDiffusion()) => 1e-7, + DiagonalEK1(order=3, smooth=false, diffusionmodel=DynamicDiffusion()) => 1e-7, + DiagonalEK1(order=3, smooth=false, initialization=ClassicSolverInit()) => 1e-7, # smoothing EK0(order=3, smooth=true) => 1e-8, EK0(order=3, smooth=true, diffusionmodel=FixedDiffusion()) => 2e-8, @@ -37,6 +41,10 @@ CONSTANT_ALGS = ( EK0(order=3, smooth=true, diffusionmodel=DynamicMVDiffusion()) => 1e-8, EK1(order=3, smooth=true) => 1e-8, EK1(order=3, smooth=true, diffusionmodel=FixedDiffusion()) => 1e-8, + DiagonalEK1(order=3, smooth=true) => 1e-7, + DiagonalEK1(order=3, smooth=true, diffusionmodel=FixedDiffusion()) => 1e-7, + DiagonalEK1(order=3, smooth=true, diffusionmodel=DynamicDiffusion()) => 1e-7, + DiagonalEK1(order=3, smooth=true, initialization=ClassicSolverInit()) => 1e-7, # Priors EK0(prior=IOUP(3, -1), smooth=true) => 2e-9, EK1(prior=IOUP(3, -1), smooth=true, diffusionmodel=FixedDiffusion()) => 1e-9, @@ -58,6 +66,10 @@ ADAPTIVE_ALGS = ( EK1(order=8) => 5e-6, EK1(order=3, initialization=ClassicSolverInit()) => 1e-5, EK1(order=3, initialization=SimpleInit()) => 1e-4, + DiagonalEK1(order=3) => 1e-4, + DiagonalEK1(order=3, diffusionmodel=FixedDiffusion()) => 1e-4, + DiagonalEK1(order=3, diffusionmodel=DynamicDiffusion()) => 1e-4, + DiagonalEK1(order=3, initialization=ClassicSolverInit()) => 1e-4, # Priors EK0(prior=IOUP(3, -1), smooth=true) => 1e-5, EK1(prior=IOUP(3, -1), smooth=true, diffusionmodel=FixedDiffusion()) => 1e-5, diff --git a/test/mass_matrix.jl b/test/mass_matrix.jl index b24ebe170..fad2f4064 100644 --- a/test/mass_matrix.jl +++ b/test/mass_matrix.jl @@ -34,11 +34,22 @@ using Test adaptive=false, dt=1e-2, ) + diagonalek1() = solve( + prob, + DiagonalEK1(smooth=false), + save_everystep=false, + dense=false, + adaptive=false, + dt=1e-2, + ) s1 = ek1() s0 = ek0() + s1diag = diagonalek1() ref = solve(prob, RadauIIA5(), abstol=1e-9, reltol=1e-6) @test s0.u[end] ≈ ref.u[end] rtol = 1e-7 + @test s1.u[end] ≈ ref.u[end] rtol = 1e-7 + @test s1diag.u[end] ≈ ref.u[end] rtol = 1e-7 @test s1.pu.Σ[1] isa PSDMatrix{<:Number,<:Matrix} @test s0.pu.Σ[1] isa PSDMatrix{<:Number,<:ProbNumDiffEq.IsometricKroneckerProduct} @@ -63,6 +74,7 @@ end 0 1 0 0 0 0 ] + M = Diagonal([1, 1, 0]) f = ODEFunction(rober, mass_matrix=M) prob = ODEProblem(f, [1.0, 0.0, 0.0], (0.0, 1e-2), (0.04, 3e7, 1e4)) @@ -78,4 +90,7 @@ end sol = solve(prob, EK1(order=3, initialization=SimpleInit())) @test sol.u[end] ≈ ref.u[end] rtol = 1e-8 + + sol = solve(prob, DiagonalEK1(order=3)) + @test sol.u[end] ≈ ref.u[end] rtol = 1e-8 end diff --git a/test/stiff_problem.jl b/test/stiff_problem.jl index de8a6009f..e94a3ad54 100644 --- a/test/stiff_problem.jl +++ b/test/stiff_problem.jl @@ -8,3 +8,7 @@ appxsol = solve(prob, RadauIIA5()) sol = solve(prob, EK1(order=3)) @test appxsol.u[end] ≈ sol.u[end] rtol = 5e-3 @test appxsol(0.5) ≈ sol(0.5).μ rtol = 5e-3 + +sol = solve(prob, DiagonalEK1(order=3)) +@test appxsol.u[end] ≈ sol.u[end] rtol = 5e-3 +@test appxsol(0.5) ≈ sol(0.5).μ rtol = 5e-3 From d7e0cc0904ff07e37aacf64f4e6216bb5c21a108 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 21:00:59 +0100 Subject: [PATCH 68/99] JuliaFormatter.jl --- src/blockdiagonals.jl | 1 - test/autodiff.jl | 1 - 2 files changed, 2 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 673516487..9714ebd37 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -272,7 +272,6 @@ for _mul! in (:mul!, :_matmul!) end end - Base.isequal(A::BlockDiag, B::BlockDiag) = length(A.blocks) == length(B.blocks) && all(map(isequal, A.blocks, B.blocks)) ==(A::BlockDiag, B::BlockDiag) = diff --git a/test/autodiff.jl b/test/autodiff.jl index a6abc357a..235478b23 100644 --- a/test/autodiff.jl +++ b/test/autodiff.jl @@ -10,7 +10,6 @@ using ForwardDiff import ODEProblemLibrary: prob_ode_fitzhughnagumo @testset "solver: $ALG" for ALG in (EK0, EK1, DiagonalEK1) - _prob = prob_ode_fitzhughnagumo prob = ODEProblem(modelingtoolkitize(_prob), _prob.u0, _prob.tspan, jac=true) function param_to_loss(p) From 9bf877057afa2fece418db2fee680ba51ae0ce12 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 22:19:50 +0100 Subject: [PATCH 69/99] Better complexity test --- src/caches.jl | 4 +-- src/derivative_utils.jl | 3 +- src/fast_linalg.jl | 12 +++++++ test/complexity.jl | 76 ++++++++++++++++------------------------- 4 files changed, 46 insertions(+), 49 deletions(-) diff --git a/src/caches.jl b/src/caches.jl index c2e8595b6..8d0b89ea5 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -178,8 +178,8 @@ function OrdinaryDiffEq.alg_cache( # Caches du = is_secondorder_ode ? similar(u.x[2]) : similar(u) - # ddu = factorized_similar(FAC, length(u), length(u)) - ddu = similar(u, length(u), length(u)) + ddu = !isnothing(f.jac_prototype) ? + f.jac_prototype : zeros(uElType, length(u), length(u)) _d = is_secondorder_ode ? 2d : d pu_tmp = if is_secondorder_ode Gaussian(similar(Array{uElType}, 2d), diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index 55391644a..c9870eb8a 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -8,7 +8,8 @@ function calc_H!(H, integ, cache) calc_H_EK0!(H, integ, cache) # @assert integ.u == @view x_pred.μ[1:(q+1):end] OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) - _matmul!(H, view(ddu, 1:d, :), cache.SolProj, -1.0, 1.0) + _ddu = size(ddu, 2) != d ? view(ddu, 1:d, :) : ddu + _matmul!(H, _ddu, cache.SolProj, -1.0, 1.0) elseif integ.alg isa DiagonalEK1 calc_H_EK0!(H, integ, cache) # @assert integ.u == @view x_pred.μ[1:(q+1):end] diff --git a/src/fast_linalg.jl b/src/fast_linalg.jl index 7a311424b..6b758ff0b 100644 --- a/src/fast_linalg.jl +++ b/src/fast_linalg.jl @@ -43,6 +43,18 @@ _matmul!(C::MSR{T}, A::Diagonal{T}, B::MSR{T}) where {T<:LinearAlgebra.BlasFloat (@.. C = A.diag * B) _matmul!(C::MSR{T}, A::Diagonal{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = (@.. C = A * B) + +_matmul!(C::MSR, A::MSR, B::Diagonal, alpha::Number, beta::Number) = + @.. C = A * B.diag' * alpha + C * beta +_matmul!(C::MSR{T}, A::MSR{T}, B::Diagonal{T}, alpha::Number, beta::Number +) where {T<:LinearAlgebra.BlasFloat} = + @.. C = A * B.diag' * alpha + C * beta +_matmul!(C::MSR, A::Diagonal, B::MSR, alpha::Number, beta::Number) = + (@.. C = A.diag * B * alpha + C * beta) +_matmul!(C::MSR{T}, A::Diagonal{T}, B::MSR{T}, alpha::Number, beta::Number +) where {T<:LinearAlgebra.BlasFloat} = + (@.. C = A.diag * B * alpha + C * beta) + _matmul!( C::Matrix{T}, A::LowerTriangular{T}, diff --git a/test/complexity.jl b/test/complexity.jl index 31482dea5..4b72224f8 100644 --- a/test/complexity.jl +++ b/test/complexity.jl @@ -7,35 +7,39 @@ using Test, SafeTestsets @testset "Scaling with ODE dimension" begin f(du, u, p, t) = mul!(du, -0.9I, u) + jac(J, u, p, t) = @simd ivdep for i in 1:size(J,1) + J[i,i] = -0.9 + end tspan = (0.0, 1.0) prob = ODEProblem(f, ones(1), tspan) NUMRUNS = 20 - @testset "Order 1 + perfect init + no smoothing" begin - time_dim(d; Alg) = begin - _prob = remake(prob, u0=ones(d)) - tmin = Inf - for _ in 1:NUMRUNS - integ = init(_prob, - Alg( - smooth=false, - order=1, - initialization=ClassicSolverInit(), - ), - dense=false, save_everystep=false, - adaptive=false, dt=1e-2, - ) - t = @elapsed solve!(integ) - tmin = min(tmin, t) - end - return tmin + time_dim(d, alg; kwargs...) = begin + _prob = remake( + prob, + u0=ones(d), + f=ODEFunction(f; jac=jac, jac_prototype=Diagonal(ones(d))) + ) + tmin = Inf + for _ in 1:NUMRUNS + integ = init(_prob, alg; adaptive=false, dt=1e-2, kwargs...) + t = @elapsed solve!(integ) + tmin = min(tmin, t) end + return tmin + end + + @testset "Order 1 + perfect init + no smoothing" begin + f(d, Alg) = time_dim( + d, Alg(smooth=false, order=1, initialization=ClassicSolverInit()); + dense=false, save_everystep=false + ) dims_ek0 = 2 .^ (8:15) - times_ek0 = [time_dim(d; Alg=EK0) for d in dims_ek0] + times_ek0 = [f(d, EK0) for d in dims_ek0] dims_ek1 = 2 .^ (2:6) - times_ek1 = [time_dim(d; Alg=EK1) for d in dims_ek1] + times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) @test_skip slope(lr_ek0)[1] ≈ 1 atol = 0.1 @@ -48,23 +52,12 @@ using Test, SafeTestsets end @testset "Order 3 + Taylor-init + no smoothing" begin - time_dim(d; Alg) = begin - _prob = remake(prob, u0=ones(d)) - tmin = Inf - for _ in 1:NUMRUNS - integ = init(_prob, Alg(smooth=false), - dense=false, save_everystep=false, - adaptive=false, dt=1e-2) - t = @elapsed solve!(integ) - tmin = min(tmin, t) - end - return tmin - end + f(d, Alg) = time_dim(d, Alg(smooth=false); dense=false, save_everystep=false) dims_ek0 = 2 .^ (8:15) - times_ek0 = [time_dim(d; Alg=EK0) for d in dims_ek0] + times_ek0 = [f(d, EK0) for d in dims_ek0] dims_ek1 = 2 .^ (2:5) - times_ek1 = [time_dim(d; Alg=EK1) for d in dims_ek1] + times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) @test_skip slope(lr_ek0)[1] ≈ 1 atol = 0.1 @@ -77,21 +70,12 @@ using Test, SafeTestsets end @testset "Order 3 with smoothing and everyting" begin - time_dim(d; Alg) = begin - _prob = remake(prob, u0=ones(d)) - tmin = Inf - for _ in 1:NUMRUNS - integ = init(_prob, Alg(), adaptive=false, dt=1e-2) - t = @elapsed solve!(integ) - tmin = min(tmin, t) - end - return tmin - end + f(d, Alg) = time_dim(d, Alg()) dims_ek0 = 2 .^ (8:13) - times_ek0 = [time_dim(d; Alg=EK0) for d in dims_ek0] + times_ek0 = [f(d, EK0) for d in dims_ek0] dims_ek1 = 2 .^ (1:4) - times_ek1 = [time_dim(d; Alg=EK1) for d in dims_ek1] + times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) @test 0.5 < slope(lr_ek0)[1] < 1.3 From 2a2533fcd4b8a1f3d36808d7d65e60e2e020c076 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 22:41:27 +0100 Subject: [PATCH 70/99] Actually do a proper test to check the scaling of the solvers --- test/complexity.jl | 53 +++++++++++++++++++++++++++++----------------- 1 file changed, 33 insertions(+), 20 deletions(-) diff --git a/test/complexity.jl b/test/complexity.jl index 4b72224f8..56baf4b9f 100644 --- a/test/complexity.jl +++ b/test/complexity.jl @@ -38,17 +38,21 @@ using Test, SafeTestsets dims_ek0 = 2 .^ (8:15) times_ek0 = [f(d, EK0) for d in dims_ek0] - dims_ek1 = 2 .^ (2:6) - times_ek1 = [f(d, EK1) for d in dims_ek1] - lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) - @test_skip slope(lr_ek0)[1] ≈ 1 atol = 0.1 + slope(lr_ek0)[1] # should be 1 @test 0.5 < slope(lr_ek0)[1] < 1.3 + dims_ek1 = 2 .^ (3:6) + times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) - @test_skip slope(lr_ek1)[1] ≈ 2 atol = 0.2 - # This is what we would actually expect, not sure what's going wrong: - @test_skip slope(lr_ek1)[1] ≈ 3 atol = 0.1 + slope(lr_ek1)[1] # shoudl be 3 + @test 2.5 < slope(lr_ek1)[1] < 3.5 + + dims_dek1 = 2 .^ (4:10) + times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] + lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) + slope(lr_dek1)[1] # should be 1 + @test 0.5 < slope(lr_dek1)[1] < 1.3 end @testset "Order 3 + Taylor-init + no smoothing" begin @@ -56,17 +60,21 @@ using Test, SafeTestsets dims_ek0 = 2 .^ (8:15) times_ek0 = [f(d, EK0) for d in dims_ek0] - dims_ek1 = 2 .^ (2:5) - times_ek1 = [f(d, EK1) for d in dims_ek1] - lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) - @test_skip slope(lr_ek0)[1] ≈ 1 atol = 0.1 + slope(lr_ek0)[1] # should be 1 @test 0.5 < slope(lr_ek0)[1] < 1.3 + dims_ek1 = 2 .^ (3:6) + times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) - @test_skip slope(lr_ek1)[1] ≈ 2 atol = 0.5 - # This is what we would actually expect, not sure what's going wrong: - @test_skip slope(lr_ek1)[1] ≈ 3 atol = 0.1 + slope(lr_ek1)[1] # should be 3 + @test 2.5 < slope(lr_ek1)[1] < 3.5 + + dims_dek1 = 2 .^ (4:10) + times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] + lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) + slope(lr_dek1)[1] # should be 1 + @test 0.5 < slope(lr_dek1)[1] < 1.3 end @testset "Order 3 with smoothing and everyting" begin @@ -74,15 +82,20 @@ using Test, SafeTestsets dims_ek0 = 2 .^ (8:13) times_ek0 = [f(d, EK0) for d in dims_ek0] - dims_ek1 = 2 .^ (1:4) - times_ek1 = [f(d, EK1) for d in dims_ek1] - lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) + slope(lr_ek0)[1] # should be 1 @test 0.5 < slope(lr_ek0)[1] < 1.3 + dims_ek1 = 2 .^ (3:6) + times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) - @test_skip slope(lr_ek1)[1] ≈ 2 atol = 0.2 - # This is what we would actually expect, not sure what's going wrong: - @test_skip slope(lr_ek1)[1] ≈ 3 atol = 0.1 + slope(lr_ek1)[1] # should be 3 + @test 2.5 < slope(lr_ek1)[1] < 3.5 + + dims_dek1 = 2 .^ (4:10) + times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] + lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) + slope(lr_dek1)[1] # should be 1 + @test 0.5 < slope(lr_dek1)[1] < 1.3 end end From 8fd54bc03122ec4df2960740159b91f3b958f27c Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Fri, 16 Feb 2024 22:43:34 +0100 Subject: [PATCH 71/99] JuliaFormatter.jl --- src/caches.jl | 3 ++- src/fast_linalg.jl | 14 ++++++++++++-- test/complexity.jl | 8 ++++---- 3 files changed, 18 insertions(+), 7 deletions(-) diff --git a/src/caches.jl b/src/caches.jl index 8d0b89ea5..ab28e0657 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -178,7 +178,8 @@ function OrdinaryDiffEq.alg_cache( # Caches du = is_secondorder_ode ? similar(u.x[2]) : similar(u) - ddu = !isnothing(f.jac_prototype) ? + ddu = + !isnothing(f.jac_prototype) ? f.jac_prototype : zeros(uElType, length(u), length(u)) _d = is_secondorder_ode ? 2d : d pu_tmp = if is_secondorder_ode diff --git a/src/fast_linalg.jl b/src/fast_linalg.jl index 6b758ff0b..b7016d677 100644 --- a/src/fast_linalg.jl +++ b/src/fast_linalg.jl @@ -46,12 +46,22 @@ _matmul!(C::MSR{T}, A::Diagonal{T}, B::Diagonal{T}) where {T<:LinearAlgebra.Blas _matmul!(C::MSR, A::MSR, B::Diagonal, alpha::Number, beta::Number) = @.. C = A * B.diag' * alpha + C * beta -_matmul!(C::MSR{T}, A::MSR{T}, B::Diagonal{T}, alpha::Number, beta::Number +_matmul!( + C::MSR{T}, + A::MSR{T}, + B::Diagonal{T}, + alpha::Number, + beta::Number, ) where {T<:LinearAlgebra.BlasFloat} = @.. C = A * B.diag' * alpha + C * beta _matmul!(C::MSR, A::Diagonal, B::MSR, alpha::Number, beta::Number) = (@.. C = A.diag * B * alpha + C * beta) -_matmul!(C::MSR{T}, A::Diagonal{T}, B::MSR{T}, alpha::Number, beta::Number +_matmul!( + C::MSR{T}, + A::Diagonal{T}, + B::MSR{T}, + alpha::Number, + beta::Number, ) where {T<:LinearAlgebra.BlasFloat} = (@.. C = A.diag * B * alpha + C * beta) diff --git a/test/complexity.jl b/test/complexity.jl index 56baf4b9f..d34708192 100644 --- a/test/complexity.jl +++ b/test/complexity.jl @@ -7,8 +7,8 @@ using Test, SafeTestsets @testset "Scaling with ODE dimension" begin f(du, u, p, t) = mul!(du, -0.9I, u) - jac(J, u, p, t) = @simd ivdep for i in 1:size(J,1) - J[i,i] = -0.9 + jac(J, u, p, t) = @simd ivdep for i in 1:size(J, 1) + J[i, i] = -0.9 end tspan = (0.0, 1.0) prob = ODEProblem(f, ones(1), tspan) @@ -19,7 +19,7 @@ using Test, SafeTestsets _prob = remake( prob, u0=ones(d), - f=ODEFunction(f; jac=jac, jac_prototype=Diagonal(ones(d))) + f=ODEFunction(f; jac=jac, jac_prototype=Diagonal(ones(d))), ) tmin = Inf for _ in 1:NUMRUNS @@ -33,7 +33,7 @@ using Test, SafeTestsets @testset "Order 1 + perfect init + no smoothing" begin f(d, Alg) = time_dim( d, Alg(smooth=false, order=1, initialization=ClassicSolverInit()); - dense=false, save_everystep=false + dense=false, save_everystep=false, ) dims_ek0 = 2 .^ (8:15) From d9501bf06787259857d20c36f5e826d250a33cd2 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 11:10:12 +0100 Subject: [PATCH 72/99] Make preconditioner computation simpler and test better --- src/preconditioning.jl | 56 ++++++++---------------------------- test/core/preconditioning.jl | 23 ++++++++++++--- 2 files changed, 31 insertions(+), 48 deletions(-) diff --git a/src/preconditioning.jl b/src/preconditioning.jl index b6f323853..ad0e494a4 100644 --- a/src/preconditioning.jl +++ b/src/preconditioning.jl @@ -9,8 +9,10 @@ function init_preconditioner(C::DenseCovariance{elType}) where {elType} return P, PI end function init_preconditioner(C::BlockDiagonalCovariance{elType}) where {elType} - P = BlockDiag([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) - PI = BlockDiag([Diagonal(ones(elType, C.q + 1)) for _ in 1:C.d]) + B = Diagonal(ones(elType, C.q + 1)) + P = BlockDiag([B for _ in 1:C.d]) + BI = Diagonal(ones(elType, C.q + 1)) + PI = BlockDiag([BI for _ in 1:C.d]) return P, PI end @@ -36,26 +38,10 @@ end end return P end - -@fastmath @inbounds function make_preconditioner!(P::IsometricKroneckerProduct, h, d, q) - val = factorial(q) / h^(q + 1 / 2) - @simd ivdep for j in 0:q - P.B.diag[j+1] = val - val /= (q - j) / h - end - return P -end - -@fastmath @inbounds function make_preconditioner!(P::BlockDiag, h, d, q) - val = factorial(q) / h^(q + 1 / 2) - @simd ivdep for j in 0:q - for M in P.blocks - M.diag[j+1] = val - end - val /= (q - j) / h - end - return P -end +make_preconditioner!(P::IsometricKroneckerProduct, h, d, q) = + (make_preconditioner!(P.B, h, 1, q); P) +make_preconditioner!(P::BlockDiag, h, d, q) = + (make_preconditioner!(blocks(P)[1], h, 1, q); P) @fastmath @inbounds function make_preconditioner_inv!(PI::Diagonal, h, d, q) val = h^(q + 1 / 2) / factorial(q) @@ -68,25 +54,7 @@ end end return PI end - -@fastmath @inbounds function make_preconditioner_inv!( - PI::IsometricKroneckerProduct, h, d, q) - val = h^(q + 1 / 2) / factorial(q) - @simd ivdep for j in 0:q - PI.B.diag[j+1] = val - val *= (q - j) / h - end - return PI -end - -@fastmath @inbounds function make_preconditioner_inv!( - PI::BlockDiag, h, d, q) - val = h^(q + 1 / 2) / factorial(q) - @simd ivdep for j in 0:q - for M in PI.blocks - M.diag[j+1] = val - end - val *= (q - j) / h - end - return PI -end +make_preconditioner_inv!(PI::IsometricKroneckerProduct, h, d, q) = + (make_preconditioner_inv!(PI.B, h, 1, q); PI) +make_preconditioner_inv!(PI::BlockDiag, h, d, q) = + (make_preconditioner_inv!(blocks(PI)[1], h, 1, q); PI) diff --git a/test/core/preconditioning.jl b/test/core/preconditioning.jl index 41502fa6b..27c02cd62 100644 --- a/test/core/preconditioning.jl +++ b/test/core/preconditioning.jl @@ -2,10 +2,6 @@ using Test using LinearAlgebra using ProbNumDiffEq import ProbNumDiffEq as PNDE -import ODEProblemLibrary: prob_ode_lotkavolterra - -prob = prob_ode_lotkavolterra -prob = remake(prob, tspan=(0.0, 10.0)) @testset "Condition numbers of A,Q" begin h = 0.1 * rand() @@ -34,3 +30,22 @@ prob = remake(prob, tspan=(0.0, 10.0)) @test cond(Matrix(Qh)) > cond(Matrix(Qh_p)) @test cond(Matrix(Qh)) > cond(Matrix(Qh_p))^2 end + +@testset "Covariance Factorizations all doing the correct thing" begin + h = rand() + d, q = 2, 3 + + function make_preconditioners(FAC) + P, PI = PNDE.init_preconditioner(FAC{Float64}(d, q)) + PNDE.make_preconditioner!(P, h, d, q) + PNDE.make_preconditioner_inv!(PI, h, d, q) + return P, PI + end + + PK, PIK = make_preconditioners(PNDE.IsometricKroneckerCovariance) + PD, PID = make_preconditioners(PNDE.DenseCovariance) + PB, PIB = make_preconditioners(PNDE.BlockDiagonalCovariance) + + @test PK == PD == PB + @test PIK == PID == PIB +end From 3609d7ecc7496fee5a30b24b47d5157ec7276d4a Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 11:13:05 +0100 Subject: [PATCH 73/99] One more prior test --- test/core/priors.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/core/priors.jl b/test/core/priors.jl index 3d8564818..f5ed0faa2 100644 --- a/test/core/priors.jl +++ b/test/core/priors.jl @@ -165,7 +165,7 @@ end make_transition_matrices!(cache, prior, h) @test AH_22_IBM ≈ cache.Ah - for Γ in (σ^2, σ^2 * Eye(d)) + for Γ in (σ^2, σ^2 * Eye(d), σ^2 * I(d)) @test QH_22_IBM ≈ Matrix(PNDE.apply_diffusion(cache.Qh, Γ)) end if FAC != PNDE.IsometricKroneckerCovariance From f08ff1ed1416e20d45276099397aff540d3b4dfe Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 12:10:51 +0100 Subject: [PATCH 74/99] Fix some tests --- src/derivative_utils.jl | 9 +-------- src/diffusions/typedefs.jl | 3 ++- src/fast_linalg.jl | 2 +- test/correctness.jl | 1 + 4 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index c9870eb8a..80365e28a 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -12,15 +12,8 @@ function calc_H!(H, integ, cache) _matmul!(H, _ddu, cache.SolProj, -1.0, 1.0) elseif integ.alg isa DiagonalEK1 calc_H_EK0!(H, integ, cache) - # @assert integ.u == @view x_pred.μ[1:(q+1):end] - # ddu_full = Matrix(ddu) - # @info "ddu" ddu_full - # error() OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) - - @unpack C_dxd = cache - copy!(C_dxd, Diagonal(ddu)) - _matmul!(H, C_dxd, cache.SolProj, -1.0, 1.0) + _matmul!(H, Diagonal(ddu), cache.SolProj, -1.0, 1.0) else error("Unknown algorithm") end diff --git a/src/diffusions/typedefs.jl b/src/diffusions/typedefs.jl index 8d754b163..8ef16d547 100644 --- a/src/diffusions/typedefs.jl +++ b/src/diffusions/typedefs.jl @@ -98,4 +98,5 @@ function initial_diffusion(diffusionmodel::FixedMVDiffusion, d, q, Eltype) ) end end -estimate_local_diffusion(::FixedMVDiffusion, integ) = local_diagonal_diffusion(integ.cache) +estimate_local_diffusion(::FixedMVDiffusion, integ) = + integ.alg isa EK0 ? local_diagonal_diffusion(integ.cache) : local_scalar_diffusion(integ.cache) diff --git a/src/fast_linalg.jl b/src/fast_linalg.jl index b7016d677..05ae8c25b 100644 --- a/src/fast_linalg.jl +++ b/src/fast_linalg.jl @@ -30,7 +30,7 @@ _matmul!( _matmul!(C::AbstractVecOrMat, A::AbstractVecOrMat, b::Number) = @.. C = A * b _matmul!(C::AbstractVecOrMat, a::Number, B::AbstractVecOrMat) = @.. C = a * B # Matrix matrix products with diagonal matrices -const MSR{T} = Union{Matrix{T},SubArray{T},Base.ReshapedArray{T}} +const MSR{T} = Union{Matrix{T},SubArray{T},Base.ReshapedArray{T},Adjoint{T,<:Matrix}} _matmul!(C::MSR, A::MSR, B::Diagonal) = @.. C = A * B.diag' _matmul!(C::MSR, A::Diagonal, B::MSR) = diff --git a/test/correctness.jl b/test/correctness.jl index 049b5fd3f..a3eacb705 100644 --- a/test/correctness.jl +++ b/test/correctness.jl @@ -60,6 +60,7 @@ ADAPTIVE_ALGS = ( EK0(order=3, diffusionmodel=DynamicMVDiffusion()) => 5e-5, EK0(order=3, initialization=ClassicSolverInit()) => 5e-5, EK0(order=3, initialization=SimpleInit()) => 1e-4, + EK0(order=3, diffusionmodel=DynamicMVDiffusion(), initialization=ClassicSolverInit()) => 4e-5, EK1(order=2) => 2e-5, EK1(order=3) => 1e-5, EK1(order=5) => 1e-6, From 786b1edb52fdaa75ad0eee9263b21e261c8bc3c5 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 12:10:56 +0100 Subject: [PATCH 75/99] Get the data likelihoods to work with DiagonalEK1 and the EK0 --- src/ProbNumDiffEq.jl | 2 ++ src/blockdiagonals.jl | 19 +++++++++++++ src/callbacks/dataupdate.jl | 33 ++++++++++++++--------- src/covariance_structure.jl | 9 +++++++ src/data_likelihoods/fenrir.jl | 17 ++++++------ src/filtering/update.jl | 2 +- test/data_likelihoods.jl | 49 +++++++++++++++++++++++----------- 7 files changed, 93 insertions(+), 38 deletions(-) diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index 3f1e7fde8..a6771bdfc 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -46,6 +46,8 @@ vecvec2mat(x) = reduce(hcat, x)' cov2psdmatrix(cov::Number; d) = PSDMatrix(sqrt(cov) * Eye(d)) cov2psdmatrix(cov::UniformScaling; d) = PSDMatrix(sqrt(cov.λ) * Eye(d)) +cov2psdmatrix(cov::Diagonal{<:Number,<:FillArrays.Fill}; d) = + (@assert size(cov, 1) == size(cov, 2) == d; cov2psdmatrix(cov.diag.value; d)) cov2psdmatrix(cov::Diagonal; d) = (@assert size(cov, 1) == size(cov, 2) == d; PSDMatrix(sqrt.(cov))) cov2psdmatrix(cov::AbstractMatrix; d) = diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 9714ebd37..240e9d10e 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -270,8 +270,27 @@ for _mul! in (:mul!, :_matmul!) end return C end + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Diagonal, alpha::Number, beta::Number) = begin + local i = 1 + map(zip(blocks(C), blocks(A))) do (Ci, Ai) + d = size(Ai, 2) + $_mul!(Ci, Ai, Diagonal(view(B.diag, i:(i+d-1))), alpha, beta) + i += d + end + return C + end + @eval $_mul!(C::BlockDiag, A::Diagonal, B::BlockDiag, alpha::Number, beta::Number) = begin + local i = 1 + map(zip(blocks(C), blocks(B))) do (Ci, Bi) + d = size(Bi, 1) + $_mul!(Ci, Diagonal(view(A.diag, i:(i+d-1))), Bi, alpha, beta) + i += d + end + return C + end end + Base.isequal(A::BlockDiag, B::BlockDiag) = length(A.blocks) == length(B.blocks) && all(map(isequal, A.blocks, B.blocks)) ==(A::BlockDiag, B::BlockDiag) = diff --git a/src/callbacks/dataupdate.jl b/src/callbacks/dataupdate.jl index c14223cab..67823fdcc 100644 --- a/src/callbacks/dataupdate.jl +++ b/src/callbacks/dataupdate.jl @@ -52,30 +52,26 @@ function DataUpdateCallback( o = length(val) @unpack x, E0, m_tmp, G1 = integ.cache - H = view(G1, 1:o, :) - if observation_matrix === I - @.. H = E0 - elseif observation_matrix isa UniformScaling - @.. H = observation_matrix.λ * E0 - else - matmul!(H, observation_matrix, E0) - end + M = observation_matrix + H = M * E0 obs_mean = _matmul!(view(m_tmp.μ, 1:o), H, x.μ) obs_mean .-= val R = cov2psdmatrix(observation_noise_cov; d=o) + R = to_factorized_matrix(integ.cache.covariance_factorization, R) # _A = x.Σ.R * H' # obs_cov = _A'_A + R - obs_cov = PSDMatrix(qr!([x.Σ.R * H'; R.R]).R) + obs_cov = PSDMatrix(make_obscov_sqrt(x.Σ.R, H, R.R)) + obs = Gaussian(obs_mean, obs_cov) @unpack x_tmp, K1, C_DxD, C_dxd, C_Dxd, C_d = integ.cache - K1 = view(K1, :, 1:o) - C_dxd = view(C_dxd, 1:o, 1:o) - C_Dxd = view(C_Dxd, :, 1:o) - C_d = view(C_d, 1:o) + K1 = K1 * M' + C_dxd = M * C_dxd * M' + C_Dxd = C_Dxd * M' + C_d = M * C_d _x = copy!(x_tmp, x) _, ll = update!(x, _x, obs, H, K1, C_Dxd, C_DxD, C_dxd, C_d; R=R) @@ -85,3 +81,14 @@ function DataUpdateCallback( end return PresetTimeCallback(data.t, affect!; save_positions, kwargs...) end + +make_obscov_sqrt(PR::AbstractMatrix, H::AbstractMatrix, RR::AbstractMatrix) = + qr!([PR * H'; RR]).R +make_obscov_sqrt( + PR::IsometricKroneckerProduct, + H::IsometricKroneckerProduct, + RR::IsometricKroneckerProduct, +) = + IsometricKroneckerProduct(PR.ldim, make_obscov_sqrt(PR.B, H.B, RR.B)) +make_obscov_sqrt(PR::BlockDiag, H::BlockDiag, RR::BlockDiag) = + BlockDiag([make_obscov_sqrt(blocks(PR)[i], blocks(H)[i], blocks(RR)[i]) for i in eachindex(blocks(PR))]) diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index 234a77688..41c003e71 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -47,6 +47,15 @@ to_factorized_matrix(::DenseCovariance, M::AbstractMatrix) = Matrix(M) to_factorized_matrix(::IsometricKroneckerCovariance, M::IsometricKroneckerProduct) = M to_factorized_matrix(C::BlockDiagonalCovariance, M::IsometricKroneckerProduct) = BlockDiag([copy(M.B) for _ in 1:C.d]) +to_factorized_matrix(C::BlockDiagonalCovariance, M::Diagonal) = + copy!(factorized_similar(C, size(M)...), M) +to_factorized_matrix( + C::IsometricKroneckerCovariance, M::Diagonal{<:Number, <:FillArrays.Fill}) = begin + out = factorized_similar(C, size(M)...) + @assert length(out.B) == 1 + out.B .= M.diag.value + out + end for FT in [:DenseCovariance, :IsometricKroneckerCovariance, :BlockDiagonalCovariance] @eval to_factorized_matrix(FAC::$FT, M::PSDMatrix) = diff --git a/src/data_likelihoods/fenrir.jl b/src/data_likelihoods/fenrir.jl index a805f9de9..05e1c3fd6 100644 --- a/src/data_likelihoods/fenrir.jl +++ b/src/data_likelihoods/fenrir.jl @@ -57,6 +57,7 @@ function fenrir_data_loglik( # Fit the ODE solution / PN posterior to the provided data; this is the actual Fenrir o = length(data.u[1]) R = cov2psdmatrix(observation_noise_cov; d=o) + R = to_factorized_matrix(integ.cache.covariance_factorization, R) LL, _, _ = fit_pnsolution_to_data!(sol, R, data; proj=observation_matrix) return LL @@ -78,12 +79,12 @@ function fit_pnsolution_to_data!( _cache = ( x_tmp=x_tmp, C_DxD=C_DxD, - C_Dxd=view(C_Dxd, :, 1:o), - C_dxd=view(C_dxd, 1:o, 1:o), - C_d=view(C_d, 1:o), - K1=view(K1, :, 1:o), - K2=view(C_Dxd, :, 1:o), - m_tmp=Gaussian(view(m_tmp.μ, 1:o), view(m_tmp.Σ, 1:o, 1:o)), + C_Dxd=C_Dxd * proj', + C_dxd=proj * C_dxd * proj', + C_d=proj * C_d, + K1=K1 * proj', + K2=C_Dxd * proj', + m_tmp=proj * m_tmp, ) x_posterior = copy(sol.x_filt) # the object to be filled @@ -136,9 +137,7 @@ function measure_and_update!(x, u, H, R::PSDMatrix, cache) z, S = cache.m_tmp _matmul!(z, H, x.μ) z .-= u - _matmul!(cache.C_Dxd, x.Σ.R, H') - _matmul!(S, cache.C_Dxd', cache.C_Dxd) - S .+= _matmul!(cache.C_dxd, R.R', R.R) + S = PSDMatrix(make_obscov_sqrt(x.Σ.R, H, R.R)) msmnt = Gaussian(z, S) return update!(x, copy!(cache.x_tmp, x), msmnt, H; R=R, cache) diff --git a/src/filtering/update.jl b/src/filtering/update.jl index 3b07b0a45..cf10c6df1 100644 --- a/src/filtering/update.jl +++ b/src/filtering/update.jl @@ -230,7 +230,7 @@ function update!( M_cache.blocks[i], C_dxd.blocks[i], view(C_d, i:i); - R, + R=isnothing(R) ? nothing : PSDMatrix(blocks(R.R)[i]) ) ll += _ll end diff --git a/test/data_likelihoods.jl b/test/data_likelihoods.jl index dacb2b6fd..2a25c09db 100644 --- a/test/data_likelihoods.jl +++ b/test/data_likelihoods.jl @@ -33,25 +33,27 @@ kwargs = ( ) @testset "Compare data likelihoods" begin @testset "$alg" for alg in ( + # EK0 + EK0(), + EK0(diffusionmodel=FixedDiffusion()), + EK0(prior=IOUP(3, -1)), + EK0(prior=Matern(3, 1.5)), + # EK1 EK1(), EK1(diffusionmodel=FixedDiffusion()), # EK1(diffusionmodel=FixedMVDiffusion(rand(2), false)), # not yet supported EK1(prior=IOUP(3, -1)), EK1(prior=Matern(3, 1.5)), EK1(prior=IOUP(3, update_rate_parameter=true)), + # DiagonalEK1 + DiagonalEK1(), + DiagonalEK1(diffusionmodel=FixedDiffusion()), + DiagonalEK1(diffusionmodel=FixedMVDiffusion(rand(2), false)), ) compare_data_likelihoods(alg; kwargs...) end end -@testset "EK0 is not (yet) supported" begin - for ll in (PNDE.dalton_data_loglik, PNDE.filtering_data_loglik) - @test_broken ll(prob, EK0(smooth=false); kwargs...) - end - @test_broken PNDE.fenrir_data_loglik( - prob, EK0(smooth=true); kwargs...) -end - @testset "Partial observations" begin H = [1 0;] data_part = (t=times, u=[H * d for d in obss]) @@ -63,6 +65,14 @@ end adaptive=false, dt=DT, dense=false, ) + @test_broken compare_data_likelihoods( + DiagonalEK1(); + observation_matrix=H, + observation_noise_cov=σ^2, + data=data_part, + adaptive=false, dt=DT, + dense=false, + ) end @testset "Observation noise types: $(typeof(Σ))" for Σ in ( @@ -70,15 +80,24 @@ end σ^2 * I, σ^2 * I(2), σ^2 * Eye(2), + Diagonal([σ^2 0; 0 2σ^2]), [σ^2 0; 0 2σ^2], (A = randn(2, 2); A'A), (PSDMatrix(randn(2, 2))), ) - compare_data_likelihoods( - EK1(); - observation_noise_cov=Σ, - data=data, - adaptive=false, dt=DT, - dense=false, - ) + @testset "$alg" for alg in (EK0(), DiagonalEK1(), EK1()) + if alg isa EK0 && !(Σ isa Number || Σ isa UniformScaling || Σ isa Diagonal{<:Number,<:FillArrays.Fill}) + continue + end + if alg isa DiagonalEK1 && !(Σ isa Number || Σ isa UniformScaling || Σ isa Diagonal) + continue + end + compare_data_likelihoods( + alg; + observation_noise_cov=Σ, + data=data, + adaptive=false, dt=DT, + dense=false, + ) + end end From 29be0c089a8609cbf696c6205830b0d566fa470c Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 12:53:51 +0100 Subject: [PATCH 76/99] Make the data likelihoods better --- src/blockdiagonals.jl | 20 ++++++++++++----- src/callbacks/dataupdate.jl | 40 ++++++++++++++++++++++++++++------ src/checks.jl | 2 +- src/data_likelihoods/fenrir.jl | 19 +++++++--------- src/derivative_utils.jl | 3 ++- src/priors/iwp.jl | 8 +++---- 6 files changed, 62 insertions(+), 30 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 240e9d10e..2202e8398 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -254,7 +254,9 @@ end for _mul! in (:mul!, :_matmul!) @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Diagonal) = begin local i = 1 - map(zip(blocks(C), blocks(A))) do (Ci, Ai) + @assert nblocks(C) == nblocks(A) + for j in eachindex(blocks(C)) + Ci, Ai = blocks(C)[j], blocks(A)[j] d = size(Ai, 2) $_mul!(Ci, Ai, Diagonal(view(B.diag, i:(i+d-1)))) i += d @@ -263,7 +265,9 @@ for _mul! in (:mul!, :_matmul!) end @eval $_mul!(C::BlockDiag, A::Diagonal, B::BlockDiag) = begin local i = 1 - map(zip(blocks(C), blocks(B))) do (Ci, Bi) + @assert nblocks(C) == nblocks(B) + for j in eachindex(blocks(C)) + Ci, Bi = blocks(C)[j], blocks(B)[j] d = size(Bi, 1) $_mul!(Ci, Diagonal(view(A.diag, i:(i+d-1))), Bi) i += d @@ -272,7 +276,9 @@ for _mul! in (:mul!, :_matmul!) end @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Diagonal, alpha::Number, beta::Number) = begin local i = 1 - map(zip(blocks(C), blocks(A))) do (Ci, Ai) + @assert nblocks(C) == nblocks(A) + for j in eachindex(blocks(C)) + Ci, Ai = blocks(C)[j], blocks(A)[j] d = size(Ai, 2) $_mul!(Ci, Ai, Diagonal(view(B.diag, i:(i+d-1))), alpha, beta) i += d @@ -280,10 +286,12 @@ for _mul! in (:mul!, :_matmul!) return C end @eval $_mul!(C::BlockDiag, A::Diagonal, B::BlockDiag, alpha::Number, beta::Number) = begin - local i = 1 - map(zip(blocks(C), blocks(B))) do (Ci, Bi) + i = 1 + @assert nblocks(C) == nblocks(B) + for j in eachindex(blocks(C)) + Ci, Bi = blocks(C)[j], blocks(B)[j] d = size(Bi, 1) - $_mul!(Ci, Diagonal(view(A.diag, i:(i+d-1))), Bi, alpha, beta) + @inbounds $_mul!(Ci, Diagonal(view(A.diag, i:(i+d-1))), Bi, alpha, beta) i += d end return C diff --git a/src/callbacks/dataupdate.jl b/src/callbacks/dataupdate.jl index 67823fdcc..b42ec2643 100644 --- a/src/callbacks/dataupdate.jl +++ b/src/callbacks/dataupdate.jl @@ -50,6 +50,7 @@ function DataUpdateCallback( val = values[idx] o = length(val) + d = integ.cache.d @unpack x, E0, m_tmp, G1 = integ.cache M = observation_matrix @@ -67,12 +68,16 @@ function DataUpdateCallback( obs = Gaussian(obs_mean, obs_cov) - @unpack x_tmp, K1, C_DxD, C_dxd, C_Dxd, C_d = integ.cache - K1 = K1 * M' - C_dxd = M * C_dxd * M' - C_Dxd = C_Dxd * M' - C_d = M * C_d - _x = copy!(x_tmp, x) + _cache = if o != d + if !(integ.alg isa EK1) + error("Partial observations only work with the EK1 right now") + end + make_obssized_cache(integ.cache; o) + else + integ.cache + end + @unpack K1, C_DxD, C_dxd, C_Dxd, C_d = _cache + _x = copy!(integ.cache.x_tmp, x) _, ll = update!(x, _x, obs, H, K1, C_Dxd, C_DxD, C_dxd, C_d; R=R) if !isnothing(loglikelihood) @@ -91,4 +96,25 @@ make_obscov_sqrt( ) = IsometricKroneckerProduct(PR.ldim, make_obscov_sqrt(PR.B, H.B, RR.B)) make_obscov_sqrt(PR::BlockDiag, H::BlockDiag, RR::BlockDiag) = - BlockDiag([make_obscov_sqrt(blocks(PR)[i], blocks(H)[i], blocks(RR)[i]) for i in eachindex(blocks(PR))]) + BlockDiag([ + make_obscov_sqrt(blocks(PR)[i], blocks(H)[i], blocks(RR)[i]) for + i in eachindex(blocks(PR)) + ]) + +function make_obssized_cache(cache; o) + @unpack K1, C_DxD, C_dxd, C_Dxd, C_d, m_tmp, x_tmp = cache + return make_obssized_cache(K1, C_DxD, C_dxd, C_Dxd, C_d, m_tmp, x_tmp; o) +end +function make_obssized_cache( + K1::M, C_DxD::M, C_dxd::M, C_Dxd::M, C_d::V, m_tmp::G, x_tmp; o, +) where {M<:Matrix,V<:Vector,G<:Gaussian} + return ( + K1=view(K1, :, 1:o), + C_dxd=view(C_dxd, 1:o, 1:o), + C_Dxd=view(C_Dxd, :, 1:o), + C_d=view(C_d, 1:o), + C_DxD=C_DxD, + m_tmp=Gaussian(view(m_tmp.μ, 1:o), view(m_tmp.Σ, 1:o, 1:o)), + x_tmp=x_tmp, + ) +end diff --git a/src/checks.jl b/src/checks.jl index 6fb213b08..819ebe454 100644 --- a/src/checks.jl +++ b/src/checks.jl @@ -16,7 +16,7 @@ function check_densesmooth(integ) error("To use `dense=true` you need to set `smooth=true`!") end if !integ.opts.save_everystep && integ.alg.smooth - error("If you do not save all values, you do not need to smooth!") + error("If you set `save_everystep=false` also set `smooth=false` in the alg!") end end function check_saveiter(integ) diff --git a/src/data_likelihoods/fenrir.jl b/src/data_likelihoods/fenrir.jl index 05e1c3fd6..4ce7f3799 100644 --- a/src/data_likelihoods/fenrir.jl +++ b/src/data_likelihoods/fenrir.jl @@ -75,17 +75,14 @@ function fit_pnsolution_to_data!( LL = zero(eltype(sol.prob.p)) o = length(data.u[1]) - @unpack x_tmp, C_dxd, C_d, K1, C_Dxd, C_DxD, m_tmp = cache - _cache = ( - x_tmp=x_tmp, - C_DxD=C_DxD, - C_Dxd=C_Dxd * proj', - C_dxd=proj * C_dxd * proj', - C_d=proj * C_d, - K1=K1 * proj', - K2=C_Dxd * proj', - m_tmp=proj * m_tmp, - ) + d = cache.d + @unpack x_tmp, m_tmp = cache + _cache = if o != d + make_obssized_cache(cache; o) + else + cache + end + @unpack K1, C_DxD, C_dxd, C_Dxd, C_d = _cache x_posterior = copy(sol.x_filt) # the object to be filled state2data_projmat = proj * cache.SolProj diff --git a/src/derivative_utils.jl b/src/derivative_utils.jl index 80365e28a..bbfae4482 100644 --- a/src/derivative_utils.jl +++ b/src/derivative_utils.jl @@ -13,7 +13,8 @@ function calc_H!(H, integ, cache) elseif integ.alg isa DiagonalEK1 calc_H_EK0!(H, integ, cache) OrdinaryDiffEq.calc_J!(ddu, integ, cache, true) - _matmul!(H, Diagonal(ddu), cache.SolProj, -1.0, 1.0) + ddu_diag = Diagonal(ddu) + _matmul!(H, ddu_diag, cache.SolProj, -1.0, 1.0) else error("Unknown algorithm") end diff --git a/src/priors/iwp.jl b/src/priors/iwp.jl index b62d5d3ac..43bc2397e 100644 --- a/src/priors/iwp.jl +++ b/src/priors/iwp.jl @@ -158,8 +158,8 @@ end function initialize_transition_matrices(FAC::IsometricKroneckerCovariance, p::IWP, dt) A, Q = preconditioned_discretize(p) P, PI = initialize_preconditioner(FAC, p, dt) - Ah = PI * A * P - Qh = PSDMatrix(Q.R * PI) + Ah = copy(A) + Qh = copy(Q) return A, Q, Ah, Qh, P, PI end function initialize_transition_matrices(FAC::DenseCovariance, p::IWP, dt) @@ -174,8 +174,8 @@ function initialize_transition_matrices(FAC::BlockDiagonalCovariance, p::IWP, dt A = to_factorized_matrix(FAC, A) Q = to_factorized_matrix(FAC, Q) P, PI = initialize_preconditioner(FAC, p, dt) - Ah = PI * A * P - Qh = PSDMatrix(Q.R * PI) + Ah = copy(A) + Qh = copy(Q) return A, Q, Ah, Qh, P, PI end From 3b53494af25bcda10753d91f00c44515bf56311e Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 13:03:15 +0100 Subject: [PATCH 77/99] JuliaFormatter.jl --- src/blockdiagonals.jl | 43 +++++++++++++++++++------------------ src/covariance_structure.jl | 12 +++++------ src/diffusions/typedefs.jl | 3 ++- src/filtering/update.jl | 2 +- test/data_likelihoods.jl | 5 ++++- 5 files changed, 35 insertions(+), 30 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 2202e8398..0506a2bb7 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -274,31 +274,32 @@ for _mul! in (:mul!, :_matmul!) end return C end - @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Diagonal, alpha::Number, beta::Number) = begin - local i = 1 - @assert nblocks(C) == nblocks(A) - for j in eachindex(blocks(C)) - Ci, Ai = blocks(C)[j], blocks(A)[j] - d = size(Ai, 2) - $_mul!(Ci, Ai, Diagonal(view(B.diag, i:(i+d-1))), alpha, beta) - i += d + @eval $_mul!(C::BlockDiag, A::BlockDiag, B::Diagonal, alpha::Number, beta::Number) = + begin + local i = 1 + @assert nblocks(C) == nblocks(A) + for j in eachindex(blocks(C)) + Ci, Ai = blocks(C)[j], blocks(A)[j] + d = size(Ai, 2) + $_mul!(Ci, Ai, Diagonal(view(B.diag, i:(i+d-1))), alpha, beta) + i += d + end + return C end - return C - end - @eval $_mul!(C::BlockDiag, A::Diagonal, B::BlockDiag, alpha::Number, beta::Number) = begin - i = 1 - @assert nblocks(C) == nblocks(B) - for j in eachindex(blocks(C)) - Ci, Bi = blocks(C)[j], blocks(B)[j] - d = size(Bi, 1) - @inbounds $_mul!(Ci, Diagonal(view(A.diag, i:(i+d-1))), Bi, alpha, beta) - i += d + @eval $_mul!(C::BlockDiag, A::Diagonal, B::BlockDiag, alpha::Number, beta::Number) = + begin + i = 1 + @assert nblocks(C) == nblocks(B) + for j in eachindex(blocks(C)) + Ci, Bi = blocks(C)[j], blocks(B)[j] + d = size(Bi, 1) + @inbounds $_mul!(Ci, Diagonal(view(A.diag, i:(i+d-1))), Bi, alpha, beta) + i += d + end + return C end - return C - end end - Base.isequal(A::BlockDiag, B::BlockDiag) = length(A.blocks) == length(B.blocks) && all(map(isequal, A.blocks, B.blocks)) ==(A::BlockDiag, B::BlockDiag) = diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl index 41c003e71..2809a3e42 100644 --- a/src/covariance_structure.jl +++ b/src/covariance_structure.jl @@ -50,12 +50,12 @@ to_factorized_matrix(C::BlockDiagonalCovariance, M::IsometricKroneckerProduct) = to_factorized_matrix(C::BlockDiagonalCovariance, M::Diagonal) = copy!(factorized_similar(C, size(M)...), M) to_factorized_matrix( - C::IsometricKroneckerCovariance, M::Diagonal{<:Number, <:FillArrays.Fill}) = begin - out = factorized_similar(C, size(M)...) - @assert length(out.B) == 1 - out.B .= M.diag.value - out - end + C::IsometricKroneckerCovariance, M::Diagonal{<:Number,<:FillArrays.Fill}) = begin + out = factorized_similar(C, size(M)...) + @assert length(out.B) == 1 + out.B .= M.diag.value + out +end for FT in [:DenseCovariance, :IsometricKroneckerCovariance, :BlockDiagonalCovariance] @eval to_factorized_matrix(FAC::$FT, M::PSDMatrix) = diff --git a/src/diffusions/typedefs.jl b/src/diffusions/typedefs.jl index 8ef16d547..1ebfb1954 100644 --- a/src/diffusions/typedefs.jl +++ b/src/diffusions/typedefs.jl @@ -99,4 +99,5 @@ function initial_diffusion(diffusionmodel::FixedMVDiffusion, d, q, Eltype) end end estimate_local_diffusion(::FixedMVDiffusion, integ) = - integ.alg isa EK0 ? local_diagonal_diffusion(integ.cache) : local_scalar_diffusion(integ.cache) + integ.alg isa EK0 ? local_diagonal_diffusion(integ.cache) : + local_scalar_diffusion(integ.cache) diff --git a/src/filtering/update.jl b/src/filtering/update.jl index cf10c6df1..4e4468480 100644 --- a/src/filtering/update.jl +++ b/src/filtering/update.jl @@ -230,7 +230,7 @@ function update!( M_cache.blocks[i], C_dxd.blocks[i], view(C_d, i:i); - R=isnothing(R) ? nothing : PSDMatrix(blocks(R.R)[i]) + R=isnothing(R) ? nothing : PSDMatrix(blocks(R.R)[i]), ) ll += _ll end diff --git a/test/data_likelihoods.jl b/test/data_likelihoods.jl index 2a25c09db..cf428109f 100644 --- a/test/data_likelihoods.jl +++ b/test/data_likelihoods.jl @@ -86,7 +86,10 @@ end (PSDMatrix(randn(2, 2))), ) @testset "$alg" for alg in (EK0(), DiagonalEK1(), EK1()) - if alg isa EK0 && !(Σ isa Number || Σ isa UniformScaling || Σ isa Diagonal{<:Number,<:FillArrays.Fill}) + if alg isa EK0 && !( + Σ isa Number || Σ isa UniformScaling || + Σ isa Diagonal{<:Number,<:FillArrays.Fill} + ) continue end if alg isa DiagonalEK1 && !(Σ isa Number || Σ isa UniformScaling || Σ isa Diagonal) From 6d835994bdc4db9f24144f16bb0d1416e6cd1d8b Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 13:05:57 +0100 Subject: [PATCH 78/99] Relax the complexity tests even more --- test/complexity.jl | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/test/complexity.jl b/test/complexity.jl index d34708192..e4d770b24 100644 --- a/test/complexity.jl +++ b/test/complexity.jl @@ -40,19 +40,19 @@ using Test, SafeTestsets times_ek0 = [f(d, EK0) for d in dims_ek0] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) slope(lr_ek0)[1] # should be 1 - @test 0.5 < slope(lr_ek0)[1] < 1.3 + @test slope(lr_ek0)[1] ≈ 1 atol=1 dims_ek1 = 2 .^ (3:6) times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) slope(lr_ek1)[1] # shoudl be 3 - @test 2.5 < slope(lr_ek1)[1] < 3.5 + @test slope(lr_ek1)[1] ≈ 3 atol = 1 dims_dek1 = 2 .^ (4:10) times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) slope(lr_dek1)[1] # should be 1 - @test 0.5 < slope(lr_dek1)[1] < 1.3 + @test slope(lr_dek1)[1] ≈ 1 atol = 1 end @testset "Order 3 + Taylor-init + no smoothing" begin @@ -62,19 +62,19 @@ using Test, SafeTestsets times_ek0 = [f(d, EK0) for d in dims_ek0] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) slope(lr_ek0)[1] # should be 1 - @test 0.5 < slope(lr_ek0)[1] < 1.3 + @test slope(lr_ek0)[1] ≈ 1 atol = 1 dims_ek1 = 2 .^ (3:6) times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) slope(lr_ek1)[1] # should be 3 - @test 2.5 < slope(lr_ek1)[1] < 3.5 + @test slope(lr_ek1)[1] ≈ 3 atol = 1 dims_dek1 = 2 .^ (4:10) times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) slope(lr_dek1)[1] # should be 1 - @test 0.5 < slope(lr_dek1)[1] < 1.3 + @test slope(lr_dek1)[1] ≈ 1 atol = 1 end @testset "Order 3 with smoothing and everyting" begin @@ -84,18 +84,18 @@ using Test, SafeTestsets times_ek0 = [f(d, EK0) for d in dims_ek0] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) slope(lr_ek0)[1] # should be 1 - @test 0.5 < slope(lr_ek0)[1] < 1.3 + @test slope(lr_ek0)[1] ≈ 1 atol = 1 dims_ek1 = 2 .^ (3:6) times_ek1 = [f(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) slope(lr_ek1)[1] # should be 3 - @test 2.5 < slope(lr_ek1)[1] < 3.5 + @test slope(lr_ek1)[1] ≈ 3 atol = 1 dims_dek1 = 2 .^ (4:10) times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) slope(lr_dek1)[1] # should be 1 - @test 0.5 < slope(lr_dek1)[1] < 1.3 + @test slope(lr_dek1)[1] ≈ 1 atol = 1 end end From be1fbefa287bc54c9da68a25377d5e6a866cbda3 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 13:16:45 +0100 Subject: [PATCH 79/99] JuliaFormatter.jl --- test/complexity.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/complexity.jl b/test/complexity.jl index e4d770b24..3620faa94 100644 --- a/test/complexity.jl +++ b/test/complexity.jl @@ -40,7 +40,7 @@ using Test, SafeTestsets times_ek0 = [f(d, EK0) for d in dims_ek0] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) slope(lr_ek0)[1] # should be 1 - @test slope(lr_ek0)[1] ≈ 1 atol=1 + @test slope(lr_ek0)[1] ≈ 1 atol = 1 dims_ek1 = 2 .^ (3:6) times_ek1 = [f(d, EK1) for d in dims_ek1] From dd7d08da063a0be2d907654d29adc9c447b3b7f4 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 16:35:52 +0100 Subject: [PATCH 80/99] Remove some unused code --- src/diffusions/calibration.jl | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/src/diffusions/calibration.jl b/src/diffusions/calibration.jl index 9bf388145..6b22b3da3 100644 --- a/src/diffusions/calibration.jl +++ b/src/diffusions/calibration.jl @@ -170,15 +170,9 @@ function local_diagonal_diffusion(cache) tmp else @warn "This is not yet implemented efficiently; TODO" - diag(H * unfactorize(Qh) * H') + diag(X_A_Xt(Qh, H)) end - # To double-check: - HQH = H * unfactorize(Qh) * H' - @assert Q_11 ≈ diag(HQH) - # Also if the solver is a EK0 and not a DiagonalEK1: - # @assert Q_11 |> unique |> length == 1 - @. local_diffusion.diag = z^2 / Q_11 return local_diffusion end From d74e3c8cbff3907726f211a321990d68ef8da338 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 16:59:10 +0100 Subject: [PATCH 81/99] =?UTF-8?q?Bisschen=20code=20upgrade=20f=C3=BCr=20di?= =?UTF-8?q?e=20data=20likelihoods?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/callbacks/dataupdate.jl | 22 ++++++++++------------ src/data_likelihoods/fenrir.jl | 6 +----- src/diffusions/apply_diffusion.jl | 3 --- 3 files changed, 11 insertions(+), 20 deletions(-) diff --git a/src/callbacks/dataupdate.jl b/src/callbacks/dataupdate.jl index b42ec2643..def5759ac 100644 --- a/src/callbacks/dataupdate.jl +++ b/src/callbacks/dataupdate.jl @@ -68,14 +68,10 @@ function DataUpdateCallback( obs = Gaussian(obs_mean, obs_cov) - _cache = if o != d - if !(integ.alg isa EK1) - error("Partial observations only work with the EK1 right now") - end - make_obssized_cache(integ.cache; o) - else - integ.cache + if o != d && !(integ.alg isa EK1) + error("Partial observations only work with the EK1 right now") end + _cache = make_obssized_cache(integ.cache; o) @unpack K1, C_DxD, C_dxd, C_Dxd, C_d = _cache _x = copy!(integ.cache.x_tmp, x) _, ll = update!(x, _x, obs, H, K1, C_Dxd, C_DxD, C_dxd, C_d; R=R) @@ -102,12 +98,14 @@ make_obscov_sqrt(PR::BlockDiag, H::BlockDiag, RR::BlockDiag) = ]) function make_obssized_cache(cache; o) - @unpack K1, C_DxD, C_dxd, C_Dxd, C_d, m_tmp, x_tmp = cache - return make_obssized_cache(K1, C_DxD, C_dxd, C_Dxd, C_d, m_tmp, x_tmp; o) + if o == cache.d + return cache + else + return make_obssized_cache(cache.covariance_factorization, cache; o) + end end -function make_obssized_cache( - K1::M, C_DxD::M, C_dxd::M, C_Dxd::M, C_d::V, m_tmp::G, x_tmp; o, -) where {M<:Matrix,V<:Vector,G<:Gaussian} +function make_obssized_cache(::DenseCovariance, cache; o) + @unpack K1, C_DxD, C_dxd, C_Dxd, C_d, m_tmp, x_tmp = cache return ( K1=view(K1, :, 1:o), C_dxd=view(C_dxd, 1:o, 1:o), diff --git a/src/data_likelihoods/fenrir.jl b/src/data_likelihoods/fenrir.jl index 4ce7f3799..01e3a3746 100644 --- a/src/data_likelihoods/fenrir.jl +++ b/src/data_likelihoods/fenrir.jl @@ -77,11 +77,7 @@ function fit_pnsolution_to_data!( o = length(data.u[1]) d = cache.d @unpack x_tmp, m_tmp = cache - _cache = if o != d - make_obssized_cache(cache; o) - else - cache - end + _cache = make_obssized_cache(cache; o) @unpack K1, C_DxD, C_dxd, C_Dxd, C_d = _cache x_posterior = copy(sol.x_filt) # the object to be filled diff --git a/src/diffusions/apply_diffusion.jl b/src/diffusions/apply_diffusion.jl index 3480e8119..397ac5160 100644 --- a/src/diffusions/apply_diffusion.jl +++ b/src/diffusions/apply_diffusion.jl @@ -3,7 +3,6 @@ Apply the diffusion to the PSD transition noise covariance `Q`, return the result. """ -apply_diffusion apply_diffusion( Q::PSDMatrix, diffusion::Number, @@ -31,7 +30,6 @@ end Apply the diffusion to the PSD transition noise covariance `Q` in place and return the result. """ -apply_diffusion! apply_diffusion!( Q::PSDMatrix, diffusion::Diagonal{T,<:FillArrays.Fill}, @@ -54,7 +52,6 @@ end Apply the diffusion to the PSD transition noise covariance `Q` and store the result in `out`. """ -apply_diffusion! apply_diffusion!( out::PSDMatrix, Q::PSDMatrix, From eb2e235d47bfacdde7b19c772aa880976a5e9bc6 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 17:11:29 +0100 Subject: [PATCH 82/99] Make the complexity tests more compact --- test/complexity.jl | 35 +++++++++++++---------------------- 1 file changed, 13 insertions(+), 22 deletions(-) diff --git a/test/complexity.jl b/test/complexity.jl index 3620faa94..a87bbc471 100644 --- a/test/complexity.jl +++ b/test/complexity.jl @@ -15,7 +15,7 @@ using Test, SafeTestsets NUMRUNS = 20 - time_dim(d, alg; kwargs...) = begin + _timer(d, alg; kwargs...) = begin _prob = remake( prob, u0=ones(d), @@ -31,71 +31,62 @@ using Test, SafeTestsets end @testset "Order 1 + perfect init + no smoothing" begin - f(d, Alg) = time_dim( + t(d, Alg) = _timer( d, Alg(smooth=false, order=1, initialization=ClassicSolverInit()); dense=false, save_everystep=false, ) dims_ek0 = 2 .^ (8:15) - times_ek0 = [f(d, EK0) for d in dims_ek0] + times_ek0 = [t(d, EK0) for d in dims_ek0] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) - slope(lr_ek0)[1] # should be 1 @test slope(lr_ek0)[1] ≈ 1 atol = 1 dims_ek1 = 2 .^ (3:6) - times_ek1 = [f(d, EK1) for d in dims_ek1] + times_ek1 = [t(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) - slope(lr_ek1)[1] # shoudl be 3 @test slope(lr_ek1)[1] ≈ 3 atol = 1 dims_dek1 = 2 .^ (4:10) - times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] + times_dek1 = [t(d, DiagonalEK1) for d in dims_dek1] lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) - slope(lr_dek1)[1] # should be 1 @test slope(lr_dek1)[1] ≈ 1 atol = 1 end @testset "Order 3 + Taylor-init + no smoothing" begin - f(d, Alg) = time_dim(d, Alg(smooth=false); dense=false, save_everystep=false) + t(d, Alg) = _timer(d, Alg(smooth=false); dense=false, save_everystep=false) dims_ek0 = 2 .^ (8:15) - times_ek0 = [f(d, EK0) for d in dims_ek0] + times_ek0 = [t(d, EK0) for d in dims_ek0] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) - slope(lr_ek0)[1] # should be 1 @test slope(lr_ek0)[1] ≈ 1 atol = 1 dims_ek1 = 2 .^ (3:6) - times_ek1 = [f(d, EK1) for d in dims_ek1] + times_ek1 = [t(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) - slope(lr_ek1)[1] # should be 3 @test slope(lr_ek1)[1] ≈ 3 atol = 1 dims_dek1 = 2 .^ (4:10) - times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] + times_dek1 = [t(d, DiagonalEK1) for d in dims_dek1] lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) - slope(lr_dek1)[1] # should be 1 @test slope(lr_dek1)[1] ≈ 1 atol = 1 end @testset "Order 3 with smoothing and everyting" begin - f(d, Alg) = time_dim(d, Alg()) + t(d, Alg) = _timer(d, Alg()) dims_ek0 = 2 .^ (8:13) - times_ek0 = [f(d, EK0) for d in dims_ek0] + times_ek0 = [t(d, EK0) for d in dims_ek0] lr_ek0 = linregress(log.(dims_ek0), log.(times_ek0)) - slope(lr_ek0)[1] # should be 1 @test slope(lr_ek0)[1] ≈ 1 atol = 1 dims_ek1 = 2 .^ (3:6) - times_ek1 = [f(d, EK1) for d in dims_ek1] + times_ek1 = [t(d, EK1) for d in dims_ek1] lr_ek1 = linregress(log.(dims_ek1), log.(times_ek1)) - slope(lr_ek1)[1] # should be 3 @test slope(lr_ek1)[1] ≈ 3 atol = 1 dims_dek1 = 2 .^ (4:10) - times_dek1 = [f(d, DiagonalEK1) for d in dims_dek1] + times_dek1 = [t(d, DiagonalEK1) for d in dims_dek1] lr_dek1 = linregress(log.(dims_dek1), log.(times_dek1)) - slope(lr_dek1)[1] # should be 1 @test slope(lr_dek1)[1] ≈ 1 atol = 1 end end From 367b0bf3faccd89c215b108e96bf641d5bf6a8eb Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 17:48:17 +0100 Subject: [PATCH 83/99] Remove some unused things, mainly to re-trigger gh actions --- Project.toml | 6 ------ src/ProbNumDiffEq.jl | 4 +--- src/priors/ltisde.jl | 36 ------------------------------------ test/Project.toml | 1 - 4 files changed, 1 insertion(+), 46 deletions(-) diff --git a/Project.toml b/Project.toml index d4969b1d5..6dc1f5565 100644 --- a/Project.toml +++ b/Project.toml @@ -8,9 +8,7 @@ ArrayAllocators = "c9d4266f-a5cb-439d-837c-c97b191379f5" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" -ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" -FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" FiniteHorizonGramians = "b59a298d-d283-4a37-9369-85a9f9a111a5" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -29,7 +27,6 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" -SpecialMatrices = "928aab9d-ef52-54ac-8ca1-acd7ca42c160" StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" @@ -54,9 +51,7 @@ DiffEqBase = "6.122" DiffEqCallbacks = "2.36" DiffEqDevTools = "2" DocStringExtensions = "0.9" -ExponentialUtilities = "1" FastBroadcast = "0.2" -FastGaussQuadrature = "0.5, 1" FillArrays = "1.9" FiniteHorizonGramians = "0.2" ForwardDiff = "0.10" @@ -75,7 +70,6 @@ RecursiveArrayTools = "2, 3" Reexport = "1" SciMLBase = "1.90, 2" SimpleUnPack = "1" -SpecialMatrices = "3" StaticArrayInterface = "1.3" Statistics = "1" StructArrays = "0.4, 0.5, 0.6" diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index a6771bdfc..d6909ffde 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -16,7 +16,7 @@ using Reexport import SciMLBase import SciMLBase: interpret_vars, getsyms, remake using OrdinaryDiffEq -using SpecialMatrices, ToeplitzMatrices +using ToeplitzMatrices using FastBroadcast using StaticArrayInterface using FunctionWrappersWrappers @@ -25,9 +25,7 @@ using TaylorSeries, TaylorIntegration using SimpleUnPack using RecursiveArrayTools using ForwardDiff -using ExponentialUtilities using Octavian -using FastGaussQuadrature import Kronecker using ArrayAllocators using FiniteHorizonGramians diff --git a/src/priors/ltisde.jl b/src/priors/ltisde.jl index 6f7f14d57..468da1da8 100644 --- a/src/priors/ltisde.jl +++ b/src/priors/ltisde.jl @@ -65,39 +65,3 @@ function matrix_fraction_decomposition( Q = Mexp[1:d, d+1:end] * A' return A, Q end - -# Previous implementation, outdated thanks to FiniteHorizonGramians.jl: -function _discretize_sqrt_with_quadraturetrick(sde::LTISDE, dt::Real) - F, L = drift(sde), dispersion(sde) - - D = size(F, 1) - d = size(L, 2) - N = D # more robust than Int(D / d) - R = similar(F, N * d, D) - method = ExpMethodHigham2005() - expcache = ExponentialUtilities.alloc_mem(F, method) - - Ah = exponential!(dt * F, method, expcache) - - chol_integrand(τ) = begin - E = exponential!((dt - τ) * F', method, expcache) - L'E - end - nodes, weights = gausslegendre(N) - b, a = dt, 0 - @. nodes = (b - a) / 2 * nodes + (a + b) / 2 - @. weights = (b - a) / 2 * weights - @simd ivdep for i in 1:N - R[(i-1)*d+1:i*d, 1:D] .= sqrt(weights[i]) .* chol_integrand(nodes[i]) - end - - M = R'R |> Symmetric - chol = cholesky!(M, check=false) - Qh_R = if issuccess(chol) - chol.U |> Matrix - else - qr!(R).R |> Matrix - end - - return Ah, Qh_R -end diff --git a/test/Project.toml b/test/Project.toml index 39a6ebb11..4705022bd 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -25,4 +25,3 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [compat] Aqua = "0.8.2" -DiffEqDevTools = "2.44.1" From c85602ad5b3527478f324a660862798901026255 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 19:04:09 +0100 Subject: [PATCH 84/99] Fix the bad getindex for BlockDiag --- src/blockdiagonals.jl | 35 ++++++++++++++++------------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index 0506a2bb7..ca81444fd 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -20,30 +20,27 @@ blocks(B::BlockDiag) = B.blocks nblocks(B::BlockDiag) = length(B.blocks) size(B::BlockDiag) = mapreduce(size, ((a, b), (c, d)) -> (a + c, b + d), blocks(B)) -function _block_indices(B::BlockDiag, i::Integer, j::Integer) - all((0, 0) .< (i, j) .<= size(B)) || throw(BoundsError(B, (i, j))) - # find the on-diagonal block `p` in column `j` - p = 0 - @inbounds while j > 0 - p += 1 - j -= size(blocks(B)[p], 2) - end - # isempty to avoid reducing over an empty collection - @views @inbounds i -= isempty(1:(p-1)) ? 0 : sum(size.(blocks(B)[1:(p-1)], 1)) - # if row `i` outside of block `p`, set `p` to place-holder value `-1` - if i <= 0 || i > size(blocks(B)[p], 2) - p = -1 - end - return p, i, j -end Base.@propagate_inbounds function Base.getindex( B::BlockDiag{T}, i::Integer, j::Integer, ) where {T} - p, i, j = _block_indices(B, i, j) - # if not in on-diagonal block `p` then value at `i, j` must be zero - @inbounds return p > 0 ? blocks(B)[p][i, end+j] : zero(T) + all((0, 0) .< (i, j) .<= size(B)) || throw(BoundsError(B, (i, j))) + + p = 1 + Si, Sj = size(blocks(B)[p]) + while p <= nblocks(B) + if i <= Si && j <= Sj + return blocks(B)[p][i, j] + elseif (i <= Si && j > Sj) || (j <= Sj && i > Si) + return zero(T) + else + i -= Si + j -= Sj + p += 1 + end + end + error("This shouldn't happen") end Base.view(::BlockDiag, idxs...) = From f3a04fd494ea48a88f19b9f2a5f2a0296d0457d5 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 19:21:15 +0100 Subject: [PATCH 85/99] Use the built-in matrix exponential --- src/priors/ltisde.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/priors/ltisde.jl b/src/priors/ltisde.jl index 468da1da8..059788ac0 100644 --- a/src/priors/ltisde.jl +++ b/src/priors/ltisde.jl @@ -60,7 +60,7 @@ function matrix_fraction_decomposition( ) d = size(drift, 1) M = [drift dispersion*dispersion'; zero(drift) -drift'] - Mexp = exponential!(dt * M) + Mexp = exp(dt * M) A = Mexp[1:d, 1:d] Q = Mexp[1:d, d+1:end] * A' return A, Q From 8b8a19a4b145db4b83eb8db1d15a90807b33345f Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sat, 17 Feb 2024 19:41:28 +0100 Subject: [PATCH 86/99] Remore parts of a test that are not implemented anymore --- test/core/priors.jl | 5 ----- 1 file changed, 5 deletions(-) diff --git a/test/core/priors.jl b/test/core/priors.jl index f5ed0faa2..688374876 100644 --- a/test/core/priors.jl +++ b/test/core/priors.jl @@ -31,11 +31,6 @@ h = 0.1 @test A1 ≈ A3 @test Matrix(Q1) ≈ Q3 - A4, Q4R = PNDE._discretize_sqrt_with_quadraturetrick( - PNDE.LTISDE(Matrix(sde.F), Matrix(sde.L)), h) - @test A1 ≈ A4 - @test Q1.R ≈ Q4R - ts = 0:0.1:1 marginals = @test_nowarn PNDE.marginalize(prior, ts) @test length(marginals) == length(ts) From ef76c67ca1f0adeef5424ed5471c0ac77bd13b7b Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 13:05:33 +0100 Subject: [PATCH 87/99] Improve coverage a bit --- src/algorithms.jl | 1 - src/fast_linalg.jl | 64 ++++++++++++++++++---------------------- test/core/fast_linalg.jl | 53 +++++++++++++++++++++++++++++++++ test/correctness.jl | 1 + test/errors_thrown.jl | 27 +++++++++++++++++ test/mass_matrix.jl | 2 ++ test/runtests.jl | 3 ++ 7 files changed, 115 insertions(+), 36 deletions(-) create mode 100644 test/core/fast_linalg.jl diff --git a/src/algorithms.jl b/src/algorithms.jl index 9a49481a4..babcf9ba0 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -55,7 +55,6 @@ function covariance_structure(::Type{Alg}, prior, diffusionmodel) where {Alg<:Ab throw(ArgumentError("Unknown algorithm type $Alg")) end end -covariance_structure(alg) = covariance_structure(typeof(alg), alg.prior, alg.diffusionmodel) """ EK0(; order=3, diff --git a/src/fast_linalg.jl b/src/fast_linalg.jl index 05ae8c25b..701274132 100644 --- a/src/fast_linalg.jl +++ b/src/fast_linalg.jl @@ -13,63 +13,57 @@ _matmul!(C, A, B) _matmul!(C, A, B) = mul!(C, A, B) _matmul!(C, A, B, a, b) = mul!(C, A, B, a, b) # Use Octavian.jl's matrix-matrix products whenever applicable +const MSR{T} = Union{Matrix{T},SubArray{T},Base.ReshapedArray{T},Adjoint{T,<:Matrix}} _matmul!( - C::AbstractMatrix{T}, - A::AbstractMatrix{T}, - B::AbstractMatrix{T}, + C::MSR{T}, + A::MSR{T}, + B::MSR{T}, alpha::Number, beta::Number, ) where {T<:LinearAlgebra.BlasFloat} = matmul!(C, A, B, alpha, beta) _matmul!( - C::AbstractMatrix{T}, - A::AbstractMatrix{T}, - B::AbstractMatrix{T}, + C::MSR{T}, + A::MSR{T}, + B::MSR{T}, ) where {T<:LinearAlgebra.BlasFloat} = matmul!(C, A, B) # Some exceptions where we'd rather broadcast with FastBroadcast.jl: # Matrix-scalar products _matmul!(C::AbstractVecOrMat, A::AbstractVecOrMat, b::Number) = @.. C = A * b _matmul!(C::AbstractVecOrMat, a::Number, B::AbstractVecOrMat) = @.. C = a * B # Matrix matrix products with diagonal matrices -const MSR{T} = Union{Matrix{T},SubArray{T},Base.ReshapedArray{T},Adjoint{T,<:Matrix}} _matmul!(C::MSR, A::MSR, B::Diagonal) = @.. C = A * B.diag' _matmul!(C::MSR, A::Diagonal, B::MSR) = (@.. C = A.diag * B) _matmul!(C::MSR, A::Diagonal, B::Diagonal) = (@.. C = A * B) -_matmul!(C::MSR{T}, A::MSR{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = - (@.. C = A * B.diag') -_matmul!(C::MSR{T}, A::Diagonal{T}, B::MSR{T}) where {T<:LinearAlgebra.BlasFloat} = - (@.. C = A.diag * B) -_matmul!(C::MSR{T}, A::Diagonal{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = - (@.. C = A * B) +# _matmul!(C::MSR{T}, A::MSR{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = +# (@.. C = A * B.diag') +# _matmul!(C::MSR{T}, A::Diagonal{T}, B::MSR{T}) where {T<:LinearAlgebra.BlasFloat} = +# (@.. C = A.diag * B) +# _matmul!(C::MSR{T}, A::Diagonal{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = +# (@.. C = A * B) _matmul!(C::MSR, A::MSR, B::Diagonal, alpha::Number, beta::Number) = @.. C = A * B.diag' * alpha + C * beta -_matmul!( - C::MSR{T}, - A::MSR{T}, - B::Diagonal{T}, - alpha::Number, - beta::Number, -) where {T<:LinearAlgebra.BlasFloat} = - @.. C = A * B.diag' * alpha + C * beta +# _matmul!( +# C::MSR{T}, +# A::MSR{T}, +# B::Diagonal{T}, +# alpha::Number, +# beta::Number, +# ) where {T<:LinearAlgebra.BlasFloat} = +# @.. C = A * B.diag' * alpha + C * beta _matmul!(C::MSR, A::Diagonal, B::MSR, alpha::Number, beta::Number) = (@.. C = A.diag * B * alpha + C * beta) -_matmul!( - C::MSR{T}, - A::Diagonal{T}, - B::MSR{T}, - alpha::Number, - beta::Number, -) where {T<:LinearAlgebra.BlasFloat} = - (@.. C = A.diag * B * alpha + C * beta) - -_matmul!( - C::Matrix{T}, - A::LowerTriangular{T}, - B::UpperTriangular{T}, -) where {T<:LinearAlgebra.BlasFloat} = mul!(C, A, B) +# _matmul!( +# C::MSR{T}, +# A::Diagonal{T}, +# B::MSR{T}, +# alpha::Number, +# beta::Number, +# ) where {T<:LinearAlgebra.BlasFloat} = +# (@.. C = A.diag * B * alpha + C * beta) """ getupperright!(A) diff --git a/test/core/fast_linalg.jl b/test/core/fast_linalg.jl new file mode 100644 index 000000000..988ce57aa --- /dev/null +++ b/test/core/fast_linalg.jl @@ -0,0 +1,53 @@ +using ProbNumDiffEq +import ProbNumDiffEq: _matmul! +using LinearAlgebra +using Test + + +@testset "T=$T" for T in (Float64, BigFloat) + A = rand(T, 2, 3) + B = rand(T, 3, 4) + C = rand(T, 2, 4) + alpha = rand(T) + beta = rand(T) + + @test _matmul!(C, A, B) == mul!(C, A, B) + @test _matmul!(C, A, B, alpha, beta) == mul!(C, A, B, alpha, beta) + + _B = copy(B) + @test _matmul!(_B, alpha, B) == mul!(_B, alpha, B) + @test _matmul!(_B, B, beta) == mul!(_B, B, beta) + + # Diagonals + D = Diagonal(rand(T, size(B, 1))) + @test _matmul!(_B, D, B) == mul!(_B, D, B) + @test _matmul!(_B, D, B, alpha, beta) == mul!(_B, D, B, alpha, beta) + D = Diagonal(rand(T, size(B, 2))) + @test _matmul!(_B, B, D) == mul!(_B, B, D) + @test _matmul!(_B, B, D, alpha, beta) == mul!(_B, B, D, alpha, beta) + CD, D1, D2 = rand(T, 3, 3), Diagonal(rand(T, 3)), Diagonal(rand(T, 3)) + @test _matmul!(CD, D1, D2) == _matmul!(CD, D1, D2) + + # Triangulars + ASQ, BSQ, CSQ = rand(T, 2, 2), rand(T, 2,2), rand(T,2,2) + ALT, AUT = LowerTriangular(ASQ), UpperTriangular(ASQ) + BLT, BUT = LowerTriangular(BSQ), UpperTriangular(BSQ) + @test _matmul!(CSQ, ALT, BSQ) == mul!(CSQ, ALT, BSQ) + @test _matmul!(CSQ, AUT, BSQ) == mul!(CSQ, AUT, BSQ) + @test _matmul!(CSQ, ASQ, BLT) == mul!(CSQ, ASQ, BLT) + @test _matmul!(CSQ, ASQ, BUT) == mul!(CSQ, ASQ, BUT) + @test _matmul!(CSQ, ALT, BUT) == mul!(CSQ, ALT, BUT) + @test _matmul!(CSQ, AUT, BLT) == mul!(CSQ, AUT, BLT) + @test _matmul!(CSQ, ALT, BLT) == mul!(CSQ, ALT, BLT) + @test _matmul!(CSQ, AUT, BUT) == mul!(CSQ, AUT, BUT) + + # Adjoints + AT = Matrix(A')' + @test _matmul!(C, AT, B) == mul!(C, AT, B) + BT = Matrix(B')' + @test _matmul!(C, A, BT) == mul!(C, A, BT) + + # Vectors + CV, BV = rand(T, 2), rand(T, 3) + @test _matmul!(CV, A, BV) == mul!(CV, A, BV) +end diff --git a/test/correctness.jl b/test/correctness.jl index a3eacb705..d43ab1175 100644 --- a/test/correctness.jl +++ b/test/correctness.jl @@ -61,6 +61,7 @@ ADAPTIVE_ALGS = ( EK0(order=3, initialization=ClassicSolverInit()) => 5e-5, EK0(order=3, initialization=SimpleInit()) => 1e-4, EK0(order=3, diffusionmodel=DynamicMVDiffusion(), initialization=ClassicSolverInit()) => 4e-5, + EK0(order=3, diffusionmodel=FixedMVDiffusion()) => 1e-4, EK1(order=2) => 2e-5, EK1(order=3) => 1e-5, EK1(order=5) => 1e-6, diff --git a/test/errors_thrown.jl b/test/errors_thrown.jl index 88c323259..9aa22d32a 100644 --- a/test/errors_thrown.jl +++ b/test/errors_thrown.jl @@ -31,3 +31,30 @@ end prior = IOUP(num_derivatives=1, rate_parameter=3, update_rate_parameter=true) @test_throws ArgumentError solve(prob, EK0(; prior)) end + +@testset "Invalid prior" begin + prob = prob_ode_lotkavolterra + @test_throws DimensionMismatch solve(prob, EK0(prior=IWP(dim=3, num_derivatives=2))) + prior = IOUP(num_derivatives=1, rate_parameter=3, update_rate_parameter=true) + @test_throws ArgumentError solve(prob, EK0(; prior)) +end + +@testset "Invalid solver configurations" begin + prob = prob_ode_lotkavolterra + + # Global calibration + observation noise doesn't work + @test_throws ArgumentError solve( + prob, EK0(pn_observation_noise=1, diffusionmodel=FixedDiffusion())) + @test_throws ArgumentError solve( + prob, EK0(pn_observation_noise=1, diffusionmodel=FixedMVDiffusion())) + + # EK1 + Multivariate diffusion doesn't work: + @test_throws ArgumentError solve( + prob, EK1(diffusionmodel=FixedMVDiffusion())) + @test_throws ArgumentError solve( + prob, EK1(diffusionmodel=DynamicMVDiffusion())) + + # Multivariate diffusion with non-diagonal diffusion model + @test_throws ArgumentError solve( + prob, EK0(diffusionmodel=FixedMVDiffusion(initial_diffusion=rand(2,2)))) +end diff --git a/test/mass_matrix.jl b/test/mass_matrix.jl index fad2f4064..76786da85 100644 --- a/test/mass_matrix.jl +++ b/test/mass_matrix.jl @@ -93,4 +93,6 @@ end sol = solve(prob, DiagonalEK1(order=3)) @test sol.u[end] ≈ ref.u[end] rtol = 1e-8 + + @test_throws ArgumentError solve(prob, EK0()) end diff --git a/test/runtests.jl b/test/runtests.jl index a36528f79..e980f2726 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -22,6 +22,9 @@ const GROUP = get(ENV, "GROUP", "All") @timedsafetestset "BlockDiagonals" begin include("core/blockdiagonals.jl") end + @timedsafetestset "FastLinalg (`_matmul!`)" begin + include("core/fast_linalg.jl") + end @timedsafetestset "Filtering" begin include("core/filtering.jl") end From 5f218dec635c0e3633558d31985c6555415bf345 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 13:12:27 +0100 Subject: [PATCH 88/99] JuliaFormatter.jl --- test/core/fast_linalg.jl | 3 +-- test/errors_thrown.jl | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/test/core/fast_linalg.jl b/test/core/fast_linalg.jl index 988ce57aa..2363f553f 100644 --- a/test/core/fast_linalg.jl +++ b/test/core/fast_linalg.jl @@ -3,7 +3,6 @@ import ProbNumDiffEq: _matmul! using LinearAlgebra using Test - @testset "T=$T" for T in (Float64, BigFloat) A = rand(T, 2, 3) B = rand(T, 3, 4) @@ -29,7 +28,7 @@ using Test @test _matmul!(CD, D1, D2) == _matmul!(CD, D1, D2) # Triangulars - ASQ, BSQ, CSQ = rand(T, 2, 2), rand(T, 2,2), rand(T,2,2) + ASQ, BSQ, CSQ = rand(T, 2, 2), rand(T, 2, 2), rand(T, 2, 2) ALT, AUT = LowerTriangular(ASQ), UpperTriangular(ASQ) BLT, BUT = LowerTriangular(BSQ), UpperTriangular(BSQ) @test _matmul!(CSQ, ALT, BSQ) == mul!(CSQ, ALT, BSQ) diff --git a/test/errors_thrown.jl b/test/errors_thrown.jl index 9aa22d32a..562f2074f 100644 --- a/test/errors_thrown.jl +++ b/test/errors_thrown.jl @@ -56,5 +56,5 @@ end # Multivariate diffusion with non-diagonal diffusion model @test_throws ArgumentError solve( - prob, EK0(diffusionmodel=FixedMVDiffusion(initial_diffusion=rand(2,2)))) + prob, EK0(diffusionmodel=FixedMVDiffusion(initial_diffusion=rand(2, 2)))) end From 8d7fd1f9fb52dfd5dfef32f1eee01639fc3b3a0f Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 16:15:05 +0100 Subject: [PATCH 89/99] Check better what observation noise works with what factorization --- src/ProbNumDiffEq.jl | 8 ++++++++ src/algorithms.jl | 20 ++++++++++++++++++++ src/blockdiagonals.jl | 18 ++++++++++++++++++ src/kronecker.jl | 6 ++++++ src/perform_step.jl | 2 +- test/core/blockdiagonals.jl | 11 ++++++++++- test/observation_noise.jl | 28 ++++++++++++++++++++++++++++ test/runtests.jl | 3 +++ 8 files changed, 94 insertions(+), 2 deletions(-) create mode 100644 test/observation_noise.jl diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index d6909ffde..1d3756455 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -52,6 +52,14 @@ cov2psdmatrix(cov::AbstractMatrix; d) = (@assert size(cov, 1) == size(cov, 2) == d; PSDMatrix(Matrix(cholesky(cov).U))) cov2psdmatrix(cov::PSDMatrix; d) = (@assert size(cov, 1) == size(cov, 2) == d; cov) +""" + add!(out, toadd) + +Add `toadd` to `out` in-place. +""" +add! +add!(out, toadd) = (out .+= toadd) + include("fast_linalg.jl") include("kronecker.jl") include("blockdiagonals.jl") diff --git a/src/algorithms.jl b/src/algorithms.jl index babcf9ba0..7bedc5eba 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -33,6 +33,26 @@ function ekargcheck( ) end end + if covariance_factorization == IsometricKroneckerCovariance && !( + pn_observation_noise isa Number + || pn_observation_noise isa UniformScaling + || pn_observation_noise isa Diagonal{<:Number,<:FillArrays.Fill}) + throw( + ArgumentError( + "The supplied `pn_observation_noise` is not compatible with the chosen `IsometricKroneckerCovariance` factorization. Try one of `BlockDiagonalCovariance` or `DenseCovariance` instead!", + ), + ) + end + if covariance_factorization == BlockDiagonalCovariance && !( + pn_observation_noise isa Number + || pn_observation_noise isa UniformScaling + || pn_observation_noise isa Diagonal) + throw( + ArgumentError( + "The supplied `pn_observation_noise` is not compatible with the chosen `BlockDiagonalCovariance` factorization. Try `DenseCovariance` instead!", + ), + ) + end end function covariance_structure(::Type{Alg}, prior, diffusionmodel) where {Alg<:AbstractEK} diff --git a/src/blockdiagonals.jl b/src/blockdiagonals.jl index ca81444fd..6ec085a01 100644 --- a/src/blockdiagonals.jl +++ b/src/blockdiagonals.jl @@ -57,6 +57,24 @@ end similar(B::BlockDiag) = BlockDiag(similar.(blocks(B))) zero(B::BlockDiag) = BlockDiag(zero.(blocks(B))) +# Sums of BlockDiag +Base.:+(A::BlockDiag, B::BlockDiag) = begin + @assert nblocks(A) == nblocks(B) + return BlockDiag([Ai + Bi for (Ai, Bi) in zip(blocks(A), blocks(B))]) +end +Base.:-(A::BlockDiag, B::BlockDiag) = begin + @assert nblocks(A) == nblocks(B) + return BlockDiag([Ai - Bi for (Ai, Bi) in zip(blocks(A), blocks(B))]) +end + +add!(out::BlockDiag, toadd::BlockDiag) = begin + @assert nblocks(out) == nblocks(toadd) + @simd ivdep for i in eachindex(blocks(out)) + add!(blocks(out)[i], blocks(toadd)[i]) + end + return out +end + # Mul with Scalar or UniformScaling Base.:*(a::Number, M::BlockDiag) = BlockDiag([a * B for B in blocks(M)]) Base.:*(M::BlockDiag, a::Number) = BlockDiag([B * a for B in blocks(M)]) diff --git a/src/kronecker.jl b/src/kronecker.jl index d05096703..16a803aae 100644 --- a/src/kronecker.jl +++ b/src/kronecker.jl @@ -103,6 +103,12 @@ Base.:+(A::IKP, B::IKP) = begin end Base.:+(U::UniformScaling, K::IKP) = IsometricKroneckerProduct(K.ldim, U + K.B) Base.:+(K::IKP, U::UniformScaling) = IsometricKroneckerProduct(K.ldim, U + K.B) + +add!(out::IsometricKroneckerProduct, toadd::IsometricKroneckerProduct) = begin + @assert out.ldim == toadd.ldim + add!(out.B, toadd.B) +end + Base.:-(U::UniformScaling, K::IKP) = IsometricKroneckerProduct(K.ldim, U - K.B) LinearAlgebra.inv(K::IKP) = IsometricKroneckerProduct(K.ldim, inv(K.B)) Base.:/(A::IKP, B::IKP) = begin diff --git a/src/perform_step.jl b/src/perform_step.jl index 4469c978b..d1c834284 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -163,7 +163,7 @@ compute_measurement_covariance!(cache) = begin _matmul!(cache.C_Dxd, cache.x_pred.Σ.R, cache.H') _matmul!(cache.measurement.Σ, cache.C_Dxd', cache.C_Dxd) if !isnothing(cache.R) - cache.measurement.Σ .+= _matmul!(cache.C_dxd, cache.R.R', cache.R.R) + add!(cache.measurement.Σ, _matmul!(cache.C_dxd, cache.R.R', cache.R.R)) end end diff --git a/test/core/blockdiagonals.jl b/test/core/blockdiagonals.jl index 3f9e8917b..a9528d712 100644 --- a/test/core/blockdiagonals.jl +++ b/test/core/blockdiagonals.jl @@ -1,4 +1,5 @@ using ProbNumDiffEq +import ProbNumDiffEq as PNDE import ProbNumDiffEq: BlockDiag, _matmul! using LinearAlgebra using BlockDiagonals @@ -65,5 +66,13 @@ D = d1 * d2 @test_throws ErrorException view(A, 1:2, 1:2) - tttm(copy!(A, Diagonal(A))) + tttm(copy!(copy(A), Diagonal(A))) + + @test tttm(A + A) ≈ AM + AM + @test tttm(A - A) ≈ AM - AM + + @test tttm(_matmul!(_A, a * I(D), A)) ≈ a * AM + + _A = copy(A) + @test tttm(PNDE.add!(_A, A)) == AM + AM end diff --git a/test/observation_noise.jl b/test/observation_noise.jl new file mode 100644 index 000000000..977fc8ac1 --- /dev/null +++ b/test/observation_noise.jl @@ -0,0 +1,28 @@ +using Test +using ProbNumDiffEq +using LinearAlgebra, FillArrays +import ODEProblemLibrary: prob_ode_lotkavolterra + +prob = prob_ode_lotkavolterra +d = length(prob.u0) +@testset "typeof(R)=$(typeof(R))" for (i, R) in enumerate(( + 0.1, + 0.1I, + 0.1Eye(d), + 0.1I(d), + Diagonal([0.1, 0.2]), + [0.1 0.01; 0.01 0.1], + PSDMatrix(0.1 * rand(d, d)), +)) + if i <= 3 + @test_nowarn solve(prob, EK0(pn_observation_noise=R)) + else + @test_broken solve(prob, EK0(pn_observation_noise=R)) + end + if i <= 5 + @test_nowarn solve(prob, DiagonalEK1(pn_observation_noise=R)) + else + @test_broken solve(prob, DiagonalEK1(pn_observation_noise=R)) + end + @test_nowarn solve(prob, EK1(pn_observation_noise=R)) +end diff --git a/test/runtests.jl b/test/runtests.jl index e980f2726..b39c91e9a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -129,6 +129,9 @@ const GROUP = get(ENV, "GROUP", "All") @timedsafetestset "Data Likelihoods" begin include("data_likelihoods.jl") end + @timedsafetestset "Observation noise" begin + include("observation_noise.jl") + end end end From 3ab45200979f9dac42fe2ad727d6bc7d2f1a7469 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 16:46:27 +0100 Subject: [PATCH 90/99] Fix the pn_observation_noise check --- src/algorithms.jl | 40 +++++++++++++++++++++------------------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/src/algorithms.jl b/src/algorithms.jl index 7bedc5eba..ba535e334 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -33,25 +33,27 @@ function ekargcheck( ) end end - if covariance_factorization == IsometricKroneckerCovariance && !( - pn_observation_noise isa Number - || pn_observation_noise isa UniformScaling - || pn_observation_noise isa Diagonal{<:Number,<:FillArrays.Fill}) - throw( - ArgumentError( - "The supplied `pn_observation_noise` is not compatible with the chosen `IsometricKroneckerCovariance` factorization. Try one of `BlockDiagonalCovariance` or `DenseCovariance` instead!", - ), - ) - end - if covariance_factorization == BlockDiagonalCovariance && !( - pn_observation_noise isa Number - || pn_observation_noise isa UniformScaling - || pn_observation_noise isa Diagonal) - throw( - ArgumentError( - "The supplied `pn_observation_noise` is not compatible with the chosen `BlockDiagonalCovariance` factorization. Try `DenseCovariance` instead!", - ), - ) + if !(isnothing(pn_observation_noise) || ismissing(pn_observation_noise)) + if covariance_factorization == IsometricKroneckerCovariance && !( + pn_observation_noise isa Number + || pn_observation_noise isa UniformScaling + || pn_observation_noise isa Diagonal{<:Number,<:FillArrays.Fill}) + throw( + ArgumentError( + "The supplied `pn_observation_noise` is not compatible with the chosen `IsometricKroneckerCovariance` factorization. Try one of `BlockDiagonalCovariance` or `DenseCovariance` instead!", + ), + ) + end + if covariance_factorization == BlockDiagonalCovariance && !( + pn_observation_noise isa Number + || pn_observation_noise isa UniformScaling + || pn_observation_noise isa Diagonal) + throw( + ArgumentError( + "The supplied `pn_observation_noise` is not compatible with the chosen `BlockDiagonalCovariance` factorization. Try `DenseCovariance` instead!", + ), + ) + end end end From a9f24d736eaa2cb4b9b68ceacb468db297089d9f Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 21:39:59 +0100 Subject: [PATCH 91/99] More BlockDiag tests --- test/core/blockdiagonals.jl | 36 +++++++++++++++++++++--------------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/test/core/blockdiagonals.jl b/test/core/blockdiagonals.jl index a9528d712..065ebff0f 100644 --- a/test/core/blockdiagonals.jl +++ b/test/core/blockdiagonals.jl @@ -11,6 +11,8 @@ D = d1 * d2 A = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) B = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) C = BlockDiag([randn(T, d1, d1) for _ in 1:d2]) + alpha = rand(T) + beta = rand(T) AM, BM, CM = @test_nowarn Matrix.((A, B, C)) @@ -46,23 +48,29 @@ D = d1 * d2 @test tttm(_matmul!(C, A', B)) ≈ _matmul!(CM, AM', BM) @test tttm(_matmul!(C, A, B')) ≈ _matmul!(CM, AM, BM') + @test tttm(mul!(C, A, B, alpha, beta)) ≈ mul!(CM, AM, BM, alpha, beta) + @test tttm(_matmul!(C, A, B, alpha, beta)) ≈ _matmul!(CM, AM, BM, alpha, beta) + @test tttm(A * B) ≈ AM * BM @test tttm(A' * B) ≈ AM' * BM @test tttm(A * B') ≈ AM * BM' - a = rand() - @test tttm(A * a) ≈ AM * a - @test tttm(a * A) ≈ a * AM - @test tttm(A * (a * I)) ≈ AM * a - @test tttm((a * I) * A) ≈ a * AM - @test tttm(rmul!(copy(A), a)) ≈ a * AM - - @test tttm((a * I(D)) * A) ≈ a * AM - @test tttm(A * (a * I(D))) ≈ AM * a - @test tttm(mul!(_A, A, a * I(D))) ≈ a * AM - @test tttm(mul!(_A, a * I(D), A)) ≈ a * AM - @test tttm(_matmul!(_A, A, a * I(D))) ≈ a * AM - @test tttm(_matmul!(_A, a * I(D), A)) ≈ a * AM + @test tttm(A * alpha) ≈ AM * alpha + @test tttm(alpha * A) ≈ alpha * AM + @test tttm(A * (alpha * I)) ≈ AM * alpha + @test tttm((alpha * I) * A) ≈ alpha * AM + @test tttm(rmul!(copy(A), alpha)) ≈ alpha * AM + @test tttm(mul!(_A, alpha, A)) ≈ alpha * AM + @test tttm(mul!(_A, A, alpha)) ≈ alpha * AM + @test tttm(_matmul!(_A, alpha, A)) ≈ alpha * AM + @test tttm(_matmul!(_A, A, alpha)) ≈ alpha * AM + + @test tttm((alpha * I(D)) * A) ≈ alpha * AM + @test tttm(A * (alpha * I(D))) ≈ AM * alpha + @test tttm(mul!(_A, A, alpha * I(D))) ≈ alpha * AM + @test tttm(mul!(_A, alpha * I(D), A)) ≈ alpha * AM + @test tttm(_matmul!(_A, A, alpha * I(D))) ≈ alpha * AM + @test tttm(_matmul!(_A, alpha * I(D), A)) ≈ alpha * AM @test_throws ErrorException view(A, 1:2, 1:2) @@ -71,8 +79,6 @@ D = d1 * d2 @test tttm(A + A) ≈ AM + AM @test tttm(A - A) ≈ AM - AM - @test tttm(_matmul!(_A, a * I(D), A)) ≈ a * AM - _A = copy(A) @test tttm(PNDE.add!(_A, A)) == AM + AM end From a2d527ac4eeec7732e7806769492f4be46ef8374 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 21:55:44 +0100 Subject: [PATCH 92/99] Add some more Kronecker tests --- src/kronecker.jl | 10 ++++++++++ test/core/kronecker.jl | 6 ++++++ 2 files changed, 16 insertions(+) diff --git a/src/kronecker.jl b/src/kronecker.jl index 16a803aae..ef1a01317 100644 --- a/src/kronecker.jl +++ b/src/kronecker.jl @@ -159,6 +159,16 @@ _matmul!( return A end +mul!(A::IKP, b::Number, C::IKP) = begin + check_matmul_sizes(A, C) + mul!(A.B, b, C.B) + return A +end +mul!(A::IKP, B::IKP, c::Number) = begin + check_matmul_sizes(A, B) + mul!(A.B, B.B, c) + return A +end _matmul!(A::IKP, b::Number, C::IKP) = begin check_matmul_sizes(A, C) _matmul!(A.B, b, C.B) diff --git a/test/core/kronecker.jl b/test/core/kronecker.jl index 982654da7..f7f52123e 100644 --- a/test/core/kronecker.jl +++ b/test/core/kronecker.jl @@ -76,6 +76,12 @@ q = 2 @test α * K1 isa PNDE.IsometricKroneckerProduct @test K1 * α ≈ α * M1 @test K1 * α isa PNDE.IsometricKroneckerProduct + _K1 = copy(K1) + @test mul!(_K1, α, K1) == α * K1 + @test mul!(_K1, K1, α) == α * K1 + @test _matmul!(_K1, K1, α) == α * K1 + @test _matmul!(_K1, α, K1) == α * K1 + # In-place Matrix-Matrix Multiplication β = -0.5 From 54c62ec647d8bdf5a5fa6aae02a223c29c962068 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 22:00:02 +0100 Subject: [PATCH 93/99] JuliaFormatter.jl --- test/core/kronecker.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/core/kronecker.jl b/test/core/kronecker.jl index f7f52123e..478c1ea6f 100644 --- a/test/core/kronecker.jl +++ b/test/core/kronecker.jl @@ -82,7 +82,6 @@ q = 2 @test _matmul!(_K1, K1, α) == α * K1 @test _matmul!(_K1, α, K1) == α * K1 - # In-place Matrix-Matrix Multiplication β = -0.5 @test mul!(K3, K1, K2) ≈ mul!(M3, M1, M2) From 7d7bf1abbe02776bebdea0455dafa09e83824e61 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Sun, 18 Feb 2024 22:58:59 +0100 Subject: [PATCH 94/99] Fix test --- test/core/kronecker.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/core/kronecker.jl b/test/core/kronecker.jl index 478c1ea6f..a10b15f8b 100644 --- a/test/core/kronecker.jl +++ b/test/core/kronecker.jl @@ -79,8 +79,8 @@ q = 2 _K1 = copy(K1) @test mul!(_K1, α, K1) == α * K1 @test mul!(_K1, K1, α) == α * K1 - @test _matmul!(_K1, K1, α) == α * K1 - @test _matmul!(_K1, α, K1) == α * K1 + @test PNDE._matmul!(_K1, K1, α) == α * K1 + @test PNDE._matmul!(_K1, α, K1) == α * K1 # In-place Matrix-Matrix Multiplication β = -0.5 From 569b9e34c52caf27a1302db2c7d29d397333d66c Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Mon, 19 Feb 2024 10:45:44 +0100 Subject: [PATCH 95/99] Make the FixedMVDiffusion work with dense matrices --- src/diffusions/apply_diffusion.jl | 19 +++++++++++++++--- test/core/diffusions.jl | 32 +++++++++++++++---------------- 2 files changed, 31 insertions(+), 20 deletions(-) diff --git a/src/diffusions/apply_diffusion.jl b/src/diffusions/apply_diffusion.jl index 397ac5160..9638ba587 100644 --- a/src/diffusions/apply_diffusion.jl +++ b/src/diffusions/apply_diffusion.jl @@ -22,7 +22,7 @@ apply_diffusion( ) where {T} = begin d = size(diffusion, 1) q = size(Q, 1) ÷ d - 1 - return PSDMatrix(Q.R * sqrt.(kron(diffusion, I(q + 1)))) + return PSDMatrix(Q.R * sqrt.(Kronecker.kronecker(diffusion, Eye(q + 1)))) end """ @@ -46,6 +46,18 @@ apply_diffusion!( end return Q end +apply_diffusion!( + Q::PSDMatrix, + diffusion::Diagonal, +) = begin + # @warn "This is not yet implemented efficiently; TODO" + d = size(diffusion, 1) + D = size(Q, 1) + q = D ÷ d - 1 + # _matmul!(Q.R, Q.R, Kronecker.kronecker(sqrt.(diffusion), Eye(q + 1))) + _matmul!(Q.R, Q.R, kron(sqrt.(diffusion), Eye(q + 1))) + return Q +end """ apply_diffusion!(out::PSDMatrix, Q::PSDMatrix, diffusion::Union{Number, Diagonal}) -> PSDMatrix @@ -80,10 +92,11 @@ apply_diffusion!( Q::PSDMatrix, diffusion::Diagonal, ) = begin - @warn "This is not yet implemented efficiently; TODO" + # @warn "This is not yet implemented efficiently; TODO" d = size(diffusion, 1) D = size(Q, 1) q = D ÷ d - 1 - _matmul!(out.R, Q.R, sqrt.(kron(diffusion, Eye(q + 1)))) + # _matmul!(out.R, Q.R, Kronecker.kronecker(sqrt.(diffusion), Eye(q + 1))) + _matmul!(out.R, Q.R, kron(sqrt.(diffusion), Eye(q + 1))) return out end diff --git a/test/core/diffusions.jl b/test/core/diffusions.jl index 7c5658160..2d7fb59de 100644 --- a/test/core/diffusions.jl +++ b/test/core/diffusions.jl @@ -13,44 +13,42 @@ T = Float64 FixedDiffusion(), FixedDiffusion(calibrate=false), FixedMVDiffusion(), - FixedMVDiffusion(; initial_diffusion=rand(2)), - FixedMVDiffusion(; initial_diffusion=Diagonal(rand(2))), - FixedMVDiffusion(; initial_diffusion=Diagonal(rand(2)), calibrate=false), + FixedMVDiffusion(; initial_diffusion=rand(d)), + FixedMVDiffusion(; initial_diffusion=Diagonal(rand(d))), + FixedMVDiffusion(; initial_diffusion=Diagonal(rand(d)), calibrate=false), ) # Test the initial diffusion - diff = PNDE.initial_diffusion(diffusionmodel, d, q, T) - @assert size(diff) == (d, d) - @assert diff isa Diagonal + diffusion = PNDE.initial_diffusion(diffusionmodel, d, q, T) + @assert size(diffusion) == (d, d) + @assert diffusion isa Diagonal if !(diffusionmodel isa FixedMVDiffusion || diffusionmodel isa DynamicMVDiffusion) - @assert diff isa Diagonal{T,<:Fill} + @assert diffusion isa Diagonal{T,<:Fill} end # Test applying the diffusion _, Q = PNDE.discretize(PNDE.IWP{T}(d, q), 0.1) Qmat = PSDMatrix(Matrix(Q.R)) - _diff = rand() * diff + _diffusion = rand() * diffusion @testset "$FAC" for FAC in ( PNDE.DenseCovariance{T}(d, q), PNDE.BlockDiagonalCovariance{T}(d, q), PNDE.IsometricKroneckerCovariance{T}(d, q), ) - if diff isa Diagonal{T,<:Vector} && FAC isa PNDE.IsometricKroneckerCovariance + if diffusion isa Diagonal{T,<:Vector} && FAC isa PNDE.IsometricKroneckerCovariance continue end _Q = PNDE.to_factorized_matrix(FAC, Q) - Qdiff = @test_nowarn PNDE.apply_diffusion(_Q, _diff) - Qmatdiff = @test_nowarn PNDE.apply_diffusion(Qmat, _diff) + Qdiff = @test_nowarn PNDE.apply_diffusion(_Q, _diffusion) + Qmatdiff = @test_nowarn PNDE.apply_diffusion(Qmat, _diffusion) @test Qdiff == Qmatdiff - if !(diff isa Diagonal{T,<:Vector} && FAC isa PNDE.DenseCovariance) - Qdiff = @test_nowarn PNDE.apply_diffusion!(copy(_Q), _diff) - @test Qdiff == Qmatdiff + Qdiff = @test_nowarn PNDE.apply_diffusion!(copy(_Q), _diffusion) + @test Qdiff == Qmatdiff - Qdiff = @test_nowarn PNDE.apply_diffusion!(copy(_Q), _Q, _diff) - @test Qdiff == Qmatdiff - end + Qdiff = @test_nowarn PNDE.apply_diffusion!(copy(_Q), _Q, _diffusion) + @test Qdiff == Qmatdiff end @testset "Calibration" begin From e67c66569d710e083c74a0e34104d98ea3b6ec0e Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Mon, 19 Feb 2024 10:45:58 +0100 Subject: [PATCH 96/99] Make the FixedMVDiffusion work with the data likelihoods --- test/data_likelihoods.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/data_likelihoods.jl b/test/data_likelihoods.jl index cf428109f..9fe370616 100644 --- a/test/data_likelihoods.jl +++ b/test/data_likelihoods.jl @@ -36,12 +36,14 @@ kwargs = ( # EK0 EK0(), EK0(diffusionmodel=FixedDiffusion()), + EK0(diffusionmodel=FixedMVDiffusion(rand(2), false)), + EK0(diffusionmodel=DynamicMVDiffusion()), EK0(prior=IOUP(3, -1)), EK0(prior=Matern(3, 1.5)), # EK1 EK1(), EK1(diffusionmodel=FixedDiffusion()), - # EK1(diffusionmodel=FixedMVDiffusion(rand(2), false)), # not yet supported + EK1(diffusionmodel=FixedMVDiffusion(rand(2), false)), EK1(prior=IOUP(3, -1)), EK1(prior=Matern(3, 1.5)), EK1(prior=IOUP(3, update_rate_parameter=true)), From d9b4a363b882ecc8e370c59fb79662a71801065a Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Mon, 19 Feb 2024 10:47:15 +0100 Subject: [PATCH 97/99] Remove some comments --- src/fast_linalg.jl | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/src/fast_linalg.jl b/src/fast_linalg.jl index 701274132..cb3a549e1 100644 --- a/src/fast_linalg.jl +++ b/src/fast_linalg.jl @@ -37,33 +37,10 @@ _matmul!(C::MSR, A::Diagonal, B::MSR) = (@.. C = A.diag * B) _matmul!(C::MSR, A::Diagonal, B::Diagonal) = (@.. C = A * B) -# _matmul!(C::MSR{T}, A::MSR{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = -# (@.. C = A * B.diag') -# _matmul!(C::MSR{T}, A::Diagonal{T}, B::MSR{T}) where {T<:LinearAlgebra.BlasFloat} = -# (@.. C = A.diag * B) -# _matmul!(C::MSR{T}, A::Diagonal{T}, B::Diagonal{T}) where {T<:LinearAlgebra.BlasFloat} = -# (@.. C = A * B) - _matmul!(C::MSR, A::MSR, B::Diagonal, alpha::Number, beta::Number) = @.. C = A * B.diag' * alpha + C * beta -# _matmul!( -# C::MSR{T}, -# A::MSR{T}, -# B::Diagonal{T}, -# alpha::Number, -# beta::Number, -# ) where {T<:LinearAlgebra.BlasFloat} = -# @.. C = A * B.diag' * alpha + C * beta _matmul!(C::MSR, A::Diagonal, B::MSR, alpha::Number, beta::Number) = (@.. C = A.diag * B * alpha + C * beta) -# _matmul!( -# C::MSR{T}, -# A::Diagonal{T}, -# B::MSR{T}, -# alpha::Number, -# beta::Number, -# ) where {T<:LinearAlgebra.BlasFloat} = -# (@.. C = A.diag * B * alpha + C * beta) """ getupperright!(A) From 71e4be472cf7e799f529a43ae4170c0f32db56c0 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Mon, 19 Feb 2024 10:51:50 +0100 Subject: [PATCH 98/99] Add `add!` to the Kronecker tests and shorten them a bit --- test/core/kronecker.jl | 45 +++++++++++++++++++----------------------- 1 file changed, 20 insertions(+), 25 deletions(-) diff --git a/test/core/kronecker.jl b/test/core/kronecker.jl index a10b15f8b..42938087d 100644 --- a/test/core/kronecker.jl +++ b/test/core/kronecker.jl @@ -14,11 +14,16 @@ q = 2 K2 = PNDE.IsometricKroneckerProduct(d, R2) M2 = Matrix(K2) + function tttm(M) # quick type test and to matrix + @test M isa PNDE.IsometricKroneckerProduct + return Matrix(M) + end + # Matrix-Matrix Operations - @test K1 * K2 ≈ M1 * M2 - @test K1 * K2 isa PNDE.IsometricKroneckerProduct - @test K1 + K2 ≈ M1 + M2 - @test K1 + K2 isa PNDE.IsometricKroneckerProduct + @test tttm(K1 * K2) ≈ M1 * M2 + @test tttm(K1 + K2) ≈ M1 + M2 + + @test tttm(PNDE.add!(copy(K1), K2)) ≈ M1 + M2 # DimensionMismatch X = PNDE.IsometricKroneckerProduct(d, rand(T, 1, 1)) @@ -31,29 +36,21 @@ q = 2 R4 = rand(T, q + 1) K4 = PNDE.IsometricKroneckerProduct(d, R4) M4 = Matrix(K4) - @test K1 * K4 ≈ M1 * M4 + @test tttm(K1 * K4) ≈ M1 * M4 @test_throws DimensionMismatch K1 + K4 # UniformScaling - @test I + K1 ≈ I + M1 - @test I + K1 isa PNDE.IsometricKroneckerProduct - @test K1 + I ≈ M1 + I - @test K1 + I isa PNDE.IsometricKroneckerProduct - @test I - K1 ≈ I - M1 - @test I - K1 isa PNDE.IsometricKroneckerProduct - @test K1 - I ≈ M1 - I - @test K1 - I isa PNDE.IsometricKroneckerProduct + @test tttm(I + K1) ≈ I + M1 + @test tttm(K1 + I) ≈ M1 + I + @test tttm(I - K1) ≈ I - M1 + @test tttm(K1 - I) ≈ M1 - I # Other LinearAlgebra - @test K1' ≈ M1' - @test K1' isa PNDE.IsometricKroneckerProduct - @test inv(K1) ≈ inv(M1) - @test inv(K1) isa PNDE.IsometricKroneckerProduct + @test tttm(K1') ≈ M1' + @test tttm(inv(K1)) ≈ inv(M1) @test det(K1) ≈ det(M1) - @test K1 / K2 ≈ M1 / M2 - @test K1 / K2 isa PNDE.IsometricKroneckerProduct - @test K1 \ K2 ≈ M1 \ M2 - @test K1 \ K2 isa PNDE.IsometricKroneckerProduct + @test tttm(K1 / K2) ≈ M1 / M2 + @test tttm(K1 \ K2) ≈ M1 \ M2 # Base K3 = PNDE.IsometricKroneckerProduct(d, copy(R2)) @@ -72,10 +69,8 @@ q = 2 # Matrix-Scalar α = 2.0 - @test α * K1 ≈ α * M1 - @test α * K1 isa PNDE.IsometricKroneckerProduct - @test K1 * α ≈ α * M1 - @test K1 * α isa PNDE.IsometricKroneckerProduct + @test tttm(α * K1) ≈ α * M1 + @test tttm(K1 * α) ≈ α * M1 _K1 = copy(K1) @test mul!(_K1, α, K1) == α * K1 @test mul!(_K1, K1, α) == α * K1 From 5f753250fd37e04097cd21ff4129ca42f9cfa050 Mon Sep 17 00:00:00 2001 From: Nathanael Bosch Date: Mon, 19 Feb 2024 11:01:01 +0100 Subject: [PATCH 99/99] Fix the failing test that I just found --- src/kronecker.jl | 1 + test/core/blockdiagonals.jl | 1 + test/core/kronecker.jl | 22 ++++++++++++---------- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/kronecker.jl b/src/kronecker.jl index ef1a01317..f18bffd74 100644 --- a/src/kronecker.jl +++ b/src/kronecker.jl @@ -107,6 +107,7 @@ Base.:+(K::IKP, U::UniformScaling) = IsometricKroneckerProduct(K.ldim, U + K.B) add!(out::IsometricKroneckerProduct, toadd::IsometricKroneckerProduct) = begin @assert out.ldim == toadd.ldim add!(out.B, toadd.B) + return out end Base.:-(U::UniformScaling, K::IKP) = IsometricKroneckerProduct(K.ldim, U - K.B) diff --git a/test/core/blockdiagonals.jl b/test/core/blockdiagonals.jl index 065ebff0f..c0969571b 100644 --- a/test/core/blockdiagonals.jl +++ b/test/core/blockdiagonals.jl @@ -81,4 +81,5 @@ D = d1 * d2 _A = copy(A) @test tttm(PNDE.add!(_A, A)) == AM + AM + @test Matrix(_A) == AM + AM end diff --git a/test/core/kronecker.jl b/test/core/kronecker.jl index 42938087d..5e8ad085c 100644 --- a/test/core/kronecker.jl +++ b/test/core/kronecker.jl @@ -14,6 +14,15 @@ q = 2 K2 = PNDE.IsometricKroneckerProduct(d, R2) M2 = Matrix(K2) + # Base + K3 = PNDE.IsometricKroneckerProduct(d, copy(R2)) + M3 = Matrix(K3) + @test similar(K1) isa PNDE.IsometricKroneckerProduct + @test copy(K1) isa PNDE.IsometricKroneckerProduct + @test copy!(K3, K1) isa PNDE.IsometricKroneckerProduct + @test K3 == K1 + @test size(K1) == size(M1) + function tttm(M) # quick type test and to matrix @test M isa PNDE.IsometricKroneckerProduct return Matrix(M) @@ -23,7 +32,9 @@ q = 2 @test tttm(K1 * K2) ≈ M1 * M2 @test tttm(K1 + K2) ≈ M1 + M2 - @test tttm(PNDE.add!(copy(K1), K2)) ≈ M1 + M2 + _K1 = copy(K1) + @test tttm(PNDE.add!(_K1, K2)) ≈ M1 + M2 + @test _K1 ≈ M1 + M2 # DimensionMismatch X = PNDE.IsometricKroneckerProduct(d, rand(T, 1, 1)) @@ -52,15 +63,6 @@ q = 2 @test tttm(K1 / K2) ≈ M1 / M2 @test tttm(K1 \ K2) ≈ M1 \ M2 - # Base - K3 = PNDE.IsometricKroneckerProduct(d, copy(R2)) - M3 = Matrix(K3) - @test similar(K1) isa PNDE.IsometricKroneckerProduct - @test copy(K1) isa PNDE.IsometricKroneckerProduct - @test copy!(K3, K1) isa PNDE.IsometricKroneckerProduct - @test K3 == K1 - @test size(K1) == size(M1) - # Base @test one(K1) isa PNDE.IsometricKroneckerProduct @test isone(one(K1).B)