2
$\begingroup$

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 :

enter image description here

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?

Improve this question
$\endgroup$

1 Answer 1

Reset to default
7
$\begingroup$

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.)

cKroneckerSquareThread = Compile[{{a, _Real, 1}}, 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]]; ChainKroneckerSquareMean = Mean[cKroneckerSquareThread[matT[[1]]]]; Table[ Module[{x, Wi, Bi}, x = matT[[i]]; ChainSum += x; ChainKroneckerSquareMean += Mean[cKroneckerSquareThread[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.

cTotalKroneckerSquareThread = 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 = Compile`GetElement[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]]; ThreadCount = "ParallelThreadNumber" /. ("ParallelOptions" /.SystemOptions["ParallelOptions"]); With[{dividers = Round[Subdivide[0, NumberOfChains, ThreadCount]]}, begin = Most[dividers] + 1; end = Rest[dividers]; ]; ChainSum = matT[[1]]; ChainKroneckerSquareSum = Total[cTotalKroneckerSquareThread[ChainSum, begin, end]]; i = 1; counter = 0; result = ConstantArray[0., {Quotient[Length[matT] - 2, di] + 1, 2}]; ++i; ++counter; x = matT[[i]]; ChainSum += x; ChainKroneckerSquareSum += Total[cTotalKroneckerSquareThread[x, begin, end]]; 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; ChainKroneckerSquareSum += Total[cTotalKroneckerSquareThread[x, begin, end]]; , {di}]; Wi = ChainKroneckerSquareSum - N[1/i] Total[cTotalKroneckerSquareThread[ChainSum, begin, end]]; 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.

So instead of

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.

Improve this answer
$\endgroup$
6
  • $\begingroup$ Thank you for your reply, please consider the above update $\endgroup$ Commented Apr 23, 2022 at 9:02
  • $\begingroup$ Yes, you're totally right: Eigenvalues is not the bottleneck here. Your real issue is that the algorithm has quadratic complexity with respect to Length[mat]. I have already found a solution for that. I will update my answer, soon. $\endgroup$ Commented Apr 23, 2022 at 9:04
  • $\begingroup$ Thanks ! I tried to test your code but I think I'm missing the definition of cKroneckerSquareThread[] am I right? $\endgroup$ Commented Apr 23, 2022 at 10:01
  • $\begingroup$ Oh sorry! There it is. $\endgroup$ Commented Apr 23, 2022 at 10:02
  • $\begingroup$ I am really impressed by the speed of your algorithm compared to the original. I will accept your answer as I have just tested it on the largest matrix I currently have available and it only took about 30 seconds. Of course if you find the time to detail your approach more I think it might be of interest to some people but in the opposite case many thanks for the time invested so far and thanks for sharing $\endgroup$ Commented Apr 23, 2022 at 10:19

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.