Fix v_a_vT and StandardBasisVector

src/helpers/algebra/common.jl

@@ -131,38 +131,14 @@ end
 """
  v_a_vT(v, a)
 
-Computes v*a*v^T with a single allocation.
+Computes v*a*v^T efficiently.
 """
-function v_a_vT(v::AbstractVector, a::Real)
- T = promote_type(eltype(v), typeof(a))
- result = zeros(T, length(v), length(v))
- mul!(result, v, v', a, one(T))
- return result
-end
-
-function v_a_vT(v, a)
- result = v * v'
- result *= a
- return result
-end
-
-"""
- v_a_vT(v1, a, v2)
+v_a_vT(v, a) = v * a * v'
+v_a_vT(v1, a, v2) = v1 * a * v2'
 
-Computes v1*a*v2^T with a single allocation.
-"""
-function v_a_vT(v1::AbstractVector, a::Real, v2::AbstractVector)
- T = promote_type(eltype(v1), typeof(a), eltype(v2))
- result = zeros(T, length(v1), length(v2))
- mul!(result, v1, v2', a, one(T))
- return result
-end
-
-function v_a_vT(v1, a, v2)
- result = v1 * v2'
- result *= a
- return result
-end
+# More efficient if `a` is a `Real`
+v_a_vT(v::AbstractVector, a::Real) = v * v' * a
+v_a_vT(v1::AbstractVector, a::Real, v2::AbstractVector) = v1 * v2' * a
 
 """
  mvbeta(x)

src/helpers/algebra/standard_basis_vector.jl

@@ -227,9 +227,8 @@ function v_a_vT(e::StandardBasisVector{T1}, a::T2) where {T1 <: Real, T2 <: Real
  N = length(e)
  I = getind(e)
  T = promote_type(T1, T2)
- Y = zeros(T, N, N)
- Y[I, I] = e.scale * a * e.scale
-
+ Y = zeros(T, N)
+ Y[I] = e.scale * a * e.scale
  # return output 
- return Y
+ return Diagonal(Y)
 end

test/algebra/test_common.jl

@@ -84,6 +84,28 @@ using LinearAlgebra
  @test ReactiveMP.mul_trace(a, b) ≈ a * b
  end
  end
+
+ @testset "v_a_vT" begin
+ import ReactiveMP: v_a_vT
+
+ for n in (1, 2, 3), v1 in [rand(n) for _ in 1:5], a in rand(5)
+ @test v_a_vT(v1, a) ≈ v1 * a * v1'
+ @test v_a_vT(v1, a) ≈ v1 * v1' * a
+ end
+
+ for n in (1, 2, 3), v1 in [rand(n) for _ in 1:5], v2 in [rand(n) for _ in 1:5], a in rand(5)
+ @test v_a_vT(v1, a, v2) ≈ v1 * a * v2'
+ @test v_a_vT(v1, a, v2) ≈ v1 * v2' * a
+ end
+
+ for n in (1, 2, 3), v in [rand(1, n) for _ in 1:5], a in [rand(n, n) for _ in 1:5]
+ @test v_a_vT(v, a) ≈ v * a * v'
+ end
+
+ for n in (1, 2, 3), v1 in [rand(1, n) for _ in 1:5], v2 in [rand(1, n) for _ in 1:5], a in [rand(n, n) for _ in 1:5]
+ @test v_a_vT(v1, a, v2) ≈ v1 * a * v2'
+ end
+ end
 end
 
 end

test/algebra/test_standard_basis_vector.jl

@@ -112,6 +112,23 @@ using LinearAlgebra
  end
  end
  end
+
+ @testset begin
+ import ReactiveMP: v_a_vT
+
+ for i in 2:5, j in 1:i, scale in (1.0, 2.0), a in rand(5)
+ e1 = StandardBasisVector(i, j, scale)
+ v1 = collect(e1)
+ @test v_a_vT(e1, a) ≈ v_a_vT(v1, a)
+
+ e2 = StandardBasisVector(i, i - j + 1, scale)
+ v2 = collect(e2)
+
+ @test v_a_vT(e1, a, e2) ≈ v_a_vT(v1, a, v2)
+ @test v_a_vT(e1, a, v2) ≈ v_a_vT(v1, a, v2)
+ @test v_a_vT(v1, a, e2) ≈ v_a_vT(v1, a, v2)
+ end
+ end
 end
 
 end