Make the CovarianceStructure object hold elType, d, q

Simplifying the code will happen after!
nathanaelbosch · Oct 28, 2023 · 8ce3162 · 8ce3162
1 parent e503a6d
commit 8ce3162
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 43 deletions.
diff --git a/src/caches.jl b/src/caches.jl
@@ -92,17 +92,18 @@ function OrdinaryDiffEq.alg_cache(
     d = is_secondorder_ode ? length(u[1, :]) : length(u)
     D = d * (q + 1)
 
-    FAC = get_covariance_structure(alg)
+    uType = typeof(u)
+    # uElType = eltype(u_vec)
+    uElType = uBottomEltypeNoUnits
+
+    FAC = get_covariance_structure(alg; elType=uElType, d, q)
     if FAC isa IsometricKroneckerCovariance && !(f.mass_matrix isa UniformScaling)
         error(
             "The selected algorithm uses an efficient Kronecker-factorized implementation which is incompatible with the provided mass matrix. Try using the `EK1` instead.",
         )
     end
 
-    uType = typeof(u)
-    # uElType = eltype(u_vec)
-    uElType = uBottomEltypeNoUnits
-    matType = typeof(factorized_similar(FAC, uElType, d, d; d, q))
+    matType = typeof(factorized_similar(FAC, d, d))
 
     # Projections
     Proj = projection(FAC, d, q, uElType)
@@ -150,38 +151,38 @@ function OrdinaryDiffEq.alg_cache(
     copy!(x0.Σ, apply_diffusion(x0.Σ, initdiff))
 
     # Measurement model related things
-    R = factorized_similar(FAC, uElType, d, d; d, q)
-    H = factorized_similar(FAC, uElType, d, D; d, q)
+    R = factorized_similar(FAC, d, d)
+    H = factorized_similar(FAC, d, D)
     v = similar(Array{uElType}, d)
-    S = PSDMatrix(factorized_zeros(FAC, uElType, D, d; d, q))
+    S = PSDMatrix(factorized_zeros(FAC, D, d))
     measurement = Gaussian(v, S)
 
     # Caches
     du = is_secondorder_ode ? similar(u[2, :]) : similar(u)
-    ddu = factorized_similar(FAC, uElType, length(u), length(u); d, q)
+    ddu = factorized_similar(FAC, length(u), length(u))
     pu_tmp = if !is_secondorder_ode # same dimensions as `measurement`
         copy(measurement)
     else # then `u` has 2d dimensions
         Gaussian(
             similar(Array{uElType}, 2d),
-            PSDMatrix(factorized_similar(FAC, uElType, D, 2d; d, q)),
+            PSDMatrix(factorized_similar(FAC, D, 2d)),
         )
     end
-    K = factorized_similar(FAC, uElType, D, d; d, q)
-    G = factorized_similar(FAC, uElType, D, D; d, q)
-    Smat = factorized_similar(FAC, uElType, d, d; d, q)
+    K = factorized_similar(FAC, D, d)
+    G = factorized_similar(FAC, D, D)
+    Smat = factorized_similar(FAC, d, d)
 
-    C_dxd = factorized_similar(FAC, uElType, d, d; d, q)
-    C_dxD = factorized_similar(FAC, uElType, d, D; d, q)
-    C_Dxd = factorized_similar(FAC, uElType, D, d; d, q)
-    C_DxD = factorized_similar(FAC, uElType, D, D; d, q)
-    C_2DxD = factorized_similar(FAC, uElType, 2D, D; d, q)
-    C_3DxD = factorized_similar(FAC, uElType, 3D, D; d, q)
+    C_dxd = factorized_similar(FAC, d, d)
+    C_dxD = factorized_similar(FAC, d, D)
+    C_Dxd = factorized_similar(FAC, D, d)
+    C_DxD = factorized_similar(FAC, D, D)
+    C_2DxD = factorized_similar(FAC, 2D, D)
+    C_3DxD = factorized_similar(FAC, 3D, D)
 
     backward_kernel = AffineNormalKernel(
-        factorized_similar(FAC, uElType, D, D; d, q),
+        factorized_similar(FAC, D, D),
         similar(Vector{uElType}, D),
-        PSDMatrix(factorized_similar(FAC, uElType, 2D, D; d, q)),
+        PSDMatrix(factorized_similar(FAC, 2D, D)),
     )
 
     u_pred = copy(u)

diff --git a/src/covariance_structure.jl b/src/covariance_structure.jl
@@ -1,8 +1,14 @@
-abstract type CovarianceStructure end
-struct IsometricKroneckerCovariance <: CovarianceStructure end
-struct DenseCovariance <: CovarianceStructure end
+abstract type CovarianceStructure{T} end
+struct IsometricKroneckerCovariance{T} <: CovarianceStructure{T}
+    d::Int64
+    q::Int64
+end
+struct DenseCovariance{T} <: CovarianceStructure{T}
+    d::Int64
+    q::Int64
+end
 
-function get_covariance_structure(alg)
+function get_covariance_structure(alg; elType, d, q)
     if (
         alg isa EK0 &&
         !(
@@ -11,29 +17,29 @@ function get_covariance_structure(alg)
         ) &&
         alg.prior isa IWP
     )
-        return IsometricKroneckerCovariance()
+        return IsometricKroneckerCovariance{elType}(d, q)
     else
-        return DenseCovariance()
+        return DenseCovariance{elType}(d, q)
     end
 end
 
-factorized_zeros(::IsometricKroneckerCovariance, elType, sizes...; d, q) = begin
+factorized_zeros(C::IsometricKroneckerCovariance{T}, sizes...) where {T} = begin
     for s in sizes
-        @assert s % d == 0
+        @assert s % C.d == 0
     end
-    return IsometricKroneckerProduct(d, Array{elType}(calloc, (s ÷ d for s in sizes)...))
+    return IsometricKroneckerProduct(C.d, Array{T}(calloc, (s ÷ C.d for s in sizes)...))
 end
-factorized_similar(::IsometricKroneckerCovariance, elType, size1, size2; d, q) = begin
+factorized_similar(C::IsometricKroneckerCovariance{T}, size1, size2) where {T} = begin
     for s in (size1, size2)
-        @assert s % d == 0
+        @assert s % C.d == 0
     end
-    return IsometricKroneckerProduct(d, similar(Matrix{elType}, size1 ÷ d, size2 ÷ d))
+    return IsometricKroneckerProduct(C.d, similar(Matrix{T}, size1 ÷ C.d, size2 ÷ C.d))
 end
 
-factorized_zeros(::DenseCovariance, elType, sizes...; d, q) =
-    Array{elType}(calloc, sizes...)
-factorized_similar(::DenseCovariance, elType, size1, size2; d, q) =
-    similar(Matrix{elType}, size1, size2)
+factorized_zeros(::DenseCovariance{T}, sizes...) where {T} =
+    Array{T}(calloc, sizes...)
+factorized_similar(::DenseCovariance{T}, size1, size2) where {T} =
+    similar(Matrix{T}, size1, size2)
 
 to_factorized_matrix(::DenseCovariance, M::AbstractMatrix) = Matrix(M)
 to_factorized_matrix(::IsometricKroneckerCovariance, M::AbstractMatrix) =

diff --git a/src/solution.jl b/src/solution.jl
@@ -94,8 +94,8 @@ function DiffEqBase.build_solution(
     D = d * (q + 1)
 
     FAC = cache.covariance_factorization
-    pu_cov = PSDMatrix(factorized_zeros(FAC, uElType, D, d; d, q))
-    x_cov = PSDMatrix(factorized_zeros(FAC, uElType, D, D; d, q))
+    pu_cov = PSDMatrix(factorized_zeros(FAC, D, d))
+    x_cov = PSDMatrix(factorized_zeros(FAC, D, D))
     pu = StructArray{Gaussian{Vector{uElType},typeof(pu_cov)}}(undef, 0)
     x_filt = StructArray{Gaussian{Vector{uElType},typeof(x_cov)}}(undef, 0)
     x_smooth = copy(x_filt)

diff --git a/test/core/filtering.jl b/test/core/filtering.jl
@@ -32,7 +32,6 @@ import ProbNumDiffEq as PNDE
 
     _fstr(F) = F ? "Kronecker" : "None"
     @testset "Factorization: $(_fstr(KRONECKER))" for KRONECKER in (false, true)
-        FAC = KRONECKER ? PNDE.IsometricKroneckerCovariance() : PNDE.DenseCovariance()
         if KRONECKER
             K = 2
             m = kron(ones(K), m)

diff --git a/test/core/preconditioning.jl b/test/core/preconditioning.jl
@@ -23,7 +23,7 @@ prob = remake(prob, tspan=(0.0, 10.0))
 
     # First test that they're both equivalent
     D = d * (q + 1)
-    P, PI = PNDE.init_preconditioner(PNDE.DenseCovariance(), d, q)
+    P, PI = PNDE.init_preconditioner(PNDE.DenseCovariance{Float64}(d, q), d, q)
     PNDE.make_preconditioner!(P, h, d, q)
     PNDE.make_preconditioner_inv!(PI, h, d, q)
     @test Ah_p ≈ P * Ah * PI

diff --git a/test/core/priors.jl b/test/core/priors.jl
@@ -119,8 +119,8 @@ end
             Qh::Any
         end
 
-        A, Q, Ah, Qh, P, PI =
-            PNDE.initialize_transition_matrices(PNDE.DenseCovariance(), prior, h)
+        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))