diff --git a/Project.toml b/Project.toml index ee0bb44ef..94bb3974d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ProbNumDiffEq" uuid = "bf3e78b0-7d74-48a5-b855-9609533b56a5" authors = ["Nathanael Bosch"] -version = "0.5.6" +version = "0.5.7" [deps] DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" @@ -9,7 +9,6 @@ DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GaussianDistributions = "43dcc890-d446-5863-8d1a-14597580bb8d" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" @@ -27,7 +26,6 @@ DiffEqBase = "6" DiffEqDevTools = "2" ForwardDiff = "0.10" GaussianDistributions = "0.5" -Octavian = "0.2,0.3" OrdinaryDiffEq = "5 - 5.69" RecipesBase = "1" RecursiveArrayTools = "2" diff --git a/src/ProbNumDiffEq.jl b/src/ProbNumDiffEq.jl index 6fa22b5e5..34b857fdd 100644 --- a/src/ProbNumDiffEq.jl +++ b/src/ProbNumDiffEq.jl @@ -35,26 +35,10 @@ using Tullio _matmul!(C, A, B) = mul!(C, A, B) _matmul!(C, A, B, a, b) = mul!(C, A, B, a, b) # Some special cases -_matmul!(C::AbstractMatrix{Float64}, A::AbstractMatrix{Float64}, B::Diagonal{Float64}) = - (C .= A .* B.diag') -_matmul!(C::AbstractMatrix{Float64}, A::Diagonal{Float64}, B::AbstractMatrix{Float64}) = - (C .= A.diag .* B) -_matmul!(C::Diagonal{Float64}, A::AbstractMatrix{Float64}, B::AbstractMatrix{Float64}) = +_matmul!(C::AbstractMatrix, A::AbstractMatrix, B::Diagonal) = (C .= A .* B.diag') +_matmul!(C::AbstractMatrix, A::Diagonal, B::AbstractMatrix) = (C .= A.diag .* B) +_matmul!(C::Diagonal, A::AbstractMatrix, B::AbstractMatrix) = @tullio C[i, i] = A[i, j] * B[j, i] -# Use Octavian's matmul! when the types are very specific -import Octavian: matmul! -_matmul!( - C::AbstractVecOrMat{Float64}, - A::AbstractVecOrMat{Float64}, - B::AbstractVecOrMat{Float64}, -) = matmul!(C, A, B) -_matmul!( - C::AbstractVecOrMat{Float64}, - A::AbstractVecOrMat{Float64}, - B::AbstractVecOrMat{Float64}, - a::Float64, - b::Float64, -) = matmul!(C, A, B, a, b) # @reexport using PSDMatrices # import PSDMatrices: X_A_Xt @@ -64,7 +48,8 @@ export SRMatrix X_A_Xt(A, X) = X * A * X' X_A_Xt!(out, A, X) = (out .= X * A * X') apply_diffusion(Q, diffusion::Diagonal) = X_A_Xt(Q, sqrt.(diffusion)) -apply_diffusion(Q::SRMatrix, diffusion::Number) = SRMatrix(sqrt.(diffusion) * Q.squareroot) +apply_diffusion(Q::SRMatrix, diffusion::Number) = + SRMatrix(sqrt.(diffusion) * Q.squareroot, diffusion * Q.mat) # All the Gaussian things @reexport using GaussianDistributions diff --git a/src/caches.jl b/src/caches.jl index f3330f7a4..2b730ae05 100644 --- a/src/caches.jl +++ b/src/caches.jl @@ -81,9 +81,9 @@ function OrdinaryDiffEq.alg_cache( alg::GaussianODEFilter, u, rate_prototype, - uEltypeNoUnits, - uBottomEltypeNoUnits, - tTypeNoUnits, + ::Type{uEltypeNoUnits}, + ::Type{uBottomEltypeNoUnits}, + ::Type{tTypeNoUnits}, uprev, uprev2, f, @@ -93,7 +93,7 @@ function OrdinaryDiffEq.alg_cache( p, calck, ::Val{IIP}, -) where {IIP} +) where {IIP,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits} initialize_derivatives = true if u isa Number @@ -113,11 +113,12 @@ function OrdinaryDiffEq.alg_cache( t0 = t uType = typeof(u) - uElType = eltype(u_vec) + # uElType = eltype(u_vec) + uElType = uBottomEltypeNoUnits matType = Matrix{uElType} # Projections - Proj = projection(d, q, Val(uElType)) + Proj = projection(d, q, uElType) E0, E1, E2 = Proj(0), Proj(1), Proj(2) @assert f isa AbstractODEFunction SolProj = f isa DynamicalODEFunction ? [Proj(1); Proj(0)] : Proj(0) @@ -127,8 +128,11 @@ function OrdinaryDiffEq.alg_cache( A, Q = ibm(d, q, uElType) - initial_variance = 1.0 * ones(uElType, D) - x0 = Gaussian(zeros(uElType, D), SRMatrix(Matrix(Diagonal(sqrt.(initial_variance))))) + initial_variance = ones(uElType, D) + x0 = Gaussian( + zeros(uElType, D), + SRMatrix(diagm(sqrt.(initial_variance)), diagm(initial_variance)), + ) # Measurement model R = zeros(uElType, d, d) @@ -141,7 +145,7 @@ function OrdinaryDiffEq.alg_cache( v, S = similar(h), similar(ddu) v = similar(h) S = - alg isa EK0 ? SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) : + # alg isa EK0 ? SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) : SRMatrix(zeros(uElType, d, D), zeros(uElType, d, d)) measurement = Gaussian(v, S) pu_tmp = diff --git a/src/filtering/smooth.jl b/src/filtering/smooth.jl index 7945b5d59..a54255773 100644 --- a/src/filtering/smooth.jl +++ b/src/filtering/smooth.jl @@ -58,12 +58,12 @@ function smooth( return x_curr_smoothed, G end -function smooth!(x_curr, x_next, Ah, Qh, integ, diffusion=1) +function smooth!(x_curr, x_next, Ah, Qh, cache, diffusion=1) # x_curr is the state at time t_n (filter estimate) that we want to smooth # x_next is the state at time t_{n+1}, already smoothed, which we use for smoothing - @unpack d, q = integ.cache - @unpack x_pred = integ.cache - @unpack C1, G1, G2, C2 = integ.cache + @unpack d, q = cache + @unpack x_pred = cache + @unpack C1, G1, G2, C2 = cache # Prediction: t -> t+1 predict_mean!(x_pred, x_curr, Ah) diff --git a/src/gaussians.jl b/src/gaussians.jl index 956161f18..749510caa 100644 --- a/src/gaussians.jl +++ b/src/gaussians.jl @@ -20,14 +20,14 @@ ndims(g::Gaussian) = ndims(g.μ) Base.:*(M, g::SRGaussian) = Gaussian(M * g.μ, X_A_Xt(g.Σ, M)) # GaussianDistributions.whiten(Σ::SRMatrix, z) = Σ.L\z -function mul!(g_out::SRGaussian, M::AbstractMatrix, g_in::SRGaussian) +function _gaussian_mul!(g_out::SRGaussian, M::AbstractMatrix, g_in::SRGaussian) _matmul!(g_out.μ, M, g_in.μ) X_A_Xt!(g_out.Σ, g_in.Σ, M) return g_out end -var(p::SRGaussian{T}) where {T} = diag(p.Σ) -std(p::SRGaussian{T}) where {T} = sqrt.(var(p)) -mean(s::SRGaussianList{T}) where {T} = s.μ -var(s::SRGaussianList{T}) where {T} = diag.(s.Σ) -std(s::SRGaussianList{T}) where {T} = map(std, s) +var(p::SRGaussian) = diag(p.Σ) +std(p::SRGaussian) = sqrt.(var(p)) +mean(s::SRGaussianList) = s.μ +var(s::SRGaussianList) = diag.(s.Σ) +std(s::SRGaussianList) = map(std, s) diff --git a/src/initialization/classicsolverinit.jl b/src/initialization/classicsolverinit.jl index 288cd47c5..121c7591a 100644 --- a/src/initialization/classicsolverinit.jl +++ b/src/initialization/classicsolverinit.jl @@ -11,10 +11,10 @@ function initial_update!(integ, cache, init::ClassicSolverInit) @unpack ddu, du, x_tmp, x_tmp2, m_tmp, K1 = cache # Initialize on u0 - condition_on!(x, Proj(0), view(u, :), m_tmp, K1, x_tmp.Σ, x_tmp2.Σ.mat) + condition_on!(x, Proj(0), view(u, :), m_tmp.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) f(du, u, p, t) integ.destats.nf += 1 - condition_on!(x, Proj(1), view(du, :), m_tmp, K1, x_tmp.Σ, x_tmp2.Σ.mat) + condition_on!(x, Proj(1), view(du, :), m_tmp.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) if q < 2 return @@ -31,7 +31,7 @@ function initial_update!(integ, cache, init::ClassicSolverInit) ForwardDiff.jacobian!(ddu, (du, u) -> f(du, u, p, t), du, u) end ddfddu = ddu * view(du, :) + view(dfdt, :) - condition_on!(x, Proj(2), ddfddu, m_tmp, K1, x_tmp.Σ, x_tmp2.Σ.mat) + condition_on!(x, Proj(2), ddfddu, m_tmp.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) if q < 3 return end @@ -88,7 +88,7 @@ function rk_init_improve(integ, cache::GaussianODEFilterCache, ts, us, dt) make_preconditioners!(cache, dt) @unpack P, PI = cache - mul!(x, P, x) + _gaussian_mul!(x, P, x) preds = [] filts = [copy(x)] @@ -114,10 +114,10 @@ function rk_init_improve(integ, cache::GaussianODEFilterCache, ts, us, dt) xf = filts[i-1] xs = filts[i] xp = preds[i-1] # Since `preds` is one shorter - smooth!(xf, xs, A, Q, integ, 1) + smooth!(xf, xs, A, Q, integ.cache, 1) end - mul!(cache.x, PI, filts[1]) + _gaussian_mul!(cache.x, PI, filts[1]) return nothing end diff --git a/src/initialization/common.jl b/src/initialization/common.jl index dac15d5ad..6ffa57803 100644 --- a/src/initialization/common.jl +++ b/src/initialization/common.jl @@ -25,14 +25,13 @@ function condition_on!( x::SRGaussian, H::AbstractMatrix, data::AbstractVector, - meascache, + Scache, Kcache, covcache, Mcache, ) - z, S = meascache + S = Scache - _matmul!(z, H, x.μ) X_A_Xt!(S, x.Σ, H) @assert isdiag(S) S_diag = diag(S) @@ -43,11 +42,12 @@ function condition_on!( _matmul!(Kcache, x.Σ.mat, H') K = Kcache ./= S_diag' - _matmul!(x.μ, K, data - z, 1.0, 1.0) # x.μ .+= K*(data - z) + datadiff = _matmul!(data, H, x.μ, -1, 1) + _matmul!(x.μ, K, datadiff, 1, 1) D = length(x.μ) - mul!(Mcache, K, H, -1.0, 0.0) + _matmul!(Mcache, K, H, -1, 0) @inbounds @simd ivdep for i in 1:D Mcache[i, i] += 1 end diff --git a/src/initialization/taylormode.jl b/src/initialization/taylormode.jl index 822d909b8..c1b4d1a65 100644 --- a/src/initialization/taylormode.jl +++ b/src/initialization/taylormode.jl @@ -6,6 +6,9 @@ function initial_update!(integ, cache, init::TaylorModeInit) D = d * (q + 1) @unpack x_tmp, x_tmp2, m_tmp, K1, K2 = cache + if size(K1, 2) != d + K1 = K1[:, 1:d] + end f_derivatives = taylormode_get_derivatives(u, f, p, t, q) integ.destats.nf += q @@ -20,15 +23,12 @@ function initial_update!(integ, cache, init::TaylorModeInit) df = df[2, :] end pmat = f.mass_matrix * Proj(o) - condition_on!( - x, - pmat, - view(df, :), - m_cache, - view(K1, :, 1:d), - x_tmp.Σ, - x_tmp2.Σ.mat, - ) + + if !(df isa AbstractVector) + df = df[:] + end + + condition_on!(x, pmat, df, m_cache.Σ, K1, x_tmp.Σ, x_tmp2.Σ.mat) end end @@ -40,8 +40,12 @@ function taylormode_get_derivatives(u, f::AbstractODEFunction{false}, p, t, q) tT = Taylor1(typeof(t), q) tT[0] = t - uT = Taylor1.(u, tT.order) - duT = zero.(Taylor1.(u, tT.order)) + # uT = Taylor1.(u, tT.order) # precompilation could be faster with the following: + uT = similar(u, Taylor1{eltype(u)}) + @inbounds @simd ivdep for i in eachindex(u) + uT[i] = Taylor1(u[i], q) + end + duT = zero(uT) uauxT = similar(uT) TaylorIntegration.jetcoeffs!(f, tT, uT, duT, uauxT, p) return [evaluate.(differentiate.(uT, i)) for i in 0:q] @@ -52,8 +56,12 @@ end function taylormode_get_derivatives(u, f::AbstractODEFunction{true}, p, t, q) tT = Taylor1(typeof(t), q) tT[0] = t - uT = Taylor1.(u, tT.order) - duT = zero.(Taylor1.(u, tT.order)) + # uT = Taylor1.(u, tT.order) + uT = similar(u, Taylor1{eltype(u)}) + @inbounds @simd ivdep for i in eachindex(u) + uT[i] = Taylor1(u[i], q) + end + duT = zero(uT) uauxT = similar(uT) TaylorIntegration.jetcoeffs!(f, tT, uT, duT, uauxT, p) return [evaluate.(differentiate.(uT, i)) for i in 0:q] diff --git a/src/integrator_utils.jl b/src/integrator_utils.jl index 35b50cb1e..4da462300 100644 --- a/src/integrator_utils.jl +++ b/src/integrator_utils.jl @@ -40,13 +40,13 @@ 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 - @simd ivdep for s in integ.sol.x_filt - copy!(s.Σ, apply_diffusion(s.Σ, mle_diffusion)) + @simd ivdep for S in integ.sol.x_filt.Σ + copy!(S, apply_diffusion(S, mle_diffusion)) end # Re-write into the solution estimates for (pu, x) in zip(integ.sol.pu, integ.sol.x_filt) - mul!(pu, integ.cache.SolProj, x) + _gaussian_mul!(pu, integ.cache.SolProj, x) end # [(su[:] .= pu) for (su, pu) in zip(integ.sol.u, integ.sol.pu.μ)] end @@ -94,13 +94,13 @@ function smooth_solution!(integ) make_preconditioners!(integ.cache, dt) P, PI = integ.cache.P, integ.cache.PI - mul!(x_tmp, P, x[i]) - mul!(x_tmp2, P, x[i+1]) - smooth!(x_tmp, x_tmp2, A, Q, integ, diffusions[i]) - mul!(x[i], PI, x_tmp) + _gaussian_mul!(x_tmp, P, x[i]) + _gaussian_mul!(x_tmp2, P, x[i+1]) + smooth!(x_tmp, x_tmp2, A, Q, integ.cache, diffusions[i]) + _gaussian_mul!(x[i], PI, x_tmp) # Save the smoothed state into the solution - mul!(integ.sol.pu[i], integ.cache.SolProj, x[i]) + _gaussian_mul!(integ.sol.pu[i], integ.cache.SolProj, x[i]) integ.sol.u[i][:] .= integ.sol.pu[i].μ end integ.sol.interp = set_smooth(integ.sol.interp) @@ -121,7 +121,7 @@ function pn_solution_endpoint_match_cur_integrator!(integ) OrdinaryDiffEq.copyat_or_push!( integ.sol.pu, integ.saveiter, - mul!(integ.cache.pu_tmp, integ.cache.SolProj, integ.cache.x), + _gaussian_mul!(integ.cache.pu_tmp, integ.cache.SolProj, integ.cache.x), ) end end @@ -148,7 +148,7 @@ function DiffEqBase.savevalues!( integ.cache.local_diffusion, ) if integ.opts.save_everystep - mul!(integ.cache.pu_tmp, integ.cache.SolProj, integ.cache.x), + _gaussian_mul!(integ.cache.pu_tmp, integ.cache.SolProj, integ.cache.x), OrdinaryDiffEq.copyat_or_push!(integ.sol.pu, integ.saveiter, integ.cache.pu_tmp) end diff --git a/src/perform_step.jl b/src/perform_step.jl index 963a7f5d8..25da7695e 100644 --- a/src/perform_step.jl +++ b/src/perform_step.jl @@ -22,7 +22,7 @@ function OrdinaryDiffEq.initialize!(integ, cache::GaussianODEFilterCache) OrdinaryDiffEq.copyat_or_push!( integ.sol.pu, integ.saveiter, - mul!(cache.pu_tmp, cache.SolProj, cache.x), + _gaussian_mul!(cache.pu_tmp, cache.SolProj, cache.x), ) return nothing end @@ -158,7 +158,9 @@ function evaluate_ode!( @unpack du1, uf, jac_config = integ.cache uf.f = OrdinaryDiffEq.nlsolve_f(f, alg) uf.t = t - uf.p = p + if !(p isa DiffEqBase.NullParameters) + uf.p = p + end OrdinaryDiffEq.jacobian!(ddu, uf, u_pred, du1, integ, jac_config) end integ.destats.njacs += 1 diff --git a/src/preconditioning.jl b/src/preconditioning.jl index 47aecfa9c..a70cb5ee6 100644 --- a/src/preconditioning.jl +++ b/src/preconditioning.jl @@ -1,4 +1,4 @@ -function init_preconditioner(d, q, elType=typeof(1.0)) +function init_preconditioner(d, q, ::Type{elType}=typeof(1.0)) where {elType} D = d * (q + 1) P = Diagonal(ones(elType, D)) PI = Diagonal(ones(elType, D)) diff --git a/src/priors.jl b/src/priors.jl index c40f96a0e..d9f060b10 100644 --- a/src/priors.jl +++ b/src/priors.jl @@ -6,7 +6,7 @@ Generate the discrete dynamics for a q-IBM model. INCLUDES AUTOMATIC PRECONDITIONING! """ -function ibm(d::Integer, q::Integer, elType=typeof(1.0)) +function ibm(d::Integer, q::Integer, ::Type{elType}=typeof(1.0)) where {elType} # Make A A_breve = zeros(elType, q + 1, q + 1) @simd ivdep for j in 1:q+1 @@ -28,9 +28,8 @@ function ibm(d::Integer, q::Integer, elType=typeof(1.0)) @inbounds Q_breve[1+row, 1+col] = val end end - QL_breve = cholesky!(Q_breve).L - QL = kron(I(d), QL_breve) - Q = SRMatrix(QL) + QL_breve = cholesky(Q_breve).L + Q = SRMatrix(kron(I(d), QL_breve), kron(I(d), Q_breve)) return A, Q end diff --git a/src/projection.jl b/src/projection.jl index 50a8b0583..e8a2e67a2 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -1,4 +1,4 @@ -function projection(d::Integer, q::Integer, ::Val{elType}=Val(typeof(1.0))) where {elType} +function projection(d::Integer, q::Integer, ::Type{elType}=typeof(1.0)) where {elType} # Proj(deriv) = kron(diagm(0 => ones(elType, d)), [i==(deriv+1) ? 1 : 0 for i in 1:q+1]') # Slightly faster version of the above: diff --git a/src/solution.jl b/src/solution.jl index 543f6fdd9..9eb8b3da9 100644 --- a/src/solution.jl +++ b/src/solution.jl @@ -63,10 +63,10 @@ function DiffEqBase.build_solution( uElType = eltype(prob.u0) D = d pu_cov = - alg isa EK0 && !(prob.f isa DynamicalODEFunction) ? - SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) : - SRMatrix(zeros(uElType, d, D)) - x_cov = SRMatrix(zeros(uElType, d, d)) + # alg isa EK0 && !(prob.f isa DynamicalODEFunction) ? + # SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) : + SRMatrix(zeros(uElType, d, D), zeros(uElType, d, d)) + x_cov = SRMatrix(zeros(uElType, d, d), zeros(uElType, d, d)) pu = StructArray{Gaussian{Vector{uElType},typeof(pu_cov)}}(undef, 0) x_filt = StructArray{Gaussian{Vector{uElType},typeof(x_cov)}}(undef, 0) x_smooth = copy(x_filt) @@ -168,7 +168,7 @@ function GaussianODEFilterPosterior(alg, u0) q = alg.order D = d * (q + 1) - Proj = projection(d, q, Val(uElType)) + Proj = projection(d, q, uElType) SolProj = u0 isa ArrayPartition ? [Proj(1); Proj(0)] : Proj(0) A, Q = ibm(d, q, uElType) diff --git a/src/squarerootmatrix.jl b/src/squarerootmatrix.jl index cfff0f656..6460663af 100644 --- a/src/squarerootmatrix.jl +++ b/src/squarerootmatrix.jl @@ -5,13 +5,12 @@ Relates to PSDMatrices.jl, and I might move this to PSDMatrices.jl in the future, but having this here allowed for easier development. """ -struct SquarerootMatrix{T<:Real,S<:AbstractMatrix{T},M<:AbstractMatrix{T}} <: - AbstractMatrix{T} +struct SquarerootMatrix{T<:Real,S<:AbstractMatrix,M<:AbstractMatrix} <: AbstractMatrix{T} squareroot::S mat::M + SquarerootMatrix(S::AbstractMatrix{T}, mat::AbstractMatrix{T}) where {T} = + new{T,typeof(S),typeof(mat)}(S, mat) end -SquarerootMatrix(S::AbstractMatrix{T}, mat::AbstractMatrix{T}) where {T} = - SquarerootMatrix{T,typeof(S),typeof(mat)}(S, mat) SquarerootMatrix(S) = SquarerootMatrix(S, S * S') Base.Matrix(M::SquarerootMatrix) = M.mat