Most efficient way to build block / banded SparseArray

Question

I have a matrix that I generate with a procedure that looks like this:

blaMX[ptsX_Integer, ptsY_Integer] :=
  Table[
   With[
    {
     ix = 1 + Floor[(i - 1)/(ptsY)], jx = 1 + Floor[(j - 1)/(ptsY)],
     iy = Mod[i, ptsY, 1], jy = Mod[j, ptsY, 1]
     },
    If[iy != jy,
      0,
      k1[ix, jx]
      ] +
     If[ix != jx,
      0,
      k2[iy, jy]
      ]
    ],
   {i, ptsX*ptsY},
   {j, ptsX*ptsY}
   ];

where I have two square matrices k1 of dimension ptsX and k2 of dimension ptsY

It generates a block-diagonal matrix with a bunch of block-bands. The basic structure looks like:

blaMX[6, 9] /. {_k1 -> 1, _k2 -> 2, c -> 3} // MatrixPlot

where each of those diagonal blocks is k2 duplicated:

blaMX[3, 6][[7 ;; 12, 7 ;; 12]] /. _k1 -> 0 // Grid

k2[1,1] k2[1,2] k2[1,3] k2[1,4] k2[1,5] k2[1,6]
k2[2,1] k2[2,2] k2[2,3] k2[2,4] k2[2,5] k2[2,6]
k2[3,1] k2[3,2] k2[3,3] k2[3,4] k2[3,5] k2[3,6]
k2[4,1] k2[4,2] k2[4,3] k2[4,4] k2[4,5] k2[4,6]
k2[5,1] k2[5,2] k2[5,3] k2[5,4] k2[5,5] k2[5,6]
k2[6,1] k2[6,2] k2[6,3] k2[6,4] k2[6,5] k2[6,6]

Each band is made up of blocks, as an unwrapping of k1 with the Band starting at {n+(i-1)*ptsY, n+(j-1)*ptsY} being k1[[i+1, j+1]] (essentially just using Quotient:

With[{n = 1},
  Table[
   blaMX[6, 9][[
      (i - 1)*9 + n, 
      (j - 1)*9 + n
      ]] /. _k2 -> 0,
   {i, 6},
   {j, 6}
   ]
  ] // Grid

Now, my generation procedure is wildly inefficient. I do ptsX^2 * ptsY^2 work, when I really only have ptsX * ptsY * (ptsX + ptsY - 1) (which is just 2 n^3 - n^2 if ptsX == ptsY == n) non-zero elements.

So, what is the best way to build a matrix of this form? My current version looks like:

makeKMat[k1_, k2_] :=
  With[{
    k1SparseRules =
     With[{ptsX = Length@k1, ptsY = Length@k2},
      Flatten[
       Table[
        Band[{(i - 1)*ptsY + 1, (j - 1)*ptsY + 1}, {i*ptsY, j*ptsY}] ->
         k1[[i, j]],
        {i, ptsX},
        {j, ptsX}
        ],
       1
       ]
      ],
    k2SparseRules =
     With[{ptsX = Length@k1, ptsY = Length@k2},
      {
       Band[{1, 1}] ->
        ConstantArray[k2, ptsX]
       }
      ]
    },
   With[{k3 = SparseArray[k2SparseRules]},
    SparseArray[k1SparseRules, Dimensions[k3]] + k3
    ]
   ];

And this works:

k1Pts = 6;
k2Pts = 9;
k1Test = Array[k1, {k1Pts, k1Pts}];
k2Test = Array[k2, {k2Pts, k2Pts}];
kMatTest = makeKMat[k1Test, k2Test];
kMatTest == blaMX[6, 9]

True

And is pretty fast, all considered:

k1Test2 = RandomReal[{}, {15, 15}];
k2Test2 = RandomReal[{}, {15, 15}];

{tt, mx} = makeKMat[k1Test2, k2Test2] // RepeatedTiming;
tt

0.0071

{tt2, lmx} = lazyMX[k1Test2, k2Test2] // RepeatedTiming;
tt2

0.338

But, I can't help but feel that this could be done more efficiently. Can someone propose a better way to build a block-banded SparseArray like this?

Some mild goal-post moving:

What if we wanted to extend this to an $N$-dimensional system? A way to implement this element-wise would be:

getCrds[i_, j_, gridDivs_] :=
  FoldList[
    With[{tot = #[[3]], new = #2},
      Append[
       Map[
        Mod[1 + Quotient[# - 1, tot], new, 1] &,
        {i, j}
        ],
       tot*new
       ]
      ] &,
    {i, j, 1},
    gridDivs
    ][[2 ;;, ;; 2]];

makeNDKMat[ks : {__List}] :=
 With[{gridDivs = Length /@ ks},
  Table[
   With[{divvy = getCrds[i, j, gridDivs]},
    Total@
     Table[
      If[AllTrue[Delete[divvy, k], Apply[Equal]],
       ks[[k, Sequence @@ divvy[[k]]]],
       0
       ],
      {k, Length@gridDivs}
      ]
    ],
   {i, Times @@ gridDivs},
   {j, Times @@ gridDivs}
   ]
  ]

For example, in 4D this looks like:

k1T = Array[k1, {2, 2}];
k2T = Array[k2, {3, 3}];
k3T = Array[k3, {4, 4}];
k4T = Array[k4, {5, 5}];
k4DLazy = 
  makeNDKMat[{k4T, k3T, k2T, k1T}] /. {_k1 -> 1, _k2 -> 2, _k3 -> 
     3, _k4 -> 4};
k4DLazy // MatrixPlot

How could we efficiently build a sparse representation of this?

Henrik Schumacher · Accepted Answer · 2020-05-30 07:17:48Z

This is definitely faster. It took me some time to figure out the combinatorics and there might still be some potential for improvement within the compiled functions.

getColumnIndices = Compile[{{m, _Integer}, {n, _Integer}, {i, _Integer}},
   Transpose[Partition[
     Partition[
      Join[
       Table[k, {k, 1, i n}],
       Flatten[Table[Table[k, {j, 1, n}], {k, i n + 1, i n + n}]],
       Table[k, {k, i n + n + 1, m n}]
       ], {1}], {n}]],
   RuntimeAttributes -> {Listable},
   Parallelization -> True,
   CompilationTarget -> "C",
   RuntimeOptions -> "Speed"
   ];

getValues = Compile[{{k1row, _Real, 1}, {k2, _Real, 2}, {i, _Integer}},
   Block[{A},
    A = Join[
      Table[Compile`GetElement[k1row, k], {l, 1, Length[k2]}, {k, 1, i}],
      k2,
      Table[Compile`GetElement[k1row, k], {l, 1, Length[k2]}, {k, i + 2, Length[k1row]}],
      2
      ];
    Do[A[[k, k + i]] += Compile`GetElement[k1row, i + 1], {k, 1, Length[k2]}];
    A
    ],
   RuntimeAttributes -> {Listable},
   Parallelization -> True,
   CompilationTarget -> "C",
   RuntimeOptions -> "Speed"
   ];

makeKMat3[k1_, k2_] := With[{m = Length[k1], n = Length[k2]},
  SparseArray @@ {Automatic, {m n, m n}, 0, {1, {
      Range[0, (m n ) (m + n - 1), m + n - 1],
      Flatten[getColumnIndices[m, n, Range[0, m - 1]], 2]
      },
     Flatten[getValues[k1, k2, Range[0, m - 1]]]
     }}
  ]

Just to give you an idea of the timings:

m = 100;
n = 6;
k1 = RandomReal[{-1, 1}, {m, m}];
k2 = RandomReal[{-1, 1}, {n, n}];
A = makeKMat2[k1, k2]; // RepeatedTiming
B = makeKMat3[k1, k2]; // RepeatedTiming
Max[Abs[A - B]]

{0.11, Null}

{0.000481, Null}

0.

Addendum

The idea for the "higher dimensional" case is similar to the idea above. I exploit that the "ColumnIndices" of the resulting matrices (when partitioned into column indices per row) stay rectangular so that column indices from the diagonal matrices can be joined to it from the left and right. Some extra index list (diagidx) is needed for book keeping of those column indices of the blocks that belong to the diagonal entries. All in all, only operations on the nonzero values and the column indices are performed. No SparseArrays are built intermediately: Even building a SparseArray from row pointers, column indices and nonzero values has still some considerable overhead, probably because there is a consistency checker involved in the backend. That's a pitty since I know that I produce consistent data. This issue is also closely related to this post by Szabolcs.

getValues2 = Compile[{{k1row, _Real, 1}, {blockvals, _Real, 2}, {diagidx, _Integer, 1}, {i, _Integer}},
   Block[{A},
    A = Join[
      Table[Compile`GetElement[k1row, k], {l, 1, Dimensions[blockvals][[1]]}, {k, 1, i}],
      blockvals,
      Table[Compile`GetElement[k1row, k], {l, 1, Dimensions[blockvals][[1]]}, {k, i + 2, Length[k1row]}],
      2];
    Do[A[[k, Compile`GetElement[diagidx, k] + i]] += Compile`GetElement[k1row, i + 1], {k, 1, Length[blockvals]}];
    A
    ],
   RuntimeAttributes -> {Listable},
   Parallelization -> True,
   CompilationTarget -> "C",
   RuntimeOptions -> "Speed"
   ];

getColumnIndices2 = Compile[{
    {blockci, _Integer, 3}, {diagci, _Integer, 
     3}, {m, _Integer}, {n, _Integer}, {i, _Integer}
    },
   If[i > 0,
    If[i < m - 1,
     Join[diagci[[All, 1 ;; i]], blockci + i n, diagci[[All, i + 1 ;; m - 1]] + (n), 2],
     Join[diagci[[All, 1 ;; i]], blockci + i n, 2]
     ],
    Join[blockci + i n, diagci[[All, i + 1 ;; m - 1]] + (n), 2]
    ],
   RuntimeAttributes -> {Listable},
   Parallelization -> True,
   CompilationTarget -> "C",
   RuntimeOptions -> "Speed"
   ];

toSparseArrayData[b_?MatrixQ] := {
  Partition[SparseArray[b]["ColumnIndices"], Dimensions[b][[2]]],
  b,
  Dimensions[b][[2]],
  Range[Dimensions[b][[2]]]
  }

toSparseArray[X_] := 
 With[{d1 = Dimensions[X[[1]]][[1]], d2 = Dimensions[X[[1]]][[2]]},
  SparseArray @@ {Automatic, {d1, d1}, 0,
    {1, {Range[0, d1 d2, d2], Flatten[X[[1]], 1]}, Flatten[X[[2]]]}}
  ]

iteration[X_, a_] := With[{
   m = Length[a],
   blockci = X[[1]],
   blockvals = X[[2]],
   n = X[[3]],
   diagidx = X[[4]]
   },
  With[{ran = Range[0, m - 1]},
   {
    Join @@ getColumnIndices2[blockci, Transpose[Partition[Partition[Range[(m - 1) n], 1], n]], m, n, ran],
    Join @@ getValues2[a, blockvals, diagidx, ran],
    m n,
    Join @@ Outer[Plus, ran, diagidx]
    }
   ]
  ]

makeKMat3ND[ks : {__List}] := 
 toSparseArray[Fold[iteration, toSparseArrayData[ks[[1]]], Rest[ks]]]

Usage example and timing test:

SeedRandom[123];
ks = Table[RandomReal[{-1, 1}, {RandomInteger[{3, 8}]}[[{1, 1}]]], 6];
A = makeKMat2ND[Reverse@ks]; // AbsoluteTiming
B = makeKMat3ND[ks]; // AbsoluteTiming
Max[Abs[A - B]]

{24.0011, Null}

{0.038756, Null}

8.88178*10^-16

This is the sparsity pattern of the resulting matrix:

It pops up in a specific version of discrete variable representation. This is the type of process you use for building the kinetic energy term in multiple dimensions for a direct product basis (which leads to a representation that is diagonal in the different coordinates, hence block-banded in the way I presented). — b3m2a1, Feb 22 '18 at 7:34
@b3m2a1 Thanks for the hint. But after some testing, I found out that MarcoB's answer there (which is actually the same J.M.'s answer here) is usually a bit more performant. — Henrik Schumacher, Jul 27 '18 at 22:09

Henrik Schumacher · Accepted Answer · 2018-03-17 12:06:14Z

12

Not faster, but maybe easier to implement

makeKMat2[k1_, k2_] := Module[{A, m, n},
  m = Length[k1];
  n = Length[k2];
  A = Map[SparseArray[Band[{1, 1}] -> #, {n, n}, 0.] &, k1, {2}];
  With[{B = SparseArray@k2}, Do[A[[i, i]] += B, {i, 1, m}]];
  ArrayFlatten[A]
  ]

Your implementation needs about 0.0048s; mine needs 0.0051s on my machine.

(b3m2a1 comment)
Note that this has a simple extension to the N-dimensional case:

makeKMat2ND[ks : {__List}] := Fold[makeKMat2, ks]

edited Mar 17 '18 at 12:06

answered Feb 14 '18 at 22:21

Henrik Schumacher

91.2k7 gold badges148 silver badges260 bronze badges

$\begingroup$ Very nice. I'm gonna be sharing this code with others so the easier it is to read for people without Mathematica experience the better. $\endgroup$
– b3m2a1
Feb 15 '18 at 7:49
$\begingroup$ Good to hear! I was also tempted to compute the "ColumnIndices" and "RowPointers" of the resulting matrix by hand since building from the "FullForm" of the matrix is usually much faster than letting Mathematica assemble it. But I am somewhat relieved that I don't have to do that. Although it should not be too complicated as each row has the same number of entries... $\endgroup$
– Henrik Schumacher
Feb 15 '18 at 8:31
$\begingroup$ One difficulty, how can I keep make sure to keep the entire thing as a SparseArray? If I feed it: makeKMat2[RandomReal[{}, {10, 10}], RandomReal[{}, {10, 10}]] it unwraps to a List for me, presumably owing to some auto-determined memory optimization, while if I feed it more regular data: makeKMat2[Table[i + j, {i, 10}, {j, 10}], Table[i + j, {i, 10}, {j, 10}]] it's happy to keep it Sparse. This becomes more of a problem as I move to 4D and beyond, which I was hoping I could do recursively with this. $\endgroup$
– b3m2a1
Feb 15 '18 at 9:27
2

$\begingroup$ When using ArrayFlatten it is important that all occuring SparseArray have the same "Background" (aka. as default values). When casting dense, real matrices to sparse ones by applying SparseArray to them, the "Background" is automatically set to 0. (machine precision) while building sparse arrays with the constructor have 0 as background by default. Adding ,0. as last argument to SparseArray cures this. I also edited the post. $\endgroup$
– Henrik Schumacher
Feb 15 '18 at 9:45
$\begingroup$ Moreover: Feel free to add as much information as other users might find useful. $\endgroup$
– Henrik Schumacher
Feb 15 '18 at 9:47

| Show 1 more comment

J. M.'s torpor · Accepted Answer · 2018-03-17 11:41:37Z

5

This is slightly faster and certainly more compact than the uncompiled implementation in Henrik's answer:

makeKMat[k1_?SquareMatrixQ, k2_?SquareMatrixQ] := Module[{m = Length[k1], n = Length[k2]},
    KroneckerProduct[k1, IdentityMatrix[n, SparseArray]] + 
    KroneckerProduct[IdentityMatrix[m, SparseArray], k2]]

and as b3m2a1 notes, the extension can be done as

makeKMatND[ks : {__?SquareMatrixQ}] := Fold[makeKMat[#2, #1] &, ks]

Using the examples in the OP:

MatrixPlot[makeKMat[ConstantArray[1, {9, 9}], ConstantArray[2, {6, 6}]]]

matrix plot of 2D case

AbsoluteTiming[tst =
               makeKMatND[{ConstantArray[4, {5, 5}], ConstantArray[3, {4, 4}],
                           ConstantArray[2, {3, 3}], ConstantArray[1, {2, 2}]}];][[1]]
   0.000406914

MatrixPlot[tst, MaxPlotPoints -> Infinity]

matrix plot of 4D case

edited Mar 17 '18 at 11:41

answered Mar 17 '18 at 11:24

J. M.'s torpor♦

117k11 gold badges364 silver badges544 bronze badges

$\begingroup$ +1 for SquareMatrixQ. I did not know that it existed! And for the KroneckerProduct trick. I learnt it from here several days later and I was really suprised that it worked so well. $\endgroup$
– Henrik Schumacher
Mar 17 '18 at 11:28
3

$\begingroup$ By sheer chance, I was rereading Van Loan's paper when I came across this old question... :D $\endgroup$
– J. M.'s torpor ♦
Mar 17 '18 at 11:36
$\begingroup$ Wow, thanks! That's a nice overview on Kronecker products. This will turn out very useful one day! $\endgroup$
– Henrik Schumacher
Mar 17 '18 at 11:55
$\begingroup$ Very cool! I had never heard of a Kronecker product (being a chemist who masquerades as a mathematician). Very cool to know that this type of operation has a full calculus behind it (and fun to see that this matrix built from Kronecker deltas is called a Kronecker product). $\endgroup$
– b3m2a1
Mar 21 '18 at 7:13
$\begingroup$ @b3m2a1 I'm a chemist pretending to be a mathematician (with varying degrees of success); to be perfectly fair, I didn't learn about Kronecker products until I was well out of school. ;) $\endgroup$
– J. M.'s torpor ♦
Mar 21 '18 at 7:19

Add a comment |

Stack Exchange Network

current community

your communities

more stack exchange communities

Most efficient way to build block / banded SparseArray

Some mild goal-post moving:

3 Answers 3

Your Answer

Not the answer you're looking for? Browse other questions tagged performance-tuning sparse-arrays or ask your own question.

Linked

Hot Network Questions

Most efficient way to build block / banded SparseArray

Some mild goal-post moving:

3 Answers 3

Your Answer

Sign up or log in

Post as a guest

Not the answer you're looking for? Browse other questions tagged performance-tuning sparse-arrays or ask your own question.

Linked

Related

Hot Network Questions