Speeding up generation of block diagonal matrix

Question

I'm struggling with the following problem. I have $48$ square matrices (full, filled with real machine precision numbers, thus are packed, all different) of size $128$. I wolud like to place them on a diagonal of sparse array of dimension $48\times 128=6144$. The method⁽¹⁾

SparseArray @ ArrayFlatten @ ReleaseHold @ DiagonalMatrix[Hold /@ matrices]
(* matrices is a list of 48 matices 128 x 128,
   e.g. matrices = RandomReal[{}, {128, 128}] & /@ Range[48] *)

is too slow (it takes ~6s on my laptop). I'm suspecting that the problem is with the ArrayFlatten function, since this produces huge matrix $6144\times 6144$ filled moslty with zeros (in some sense it unpacks sparse array). Is there any way to do the same but much faster (more efficient)? In a fraction of a second (I'm optimistic)? I've looked at "SparseArray`" context but without much success (SparseArray`VectorToDiagonalSparseArray seems to be equivalent to DiagonalMatrix so accepts only vectiors/lists). (Specific numbers given here are just for tests, in the end I would like to increase size of my problem, but then it of course gets even worse.)

After posting this question I've found the code on MathWorld which gives me the result in ~3.63s. Code by ybeltukov SparseArray[Band@{1, 1} -> matrices] is even faster (~2.48s) but still far from being ideal.

Update: I've checked that asymptotically execution time scales as (based on AbsoluteTiming):

• $m^{2}n^{2}$ for BlockDiagonalMatrix
• $m^{2}n^{1}$ for recent version of blockArray by ybeltukov

where: $n$ is a number of matrices/blocks and $m$ is a size of a single matix/block.

Answer 1

You are right, it can be done in a fraction of second. One can explicitly construct an array of indexes

blockArray[mat_] := SparseArray[
  Tuples[Range@# - {1, 0, 0}].{Rest@#, {1, 0}, {0, 1}} &@Dimensions@mat -> 
   Flatten@mat]

Timings:

matrices = RandomReal[1, {48, 128, 128}];

s1 = 
   SparseArray@
    ArrayFlatten@ReleaseHold@DiagonalMatrix[Hold /@ matrices]; // RepeatedTiming
(* {7.56, Null} *)

s2 = SparseArray[Band@{1, 1} -> matrices]; // RepeatedTiming
(* {4.03, Null} *)

s3 = blockArray[matrices]; // RepeatedTiming
(* {0.097, Null} *)

TrueQ[s1 == s2 == s3]
(* True *)

For further acceleration you can create the internal structure of the SparseArray directly

c = Compile[{{b, _Integer}, {h, _Integer}, {w, _Integer}},
   Partition[Flatten@Table[k + i w, {i, 0, b - 1}, {j, h}, {k, w}], 1], 
   CompilationTarget -> "C", RuntimeOptions -> "Speed"];

blockArray2[mat_] := 
 SparseArray @@ {Automatic, # {##2}, 
     0, {1, {Range[0, 1 ##, #3], c@##}, Flatten@mat}} & @@ Dimensions@mat

s4 = blockArray2[matrices]; // RepeatedTiming
(* {0.031, Null} *)

s3 == s4
(* True *)

Answer 2

There is an undocumented built-in solution:

rules = {#, #} -> RandomReal[{}, {128, 128}] & /@ Range[48];
SparseArray`SparseBlockMatrix[rules]; // RepeatedTiming

(* {0.042, Null} *)

Answer 3

Version 13.1 introduced BlockDiagonalMatrix (labeled as Experimental).

On my computer, it performs similarly to ybeltukov's blockArray (using the setup from the answer):

s3 = blockArray[matrices]; // RepeatedTiming
(*{0.0282592,Null}*)

s4 = BlockDiagonalMatrix@matrices; // RepeatedTiming
(* {0.0281052,Null}*)

s3 == s4
(* True *)

Answer 4

Your ArrayFlatten method can be made as fast as the other methods with two small tweaks: applying SparseArray to the list of dense matrices rather than at the end, and using N on the diagonal matrix to turn the off-diagonal zeros into approximate reals.

matrices = RandomReal[{}, {128, 128}] & /@ Range[48];

mat1 = SparseArray @ ArrayFlatten @ ReleaseHold @ DiagonalMatrix[Hold /@ matrices]; // RepeatedTiming
(* {3.72, Null} *)

mat2 = ArrayFlatten @ ReleaseHold @ N@DiagonalMatrix[Hold /@ SparseArray@matrices]; // RepeatedTiming
(* {0.045, Null} *)

mat1 == mat2
(* True *)