From 267f5649f5911c3618d1a365d0fd3274d302f8db Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 13:52:48 +0200
Subject: [PATCH 01/17] feat: add ArrowheadMatrix

---
 Project.toml                            |   5 +-
 src/BayesBase.jl                        |   1 +
 src/algebra/arrowheadmatrix.jl          | 124 ++++++++++++++++
 test/algebra/algebrasetup_setuptests.jl |   2 +
 test/algebra/arrowheadmatrix_tests.jl   | 187 ++++++++++++++++++++++++
 5 files changed, 317 insertions(+), 2 deletions(-)
 create mode 100644 src/algebra/arrowheadmatrix.jl
 create mode 100644 test/algebra/algebrasetup_setuptests.jl
 create mode 100644 test/algebra/arrowheadmatrix_tests.jl

diff --git a/Project.toml b/Project.toml
index a9138a8..82afceb 100644
--- a/Project.toml
+++ b/Project.toml
@@ -18,6 +18,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 TinyHugeNumbers = "783c9a47-75a3-44ac-a16b-f1ab7b3acf04"
 
 [compat]
+BenchmarkTools = "1.5.0"
 Distributions = "0.25"
 DomainSets = "0.5.2, 0.6, 0.7"
 LinearAlgebra = "1.9"
@@ -35,12 +36,12 @@ julia = "1.9"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CpuId = "adafc99b-e345-5852-983c-f28acb93d879"
-JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 
 [targets]
 test = ["Aqua", "CpuId", "JET", "Test", "ReTestItems", "LinearAlgebra", "StableRNGs", "HCubature"]
diff --git a/src/BayesBase.jl b/src/BayesBase.jl
index f805ea5..65ecd23 100644
--- a/src/BayesBase.jl
+++ b/src/BayesBase.jl
@@ -96,5 +96,6 @@ include("densities/samplelist.jl")
 include("densities/mixture.jl")
 include("densities/factorizedjoint.jl")
 include("densities/contingency.jl")
+include("algebra/arrowheadmatrix.jl")
 
 end
diff --git a/src/algebra/arrowheadmatrix.jl b/src/algebra/arrowheadmatrix.jl
new file mode 100644
index 0000000..9b72ab4
--- /dev/null
+++ b/src/algebra/arrowheadmatrix.jl
@@ -0,0 +1,124 @@
+export ArrowheadMatrix, InvArrowheadMatrix
+
+import Base: getindex
+import LinearAlgebra: mul!
+import Base: size, *, \, inv, convert, Matrix
+
+struct ArrowheadMatrix{O, T, Z, P} <: AbstractMatrix{O}
+    α::T         
+    z::Z         
+    D::P
+end
+function ArrowheadMatrix(a::T, z::Z, d::D) where {T,Z,D}
+    O = promote_type(typeof(a), eltype(z), eltype(d))
+    return ArrowheadMatrix{O, T, Z, D}(a, z, d)
+end
+
+function size(A::ArrowheadMatrix)
+    n = length(A.D) + 1
+    return (n, n)
+end
+
+function Base.convert(::Type{Matrix}, A::ArrowheadMatrix{O}) where {O}
+    n = length(A.z)
+    M = zeros(O, n + 1, n + 1)
+    for i in 1:n
+        M[i, i] = A.D[i]
+    end
+    M[1:n, n + 1] .= A.z
+    M[n + 1, 1:n] .= A.z
+    M[n + 1, n + 1] = A.α
+    return M
+end
+
+function LinearAlgebra.mul!(y, A::ArrowheadMatrix{T}, x::AbstractVector{T}) where T
+    n = length(A.z)
+    if length(x) != n + 1
+        throw(DimensionMismatch())
+    end
+    @inbounds @views begin 
+        y[1:n] = A.D .* x[1:n] + A.z * x[n + 1]
+        y[n + 1] = dot(A.z, x[1:n]) + A.α * x[n + 1]
+    end
+    return y
+end
+
+function linsolve!(y, A::ArrowheadMatrix, b::AbstractVector)
+    n = length(A.z)
+    @assert length(b) == n + 1 "Dimension mismatch."
+
+    z = A.z
+    D = A.D
+    α = A.α
+
+    if any(iszero, A.D)
+        throw(DomainError("Matrix is singular"))
+    end
+    @inbounds @views begin
+        s = dot(z ./ D, b[1:n])
+        t = dot(z ./ D, z)
+        denom = α - t
+    
+        if denom == 0
+            throw(DomainError("Matrix is singular"))
+        end
+
+        y[n+1] = (b[n + 1] - s) / denom
+        y[1:n] .= (b[1:n] - z * y[n+1]) ./ D
+    end
+    return y
+end
+
+function Base.:\(A::ArrowheadMatrix, b::AbstractVector{T}) where T
+    y = similar(b)
+    return linsolve!(y, A, b)
+end
+
+function LinearAlgebra.ldiv!(x::AbstractVector{T}, A::ArrowheadMatrix, b::AbstractVector{T}) where T
+    return linsolve!(x, A, b)
+end
+
+struct InvArrowheadMatrix{O, T, Z, P} <: AbstractMatrix{O}
+    A::ArrowheadMatrix{O, T, Z, P}
+end
+
+inv(A::ArrowheadMatrix) = InvArrowheadMatrix(A)
+
+function size(A_inv::InvArrowheadMatrix)
+    size(A_inv.A)
+end
+
+function LinearAlgebra.mul!(y, A_inv::InvArrowheadMatrix{T}, x::AbstractVector{T}) where T
+    A = A_inv.A
+    return linsolve!(y, A, x)
+end
+
+function Base.:\(A_inv::InvArrowheadMatrix{T}, x::AbstractVector{T}) where T
+    A = A_inv.A
+    return A * x
+end
+
+function Base.convert(::Type{Matrix}, A_inv::InvArrowheadMatrix{T}) where T
+    A = A_inv.A
+    n = length(A.z)
+    z = A.z
+    D = A.D
+    α = A.α
+
+    # Compute t = dot(z ./ D, z)
+    t = dot(z ./ D, z)
+    denom = α - t
+    @assert denom != 0 "Matrix is singular."
+
+    # Compute u = [ (z ./ D); -1 ]
+    u = [ z ./ D; -1.0 ]
+
+    # Compute the inverse diagonal elements
+    D_inv = 1.0 ./ D
+
+    # Construct the inverse matrix
+    M = zeros(T, n + 1, n + 1)
+    M[1:n, 1:n] .= Diagonal(D_inv)
+    M .+= (u * u') / denom
+    return M
+end
\ No newline at end of file
diff --git a/test/algebra/algebrasetup_setuptests.jl b/test/algebra/algebrasetup_setuptests.jl
new file mode 100644
index 0000000..f616f8d
--- /dev/null
+++ b/test/algebra/algebrasetup_setuptests.jl
@@ -0,0 +1,2 @@
+using BenchmarkTools, LinearAlgebra
+using BayesBase
\ No newline at end of file
diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
new file mode 100644
index 0000000..a042799
--- /dev/null
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -0,0 +1,187 @@
+
+
+@testitem "ArrowheadMatrix: Construction and Properties" begin
+    include("algebrasetup_setuptests.jl")
+    α = 2.0
+    z = [1.0, 2.0, 3.0]
+    D = [4.0, 5.0, 6.0]
+    A = ArrowheadMatrix(α, z, D)
+    @test size(A) == (4, 4)
+end
+
+@testitem "ArrowheadMatrix: Multiplication with Vector" begin
+    include("algebrasetup_setuptests.jl")
+    for n in 2:2
+        α = randn()
+        z = randn(n)
+        D = randn(n)
+        A = ArrowheadMatrix(α, z, D)
+        
+        x = randn(n+1)
+        y = A * x
+        
+        dense_A = [Diagonal(D) z; z' α]
+        converted_A = convert(Matrix, A)
+        @test dense_A ≈ converted_A
+
+        y_expected = dense_A * x
+        @test y ≈ y_expected
+    end
+end
+
+@testitem "ArrowheadMatrix: Solving Linear System" begin
+    include("algebrasetup_setuptests.jl")
+    for n in 2:20
+        α = randn()^2 .+ 1
+        z = randn(n)
+        D = randn(n).^2 .+ 1
+        A = ArrowheadMatrix(α, z, D)
+        
+        x = randn(n+1)
+        y = A \ x
+        
+        dense_A = convert(Matrix, A)
+        y_expected = dense_A \ x
+
+        @test y ≈ y_expected
+    end
+end
+
+@testitem "ArrowheadMatrix: Handling Singular Matrices" begin
+    include("algebrasetup_setuptests.jl")
+
+    α = 0.0
+    z = [0.0, 0.0, 0.0]
+    D = [1.0, 5.0, 6.0]
+    A = ArrowheadMatrix(α, z, D)
+    
+    b = [7.0, 8.0, 9.0, 10.0]
+
+    @test_throws DomainError A \ b
+
+    α = 0.0
+    z = [1.0, 0.0, 0.0]
+    D = [0.0, 5.0, 6.0]
+    A = ArrowheadMatrix(α, z, D)
+    b = [7.0, 8.0, 9.0, 10.0]
+
+    @test_throws DomainError A \ b
+end
+
+
+@testitem "InvArrowheadMatrix: Construction and Properties" begin
+    include("algebrasetup_setuptests.jl")
+
+    α = 2.0
+    z = [1.0, 2.0, 3.0]
+    D = [4.0, 5.0, 6.0]
+    A = ArrowheadMatrix(α, z, D)
+    A_inv = inv(A)
+    @test size(A_inv) == (4, 4)
+end
+
+@testitem "InvArrowheadMatrix: Multiplication with Vector" begin
+    α = 2.0
+    z = [1.0, 2.0, 3.0]
+    D = [4.0, 5.0, 6.0]
+    A = ArrowheadMatrix(α, z, D)
+    b = [7.0, 8.0, 9.0, 10.0]
+    A_inv = inv(A)
+    b = [7.0, 8.0, 9.0, 10.0]
+    x = A_inv * b
+
+    x_expected = A \ b
+    @test x ≈ x_expected
+end
+
+@testitem "InvArrowheadMatrix: Division with Vector" begin
+    α = 2.0
+    z = [1.0, 2.0, 3.0]
+    D = [4.0, 5.0, 6.0]
+    A = ArrowheadMatrix(α, z, D)
+    A_inv = inv(A)
+    x = [1.0, 2.0, 3.0, 4.0]
+    y = A_inv \ x
+
+    y_expected = A * x
+    @test y == y_expected
+end
+
+@testitem "InvArrowheadMatrix: Conversion to Dense Matrix" begin
+    include("algebrasetup_setuptests.jl")
+
+    α = 2.0
+    z = [1.0, 2.0, 3.0]
+    D = [4.0, 5.0, 6.0]
+    A = ArrowheadMatrix(α, z, D)
+    A_inv = inv(A)
+
+    A_inv_dense = convert(Matrix, A_inv)
+    A_dense = convert(Matrix, A)
+    
+    # Verify that A_inv_dense * A_dense ≈ Identity matrix
+    I_approx = A_inv_dense * A_dense
+    I_n = Matrix{Float64}(I, size(A_dense))
+    @test I_approx ≈ I_n
+end
+
+
+@testitem "ArrowheadMatrix: division vs ldiv!" begin
+
+    include("algebrasetup_setuptests.jl")
+
+    for n in [10, 20]
+        α = rand()^2 + 1.0  # Ensure α is not too close to zero
+        z = randn(n)
+        D = rand(n).^2 .+ 1.0  # Ensure D elements are not too close to zero
+        A = ArrowheadMatrix(α, z, D)
+        
+        b = randn(n+1)
+        x1 = A \ b
+
+        x2 = similar(b)
+        LinearAlgebra.ldiv!(x2, A, b)
+        @test x1 ≈ x2
+        
+        allocs = @allocations LinearAlgebra.ldiv!(x2, A, b)
+        @test allocs == 0
+    end
+end
+
+@testitem "ArrowheadMatrix: Performance comparison with dense matrix" begin
+    using BenchmarkTools
+    include("algebrasetup_setuptests.jl")
+
+    for n in [10, 100, 1000]
+        α = rand()^2 + 1.0  # Ensure α is not too close to zero
+        z = randn(n)
+        D = rand(n).^2 .+ 1.0 # Ensure D elements are not too close to zero
+        A_arrow = ArrowheadMatrix(α, z, D)
+        
+        # Create equivalent dense matrix
+        A_dense = [Diagonal(D) z; z' α]
+        
+        b = randn(n+1)
+        
+        # Warm-up runs
+        _ = A_arrow \ b
+        _ = A_dense \ b
+        
+        # Benchmark ArrowheadMatrix division
+        time_arrow = @btime $A_arrow \ $b
+        allocs_arrow = @allocations A_arrow \ b
+        
+        # Benchmark dense matrix division
+        time_dense = @btime $A_dense \ $b
+        allocs_dense = @allocations A_dense \ b
+        
+        # Check if ArrowheadMatrix is more efficient
+        @test time_arrow < time_dense
+        @test allocs_arrow < allocs_dense
+        
+        # Check correctness
+        x_arrow = A_arrow \ b
+        x_dense = A_dense \ b
+        @test x_arrow ≈ x_dense
+    end
+end
\ No newline at end of file

From 4c1ec7fa45afda0cc2a56618e6ed9e403506004f Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 13:57:45 +0200
Subject: [PATCH 02/17] fix: remove LinearAlgebra from test deps

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 82afceb..e6c5244 100644
--- a/Project.toml
+++ b/Project.toml
@@ -44,4 +44,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 
 [targets]
-test = ["Aqua", "CpuId", "JET", "Test", "ReTestItems", "LinearAlgebra", "StableRNGs", "HCubature"]
+test = ["Aqua", "CpuId", "JET", "Test", "ReTestItems", "StableRNGs", "HCubature"]

From c97b255ec1b9b92087b9e3b57d04cf01d62ac685 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 13:58:36 +0200
Subject: [PATCH 03/17] fix: remove BC from compat

---
 Project.toml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index e6c5244..da2b7ee 100644
--- a/Project.toml
+++ b/Project.toml
@@ -18,7 +18,6 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 TinyHugeNumbers = "783c9a47-75a3-44ac-a16b-f1ab7b3acf04"
 
 [compat]
-BenchmarkTools = "1.5.0"
 Distributions = "0.25"
 DomainSets = "0.5.2, 0.6, 0.7"
 LinearAlgebra = "1.9"

From 4dfbdc5d8ebc764ad1b87b59ff37bfb8e386e6db Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 14:02:14 +0200
Subject: [PATCH 04/17] test(fix): add BenchmarkTools into test deps

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index da2b7ee..658ca6d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -43,4 +43,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 
 [targets]
-test = ["Aqua", "CpuId", "JET", "Test", "ReTestItems", "StableRNGs", "HCubature"]
+test = ["Aqua", "BenchmarkTools", "CpuId", "JET", "Test", "ReTestItems", "StableRNGs", "HCubature"]

From d2d167d061447d88826ba0932fbb742a51ebe42d Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 14:20:03 +0200
Subject: [PATCH 05/17] fix: stupid optimization

---
 src/algebra/arrowheadmatrix.jl | 34 ++++++++++++++++++++++------------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/src/algebra/arrowheadmatrix.jl b/src/algebra/arrowheadmatrix.jl
index 9b72ab4..df70f39 100644
--- a/src/algebra/arrowheadmatrix.jl
+++ b/src/algebra/arrowheadmatrix.jl
@@ -51,21 +51,31 @@ function linsolve!(y, A::ArrowheadMatrix, b::AbstractVector)
     D = A.D
     α = A.α
 
-    if any(iszero, A.D)
+    if any(iszero, D)
         throw(DomainError("Matrix is singular"))
     end
-    @inbounds @views begin
-        s = dot(z ./ D, b[1:n])
-        t = dot(z ./ D, z)
-        denom = α - t
-    
-        if denom == 0
-            throw(DomainError("Matrix is singular"))
-        end
-
-        y[n+1] = (b[n + 1] - s) / denom
-        y[1:n] .= (b[1:n] - z * y[n+1]) ./ D
+
+    s = zero(eltype(y))
+    @inbounds for i in 1:n
+        s += (z[i] / D[i]) * b[i]
+    end
+
+    t = zero(eltype(y))
+    @inbounds for i in 1:n
+        t += (z[i] / D[i]) * z[i]
     end
+
+    denom = α - t
+    if denom == 0
+        throw(DomainError("Matrix is singular"))
+    end
+
+    y[n+1] = (b[n + 1] - s) / denom
+
+    @inbounds for i in 1:n
+        y[i] = (b[i] - z[i] * y[n+1]) / D[i]
+    end
+
     return y
 end
 

From a37543e543523ce2e1b926e178a2aeeb9f0c4da1 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 14:30:49 +0200
Subject: [PATCH 06/17] optim: use @simd

---
 src/algebra/arrowheadmatrix.jl | 36 ++++++++++++++++++++++------------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/src/algebra/arrowheadmatrix.jl b/src/algebra/arrowheadmatrix.jl
index df70f39..4cc3ef7 100644
--- a/src/algebra/arrowheadmatrix.jl
+++ b/src/algebra/arrowheadmatrix.jl
@@ -43,26 +43,34 @@ function LinearAlgebra.mul!(y, A::ArrowheadMatrix{T}, x::AbstractVector{T}) wher
     return y
 end
 
-function linsolve!(y, A::ArrowheadMatrix, b::AbstractVector)
+function linsolve!(y::AbstractVector{T}, A::ArrowheadMatrix{T}, b::AbstractVector{T}) where T
     n = length(A.z)
     @assert length(b) == n + 1 "Dimension mismatch."
+    @assert length(y) == n + 1 "Dimension mismatch."
 
     z = A.z
     D = A.D
     α = A.α
 
-    if any(iszero, D)
-        throw(DomainError("Matrix is singular"))
-    end
-
-    s = zero(eltype(y))
+    # Check for zeros in D to avoid division by zero
     @inbounds for i in 1:n
-        s += (z[i] / D[i]) * b[i]
+        if D[i] == 0
+            throw(DomainError("Matrix is singular"))
+        end
     end
 
-    t = zero(eltype(y))
-    @inbounds for i in 1:n
-        t += (z[i] / D[i]) * z[i]
+    s = zero(T)
+    t = zero(T)
+
+    # Compute s and t in a single loop to avoid recomputing z[i] / D[i]
+    @inbounds @simd for i in 1:n
+        zi = z[i]
+        Di = D[i]
+        z_div_D = zi / Di
+        bi = b[i]
+
+        s += z_div_D * bi       # Accumulate s
+        t += z_div_D * zi       # Accumulate t
     end
 
     denom = α - t
@@ -70,10 +78,12 @@ function linsolve!(y, A::ArrowheadMatrix, b::AbstractVector)
         throw(DomainError("Matrix is singular"))
     end
 
-    y[n+1] = (b[n + 1] - s) / denom
+    yn1 = (b[n + 1] - s) / denom
+    y[n + 1] = yn1
 
-    @inbounds for i in 1:n
-        y[i] = (b[i] - z[i] * y[n+1]) / D[i]
+    # Compute y[1:n]
+    @inbounds @simd for i in 1:n
+        y[i] = (b[i] - z[i] * yn1) / D[i]
     end
 
     return y

From a17a7b74c074201ef6500f442cd9112229d05138 Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Tue, 8 Oct 2024 14:45:06 +0200
Subject: [PATCH 07/17] Add exception for sampletype for Gamma

---
 src/promotion.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/promotion.jl b/src/promotion.jl
index 1a66af5..9db9a17 100644
--- a/src/promotion.jl
+++ b/src/promotion.jl
@@ -156,6 +156,9 @@ sampletype(::Type{Univariate}, distribution) = eltype(distribution)
 sampletype(::Type{Multivariate}, distribution) = Vector{eltype(distribution)}
 sampletype(::Type{Matrixvariate}, distribution) = Matrix{eltype(distribution)}
 
+# Exceptions
+sampletype(::Gamma{T}) where {T} = T
+
 """
     samplefloattype(distribution)
 

From 99527efa48ac5226d2d67ed312d2e3f3d9fe4d50 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 14:54:04 +0200
Subject: [PATCH 08/17] fix: add docs to ArrowheadMatrix

---
 src/algebra/arrowheadmatrix.jl | 100 ++++++++++++++++++++++++++++++++-
 1 file changed, 99 insertions(+), 1 deletion(-)

diff --git a/src/algebra/arrowheadmatrix.jl b/src/algebra/arrowheadmatrix.jl
index 4cc3ef7..2ee9009 100644
--- a/src/algebra/arrowheadmatrix.jl
+++ b/src/algebra/arrowheadmatrix.jl
@@ -3,7 +3,51 @@ export ArrowheadMatrix, InvArrowheadMatrix
 import Base: getindex
 import LinearAlgebra: mul!
 import Base: size, *, \, inv, convert, Matrix
-
+"""
+    ArrowheadMatrix{O, T, Z, P} <: AbstractMatrix{O}
+
+A structure representing an arrowhead matrix, which is a special type of sparse matrix.
+
+# Fields
+- `α::T`: The scalar value at the bottom-right corner of the matrix.
+- `z::Z`: A vector representing the last row/column (excluding the corner element).
+- `D::P`: A vector representing the diagonal elements (excluding the corner element).
+
+# Constructors
+    ArrowheadMatrix(a::T, z::Z, d::D) where {T,Z,D}
+
+Constructs an `ArrowheadMatrix` with the given α, z, and D values. The output type `O`
+is automatically determined as the promoted type of all input elements.
+
+# Operations
+- Matrix-vector multiplication: `A * x` or `mul!(y, A, x)`
+- Linear system solving: `A \\ b` or `ldiv!(x, A, b)`
+- Conversion to dense matrix: `convert(Matrix, A)`
+- Inversion: `inv(A)` (returns an `InvArrowheadMatrix`)
+
+# Examples
+```julia
+α = 2.0
+z = [1.0, 2.0, 3.0]
+D = [4.0, 5.0, 6.0]
+A = ArrowheadMatrix(α, z, D)
+
+# Matrix-vector multiplication
+x = [1.0, 2.0, 3.0, 4.0]
+y = A * x
+
+# Solving linear system
+b = [7.0, 8.0, 9.0, 10.0]
+x = A \\ b
+
+# Convert to dense matrix
+dense_A = convert(Matrix, A)
+```
+
+# Notes
+- The matrix is singular if α - dot(z ./ D, z) = 0 or if any element of D is zero.
+- For best performance, use `ldiv!` for solving linear systems when possible.
+"""
 struct ArrowheadMatrix{O, T, Z, P} <: AbstractMatrix{O}
     α::T         
     z::Z         
@@ -98,6 +142,60 @@ function LinearAlgebra.ldiv!(x::AbstractVector{T}, A::ArrowheadMatrix, b::Abstra
     return linsolve!(x, A, b)
 end
 
+"""
+    InvArrowheadMatrix{O, T, Z, P} <: AbstractMatrix{O}
+
+A wrapper structure representing the inverse of an `ArrowheadMatrix`.
+
+This structure doesn't explicitly compute or store the inverse matrix.
+Instead, it stores a reference to the original `ArrowheadMatrix` and
+implements efficient operations that leverage the special structure
+of the arrowhead matrix.
+
+# Fields
+- `A::ArrowheadMatrix{O, T, Z, P}`: The original `ArrowheadMatrix` being inverted.
+
+# Constructors
+    InvArrowheadMatrix(A::ArrowheadMatrix{O, T, Z, P})
+
+Constructs an `InvArrowheadMatrix` by wrapping the given `ArrowheadMatrix`.
+
+# Operations
+- Matrix-vector multiplication: `A_inv * x` or `mul!(y, A_inv, x)`
+  (Equivalent to solving the system A * y = x)
+- Linear system solving: `A_inv \\ x`
+  (Equivalent to multiplication by the original matrix: A * x)
+- Conversion to dense matrix: `convert(Matrix, A_inv)`
+  (Computes and returns the actual inverse as a dense matrix)
+
+# Examples
+```julia
+α = 2.0
+z = [1.0, 2.0, 3.0]
+D = [4.0, 5.0, 6.0]
+A = ArrowheadMatrix(α, z, D)
+A_inv = inv(A)  # Returns an InvArrowheadMatrix
+
+# Multiplication (equivalent to solving A * y = x)
+x = [1.0, 2.0, 3.0, 4.0]
+y = A_inv * x
+
+# Division (equivalent to multiplying by A)
+b = [5.0, 6.0, 7.0, 8.0]
+x = A_inv \\ b
+
+# Convert to dense inverse matrix
+dense_inv_A = convert(Matrix, A_inv)
+```
+
+# Notes
+- The inverse exists only if the original `ArrowheadMatrix` is non-singular.
+- Operations with `InvArrowheadMatrix` do not explicitly compute the inverse,
+  but instead solve the corresponding system with the original matrix.
+
+# See Also
+- [`ArrowheadMatrix`](@ref): The original arrowhead matrix structure.
+"""
 struct InvArrowheadMatrix{O, T, Z, P} <: AbstractMatrix{O}
     A::ArrowheadMatrix{O, T, Z, P}
 end

From 0f10b085c0bf414ec88458430fb41ef69eaa9be5 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 15:14:41 +0200
Subject: [PATCH 09/17] fix: check that ArrowheadMatrix use quadratically less
 memory

---
 test/algebra/arrowheadmatrix_tests.jl | 35 +++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index a042799..b8ae366 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -184,4 +184,39 @@ end
         x_dense = A_dense \ b
         @test x_arrow ≈ x_dense
     end
+end
+
+@testitem "ArrowheadMatrix: Memory allocation comparison with dense matrix" begin
+    using Test
+    include("algebrasetup_setuptests.jl")
+
+    function memory_size(x)
+        return Base.summarysize(x)
+    end
+
+    sizes = [10, 100, 1000, 10000]
+    arrow_mem = zeros(Int, length(sizes))
+    dense_mem = zeros(Int, length(sizes))
+
+    for (i, n) in enumerate(sizes)
+        α = rand()^2 + 1.0
+        z = randn(n)
+        D = rand(n).^2 .+ 1.0
+
+        A_arrow = ArrowheadMatrix(α, z, D)
+        A_dense = [Diagonal(D) z; z' α]
+
+
+        arrow_mem[i] = memory_size(A_arrow)
+        dense_mem[i] = memory_size(A_dense)
+        @test arrow_mem[i] < dense_mem[i]
+    end
+
+    mem_ratio = dense_mem ./ arrow_mem
+
+    for i in 2:length(sizes)
+        ratio_growth = mem_ratio[i] / mem_ratio[i-1]
+        size_growth = sizes[i] / sizes[i-1]
+        @test isapprox(ratio_growth, size_growth, rtol=0.5)
+    end
 end
\ No newline at end of file

From 95021c96905eaa4983c1a25866311678e07a3fa2 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 8 Oct 2024 16:12:53 +0200
Subject: [PATCH 10/17] docs: add bindings for Arrowhead docs

---
 docs/src/index.md | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/docs/src/index.md b/docs/src/index.md
index 3fc733e..b725ea6 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -100,6 +100,12 @@ BayesBase.weightedmean_cov
 BayesBase.weightedmean_invcov
 ```
 
+## [Extra matrix structures](@id matrix-structures)
+```@docs
+BayesBase.ArrowheadMatrix
+BayesBase.InvArrowheadMatrix
+```
+
 ## [Helper utilities](@id library-helpers)
 
 ```@docs

From 73607a66012f7608fbc81181762c3fe5524172e8 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Wed, 9 Oct 2024 14:54:20 +0200
Subject: [PATCH 11/17] test: @btime -> @benchmark

---
 test/algebra/arrowheadmatrix_tests.jl | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index b8ae366..9f2c5ea 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -163,23 +163,18 @@ end
         
         b = randn(n+1)
         
-        # Warm-up runs
         _ = A_arrow \ b
         _ = A_dense \ b
         
-        # Benchmark ArrowheadMatrix division
-        time_arrow = @btime $A_arrow \ $b
+        time_arrow = @benchmark $A_arrow \ $b;
         allocs_arrow = @allocations A_arrow \ b
         
-        # Benchmark dense matrix division
-        time_dense = @btime $A_dense \ $b
+        time_dense = @benchmark $A_dense \ $b;
         allocs_dense = @allocations A_dense \ b
         
-        # Check if ArrowheadMatrix is more efficient
-        @test time_arrow < time_dense
+        @test minimum(time_arrow.times) < minimum(time_dense.times)/n
         @test allocs_arrow < allocs_dense
         
-        # Check correctness
         x_arrow = A_arrow \ b
         x_dense = A_dense \ b
         @test x_arrow ≈ x_dense

From 01989489e7730a90c56d46aea5597382700f16e7 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Wed, 9 Oct 2024 16:31:33 +0200
Subject: [PATCH 12/17] feat: add show methods for ArrowheadMatrix

---
 Project.toml                          | 11 ++++++++--
 src/algebra/arrowheadmatrix.jl        | 31 +++++++++++++++++++++++++++
 test/algebra/arrowheadmatrix_tests.jl |  4 +++-
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 658ca6d..a394ca7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -17,9 +17,16 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 TinyHugeNumbers = "783c9a47-75a3-44ac-a16b-f1ab7b3acf04"
 
+[weakdeps]
+FastCholesky = "2d5283b6-8564-42b6-bb00-83ed8e915756"
+
+[extensions]
+FastCholeskyExt = "FastCholesky"
+
 [compat]
 Distributions = "0.25"
 DomainSets = "0.5.2, 0.6, 0.7"
+FastCholesky = "1.3.1"
 LinearAlgebra = "1.9"
 LoopVectorization = "0.12"
 Random = "1.9"
@@ -34,13 +41,13 @@ julia = "1.9"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CpuId = "adafc99b-e345-5852-983c-f28acb93d879"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 
 [targets]
-test = ["Aqua", "BenchmarkTools", "CpuId", "JET", "Test", "ReTestItems", "StableRNGs", "HCubature"]
+test = ["Aqua", "BenchmarkTools", "CpuId", "FastCholesky", "JET", "Test", "ReTestItems", "StableRNGs", "HCubature"]
diff --git a/src/algebra/arrowheadmatrix.jl b/src/algebra/arrowheadmatrix.jl
index 2ee9009..17000cb 100644
--- a/src/algebra/arrowheadmatrix.jl
+++ b/src/algebra/arrowheadmatrix.jl
@@ -58,6 +58,29 @@ function ArrowheadMatrix(a::T, z::Z, d::D) where {T,Z,D}
     return ArrowheadMatrix{O, T, Z, D}(a, z, d)
 end
 
+function show(io::IO, ::MIME"text/plain", A::ArrowheadMatrix)
+    n = length(A.D) + 1
+    println(io, n, "×", n, " ArrowheadMatrix{", eltype(A), "}:")
+    
+    for i in 1:n-1
+        for j in 1:n-1
+            if i == j
+                print(io, A.D[i])
+            else
+                print(io, "⋅")
+            end
+            print(io, "  ")
+        end
+        println(io, A.z[i])
+    end
+    
+    # Print the last row
+    for i in 1:n-1
+        print(io, A.z[i], "  ")
+    end
+    println(io, A.α)
+end
+
 function size(A::ArrowheadMatrix)
     n = length(A.D) + 1
     return (n, n)
@@ -200,6 +223,14 @@ struct InvArrowheadMatrix{O, T, Z, P} <: AbstractMatrix{O}
     A::ArrowheadMatrix{O, T, Z, P}
 end
 
+function show(io::IO, ::MIME"text/plain", A_inv::InvArrowheadMatrix)
+    n = size(A_inv.A, 1)
+    println(io, n, "×", n, " InvArrowheadMatrix{", eltype(A_inv), "}:")
+    println(io, "Inverse of:")
+    show(io, MIME"text/plain"(), A_inv.A)
+end
+
+
 inv(A::ArrowheadMatrix) = InvArrowheadMatrix(A)
 
 function size(A_inv::InvArrowheadMatrix)
diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index 9f2c5ea..1d2b004 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -11,7 +11,7 @@ end
 
 @testitem "ArrowheadMatrix: Multiplication with Vector" begin
     include("algebrasetup_setuptests.jl")
-    for n in 2:2
+    for n in 2:20
         α = randn()
         z = randn(n)
         D = randn(n)
@@ -163,6 +163,7 @@ end
         
         b = randn(n+1)
         
+        # warm-up runs
         _ = A_arrow \ b
         _ = A_dense \ b
         
@@ -172,6 +173,7 @@ end
         time_dense = @benchmark $A_dense \ $b;
         allocs_dense = @allocations A_dense \ b
         
+        # ours at least n times faster where n is dimensionality
         @test minimum(time_arrow.times) < minimum(time_dense.times)/n
         @test allocs_arrow < allocs_dense
         

From c270d1ff315b1ea0433fa12e3441e45983084a04 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Thu, 10 Oct 2024 09:44:42 +0200
Subject: [PATCH 13/17] feat: add cholinv interface for ArrowheadMatrix

---
 ext/FastCholeskyExt.jl                | 10 ++++++++
 test/algebra/arrowheadmatrix_tests.jl | 37 +++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 ext/FastCholeskyExt.jl

diff --git a/ext/FastCholeskyExt.jl b/ext/FastCholeskyExt.jl
new file mode 100644
index 0000000..10af35e
--- /dev/null
+++ b/ext/FastCholeskyExt.jl
@@ -0,0 +1,10 @@
+module FastCholeskyExt
+
+    using FastCholesky
+    using BayesBase
+
+    function FastCholesky.cholinv(input::ArrowheadMatrix)
+        return inv(input)
+    end
+
+end
\ No newline at end of file
diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index 1d2b004..02f9bdc 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -183,6 +183,43 @@ end
     end
 end
 
+
+@testitem "ArrowheadMatrix: Performance comparison with dense matrix" begin
+    using BenchmarkTools
+    using FastCholesky
+    include("algebrasetup_setuptests.jl")
+
+    for n in [10, 100, 1000]
+        α = rand()^2 + 1.0  # Ensure α is not too close to zero
+        z = randn(n)
+        D = rand(n).^2 .+ 1.0 # Ensure D elements are not too close to zero
+        A_arrow = ArrowheadMatrix(α, z, D)
+        
+        # Create equivalent dense matrix
+        A_dense = [Diagonal(D) z; z' α]
+        
+        b = randn(n+1)
+        
+        # warm-up runs
+        _ = cholinv(A_arrow) \ b
+        _ = cholinv(A_dense) \ b
+        
+        time_arrow = @benchmark cholinv($A_arrow) * $b;
+        allocs_arrow = @allocations cholinv(A_arrow) * b
+        
+        time_dense = @benchmark cholinv($A_dense) * $b;
+        allocs_dense = @allocations cholinv(A_dense) * b
+        
+        # ours at least n times faster where n is dimensionality
+        @test minimum(time_arrow.times) < minimum(time_dense.times)/n
+        @test allocs_arrow < allocs_dense
+        
+        x_arrow = A_arrow \ b
+        x_dense = A_dense \ b
+        @test x_arrow ≈ x_dense
+    end
+end
+
 @testitem "ArrowheadMatrix: Memory allocation comparison with dense matrix" begin
     using Test
     include("algebrasetup_setuptests.jl")

From 033d708a9b35868999927264209919929756fa6f Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Thu, 10 Oct 2024 09:47:17 +0200
Subject: [PATCH 14/17] test: rename

---
 test/algebra/arrowheadmatrix_tests.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index 02f9bdc..6264ee7 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -208,7 +208,7 @@ end
         allocs_arrow = @allocations cholinv(A_arrow) * b
         
         time_dense = @benchmark cholinv($A_dense) * $b;
-        allocs_dense = @allocations cholinv(A_dense) * b
+        allocs_dense = @allocations cholinv(A_dense)  b
         
         # ours at least n times faster where n is dimensionality
         @test minimum(time_arrow.times) < minimum(time_dense.times)/n
@@ -220,7 +220,7 @@ end
     end
 end
 
-@testitem "ArrowheadMatrix: Memory allocation comparison with dense matrix" begin
+@testitem "ArrowheadMatrix: Memory allocation comparison with cholinv" begin
     using Test
     include("algebrasetup_setuptests.jl")
 

From aa20c2789b36abc8d94c3f925f8f8d502a733404 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Thu, 10 Oct 2024 09:50:41 +0200
Subject: [PATCH 15/17] test: typo

---
 test/algebra/arrowheadmatrix_tests.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index 6264ee7..5c084d5 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -184,7 +184,7 @@ end
 end
 
 
-@testitem "ArrowheadMatrix: Performance comparison with dense matrix" begin
+@testitem "ArrowheadMatrix: Performance comparison with cholinv" begin
     using BenchmarkTools
     using FastCholesky
     include("algebrasetup_setuptests.jl")
@@ -208,7 +208,7 @@ end
         allocs_arrow = @allocations cholinv(A_arrow) * b
         
         time_dense = @benchmark cholinv($A_dense) * $b;
-        allocs_dense = @allocations cholinv(A_dense)  b
+        allocs_dense = @allocations cholinv(A_dense) * b
         
         # ours at least n times faster where n is dimensionality
         @test minimum(time_arrow.times) < minimum(time_dense.times)/n
@@ -220,7 +220,7 @@ end
     end
 end
 
-@testitem "ArrowheadMatrix: Memory allocation comparison with cholinv" begin
+@testitem "ArrowheadMatrix: Memory allocation comparison with dense matrix" begin
     using Test
     include("algebrasetup_setuptests.jl")
 

From c5a0685c927950438964b8f1b5d5852c974bbf03 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Thu, 10 Oct 2024 10:23:35 +0200
Subject: [PATCH 16/17] fix: use the same errors that LinearAlgebra

---
 src/algebra/arrowheadmatrix.jl        |  12 ++-
 test/algebra/arrowheadmatrix_tests.jl | 141 +++++++++++++++++++++-----
 2 files changed, 126 insertions(+), 27 deletions(-)

diff --git a/src/algebra/arrowheadmatrix.jl b/src/algebra/arrowheadmatrix.jl
index 17000cb..8628bc9 100644
--- a/src/algebra/arrowheadmatrix.jl
+++ b/src/algebra/arrowheadmatrix.jl
@@ -1,5 +1,7 @@
 export ArrowheadMatrix, InvArrowheadMatrix
 
+
+import LinearAlgebra: SingularException
 import Base: getindex
 import LinearAlgebra: mul!
 import Base: size, *, \, inv, convert, Matrix
@@ -112,8 +114,10 @@ end
 
 function linsolve!(y::AbstractVector{T}, A::ArrowheadMatrix{T}, b::AbstractVector{T}) where T
     n = length(A.z)
-    @assert length(b) == n + 1 "Dimension mismatch."
-    @assert length(y) == n + 1 "Dimension mismatch."
+
+    if length(b) != n + 1
+        throw(DimensionMismatch(1))
+    end
 
     z = A.z
     D = A.D
@@ -122,7 +126,7 @@ function linsolve!(y::AbstractVector{T}, A::ArrowheadMatrix{T}, b::AbstractVecto
     # Check for zeros in D to avoid division by zero
     @inbounds for i in 1:n
         if D[i] == 0
-            throw(DomainError("Matrix is singular"))
+            throw(SingularException(1))
         end
     end
 
@@ -142,7 +146,7 @@ function linsolve!(y::AbstractVector{T}, A::ArrowheadMatrix{T}, b::AbstractVecto
 
     denom = α - t
     if denom == 0
-        throw(DomainError("Matrix is singular"))
+        throw(SingularException("matrix is singular"))
     end
 
     yn1 = (b[n + 1] - s) / denom
diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index 5c084d5..e600408 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -47,28 +47,6 @@ end
     end
 end
 
-@testitem "ArrowheadMatrix: Handling Singular Matrices" begin
-    include("algebrasetup_setuptests.jl")
-
-    α = 0.0
-    z = [0.0, 0.0, 0.0]
-    D = [1.0, 5.0, 6.0]
-    A = ArrowheadMatrix(α, z, D)
-    
-    b = [7.0, 8.0, 9.0, 10.0]
-
-    @test_throws DomainError A \ b
-
-    α = 0.0
-    z = [1.0, 0.0, 0.0]
-    D = [0.0, 5.0, 6.0]
-    A = ArrowheadMatrix(α, z, D)
-    b = [7.0, 8.0, 9.0, 10.0]
-
-    @test_throws DomainError A \ b
-end
-
-
 @testitem "InvArrowheadMatrix: Construction and Properties" begin
     include("algebrasetup_setuptests.jl")
 
@@ -253,4 +231,121 @@ end
         size_growth = sizes[i] / sizes[i-1]
         @test isapprox(ratio_growth, size_growth, rtol=0.5)
     end
-end
\ No newline at end of file
+end
+
+@testitem "ArrowheadMatrix: Error handling comparison with dense matrix" begin
+    include("algebrasetup_setuptests.jl")
+
+    function test_error_consistency(A_arrow, A_dense, operation)
+        arrow_error = nothing
+        dense_error = nothing
+        
+        try
+            operation(A_arrow)
+        catch e
+            arrow_error = e
+        end
+
+        try
+            operation(A_dense)
+        catch e
+            dense_error = e
+        end
+
+        if isnothing(arrow_error) && isnothing(dense_error)
+            @test true  # Both succeeded, no error
+        elseif !isnothing(arrow_error) && !isnothing(dense_error)
+            @test typeof(arrow_error) == typeof(dense_error)  # Same error type
+        else
+            @test false  # One threw an error while the other didn't
+        end
+    end
+
+    for n in [3, 10]
+        α = randn()
+        z = randn(n)
+        D = randn(n)
+        A_arrow = ArrowheadMatrix(α, z, D)
+        A_dense = [Diagonal(D) z; z' α]
+
+        # Test invalid dimension for multiplication
+        invalid_vector = randn(n+2)
+        test_error_consistency(A_arrow, A_dense, A -> A * invalid_vector)
+
+        # Test multiplication with matrix of incorrect size
+        invalid_matrix = randn(n+2, n)
+        test_error_consistency(A_arrow, A_dense, A -> A * invalid_matrix)
+
+        # Test singularity in linear solve
+        singular_α = 0.0
+        singular_z = zeros(n)
+        singular_D = vcat(0.0, ones(n-1))
+        A_arrow_singular = ArrowheadMatrix(singular_α, singular_z, singular_D)
+        A_dense_singular = [Diagonal(singular_D) singular_z; singular_z' singular_α]
+        b = randn(n+1)
+        test_error_consistency(A_arrow_singular, A_dense_singular, A -> A \ b)
+
+        # Test linear solve with vector of incorrect size
+        invalid_b = randn(n+2)
+        test_error_consistency(A_arrow, A_dense, A -> A \ invalid_b)
+    end
+end
+# @testitem "ArrowheadMatrix: Error handling comparison with dense matrix for \ and * operations" begin
+#     include("algebrasetup_setuptests.jl")
+#     using LinearAlgebra
+
+#     function test_error_consistency(A_arrow, A_dense, operation)
+#         arrow_error = nothing
+#         dense_error = nothing
+        
+#         try
+#             operation(A_arrow)
+#         catch e
+#             arrow_error = e
+#         end
+
+#         try
+#             operation(A_dense)
+#         catch e
+#             dense_error = e
+#         end
+
+#         if isnothing(arrow_error) && isnothing(dense_error)
+#             @test true  # Both succeeded, no error
+#         elseif !isnothing(arrow_error) && !isnothing(dense_error)
+#             @test typeof(arrow_error) == typeof(dense_error)  # Same error type
+#         else
+#             @test false  # One threw an error while the other didn't
+#         end
+#     end
+
+#     # Test cases
+#     for n in [3, 10]
+#         α = randn()
+#         z = randn(n)
+#         D = randn(n)
+#         A_arrow = ArrowheadMatrix(α, z, D)
+#         A_dense = [Diagonal(D) z; z' α]
+
+#         # Test invalid dimension for multiplication
+#         invalid_vector = randn(n+2)
+#         test_error_consistency(A_arrow, A_dense, A -> A * invalid_vector)
+
+#         # Test multiplication with matrix of incorrect size
+#         invalid_matrix = randn(n+2, n)
+#         test_error_consistency(A_arrow, A_dense, A -> A * invalid_matrix)
+
+#         # Test singularity in linear solve
+#         singular_α = 0.0
+#         singular_z = zeros(n)
+#         singular_D = vcat(0.0, ones(n-1))
+#         A_arrow_singular = ArrowheadMatrix(singular_α, singular_z, singular_D)
+#         A_dense_singular = [Diagonal(singular_D) singular_z; singular_z' singular_α]
+#         b = randn(n+1)
+#         test_error_consistency(A_arrow_singular, A_dense_singular, A -> A \ b)
+
+#         # Test linear solve with vector of incorrect size
+#         invalid_b = randn(n+2)
+#         test_error_consistency(A_arrow, A_dense, A -> A \ invalid_b)
+#     end
+# end
\ No newline at end of file

From 80c29f4cb33f3d96bf20c4848b84c7f7b8a0d7ac Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Thu, 10 Oct 2024 10:35:19 +0200
Subject: [PATCH 17/17] test(fix): correct DimensionMismatch

---
 src/algebra/arrowheadmatrix.jl        |  6 +--
 test/algebra/arrowheadmatrix_tests.jl | 59 ---------------------------
 2 files changed, 3 insertions(+), 62 deletions(-)

diff --git a/src/algebra/arrowheadmatrix.jl b/src/algebra/arrowheadmatrix.jl
index 8628bc9..e5c276d 100644
--- a/src/algebra/arrowheadmatrix.jl
+++ b/src/algebra/arrowheadmatrix.jl
@@ -112,11 +112,11 @@ function LinearAlgebra.mul!(y, A::ArrowheadMatrix{T}, x::AbstractVector{T}) wher
     return y
 end
 
-function linsolve!(y::AbstractVector{T}, A::ArrowheadMatrix{T}, b::AbstractVector{T}) where T
+function linsolve!(y::AbstractVector{T2}, A::ArrowheadMatrix{T}, b::AbstractVector{T2}) where {T, T2}
     n = length(A.z)
 
     if length(b) != n + 1
-        throw(DimensionMismatch(1))
+        throw(DimensionMismatch())
     end
 
     z = A.z
@@ -146,7 +146,7 @@ function linsolve!(y::AbstractVector{T}, A::ArrowheadMatrix{T}, b::AbstractVecto
 
     denom = α - t
     if denom == 0
-        throw(SingularException("matrix is singular"))
+        throw(SingularException(1))
     end
 
     yn1 = (b[n + 1] - s) / denom
diff --git a/test/algebra/arrowheadmatrix_tests.jl b/test/algebra/arrowheadmatrix_tests.jl
index e600408..c8b629b 100644
--- a/test/algebra/arrowheadmatrix_tests.jl
+++ b/test/algebra/arrowheadmatrix_tests.jl
@@ -290,62 +290,3 @@ end
         test_error_consistency(A_arrow, A_dense, A -> A \ invalid_b)
     end
 end
-# @testitem "ArrowheadMatrix: Error handling comparison with dense matrix for \ and * operations" begin
-#     include("algebrasetup_setuptests.jl")
-#     using LinearAlgebra
-
-#     function test_error_consistency(A_arrow, A_dense, operation)
-#         arrow_error = nothing
-#         dense_error = nothing
-        
-#         try
-#             operation(A_arrow)
-#         catch e
-#             arrow_error = e
-#         end
-
-#         try
-#             operation(A_dense)
-#         catch e
-#             dense_error = e
-#         end
-
-#         if isnothing(arrow_error) && isnothing(dense_error)
-#             @test true  # Both succeeded, no error
-#         elseif !isnothing(arrow_error) && !isnothing(dense_error)
-#             @test typeof(arrow_error) == typeof(dense_error)  # Same error type
-#         else
-#             @test false  # One threw an error while the other didn't
-#         end
-#     end
-
-#     # Test cases
-#     for n in [3, 10]
-#         α = randn()
-#         z = randn(n)
-#         D = randn(n)
-#         A_arrow = ArrowheadMatrix(α, z, D)
-#         A_dense = [Diagonal(D) z; z' α]
-
-#         # Test invalid dimension for multiplication
-#         invalid_vector = randn(n+2)
-#         test_error_consistency(A_arrow, A_dense, A -> A * invalid_vector)
-
-#         # Test multiplication with matrix of incorrect size
-#         invalid_matrix = randn(n+2, n)
-#         test_error_consistency(A_arrow, A_dense, A -> A * invalid_matrix)
-
-#         # Test singularity in linear solve
-#         singular_α = 0.0
-#         singular_z = zeros(n)
-#         singular_D = vcat(0.0, ones(n-1))
-#         A_arrow_singular = ArrowheadMatrix(singular_α, singular_z, singular_D)
-#         A_dense_singular = [Diagonal(singular_D) singular_z; singular_z' singular_α]
-#         b = randn(n+1)
-#         test_error_consistency(A_arrow_singular, A_dense_singular, A -> A \ b)
-
-#         # Test linear solve with vector of incorrect size
-#         invalid_b = randn(n+2)
-#         test_error_consistency(A_arrow, A_dense, A -> A \ invalid_b)
-#     end
-# end
\ No newline at end of file