Simplify the whole I code a bit

nathanaelbosch · Jun 13, 2023 · a58030d · a58030d
1 parent 11d6ded
commit a58030d
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 9 deletions.
diff --git a/src/caches.jl b/src/caches.jl
@@ -92,7 +92,7 @@ function OrdinaryDiffEq.alg_cache(
     D = d * (q + 1)
 
     KRONECKER = alg isa EK0
-    Id = _mul_stable_I(d)
+    Id = _I(d)
 
     uType = typeof(u)
     # uElType = eltype(u_vec)

diff --git a/src/initialization/common.jl b/src/initialization/common.jl
@@ -104,7 +104,7 @@ function condition_on!(
     end
 
     d, q1 = size(H.A, 1), size(x.Σ.R.B, 1)
-    _I = kronecker(_mul_stable_I(d), I(q1))
+    _I = kronecker(I(d), I(q1))
     KH = K*H
     @assert _I.A == KH.A
     @. KH.B = _I.B - KH.B

diff --git a/src/kronecker.jl b/src/kronecker.jl
@@ -18,8 +18,23 @@ end
 copy(A::KP) = kronecker((A.A), copy(A.B))
 
 """
-    _mul_stable_I(d) = I(d) * I(d)
+    _I(d) = I(d) * I(d)
 
 Create an identity matrix that does not change its type when multiplied by another identity matrix.
+
+# Examples
+```julia-repl
+julia> I(2)|> typeof
+Diagonal{Bool, Vector{Bool}}
+
+julia> I(2) * I(2) |> typeof
+Diagonal{Bool, BitVector}
+
+julia> _I(2) |> typeof
+Diagonal{Bool, BitVector}
+
+julia> _I(2) * _I(2) |> typeof
+Diagonal{Bool, BitVector}
+```
 """
-_mul_stable_I(d) = I(d) * I(d)
+_I(d) = I(d) * I(d)
diff --git a/src/preconditioning.jl b/src/preconditioning.jl
@@ -1,5 +1,5 @@
 function init_preconditioner(d, q, ::Type{elType}=typeof(1.0)) where {elType}
-    Id = _mul_stable_I(d)
+    Id = _I(d)
     P = kronecker(Id, Diagonal(ones(elType, q + 1)))
     PI = kronecker(Id, Diagonal(ones(elType, q + 1)))
     return P, PI

diff --git a/src/priors/iwp.jl b/src/priors/iwp.jl
@@ -70,7 +70,7 @@ function preconditioned_discretize(iwp::IWP)
     QR_breve = Q_breve.R |> Matrix
 
     d = iwp.wiener_process_dimension
-    Id = _mul_stable_I(d)
+    Id = _I(d)
     A = kronecker(Id, A_breve)
     QR = kronecker(Id, QR_breve)
     Q = PSDMatrix(QR)
@@ -100,7 +100,7 @@ end
 function discretize(p::IWP, dt::Real)
     A_breve, Q_breve = discretize_1d(p, dt)
     d = p.wiener_process_dimension
-    Id = _mul_stable_I(d)
+    Id = _I(d)
     A = kronecker(Id, A_breve)
     QR = kronecker(Id, Q_breve.R)
     Q = PSDMatrix(QR)

diff --git a/src/projection.jl b/src/projection.jl
@@ -1,5 +1,5 @@
 function projection(d::Integer, q::Integer, ::Type{elType}=typeof(1.0)) where {elType}
-    Id = _mul_stable_I(d)
+    Id = _I(d)
     Proj(deriv) = kronecker(Id, [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
@@ -93,7 +93,7 @@ function DiffEqBase.build_solution(
     uElType = eltype(prob.u0)
     D = d
     KRONECKER = true
-    Id = _mul_stable_I(d)
+    Id = _I(d)
     pu_cov = if KRONECKER
         PSDMatrix(kronecker(Id, zeros(uElType, D ÷ d + 1)))
     else