Overload LinearAlgebra.diagm for vectors

Author: asterycs

src/ricci.jl

@@ -781,6 +781,22 @@ function Base.adjoint(arg::Union{Variable,Literal,KrD,Zero})
     return evaluate(e)
 end
 
+function LinearAlgebra.diagm(v::Tensor)
+    indices = get_free_indices(v)
+
+    if length(indices) != 1
+        throw(
+            DomainError(indices, "Input is not a vector, cannot create a diagonal matrix"),
+        )
+    end
+
+    vector_index = only(indices)
+    next_letter = get_next_letter(v)
+    id = KrD(vector_index, flip_to(vector_index, next_letter))
+
+    return BinaryOperation{Mult}(id, v)
+end
+
 function script(index::Lower)
     @assert index.letter >= 0
     text = []

src/simplify.jl

@@ -135,33 +135,46 @@ function simplify(::Mult, arg1::Tensor, arg2::BinaryOperation{Mult})
 end
 
 function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
-    if is_diag(arg1) && !is_elementwise_multiplication(arg1, arg2)
+    op = BinaryOperation{Mult}(arg1, arg2)
+
+    if is_diag(arg1) &&
+       !is_elementwise_multiplication(arg1, arg2) &&
+       length(get_free_indices(op)) == 1
         d = get_diag_delta(arg1)
 
         @assert !isnothing(d)
 
         target_indices = eliminate_indices(vcat(get_free_indices(arg1), get_indices(arg2)))
         factors = collect_factors(arg1)
+        vector_factors = filter(f -> f != d, factors)
         reshaped = []
 
         for f ∈ factors
-            if f isa KrD
+            if isequal(f, d)
                 continue
             end
 
             free_ids = get_free_indices(f)
+
             if isempty(free_ids)
                 push!(reshaped, f)
-            elseif length(free_ids) == 1
+            elseif length(free_ids) == 1 || length(free_ids) == 2
                 @assert length(target_indices) == 1
 
-                current_idx = intersect(free_ids, get_free_indices(d))
-                f = update_index(
-                    f,
-                    only(current_idx),
-                    only(target_indices);
-                    allow_shape_change = true,
-                )
+                vector_index =
+                    only(get_free_indices(to_binary_operation(Mult(), vector_factors)))
+
+                current_idx = intersect(free_ids, [vector_index])
+
+                if !isempty(current_idx)
+                    f = update_index(
+                        f,
+                        vector_index,
+                        only(target_indices);
+                        allow_shape_change = true,
+                    )
+                end
+
                 push!(reshaped, f)
             else
                 @assert false "Not implemented, please open an issue with your input"
@@ -185,7 +198,7 @@ function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
         return to_binary_operation(Mult(), reshaped)
     end
 
-    return BinaryOperation{Mult}(arg1, arg2)
+    return op
 end
 
 function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::BinaryOperation{Mult})

test/RicciTest.jl

@@ -9,7 +9,7 @@ using Test
 using DiffMatic: Variable, Literal, KrD, Zero
 using DiffMatic: Upper, Lower
 
-using LinearAlgebra: norm, tr, I
+using LinearAlgebra: norm, tr, I, diagm
 
 dc = DiffMatic
 
@@ -48,6 +48,14 @@ end
     @test matrix(4.2) == Literal(4.2, Upper(1), Lower(2))
 end
 
+@testset "Diagonal matrix constructor" begin
+    x = Variable("x", Upper(1))
+    xt = Variable("x", Lower(1))
+
+    @test diagm(x) == dc.BinaryOperation{dc.Mult}(KrD(Upper(1), Lower(2)), x)
+    @test diagm(xt) == dc.BinaryOperation{dc.Mult}(KrD(Lower(1), Upper(2)), xt)
+end
+
 @testset "index equality operator" begin
     left = Lower(3)
     right = Lower(3)

test/StdTest.jl

@@ -240,6 +240,24 @@ end
     @test to_std(mul(y, A)) == "Adiag(y)"
 end
 
+@testset "to_std output is correct with diagonal matrix-matrix multiplication" begin
+    A = Variable("A", Upper(1), Lower(2))
+    x = Variable("x", Upper(3))
+    y = Variable("y", Upper(2))
+
+    d = KrD(Upper(1), Lower(3))
+
+    function mul(l, r)
+        return dc.BinaryOperation{dc.Mult}(l, r)
+    end
+
+    @test to_std(mul(mul(A, y), KrD(Upper(1), Lower(3)))) == "diag(Ay)I"
+    @test to_std(mul(KrD(Upper(1), Lower(3)), mul(A, y))) == "diag(Ay)I"
+    @test to_std(mul(mul(mul(A, y), KrD(Upper(1), Lower(3))), x)) == "diag(x)Ay"
+    @test to_std(mul(x, mul(KrD(Upper(1), Lower(3)), mul(A, y)))) == "diag(x)Ay"
+    @test to_std(mul(x, mul(mul(A, y), KrD(Upper(1), Lower(3))))) == "diag(x)Ay"
+end
+
 @testset "to_std output is correct with vector-vector element wise multiplication" begin
     x = Variable("x", Upper(1))
     y = Variable("y", Upper(1))