[FIXED] Numpy: Efficiently create this Matrix (N,3) base values of another list and repeating them

Issue

How can I create the matrix

[[a, 0, 0],
 [0, a, 0],
 [0, 0, a],
 [b, 0, 0],
 [0, b, 0],
 [0, 0, b],
 ...]

from the vector

[a, b, ...]

efficiently?

There must be a better solution than

np.squeeze(np.reshape(np.tile(np.eye(3), (len(foo), 1, 1)) * np.expand_dims(foo, (1, 2)), (1, -1, 3)))

right?

Solution

You can create a zero array in advance, and then quickly assign values by slicing:

def concated_diagonal(ar, col):
    ar = np.asarray(ar).ravel()
    size = ar.size
    ret = np.zeros((col * size, col), ar.dtype)
    for i in range(col):
        ret[i::col, i] = ar
    return ret

Test:

>>> concated_diagonal([1, 2, 3], 3)
array([[1, 0, 0],
       [0, 1, 0],
       [0, 0, 1],
       [2, 0, 0],
       [0, 2, 0],
       [0, 0, 2],
       [3, 0, 0],
       [0, 3, 0],
       [0, 0, 3]])

Note that because the number of columns you require is small, the impact of the relatively slow Python level for loop is acceptable:

%timeit concated_diagonal(np.arange(1_000_000), 3)
17.1 ms ± 84.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answered By - Mechanic Pig