Add the DiagonalEK1

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,15 +38,18 @@ 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"
 
 [compat]
 ArrayAllocators = "0.3"
+BlockDiagonals = "0.1"
 DiffEqBase = "6.122"
 DiffEqCallbacks = "2.36"
 DiffEqDevTools = "2"

ext/BlockDiagonalsExt.jl

@@ -0,0 +1,8 @@
+module BlockDiagonalsExt
+
+import ProbNumDiffEq: ProbNumDiffEqBlockDiagonal, blocks
+import BlockDiagonals: BlockDiagonal
+
+BlockDiagonal(M::ProbNumDiffEqBlockDiagonal) = BlockDiagonal(blocks(M))
+
+end

src/ProbNumDiffEq.jl

@@ -2,7 +2,8 @@ __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!
@@ -45,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) =
@@ -53,7 +56,9 @@ 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")
+export IsometricKroneckerCovariance, DenseCovariance, BlockDiagonalCovariance
 
 abstract type AbstractODEFilterCache <: OrdinaryDiffEq.OrdinaryDiffEqCache end
 
@@ -65,14 +70,16 @@ 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")
 export TaylorModeInit, ClassicSolverInit, SimpleInit, ForwardDiffInit
 
 include("algorithms.jl")
-export EK0, EK1
+export EK0, EK1, DiagonalEK1
 export ExpEK, RosenbrockExpEK
 
 include("alg_utils.jl")

src/alg_utils.jl

@@ -4,19 +4,25 @@
 ############################################################################################
 
 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

src/algorithms.jl

@@ -3,7 +3,13 @@
 ########################################################################################
 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(
@@ -12,16 +18,44 @@ function ekargcheck(alg; diffusionmodel, pn_observation_noise, kwargs...)
             ),
         )
     end
-    if (
-        (diffusionmodel isa FixedMVDiffusion && diffusionmodel.calibrate) ||
-        diffusionmodel isa DynamicMVDiffusion) && alg == EK1
-        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`.",
-            ),
-        )
+    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`.",
+                ),
+            )
+        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
+
+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
+            end
+        else
+            # This is not great as other priors can be Kronecker too; TODO
+            return DenseCovariance
+        end
+    elseif Alg <: DiagonalEK1
+        return BlockDiagonalCovariance
+    elseif Alg <: EK1
+        return DenseCovariance
+    else
+        throw(ArgumentError("Unknown algorithm type $Alg"))
     end
 end
+covariance_structure(alg) = covariance_structure(typeof(alg), alg.prior, alg.diffusionmodel)
 
 """
     EK0(; order=3,
@@ -58,22 +92,25 @@ 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
-        ekargcheck(EK0; diffusionmodel, pn_observation_noise)
-        new{PT,DT,IT,RT}(
-            prior, diffusionmodel, smooth, initialization, pn_observation_noise)
+        covariance_factorization::CF=covariance_structure(EK0, prior, diffusionmodel),
+    ) where {PT,DT,IT,RT,CF} = begin
+        ekargcheck(EK0; diffusionmodel, pn_observation_noise, covariance_factorization)
+        new{PT,DT,IT,RT,CF}(
+            prior, diffusionmodel, smooth, initialization, pn_observation_noise,
+            covariance_factorization)
     end
 end
 
@@ -117,12 +154,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),
@@ -135,8 +173,57 @@ 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
-        ekargcheck(EK1; diffusionmodel, pn_observation_noise)
+        covariance_factorization::CF=covariance_structure(EK1, prior, diffusionmodel),
+    ) where {PT,DT,IT,RT,CF} = begin
+        ekargcheck(EK1; diffusionmodel, pn_observation_noise, covariance_factorization)
+        new{
+            _unwrap_val(chunk_size),
+            _unwrap_val(autodiff),
+            diff_type,
+            _unwrap_val(standardtag),
+            _unwrap_val(concrete_jac),
+            PT,
+            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,CF} <: AbstractEK
+    prior::PT
+    diffusionmodel::DT
+    smooth::Bool
+    initialization::IT
+    pn_observation_noise::RT
+    covariance_factorization::CF
+    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,
+        covariance_factorization::CF=covariance_structure(
+            DiagonalEK1,
+            prior,
+            diffusionmodel,
+        ),
+    ) where {PT,DT,IT,RT,CF} = begin
+        ekargcheck(DiagonalEK1; diffusionmodel, pn_observation_noise, covariance_factorization)
         new{
             _unwrap_val(chunk_size),
             _unwrap_val(autodiff),
@@ -147,12 +234,14 @@ 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
@@ -236,7 +325,12 @@ 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