Function to convert TensorContract[_TensorProduct, indices] into equivalent Dot + Tr version

Question

Mathematica can use either Dot + Tr to represent some tensors, or TensorContract + TensorProduct. I believe that the TensorContract + TensorProduct representation, while verbose, is more powerful for a couple reasons:

1. It can represent a wider variety of tensors, e.g., TensorContract[TensorProduct[a, b], {{1, 4}, {2, 5}, {3, 6}}] where a and b are rank 3 tensors doesn't have an equivalent Dot + Tr representation (at least, I can't think of one).
2. TensorReduce can in some cases reduce pure TensorContract + TensorProduct expressions better than the equivalent Dot + Tr expressions.

Because of the above, it would be convenient to have a function that converted a Dot + Tr representation into a TensorContract + TensorProduct representation. Another reason why it would be nice to have such a function is that TensorReduce of a pure TensorContract + TensorProduct often works much better than TensorReduce of a mixture of a Dot + Tr and TensorContract + TensorProduct representation.

Pure vs mixed

Here is an example where TensorReduce works better with pure TensorContract representations instead of mixed representations:

TensorReduce[
    r.R - TensorContract[TensorProduct[R, r], {{1, 2}}],
    Assumptions -> (r|R) \[Element] Vectors[3]
]

TensorReduce[
    TensorContract[TensorProduct[r, R], {{1, 2}}] - TensorContract[TensorProduct[R, r], {{1, 2}}],
    Assumptions -> (r|R) \[Element] Vectors[3]
]       

r.R - TensorContract[r[TensorProduct]R, {{1, 2}}]

0

ToTensor

The following function can be used to convert at Dot + Tr representation into a TensorContract + TensorProduct representation:

ToTensor[expr_] := expr /. {Dot->dot, Tr->tr}

dot[a__] := With[{indices = Accumulate@Map[TensorRank]@{a}},
    TensorContract[TensorProduct[a], {#, # + 1} & /@ Most[indices]]
]

tr[a_] /; TensorRank[a] == 2 := TensorContract[a, {{1, 2}}]
tr[a_, Plus, 2] := TensorContract[a, {{1, 2}}]
tr[a___] := Tr[a]

FromTensor

It would be nice to have a function that converts a TensorContract + TensorProduct representation into a Dot + Tr representation, if possible. Let's call such a function FromTensor. Then, a TensorSimplify function that does something like FromTensor @ TensorReduce @ ToTensor @ expr could be defined that is as powerful as a simple TensorReduce, but allows one to work with Dot + Tr or mixed representations.

Examples

The kinds of TensorContract + TensorProduct representations that should be converted into a Dot + Tr representation include at least the following, where a and b are vectors, and m and n are matrices:

1. Tr[m.n] ⇔ TensorContract[TensorProduct[m, n], {{1, 4}, {2,3}}]
2. m.n ⇔ TensorContract[TensorProduct[m, n], {{2, 3}}]
3. a.m.n ⇔ TensorContract[TensorProduct[a, m, n], {{1, 2}, {3, 4}}]
4. a.m.n.b ⇔ TensorContract[TensorProduct[a, m, n, b], {{1, 2}, {3, 4}, {5, 6}}]

Some other similar examples:

1. a.Transpose[n].Transpose[m] ⇔ TensorContract[TensorProduct[a, m, n], {{1, 5}, {4, 3}}]
2. Tr[Transpose[m].n] ⇔ TensorContract[TensorProduct[m, n], {{1, 3}, {2, 4}}]

There may be other equivalent representations.

So, my question is, can somebody write such a FromTensor function?

(I have written such a function, but I am unhappy with it. I'm hopeful that someone can write a better one. I will post my version as an answer at some point, but for now I'm curious what other independent answers are possible)

Answer 1

Update

I have put the package on GitHub. One can install the paclet using:

PacletInstall[
    "TensorSimplify", 
    "Site" -> "http://raw.githubusercontent.com/carlwoll/TensorSimplify/master"
]

and then load it with

<<TensorSimplify`

The package includes 4 functions. One (FromTensor) is described below. The other three are ToTensor, IdentityReduce and TensorSimplify. ToTensor converts Dot/Tr to TensorContract objects, IdentityReduce simplifies identity tensors (typically inactive IdentityMatrix objects) and TensorSimplify converts to TensorContract objects, then uses both TensorReduce and IdentityReduce.

As promised, I will present my current code for FromTensor. This function will shortly be part of a package on GitHub, augmented with code to handle symbolic identity tensors.

FromTensor[expr_] := expr /. TensorContract->tc

tc[a_TensorProduct, i_] := Module[{res = itc[a, i]},
    res /; res =!= $Failed
]
tc[a_, {{1, 2}}] /; TensorRank[a] == 2 := Tr[Replace[a, (Transpose|TensorTranspose)[m_, {2, 1} | PatternSequence[]]-> m]]
tc[a__] := TensorContract[a]

itc[a_TensorProduct, i_] := Module[
    {indices, rnk, s=0, ends, g, nodes, info, tlist, res},
    indices = tensorIndices[a];
    rnk = TensorRank @ TensorContract[a,i];

    (* 
     * Determine ends of the contraction chain.
     * For Tr, remove one set of indices, and find contraction 
     * chain of remaining indices
     *)

    ends = Switch[{rnk, Sort@Tally[Length/@indices]},
        {0, {{2,_}}}, Complement[Range@TensorRank[a], Flatten@Most@i],
        {2, {{2,_}}}, Complement[Range@TensorRank[a],Flatten@i],
        {1, {{1,1},{2,_}}}, {0, First@Complement[Range@TensorRank[a],Flatten@i]},
        {0, {{1,2},{2,_}|PatternSequence[]}}, {0,-1},
        _,Return[$Failed]
    ];

    (* find contraction chain. Augment vectors with 0 | -1 so that each node is a pair *)
    g = FindPath[
            Graph @ Join[
            Cases[indices, p:{_,_} :> UndirectedEdge@@p],
            Cases[indices,{p_} :> UndirectedEdge[s--, p]],
            UndirectedEdge @@@ i
        ],
        First@ends,
        Last@ends,
        {2 (Length[i] - Boole[rnk == 0 && Min[ends]>0])+ 1}
    ];
    (* unable to find a single contraction containing all tensors *)
    If[g === {}, Return[$Failed, Module]];

    (* find node (tensor) indices in the contraction chain *)
    nodes = DeleteCases[Partition[First@g, 2, 2], 0|-1, Infinity];

    (* determine tensors corresponding to indices, and whether to transpose tensor *)
    info=Table[
        Query[Select[MemberQ[n]], MatchQ[{n,___}]][indices],
        {n, nodes[[All,1]]}
    ];

    (* standardize Transpose *)
    tlist = Replace[
        List@@a,
        (TensorTranspose | Transpose)[m_, {2, 1}] -> Transpose[m],
        {1}
    ]; 

    (* create equivalent Dot product *)
    res = Dot @@ MapThread[
        If[#2, #1, Transpose[#1]]&,
        {
            tlist[[Flatten@Keys[info]]],
            Flatten@Values[info]
        }
    ];
    res = Replace[res, Transpose[Transpose[m_]] :> m, {1}];

    (* For 0-rank outputs, determine whether the normal or "transposed" version has fewer Transpose's *)
    Which[
        rnk > 0,
        res,

        TensorRank[res] > 0,
        If[Count[res, _Transpose] > Length[a]/2,
            Tr @ Replace[Reverse[res], {Transpose[m_]:>m, m_:>Transpose[m]}, {1}],
            Tr @ res
        ],

        Count[res,_Transpose] > Length[a]/2-1,
        res = Reverse[res];
        res[[2 ;; -2]] = Replace[List @@ res[[2 ;; -2]], {Transpose[m_]:>m, m_:>Transpose[m]}, {1}];
        res,

        True,
        res
    ]
]

(* tensorIndices returns a list of node -> indices rules *)
tensorIndices[Verbatim[TensorProduct][t__]] := With[{r=Accumulate @* Map[TensorRank] @ {1,t}},
    If[MatchQ[r, {__Integer}],
        Association @ Thread @ Rule[
            Range@Length[{t}], 
            Range[1+Most[r], Rest[r]]
        ],
        $Failed
    ]
]

Here are some examples of FromTensor usage:

$Assumptions = Element[a|b, Vectors[n]] && Element[M|T|A|B, Matrices[{n,n}]];

FromTensor @ TensorContract[TensorProduct[M,T,Transpose[A],B], {{1,4},{3,6},{5,7},{8,2}}]
FromTensor @ TensorContract[TensorProduct[a,M,T], {{1,3},{2,4}}]
FromTensor @ TensorContract[M, {{1,2}}]
FromTensor @ TensorContract[TensorProduct[M,T], {{2,3}}]
FromTensor @ TensorContract[TensorProduct[a,b], {{1,2}}]
FromTensor @ TensorContract[TensorProduct[a,Transpose[M],T,b], {{1,3}, {2,5}, {4,6}}]

Tr[Transpose[A].T.M.Transpose[B]]

a.Transpose[M].T

Tr[M]

M.T

a.b

a.M.Transpose[T].b

One can use random inputs to test FromTensor. For example, here is a list of rules:

rules = {
    A -> RandomReal[1, {3,3}],
    B -> RandomReal[1, {3,3}],
    M -> RandomReal[1, {3,3}],
    T -> RandomReal[1, {3,3}],
    a -> RandomReal[1, 3],
    b -> RandomReal[1, 3]
};

And here are tests of some of the previous examples:

tensor = TensorContract[TensorProduct[M,T,Transpose[A],B], {{1,4},{3,6},{5,7},{8,2}}];
tensor /. rules
FromTensor @ tensor /. rules

tensor = TensorContract[TensorProduct[a,M,T], {{1,3},{2,4}}];
tensor /. rules
FromTensor @ tensor /. rules

tensor = TensorContract[TensorProduct[a,Transpose[M],T,b], {{1,3}, {2,5}, {4,6}}];
tensor /. rules
FromTensor @ tensor /. rules

4.17922

4.17922

{1.4244, 0.45804, 1.00372}

{1.4244, 0.45804, 1.00372}

2.40798

2.40798

If the TensorContract object consists of multiple distinct contractions, e.g., TensorContract[TensorProduct[M, T], {{1,2}, {3,4}}], then FromTensor will not perform a conversion. In this case, one should use TensorReduce first. Here is an example:

tensor = TensorContract[TensorProduct[a,M,T],{{1,3},{4,5}}];

FromTensor @ tensor
FromTensor @ TensorReduce @ tensor

TensorContract[TensorProduct[a, M, T], {{1, 3}, {4, 5}}]

a.Transpose[M] Tr[T]

This version of FromTensor doesn't work with individual tensors higher than rank 2.