diff --git a/src/linear_algebra_utilities/selection_matrices.jl b/src/linear_algebra_utilities/selection_matrices.jl new file mode 100644 index 000000000..0f6055e4d --- /dev/null +++ b/src/linear_algebra_utilities/selection_matrices.jl @@ -0,0 +1,56 @@ +struct SelectionMatrix{T<:Number} <: AbstractMatrix{T} + dims_to_select::Vector{Int64} + input_size::Int64 +end +Base.size(A::SelectionMatrix) = (length(A.dims_to_select), A.input_size) +SelectionMatrix(dims_to_select, input_size) = + SelectionMatrix{Bool}(dims_to_select, input_size) + +# Base.@propagate_inbounds function Base.getindex( +# M::SelectionMatrix{T}, +# i::Integer, +# j::Integer, +# ) where {T} +# all((0, 0) .< (i, j) .<= size(M)) || throw(BoundsError(M, (i, j))) + +# if j != M.dims_to_select[i] +# return zero(T) +# else +# return one(T) +# end +# end + +Matrix(A::SelectionMatrix{T}) where {T} = begin + M = zeros(T, size(A)...) + for (i, d) in enumerate(A.dims_to_select) + M[i, d] = one(T) + end + return M +end + +Base.:*(M::SelectionMatrix, v::AbstractVector) = + v[M.dims_to_select] +Base.:*(M::SelectionMatrix, A::AbstractMatrix) = + A[M.dims_to_select, :] +Base.:*(M::Adjoint{<:Any,<:SelectionMatrix}, v::AbstractVector) = begin + out = zeros(M.parent.input_size) + out[M.parent.dims_to_select] = v + return out +end +Base.:*(M::Adjoint{<:Any,<:SelectionMatrix}, A::AbstractMatrix) = begin + out = similar(A, M.parent.input_size, size(A, 2)) + out .= 0 + out[M.parent.dims_to_select, :] .= A + return out +end +Base.:*(A::AbstractMatrix, M::Adjoint{<:Any,<:SelectionMatrix}) = (M' * A')' +Base.:*(v::AbstractVector, M::Adjoint{<:Any,<:SelectionMatrix}) = (M' * v')' +Base.:*(A::Adjoint{<:Any,<:AbstractMatrix}, M::Adjoint{<:Any,<:SelectionMatrix}) = + (M' * A')' +Base.:*(v::Adjoint{<:Any,<:AbstractVector}, M::Adjoint{<:Any,<:SelectionMatrix}) = + (M' * v')' + +Base.:*(M::SelectionMatrix, A::BlocksOfDiagonals) = begin + @assert size(M, 2) == size(A, 1) + @assert M.input_size == length(A.blocks) +end diff --git a/test/core/markov_kernels.jl b/test/core/markov_kernels.jl new file mode 100644 index 000000000..0e3580ee9 --- /dev/null +++ b/test/core/markov_kernels.jl @@ -0,0 +1,67 @@ +using ProbNumDiffEq +import ProbNumDiffEq as PNDE +import ProbNumDiffEq: AffineNormalKernel +using LinearAlgebra, Statistics +using FillArrays +using Test + +d_in, d_out = 3, 2 +d_out = d_in +D = 3 + +T = Float64 +@testset "T=$T" for T in (Float64, BigFloat) + @testset "FAC=$_FAC" for _FAC in ( + PNDE.IsometricKroneckerCovariance, + PNDE.BlockDiagonalCovariance, + PNDE.DenseCovariance, + ) + _FAC = PNDE.DenseCovariance + FAC = _FAC{T}(d_in, D - 1) + + _A = PNDE.RightIsometricKroneckerProduct(D, rand(T, d_out, d_in)) + _b = rand(T, d_out * D) + _C = PSDMatrix(PNDE.RightIsometricKroneckerProduct(D, rand(T, d_out, d_out))) + + A = PNDE.to_factorized_matrix(FAC, _A) + for b in (_b, Zeros(T, size(_b)...)), + C in (PNDE.to_factorized_matrix(FAC, _C), PSDMatrix(Zeros(T, size(_C)...))) + + K = AffineNormalKernel(A, b, C) + + # definition + x = rand(T, d_in * D) + @test K(x) == Gaussian(A * x + b, C) + + # base + @test_nowarn _, _, _ = K + @test (@test_nowarn similar(K) isa AffineNormalKernel) + @test copy(K) == AffineNormalKernel(copy(A), copy(b), copy(C)) + _K = similar(K) + @test copy!(_K, K) == K + + # marginalize + d_ov = 2 * d_in + x = Gaussian(rand(T, d_in * D), PSDMatrix(rand(T, d_ov * D, d_in * D))) + y = @test_nowarn PNDE.marginalize(x, K) + @test y isa Gaussian + @test mean(y) ≈ A * mean(x) + b + @test PNDE.unfactorize(cov(y)) ≈ + A * PNDE.unfactorize(cov(x)) * A' + PNDE.unfactorize(C) + + # forward-backward marginalization returns the original! + Kback = @test_nowarn PNDE.compute_backward_kernel(y, x, K) + xback = @test_nowarn PNDE.marginalize(y, Kback) + @test mean(xback) ≈ mean(x) + @test PNDE.unfactorize(cov(xback)) ≈ PNDE.unfactorize(cov(x)) + + # marginalize in-place + y = similar(y) + cachemat = rand(T, size(cov(x).R, 1) + size(K.C.R, 1), size(A, 1)) + @test_nowarn PNDE.marginalize!(y, x, K; cachemat) + @test mean(y) ≈ A * mean(x) + b + @test PNDE.unfactorize(cov(y)) ≈ + A * PNDE.unfactorize(cov(x)) * A' + PNDE.unfactorize(C) + end + end +end diff --git a/test/core/selection_matrices.jl b/test/core/selection_matrices.jl new file mode 100644 index 000000000..4c928eaa4 --- /dev/null +++ b/test/core/selection_matrices.jl @@ -0,0 +1,27 @@ +using ProbNumDiffEq +import ProbNumDiffEq as PNDE +import ProbNumDiffEq: SelectionMatrix +using LinearAlgebra +using Test + +d_out = 2 +d_in = 3 +d_lat = 4 + +@testset "T=$T" for T in (Float64, BigFloat) + E = SelectionMatrix([2, 1], d_in) + M = Matrix(E) + + v = rand(d_in) + @test E * v == M * v + X = rand(d_in, d_lat) + @test E * X == M * X + @test X' * E' == X' * M' + Xt = Matrix(X') + @test Xt * E' == Xt * M' + + v = rand(d_out) + @test E' * v == M' * v + X = rand(d_out, d_lat) + @test E' * X == M' * X +end