Fast vector - "diagonal" matrix multiplication

Question

Let $\mathbf{1}\in\mathbb{R}^d$ be a vector with all elements equal to $1$. Define: $$\mathbf{D} = \mathrm{diag}(\mathbf{1}^\top,\mathbf{1}^\top,\ldots,\mathbf{1}^\top) = \begin{bmatrix} 1 \cdots 1 & & \\ & 1\cdots 1 & \\ & & 1\cdots 1 \end{bmatrix} \in\mathbb{R}^{d\times d^2}$$

I would like to compute $\mathbf{D}\cdot \mathbf{x}$ for any vector $\mathbf{x}\in\mathbb{R}^{d^2}$. Could anybody suggest me a way to efficiently compute this product? (Since $\mathbf{D}$ has a very special structure, I guess there should be such a way).

Thank you in advance.

Answer 1

This can be interpreted as summing over an index of a tensor when the vector $x$ is reshaped into a box of numbers instead of a list. In particular, if $X$ is the $d\text{-by-}d$ folded version of $x$, then the operation you are doing is, \begin{align} Dx &= \mathrm{vec}\left((I \otimes \mathbf{1})\mathrm{vec}(X)\right) \\ &= \mathrm{vec}(\mathbf{1}^T X I) \\ &= \mathrm{vec}(\mathbf{1}^T X). \end{align}

The matlab code to do this is surprisingly simple. It is,

sum(reshape(x,d,d))'

Here's an example of it in action - you can see that it far outperforms the standard dense multiply, sparse matrix multiply, and for loop versions:

>> onesmatrixquestion
dense matrix multiply
Elapsed time is 0.000873 seconds.
sparse matrix multiply
Elapsed time is 0.000115 seconds.
for loop version
Elapsed time is 0.000154 seconds.
tensorized version
Elapsed time is 0.000018 seconds.

---

Here's the code that generated those timing results:

%onesmatrixquestion.m
d = 100;
onevec = ones(1,d);
D = kron(eye(d,d),onevec);
Dsparse = sparse(D);

x = randn(d^2,1);

disp('dense matrix multiply')
tic
aa = D*x;
toc

disp('sparse matrix multiply')
tic
bb = Dsparse*x;
toc

disp('for loop version')
tic
cc = zeros(d,1);
ind = 1;
for kk=1:d
    for jj=1:d
        cc(kk) = cc(kk) + x(ind);
        ind = ind + 1;
    end
end
toc

disp('tensorized version')
tic
dd = sum(reshape(x,d,d))';
toc

if (norm(aa - bb) > 1e-9 || norm(bb - cc) > 1e-9 ...
        || norm(cc - dd) > 1e-9)
    disp('error: different methods give different results')
end

Answer 2

This is equivalent to computing sums of consecutive contiguous subvectors of $\mathbf{x}$. You won't do much better than simple hand-coded nested loops if you have an automatically vectorizing compiler since you will probably be memory bandwidth limited for large $d$.

Answer 3

Edit: removed confusing code

---

Yes there is a shortcut, here's some Python code:

import numpy as np

d = 3


D = np.repeat(np.identity(d), d, axis=1)
print('D:')
print D
print


print('x:')
x = np.arange(d*d, dtype=float)
print x
print


print('D * x:')
print D.dot(x)
print


print('shortcut:')
print x.reshape((d, d)).sum(axis=1)
print

output:

D:
[[ 1.  1.  1.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  1.  1.  1.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  1.  1.  1.]]


x:
[ 0.  1.  2.  3.  4.  5.  6.  7.  8.]


D * x:
[  3.  12.  21.]


shortcut:
[  3.  12.  21.]