Getting formulas like $E[\operatorname{Tr}(A^3 (A^T)^3]=n^4+10n^2+4n$

Question

Suppose $A$ is an random $n\times n$ matrix with standard normal entries. A user was able to obtain the following formula by enumerating all diagrams by hand: $$E[\operatorname{Tr}(A^3 (A^T)^3]=n^4+10n^2+4n$$

There's also this formula $$E[AA'AA'AA']=4n^2+6n^3+5n^4$$

Then

$$E[Tr[A^6 (A^T)^6]] = ???$$

How could I use Mathematica to discover these formulas automatically?

A user "Ben" provided magically compact Mathematica code to obtain the following formula for independent $A,B$: $$E||ABB^TA^T||^2_F=3 n^3 \left(n^2+n+1\right)$$

However it doesn't work out of the box for the problem at hand, perhaps someone fluent in both Mathematica and Wicks contractions can fix this -

Wick[x_Times] := 
 Block[{L = List @@ x}, 
  Sum[c[L[[1]], L[[i]]] Wick[Times @@ (Drop[L, {i}][[2 ;;]])], {i, 2, 
    Length[L]}]]
Wick[1] := 1
c[A_[i_, j_], A_[k_, l_]] := [Delta][i, k] [Delta][j, l]
c[__] := 0
rules = {[Delta][a1_, a1_] -> n,
   [Delta][a1_, a2_]^2 -> n,
   [Delta][a1_, a2_] a3_ :> (a3 /. a1 -> a2) /; Not@FreeQ[a3, a1]};
res = Wick[
    A[i1, i2] B[i2, i3] B[i4, i3] A[i5, i4] A[i5, i6] B[i6, i7] B[i8, 
      i7] A[i1, i8]] //. rules // Simplify
Table[{n, res}, {n, 1, 5}] // TableForm

For $q=7$ I'm getting 41312 n + 58312 n^2 + 25512 n^3 + 9046 n^4 + 748 n^5 + 204 n^6 + n^8. I'll add a speculative answer about this a bit later. — JimB, Feb 22 '24 at 00:17
If there is a solution for $q>9$ that can be executed without a quantum computer, I think it will be one that efficiently filters out all of the products whose expectations end up being zero. — JimB, Feb 28 '24 at 16:26
@JimB the Wicks contraction approach should be feasible for q>9, the issue right now is that there's bug in the code which makes it give formulas for two matrix expressions like $\operatorname{tr}(ABB'A'ABB'A)$, and not for ones involving just one matrix — Yaroslav Bulatov, Feb 28 '24 at 17:08
You might want to add the requirement stating the values of $q$ that are needed if $q>9$ is what is desired. I'd argue that all of the answers so far are attempts at efficient complete enumeration which are probably limited to $q<9$ (at least for getting the equation $n$ given $q$). — JimB, Feb 28 '24 at 17:28
@JimB ultimately I want generic code that can do arbitrary trace formula for Gaussian A in any position, like perhaps Tr((AAA')(AAA')') — Yaroslav Bulatov, Feb 28 '24 at 18:22
For whatever it's worth the formula for $q=9$ is $n^{10}+500 n^8+2952 n^7+70354 n^6+423642 n^5+2609760 n^4+7728888 n^3+13980320 n^2+9643008 n$. I'll be updating my answer with faster code certainly won't give general answers as you desire but should help with brute-force approaches. — JimB, Mar 08 '24 at 15:48

score 8 · Answer 1 · answered Feb 18 '24 at 06:37

$Version

(* "14.0.0 for Mac OS X ARM (64-bit) (December 13, 2023)" *)

Clear["Global`*"]

First problem:

f1[n_Integer?Positive] := f1[n] =
  Module[{A, distSpec},
   A = Array[a, {n, n}];
   distSpec = # \[Distributed] NormalDistribution[] & /@ 
     Variables[A];
   Expectation[
    Tr[A . A . A . A\[Transpose] . A\[Transpose] . A\[Transpose]], 
    distSpec]]
seq1 = f1 /@ Range[6]
(* {15, 64, 183, 432, 895, 1680} *)
FindSequenceFunction[seq1, n]
(* 4 n + 10 n^2 + n^4 *)

Second problem:

f2[n_Integer?Positive] := f2[n] =
  Module[{A, distSpec},
   A = Array[a, {n, n}];
   distSpec = # \[Distributed] NormalDistribution[] & /@ 
     Variables[A];
   Expectation[
    Tr[A . A\[Transpose] . A . A\[Transpose] . A . A\[Transpose]], 
    distSpec]]
seq2 = f2 /@ Range[6]
(* {15, 144, 603, 1728, 3975, 7920} *)
FindSequenceFunction[seq2, n]
(* 4 n^2 + 6 n^3 + 5 n^4 *)

thanks, that works, now I'm curious about larger formulas – Yaroslav Bulatov Feb 18 '24 at 17:19 — Yaroslav Bulatov, Feb 18 '24 at 17:19

Roman · Answer 2 · 2024-02-19T06:53:28.240

Here's a relatively fast algorithm for calculating $F_{q}(n)=E\left[\text{Tr}\left(A^q(A^T)^q\right)\right]$:

f[A_, q_] := Tr[MatrixPower[A, q] . MatrixPower[Transpose[A], q]]
Clear[F];
F[n_Integer?Positive, q_Integer?Positive] := F[n, q] =
  Expand[f[Array[a, {n, n}], q]] /. {a[__]^(m_?EvenQ) -> (m - 1)!!, a[__] -> 0}

Thanks to @JimB for suggesting a faster version:

f[A_, q_] := Length[A] (# . # &[MatrixPower[A, q][[1]]])

From this we can build formulas with FindSequenceFunction.

For $q\ge6$ it's tricky to go to high-enough $n$; but getting some inspiration from the lower-$q$ formulas and the OEIS we can suppose it's going to be $n\left(n^q+O(n^{q-2})\right)$:

Table[F[n, 1] == n^2, {n, 10}]
(*    {True, True, True, True, True, True, True, True, True, True}    *)
Table[F[n, 2] == 2 n + n^3, {n, 9}]
(*    {True, True, True, True, True, True, True, True, True}    *)
Table[F[n, 3] == 4 n + 10 n^2 + n^4, {n, 8}]
(*    {True, True, True, True, True, True, True, True}    *)
Table[F[n, 4] == 48 n + 28 n^2 + 28 n^3 + n^5, {n, 7}]
(*    {True, True, True, True, True, True, True}    *)
Table[F[n, 5] == 304 n + 468 n^2 + 110 n^3 + 62 n^4 + n^6, {n, 6}]
(*    {True, True, True, True, True, True}    *)
Table[F[n, 6] == 3656 n + 3880 n^2 + 2426 n^3 + 314 n^4 + 118 n^5 + n^7, {n, 5}]
(*    {True, True, True, True, True}    *)

+1 Some additional speed can be achieved by noting that f[A_, q_] := Dimensions[A][[1]] (MatrixPower[A, q] . MatrixPower[Transpose[A], q])[[1, 1]] gives the same expected value. This follows from the expected values of each diagonal element of $A^q(A^T)^q$ are the same. — JimB, Feb 19 '24 at 02:56
@roman Thanks for putting these formulas together, they were useful for debugging the diagrammatic visualization here — Yaroslav Bulatov, Mar 28 '24 at 17:21

score 1 · Answer 3 · answered Feb 22 '24 at 05:51

Here is an approach that uses several simplifications to speed up the calculations.

Define parameters and matrices - with some of the code copied directly from @Roman 's answer.

n = 3;
q = 5;
e = {a[__]^(m_?EvenQ) -> (m - 1)!!, a[__] -> 0};
Aq = MatrixPower[Array[a, {n, n}], q];
ATq = MatrixPower[Transpose[Array[a, {n, n}]], q];

The direct approach by @Roman that works for relatively small values of $n$ and $q$ is the following:

(Tr[Aq . ATq] // Expand) /. e
(* 13845 *)

That can be shortened a bit by noting that each diagonal element has the same expectation.

Table[((Aq . ATq)[[i, i]] // Expand) /. e , {i, n}]
(* {4615, 4615, 4615} *)

And that can be shortened a bit more by noting that the first element of Aq.ATq is just sum of the squares of the elements in the first row of Aq:

(Aq . ATq)[[1, 1]] - Total[Aq[[1]]^2] // Expand
(* 0 *)

Further simplification is that when $n > 2$, elements Aq[[1,2;;]] all have the same expected values.

TableForm[Transpose[{Range[n], ((Aq[[1]]^2 // Expand) /. e)}], 
 TableHeadings -> {None, {"i", "E(Aq[[1,i]])"}}]

Putting all of the simplifications together we have the following:

F[n_?IntegerQ, q_?IntegerQ] := Module[{Aq},
Aq = MatrixPower[Array[a, {n, n}], q];
  If[n == 1, (2 q - 1)!!,
   n  ((Aq[[1, 1]]^2 // Expand) /. e) + 
    n (n - 1) ((Aq[[1, 2]]^2 // Expand) /. e)]
  ]
F[3, 5]
(* 13845 *)

Using this approach I was able to find a formula when $q=7$:

41312 n + 58312 n^2 + 25512 n^3 + 9046 n^4 + 748 n^5 + 204  n^6 + n^8

But I did so under the assumption that the coefficient of $n^8$ is 1 as I didn't want to wait around for 'F[7, 7]` to finish which would be required to determine all 7 parameters.

Unfortunately while the above simplifications can speed things up considerably, this approach still doesn't expand the workable combinations of $n$ and $q$ by very much.

@YarsolovBulatov Just to be clear: I don't consider what I wrote above to be an answer. It only shaves off a bit of time from the other responses. I have looked at how a sampling approach might be used (by sampling a fraction of the millions of terms that can result when $n$ or $q$ is large. But most (99% ?) evaluate to zero when taking the expected value and no apparent way to pre-identify more than about half of those. — JimB, Feb 22 '24 at 21:42
higher moments of Gaussian distributions are notoriously hard to find by random sampling. To estimate how hard this is, you can compute the expected standard error on your expectation value, then estimate how many samples you’ll need to get this standard error to a good size. — Roman, Feb 23 '24 at 08:43
@Roman. Correct. I was treating the individual products ($a_{1,2}^4 a_{3,2}^3$, etc.) as elements of a finite population. The expected values of those terms are just how you wrote with a[__]^(m_?EvenQ) -> (m-1)!!. Because nearly all of those expectations are zero, simple random sampling from that population is very inefficient. I tried separating the squared terms from the rest of the cross-products but still the variance in the expectations of the cross-products was extremely large. I have one more stratification to try this weekend. — JimB, Feb 23 '24 at 16:30

score 1 · Answer 4 · answered Feb 25 '24 at 22:43

Using the following adjustment of the other answers, it was pretty fast to obtain all but 1 value needed for the case q == 8

dtp = Developer`ToPackedArray;
F2[n_, q_] := F2[n, q] = Which[MatchQ[n, 0 | 1], n (2 q - 1)!!, MatchQ[q, 0 | 1], n^(q + 1), True,
    Module[{Q = List @@@ Expand[MatrixPower[Array[a, {n, n}], q][[1, ;; 2]]], f, R, map, odd},
            Table[f = MapAt[If[ListQ[#], #, {#}] &, t /. Times -> List, {1}];
                map = AssociationThread[R = Range[Length[t]], 0];
                odd = dtp[Position[f, #, {2}]][[All, 1]] & /@ Alternatives @@@ Outer[Power, Variables[f], Range[1, q, 2]];
                Do[If[KeyExistsQ[map, i], If[# =!= {}, map[i] = #, map[i] =.] &[If[map[i] =!= 0, map[i], Drop[R, i]] ⋂ #]], {i, #}] & /@ odd;
                Do[If[KeyExistsQ[map, i], If[# =!= {}, map[i] = #, map[i] =.] &[Complement[If[map[i] =!= 0, map[i], Drop[R, i]], #]]], {i, Complement[Keys[map], #]}] & /@ odd;
                Total[t t /. _^m_ :> (m - 1)!!] + 2 Total[KeyValueMap[t[[#2]] t[[#]] &, map] /. _^m_ :> (m - 1)!!, {1, 2}],
            {t, Q}] . {1, (n - 1)} n]]

Solving with one missing value:

form = With[{q = 8}, # /. Solve[Table[F2[n, q] == #, {n, q - 2}]] &[
    Sum[Which[i === q + 1, 1, i === q, 0, True, c[i]] n^i, {i, q + 1}]]]

And notice that when using Ratios on the sequence of values from the previous polynomials evaluated at n == -2 you simply get {-2, -3, -4, -5, -6, -7, -8}. Assuming the next value will be -9 implies

Solve[-725760 == form /. n -> -2]
(* {{c[1] -> 609600}} *)
form /. c[1] -> 609600
(* {609600 n + 800624 n^2 + 470832 n^3 + 116900 n^4 + 27184 n^5 + 1556 n^6 + 328 n^7 + n^9} *)

+1 It is beyond my comprehension but definitely faster. And a very useful observation about using Ratios. — JimB, Feb 28 '24 at 03:33

Getting formulas like $E[\operatorname{Tr}(A^3 (A^T)^3]=n^4+10n^2+4n$

4 Answers4

Linked