Mapping a[i+j] array to A[i,j] matrix

Question

I have an array a, and want to create a new matrix A for which A[i,j] = a[i+j]. Example:

import numpy as np

a = np.random.rand(3)
A = np.zeros((2,2));

for i in range(2):
    for j in range(2):
        A[i,j] = a[i+j]

Is there a way of doing this without a for loop? (with numpy)

Answer 1

Using stride_tricks.as_strided

This would be a perfect use case for stride_tricks:

from np.lib.stride_tricks import as_strided

Set the strides as (8, 8) (i.e. (1, 1) in terms of slots). This way we essentially map the resulting array A as i, j -> k = i + j. A more detailed description would be: we map every i, j pair to a natural number k defined by the strides as k = i*s_i + j*s_j, where s_i and s_j are the strides set to 1s, i.e. k = i + j. Thus ending up with the desired result: A[i, j] = a[k] = a[i + j].

>>> a
array([0.53954179, 0.51789927, 0.33982179])

>>> A = as_strided(a, shape=(2,2), strides=(8,8))
array([[0.53954179, 0.51789927],
       [0.51789927, 0.33982179]])

Additional considerations

A more general solution would be to get the shapes as well as the strides from a's metadata.

• The shape of A is given by (len(a)//2+1,)*2.

• As noted by @Daniel F, the memory slot size does not always equal to 8, this indeed depends on the dtype of your array. It would be better to define strides from a's strides instead: a.strides*2.

This comes down to:

>>> A = as_strided(a, shape=(len(a)//2+1,)*2, strides=a.strides*2)

---

Using grid indexing

Alternatively, you can construct a grid of coordinates (you can do so using itertools.product) then copy the appropriate values from a to A:

i, j = np.array(list(product(range(2), range(2)))).T

Then initialize A and copy:

>>> A = np.zeros((2,2))
>>> A[i, j] = a[i + j]

This will however double the memory usage, compared to the as_strided method.

Answer 2

I think the following is a bit more readable than using stride_tricks. tril_indices,triu_indices gives you the indices of the lower/upper triangle. So you can write to them directly.

A = np.zeros((2,2))

A[np.tril_indices(2)] = a
A[np.triu_indices(2)] = a

And if you're worried it might me inefficient because we write the diagonal twice. We can write the uppwer triangle and copy the upper non diagnonal values to the lower non diagonal values:

A = np.zeros((2,2))
A[np.tril_indices(2)] = a
A[np.triu_indices(2,1)] = B[np.tril_indices(2,-1)]

Answer 3

In [40]: a=np.random.randint(1,100, (9,))
In [41]: a
Out[41]: array([ 6, 35, 69, 60, 63, 51, 72, 57, 22])

broadcastable indices work nicely:

In [42]: i,j=np.arange(3)[:,None], np.arange(3)
In [43]: i,j
Out[43]: 
(array([[0],
        [1],
        [2]]),
 array([0, 1, 2]))
In [44]: i+j
Out[44]: 
array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4]])

Applying that to the 1d array:

In [45]: a[i+j].reshape(3,3)
Out[45]: 
array([[ 6, 35, 69],
       [35, 69, 60],
       [69, 60, 63]])

they could also be used for assignment

A[i,j] = a[i+j]

In [46]: _[i,j]
Out[46]: 
array([[ 6, 35, 69],
       [35, 69, 60],
       [69, 60, 63]])

as_strided works fine, but is harder to understand, and subject to error. Recent versions encourage us to use sliding_window_view.

In [47]: np.lib.stride_tricks.as_strided(a, shape=(3,3), strides=(8,8))
Out[47]: 
array([[ 6, 35, 69],
       [35, 69, 60],
       [69, 60, 63]])

As long as the strided array is just 'read', used as a view it is faster, but subsequent operations may require a copy.

In fact, for this small case, the view generation is not faster:

In [48]: timeit np.lib.stride_tricks.as_strided(a, shape=(3,3), strides=(8,8))
12.4 µs ± 43.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [49]: %%timeit
    ...: i,j=np.arange(3)[:,None], np.arange(3)
    ...: a[i+j].reshape(3,3)
    ...: 
    ...: 
8.9 µs ± 189 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

These also produce the required i,j:

np.ix_(range(3),range(3))
np.array(list(itertools.product(range(3), range(3)))).T