Collapse traces by default

Author: asterycs

src/forward.jl

@@ -253,10 +253,6 @@ function evaluate(::Mult, arg1::KrD, arg2::BinaryOperation{Mult})
 end
 
 function evaluate(::Mult, arg1::BinaryOperation{Mult}, arg2::KrD)
-    if is_trace(BinaryOperation{Mult}(arg1, arg2))
-        return BinaryOperation{Mult}(arg1, arg2)
-    end
-
     ci = indices_in_common(arg1.arg1, arg1.arg2)
 
     if !isempty(ci) && !is_trace(arg2)
@@ -339,10 +335,6 @@ function evaluate(::Mult, arg1::UnaryOperation, arg2::KrD)
 end
 
 function evaluate(::Mult, arg1::KrD, arg2::UnaryOp) where {UnaryOp<:UnaryOperation}
-    if is_trace(BinaryOperation{Mult}(arg1, arg2))
-        return BinaryOperation{Mult}(arg1, arg2)
-    end
-
     if can_contract(evaluate(arg1), evaluate(arg2.arg))
         return UnaryOp(evaluate(Mult(), evaluate(arg1), evaluate(arg2.arg)))
     end
@@ -374,10 +366,6 @@ function _multiply_with_krd(arg1::Union{Monomial,KrD}, arg2::KrD)
         return BinaryOperation{Mult}(arg1, arg2)
     end
 
-    if is_trace(BinaryOperation{Mult}(arg1, arg2))
-        return BinaryOperation{Mult}(arg1, arg2)
-    end
-
     if is_trace(arg1) || is_trace(arg2)
         return BinaryOperation{Mult}(arg1, arg2)
     end
@@ -410,10 +398,6 @@ function evaluate(
     arg1::BinaryOperation{Op},
     arg2::KrD,
 ) where {Op<:AdditiveOperation}
-    if is_trace(BinaryOperation{Mult}(arg1, arg2))
-        return BinaryOperation{Mult}(arg1, arg2)
-    end
-
     return evaluate(
         Op(),
         evaluate(Mult(), evaluate(arg1.arg1), evaluate(arg2)),
@@ -426,10 +410,6 @@ function evaluate(
     arg1::KrD,
     arg2::BinaryOperation{Op},
 ) where {Op<:AdditiveOperation}
-    if is_trace(BinaryOperation{Mult}(arg1, arg2))
-        return BinaryOperation{Mult}(arg1, arg2)
-    end
-
     return evaluate(
         Op(),
         evaluate(Mult(), arg1, evaluate(arg2.arg1)),

src/ricci.jl

@@ -232,7 +232,7 @@ function tr(arg::Tensor)
         throw(de)
     end
 
-    return BinaryOperation{Mult}(arg, KrD(flip(free_ids[2]), flip(free_ids[1])))
+    return evaluate(BinaryOperation{Mult}(arg, KrD(flip(free_ids[2]), flip(free_ids[1]))))
 end
 
 function Base.sum(arg::Tensor)

src/std.jl

@@ -214,7 +214,9 @@ function _to_std_string(arg::Monomial)
     ids = get_indices(arg)
 
     if length(ids) == 2
-        if typeof(ids[1]) == Upper && typeof(ids[2]) == Lower
+        if flip(ids[1]) == ids[2]
+            return "tr(" * arg.id * ")"
+        elseif typeof(ids[1]) == Upper && typeof(ids[2]) == Lower
             return arg.id
         elseif typeof(ids[1]) == Lower && typeof(ids[2]) == Upper
             return arg.id * "ᵀ"
@@ -300,14 +302,16 @@ function get_contra_covariant_matrix(arg1::Tensor, arg2::Tensor)
     arg2_letters = [i.letter for i ∈ get_free_indices(arg2)]
     common_letter = intersect(arg1_letters, arg2_letters)
 
-    @assert length(common_letter) == 1
-
-    arg1_filt = filter(i->i.letter == first(common_letter), arg1_ids)
-    arg2_filt = filter(i->i.letter == first(common_letter), arg2_ids)
+    if length(common_letter) == 1
+        arg1_filt = filter(i->i.letter == first(common_letter), arg1_ids)
+        arg2_filt = filter(i->i.letter == first(common_letter), arg2_ids)
 
-    if typeof(first(arg1_filt)) == Lower && typeof(first(arg2_filt)) == Upper
-        return (arg2, arg1)
-    elseif typeof(first(arg1_filt)) == Upper && typeof(first(arg2_filt)) == Lower
+        if typeof(first(arg1_filt)) == Lower && typeof(first(arg2_filt)) == Upper
+            return (arg2, arg1)
+        elseif typeof(first(arg1_filt)) == Upper && typeof(first(arg2_filt)) == Lower
+            return (arg1, arg2)
+        end
+    elseif length(common_letter) == 2 # is a trace
         return (arg1, arg2)
     end
 
@@ -377,6 +381,10 @@ function _to_std_string(arg::BinaryOperation{Mult})
         end
     end
 
+    if is_trace(arg)
+        return "tr(" * _to_std_string(arg.arg1) * _to_std_string(arg.arg2) * ")"
+    end
+
     if isempty(target_indices) && (typeof(terms[1]) == KrD || typeof(terms[2]) == KrD)
         tensor = if typeof(first(terms)) == KrD
             last(terms)
@@ -386,9 +394,7 @@ function _to_std_string(arg::BinaryOperation{Mult})
 
         tensor_free_ids = get_free_indices(tensor)
 
-        if length(tensor_free_ids) == 2
-            return "tr(" * to_std_string(tensor) * ")"
-        elseif length(tensor_free_ids) == 1
+        if length(tensor_free_ids) == 1
             return "sum(" * _to_std_string(tensor) * ")"
         end
 

test/ForwardTest.jl

@@ -539,10 +539,7 @@ end
     A = Monomial("A", Upper(1), Lower(2))
     B = Monomial("B", Upper(2), Lower(3))
 
-    @test dc.evaluate(tr(A)) == BinaryOperation{dc.Mult}(
-        Monomial("A", Upper(1), Lower(2)),
-        KrD(Upper(2), Lower(1)),
-    )
+    @test dc.evaluate(tr(A)) == Monomial("A", Upper(2), Lower(2))
     @test equivalent(
         dc.evaluate(tr(A * B)),
         dc.BinaryOperation{dc.Mult}(A, Monomial("B", Upper(2), Lower(1))),

test/StdTest.jl

@@ -109,7 +109,8 @@ end
         return dc.BinaryOperation{dc.Add}(l, r)
     end
 
-    trA = mul(Monomial("A", Upper(1), Lower(2)), KrD(Upper(2), Lower(1)))
+    trA = Monomial("A", Upper(2), Lower(2))
+
     A = Monomial("A", Upper(3), Lower(4))
     B = Monomial("B", Upper(4), Lower(5))
     x = Monomial("x", Upper(4))
@@ -119,23 +120,19 @@ end
     @test to_std_string(mul(A, trA)) == "tr(A)A"
     @test to_std_string(mul(mul(trA, A), x)) == "tr(A)Ax"
 
-    trAB = mul(
-        mul(Monomial("A", Upper(1), Lower(2)), Monomial("B", Upper(2), Lower(3))),
-        KrD(Upper(3), Lower(1)),
-    )
+    trAB = mul(Monomial("A", Upper(1), Lower(2)), Monomial("B", Upper(2), Lower(1)))
+
     @test to_std_string(trAB) == "tr(AB)"
     @test to_std_string(mul(trAB, A)) == "tr(AB)A"
     @test to_std_string(mul(trAB, B)) == "tr(AB)B"
     @test to_std_string(mul(mul(trAB, A), x)) == "tr(AB)Ax"
 
-    trApB = mul(
-        add(Monomial("A", Upper(1), Lower(2)), Monomial("B", Upper(1), Lower(2))),
-        KrD(Upper(2), Lower(1)),
-    )
-    @test to_std_string(trApB) == "tr(A + B)"
-    @test to_std_string(mul(trApB, A)) == "tr(A + B)A"
-    @test to_std_string(mul(trApB, B)) == "tr(A + B)B"
-    @test to_std_string(mul(mul(trApB, A), x)) == "tr(A + B)Ax"
+    trApB = add(Monomial("A", Upper(2), Lower(2)), Monomial("B", Upper(2), Lower(2)))
+
+    @test to_std_string(trApB) == "tr(A) + tr(B)"
+    @test to_std_string(mul(trApB, A)) == "(tr(A) + tr(B))A"
+    @test to_std_string(mul(trApB, B)) == "(tr(A) + tr(B))B"
+    @test to_std_string(mul(mul(trApB, A), x)) == "(tr(A) + tr(B))Ax"
 end
 
 @testset "to_std_string output is correct with all covariant bilinar form-vector contraction" begin