# Speed Up the Gelman-Rubin-History function

A bit of background information

I have created a function to calculate the convergence of the Gelman Rubin diagnostic (useful in Bayesian analysis). The function seems to work as I want, however I have serious speed problems when the matrix A becomes larger, hence the question.

I have created a small matrix for the example (it is already big I know) :

A={{{18127., 13492.2, 17020.3, 14306.6, 20087.8, 15322.7, 18611.3,
14087.3, 18699., 15339.8}, {40095.6, 81848.7, 41145.4, 63459.7,
38075.1, 52870.1, 84746.4, 94331.2, 46547.7, 53182.5}, {105920.,
59366., 113773., 105011., 63431.1, 88142.4, 84235.7, 80818.9,
75614.5, 109016.}, {0.621616, 0.598056, 0.562754, 0.592478,
0.57167, 0.453746, 0.529772, 0.476047, 0.556064,
0.460045}, {26.1633, 10.2077, 35.2138, 24.4782, 10.9868, 37.6809,
16.1979, 11.3945, 36.6485, 20.4037}}, {{17874.9, 13492.2, 17066.7,
14426.7, 21057.5, 15322.9, 20614.6, 14087.3, 18699.,
15339.8}, {43479.8, 81848.7, 41098.9, 60746.7, 31639.3, 52874.1,
80502.1, 94331.2, 46547.7, 53182.5}, {104353., 59366., 113010.,
111745., 64028.7, 88410.6, 85748.7, 80818.9, 75614.5,
109016.}, {0.612533, 0.598056, 0.562889, 0.591655, 0.56779,
0.453827, 0.553562, 0.476047, 0.556064, 0.460045}, {25.2417,
10.2077, 34.8469, 26.5836, 11.1013, 37.4588, 18.325, 11.3945,
36.6485, 20.4037}}, {{17780., 13880.6, 19410.8, 14426.7, 20276.7,
15322.9, 20988.3, 14087.3, 18699., 15479.1}, {43832.8, 77421.4,
23653.9, 60746.7, 35067.1, 52874.1, 87127.1, 94331.2, 46547.7,
52648.9}, {105556., 65194.3, 114134., 111745., 69647.9, 88410.6,
87726.1, 80818.9, 75614.5, 108819.}, {0.613867, 0.594235, 0.537348,
0.591655, 0.570601, 0.453827, 0.553074, 0.476047, 0.556064,
0.468825}, {25.6633, 12.8847, 42.1838, 26.5836, 12.9246, 37.4588,
14.7498, 11.3945, 36.6485, 20.7039}}, {{17780., 9790.14, 19410.8,
13485.3, 20477.9, 15322.9, 20988.3, 13455.1, 18699.,
15000.1}, {43832.8, 71835.8, 23653.9, 64879.4, 33602.8, 52874.1,
87127.1, 113265., 46547.7, 55586.9}, {105556., 52227.3, 114134.,
118519., 68886.9, 88410.6, 87726.1, 68100.3, 75614.5,
111149.}, {0.613867, 0.617923, 0.537348, 0.595043, 0.572455,
0.453827, 0.553074, 0.479327, 0.556064, 0.499145}, {25.6633,
11.8113, 42.1838, 28.7817, 9.89854, 37.4588, 14.7498, 7.16584,
36.6485, 22.6444}}, {{16787.6, 9790.14, 19410.8, 12760.4, 20477.9,
15304.1, 17627.4, 14013.4, 20240.4, 13360.3}, {47595.9, 71835.8,
23653.9, 69326., 33602.8, 53481.3, 82537.7, 95212.8, 45702.4,
62917.6}, {119044., 52227.3, 114134., 122046., 68886.9, 88206.4,
77071.8, 74171.5, 40592.9, 103906.}, {0.6291, 0.617923, 0.537348,
0.640932, 0.572455, 0.454083, 0.572537, 0.471704, 0.497167,
0.379928}, {31.4623, 11.8113, 42.1838, 31.7186, 9.89854, 37.1542,
13.8679, 16.2211, 40.8306, 14.555}}, {{16787.6, 11441.1, 17257.,
13894.9, 20477.9, 15304.1, 17303.4, 14013.4, 20240.4,
12897.5}, {47595.9, 65929.8, 38444.8, 63204.7, 33602.8, 53481.3,
83674., 95212.8, 45702.4, 40029.9}, {119044., 54800.7, 116696.,
121201., 68886.9, 88206.4, 76811.8, 74171.5, 40592.9,
124980.}, {0.6291, 0.610899, 0.570894, 0.637599, 0.572455,
0.454083, 0.563498, 0.471704, 0.497167, 0.314885}, {31.4623,
11.5158, 38.7946, 31.6464, 9.89854, 37.1542, 14.0788, 16.2211,
40.8306, 13.3743}}, {{16194.8, 12975.2, 17283.4, 12940.3, 20477.9,
15874.9, 17386.6, 14070.8, 20016.8, 14001.1}, {49629., 58680.1,
38167.1, 68839.5, 33602.8, 51288.2, 82598.9, 97232.4, 45919.5,
44202.9}, {111920., 71126.9, 116815., 125052., 68886.9, 86074.9,
75786.3, 74339.8, 44535.6, 110559.}, {0.627082, 0.600347, 0.569499,
0.641965, 0.572455, 0.467143, 0.56162, 0.464593, 0.504349,
0.371321}, {29.2505, 18.7111, 39.0316, 32.6408, 9.89854, 34.1472,
14.8363, 16.0834, 40.1905, 21.074}}, {{16216.9, 12975.2, 17612.9,
14628.4, 19221.2, 15874.9, 17631.4, 13637.1, 20016.8,
14001.1}, {49683.7, 58680.1, 39101.7, 63371.8, 39998.3, 51288.2,
76913.4, 81970.9, 45919.5, 44202.9}, {111934., 71126.9, 108102.,
105845., 76826.9, 86074.9, 75925.1, 73067.9, 44535.6,
110559.}, {0.629663, 0.600347, 0.561645, 0.609136, 0.580335,
0.467143, 0.564117, 0.518333, 0.504349, 0.371321}, {29.333,
18.7111, 39.1713, 34.4419, 15.1714, 34.1472, 14.881, 17.1236,
40.1905, 21.074}}, {{15941.1, 12975.2, 17768.2, 14379., 19200.5,
15606.4, 19784.9, 13637.1, 20016.8, 14001.1}, {44374.8, 58680.1,
39188.3, 65292.8, 40480., 51972.8, 84607.5, 81970.9, 45919.5,
44202.9}, {118954., 71126.9, 105082., 105905., 76815.1, 84690.5,
56074.9, 73067.9, 44535.6, 110559.}, {0.642443, 0.600347, 0.56345,
0.612765, 0.580124, 0.47948, 0.531908, 0.518333, 0.504349,
0.371321}, {32.1507, 18.7111, 36.8538, 34.2503, 15.1676, 32.7175,
2.0564, 17.1236, 40.1905, 21.074}}, {{15941.1, 12975.2, 17581.4,
14379., 19565.1, 15606.4, 19784.9, 14653.2, 20016.8,
15776.6}, {44374.8, 58680.1, 39948., 65292.8, 40599.9, 51972.8,
84607.5, 65390.7, 45919.5, 42092.9}, {118954., 71126.9, 103758.,
105905., 70746.5, 84690.5, 56074.9, 93304.3, 44535.6,
107187.}, {0.642443, 0.600347, 0.564888, 0.612765, 0.583555,
0.47948, 0.531908, 0.573067, 0.504349, 0.467311}, {32.1507,
18.7111, 36.1466, 34.2503, 10.4423, 32.7175, 2.0564, 23.7507,
40.1905, 28.5485}}}

The dimensions of A can be given by

Dimensions[A]
(*{10, 5, 10}*)

Which represent 10 steps, 5 variables and 10 Markov chains. In other simulations, the size of the matrix can easily reach (for instance) {5000,10,500}.

My code already uses parallelization notions like ParallelTable[] but I am open to any proposal to improve the functions.

The code :

W[mat_, i_] := Module[{NumberOfChains},
NumberOfChains = Dimensions[mat][[3]];
1/NumberOfChains*Sum[Covariance[mat[[1 ;; i, All, j]]], {j, 1., NumberOfChains}]]

ChainMean[mat_, i_] := Transpose@Mean[mat[[1 ;; i, All, All]]]

B[mat_, i_] := Module[{NumberOfChains},
NumberOfChains = Dimensions[mat][[3]];
Covariance[ChainMean[mat, i]]]

GelmanRubin[mat_, di_] := Module[{NumberOfChains},
NumberOfChains = Dimensions[mat][[3]];
ParallelEvaluate[Off[Inverse::luc]];
ParallelTable[
Transpose@{i,i/(i + 1.) + (NumberOfChains + 1.0)/NumberOfChains*Max@Eigenvalues[Inverse[W[mat, i]] . B[mat, i]]}, {i, 2., Dimensions[mat][[1]], di}]]

And the plotting of the results :

GelmanRubin[A, 1];
ListPlot[%, Joined -> True]

Which gives :

This is typically the kind of curve we expect

UPDATE after Henrik Schumacher response

With the proposal below I made another test on a larger matrix :

Dimensions[B]
(*{1000,11,500}*)

and compared the calculation times :

(*With increment dStep=10 (100 computation to run)*)
GelmanRubin[B, 10] // AbsoluteTiming (*Old*)
GelmanRubin[B, 10] // AbsoluteTiming (*New*)

which gives :

(*112.989*)
(*110.036*)

I actually think that the problem does not come from the calculation of the eigenvalues since the size of the matrix depends on the second component of the vector Dimensions[B] = 11, which is probably not a problem for mathematica.

However, I think the problem is maybe with the sum of the W function, so I tried the following:

ParallelTable[B[B, i], {i, 2, Dimensions[B][[1]], 10}] // AbsoluteTiming
(*1.91474*)
ParallelTable[W[B, i], {i, 2, Dimensions[B][[1]], 10}] // AbsoluteTiming
(*80.6489*)(*OK the W function slow down the computation*)

I have tried to compile the functions but I have matrix rank errors, I have never compiled a function so far. Maybe a clue?

Update:

In the meantime, OP and I independently found out that Eigenvalues is not the bottleneck here as the dimensions of the matrices W[mat, i] and B[mat, i] are quite small in the problem instances of interest.

The real bottleneck stems from the fact that GelmanRubin[mat,1] has quadratic complexity in Length[mat]. The reason is that GelmanRubin[mat,1] calls W[mat, i]] in a loop for i from 2 to Lenght[mat]. And W[mat, i]] has complexity linear in i because it contains

Covariance[mat[[1 ;; i, All, j]]]

The solution is to use a recursive formala for the covariance matrix that allows us to compute the covariance of mat[[1 ;; i, All, j]] from some auxiliary data related to mat[[1 ;; i-1, All, j]]. This post explains this strategy.

I think the function GelmanRubin2 below accomplishes this. It was a bit fiddly to get all the updates right, so I am not totally sure. (And I'd rather don't want to explain how this works exactly.)

Transpose[List[a]] . List[a],
RuntimeAttributes -> {Listable},
Parallelization -> True
];

GelmanRubin2[matT_] :=
Module[{NumberOfChains, d, ChainSum, ChainKroneckerSquareMean},
NumberOfChains = Dimensions[matT][[2]];
d = Dimensions[matT][[3]];
ChainSum = matT[[1]];

Table[
Module[{x, Wi, Bi},
x = matT[[i]];
ChainSum += x;
Wi = N[1/(i - 1)] ChainKroneckerSquareMean - N[1/(i (i - 1))] Mean[cKroneckerSquareThread[ChainSum]];
Bi = N[1/i^2] Covariance[ChainSum];
{i, i/(i + 1.) + (1. + 1./NumberOfChains) Eigenvalues[{Bi, Wi},
1][[1]]}
]
, {i, 2, Length[matT]}]
]

For performance reasons, it wants to be called upon some slightly modified input data, though. Here is usage example that shows (i) how the inputs have to be transformed, (ii) that the new method works correctly, and (iii) that the new method faster for large values of Length[mat].

mat = RandomReal[{-10, 10}, {5000, 10, 50}];
matT = Transpose[mat, {1, 3, 2}];
result = GelmanRubin[mat, 1]; // AbsoluteTiming // First
result2 = GelmanRubin2[matT]; // AbsoluteTiming // First
Max[Abs[result - result2]]
28.6856

1.41671

3.9968*10^-15

Be cautious though: The recursive update formula for the covariance is a bit more sensitive to loss of precision. However I doubt that this will actually affect the results in practical use cases.

#Edit

Another drawback of the new method is that its recursive nature does not allow us to parallelize it in the i variable. We can still parallelize a bit in the j variable, though. The function GelmanRubin3 below tries to improve this parallelization; it achieves some performance gain for large values of NumberOfChains. Moreover, it reintroduces the stepping parameter di. Of course, all this comes of the cost of legibility.

Compile[{{a, _Real, 2}, {begin, _Integer}, {end, _Integer}},
Block[{d, A, x, ai},
d = Dimensions[a][[2]];
A = Table[0., {d}, {d}];
Do[
Do[
ai = CompileGetElement[a, k, i];
Do[
A[[i, j]] += ai Compile`GetElement[a, k, j];
, {j, i, d}]
, {i, 1, d}]
, {k, begin, end}];
Do[
Do[
A[[i, j]] = A[[j, i]];
, {j, 1, i - 1}]
, {i, 1, d}];
A
],
CompilationTarget -> "C",
RuntimeAttributes -> {Listable},
Parallelization -> True,
RuntimeOptions -> "Speed"
];

GelmanRubin3[matT_, di_] :=
Module[{NumberOfChains, d, ChainSum, ChainKroneckerSquareSum, i, x, Wi, Bi, result, counter, ThreadCount, begin, end},
NumberOfChains = Dimensions[matT][[2]];
d = Dimensions[matT][[3]];
begin = Most[dividers] + 1;
end = Rest[dividers];
];

ChainSum = matT[[1]];

i = 1;
counter = 0;
result = ConstantArray[0., {Quotient[Length[matT] - 2, di] + 1, 2}];

++i;
++counter;
x = matT[[i]];
ChainSum += x;
Wi = ChainKroneckerSquareSum - N[1/i] Total[cTotalKroneckerSquareThread[ChainSum, begin, end]];
Bi = Covariance[ChainSum];
result[[1, 1]] = N[i];
result[[1, 2]] = i/(i + 1.) + N[(i - 1)/i^2] (NumberOfChains + 1.) Eigenvalues[{Bi, Wi}, 1][[1]];

Do[
++counter;
Do[
++i;
x = matT[[i]];
ChainSum += x;
, {di}];

Wi = ChainKroneckerSquareSum -
Bi = Covariance[ChainSum];

result[[counter, 1]] = N[i];
result[[counter, 2]] = i/(i + 1.) + N[(i - 1)/i^2] (NumberOfChains + 1.) Eigenvalues[{Bi, Wi}, 1][[1]];

, {chunkindex, 2 + di, Length[matT], di}];

result
]

Old post; it's correct in principle, but it does not address the real bottleneck:

Maybe this helps. First, I observed that you seem to compute generalized eigenvalues.

Eigenvalues[Inverse[W[mat, i]] . B[mat, i]]

we can use

Eigenvalues[{B[mat, i], W[mat, i]}]

which in particular avoids the with Inverse.

But you actually want the greatest eigenvalue. So no need to compute the others! This can be done like this:

Eigenvalues[{B[mat, i], W[mat, i]}, 1][[1]]

If the matrices are large, then Arnoldi's method might perform a bit better:

Eigenvalues[{B[mat, i], W[mat, i]}, 1, Method -> "Arnoldi"][[1]]

Note the this assumes that all eigenvalues are positive! Eigenvalues orders the eigenvalues by absolute value from greatest to smallest.